From db02f3f02aacf21ca2f5878a1012feb54d7af447 Mon Sep 17 00:00:00 2001
From: Tamas Bela Feher <tfeher@nvidia.com>
Date: Thu, 30 May 2024 21:21:48 +0200
Subject: [PATCH] Fix OOM handling

---
 cpp/src/neighbors/detail/cagra/cagra_build.cuh | 9 ++++++---
 cpp/src/neighbors/nn_descent.cu                | 4 ++++
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh
index a50ac2e65..43e720077 100644
--- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh
+++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh
@@ -480,9 +480,12 @@ index<T, IdxT> build(
     RAFT_LOG_DEBUG("Insufficient GPU memory to construct CAGRA index with dataset on GPU");
     // We just add the graph. User is expected to update dataset separately (e.g allocating in
     // managed memory).
-    index<T, IdxT> idx(res, params.metric);
-    idx.update_graph(res, raft::make_const_mdspan(cagra_graph.view()));
-    return idx;
+  } catch (raft::logic_error& e) {
+    // The memory error can also manifest as logic_error.
+    RAFT_LOG_DEBUG("Insufficient GPU memory to construct CAGRA index with dataset on GPU");
   }
+  index<T, IdxT> idx(res, params.metric);
+  idx.update_graph(res, raft::make_const_mdspan(cagra_graph.view()));
+  return idx;
 }
 }  // namespace cuvs::neighbors::cagra::detail
diff --git a/cpp/src/neighbors/nn_descent.cu b/cpp/src/neighbors/nn_descent.cu
index 3fa87d7ab..548ed00f3 100644
--- a/cpp/src/neighbors/nn_descent.cu
+++ b/cpp/src/neighbors/nn_descent.cu
@@ -18,6 +18,7 @@
 #include <cuvs/neighbors/nn_descent.hpp>
 #include <raft/core/logger.hpp>
 
+using namespace raft;
 namespace cuvs::neighbors::nn_descent {
 
 /**
@@ -53,6 +54,9 @@ bool has_enough_device_memory(raft::resources const& res,
   } catch (std::bad_alloc& e) {
     RAFT_LOG_DEBUG("Insufficient memory for NN descent");
     return false;
+  } catch (raft::logic_error& e) {
+    RAFT_LOG_DEBUG("Insufficient memory for NN descent (logic error)");
+    return false;
   }
 }