Fix Memory cleanup issues for GPU and CPU index (#44)

Signed-off-by: Rajvaibhav Rahane <[email protected]>

Author: Rajrahane

remote_vector_index_builder/core/common/models/index_builder/faiss/faiss_gpu_index_cagra_builder.py

@@ -195,14 +195,15 @@ def build_gpu_index(
             raise Exception(f"Failed to create faiss GPU index config: {str(e)}") from e
 
         try:
-            gpu_resources = faiss.StandardGpuResources()
-
             # Configure the distance metric
             metric = configure_metric(space_type)
 
             # Create GPU CAGRA index with specified configuration
             faiss_gpu_index = faiss.GpuIndexCagra(
-                gpu_resources, dataset_dimension, metric, faiss_gpu_index_config
+                faiss.StandardGpuResources(),
+                dataset_dimension,
+                metric,
+                faiss_gpu_index_config,
             )
 
             # Create ID mapping layer to preserve document IDs
@@ -218,7 +219,9 @@ def build_gpu_index(
         except Exception as e:
             if faiss_gpu_index is not None:
                 faiss_gpu_index.thisown = True
-                del faiss_gpu_index
+                faiss_gpu_index.__swig_destroy__(faiss_gpu_index)
             if faiss_index_id_map is not None:
-                del faiss_index_id_map
+                faiss_index_id_map.thisown = True
+                faiss_index_id_map.own_fields = False
+                faiss_index_id_map.__swig_destroy__(faiss_index_id_map)
             raise Exception(f"Failed to create faiss GPU index: {str(e)}") from e

remote_vector_index_builder/core/common/models/index_builder/faiss/faiss_index_hnsw_cagra_builder.py

@@ -78,17 +78,21 @@ def convert_gpu_to_cpu_index(
             cpu_index.base_level_only = self.base_level_only
 
             # Copy GPU index to CPU index
-            gpu_index = faiss_gpu_build_index_output.gpu_index
-            gpu_index.copyTo(cpu_index)
+            faiss_gpu_build_index_output.gpu_index.copyTo(cpu_index)
+
+            # Remove reference of GPU Index from the IndexIDMap
+            faiss_gpu_build_index_output.index_id_map.index = None
 
             # Update the ID map index with the CPU index
             index_id_map = faiss_gpu_build_index_output.index_id_map
+
+            # Remove reference of the IndexIDMap from the GPU Build Index Output before cleanup
             faiss_gpu_build_index_output.index_id_map = None
+
             index_id_map.index = cpu_index
 
             # Free memory taken by GPU Index
-            del gpu_index
-            del faiss_gpu_build_index_output
+            faiss_gpu_build_index_output.cleanup()
 
             return FaissCpuBuildIndexOutput(
                 cpu_index=cpu_index, index_id_map=index_id_map
@@ -121,7 +125,7 @@ def write_cpu_index(
                 cpu_build_index_output.index_id_map, cpu_index_output_file_path
             )
             # Free memory taken by CPU Index
-            del cpu_build_index_output
+            cpu_build_index_output.cleanup()
         except IOError as io_error:
             raise Exception(
                 f"Failed to write index to file {cpu_index_output_file_path}: {str(io_error)}"

remote_vector_index_builder/core/common/models/index_builder/response/faiss_cpu_build_index_output.py

@@ -21,20 +21,40 @@ class FaissCpuBuildIndexOutput:
     index_id_map: faiss.IndexIDMap
 
     def __del__(self):
+        self.cleanup()
+
+    def cleanup(self):
         """
-        Destructor to clean up FAISS resources.
+        Clean up method for FAISS resources.
         Ensures memory is properly freed when the object is destroyed.
 
         The method handles cleanup by
         explicitly deleting the internal Index and Vectors data
         """
         try:
+            # Cleanup the CPU Index explicitly after it is succussfully persisted to disc
+            # OR if the orchestrator fails after creating the CPU Index
+            # and before the CPU Index is persisted to disc
             if self.cpu_index:
                 # Delete the internal Index
-                del self.cpu_index
+                self.cpu_index.thisown = True
+
+                cpu_index = self.cpu_index
+                self.cpu_index = None
+                cpu_index.__swig_destroy__(cpu_index)
+
+            # Delete the IndexIDMap not containing any underlying Index
+            # This block runs after the CPU Index is cleaned up
+            # if the CPU Index is succussfully persisted to disc
+            # OR if the orchestrator fails after creating the CPU Index
+            # and before the CPU Index is persisted to disc
             if self.index_id_map:
                 # Delete the vectors, vector ids stored as part of the IndexIDMap
-                self.index_id_map.own_fields = True
-                del self.index_id_map
+                self.index_id_map.own_fields = False
+                self.index_id_map.thisown = True
+
+                index_id_map = self.index_id_map
+                self.index_id_map = None
+                index_id_map.__swig_destroy__(index_id_map)
         except Exception as e:
             print(f"Error during cleanup of FaissCpuBuildIndexOutput: {e}")

remote_vector_index_builder/core/common/models/index_builder/response/faiss_gpu_build_index_output.py

@@ -20,21 +20,42 @@ class FaissGpuBuildIndexOutput:
     index_id_map: faiss.IndexIDMap
 
     def __del__(self):
+        self.cleanup()
+
+    def cleanup(self):
         """
-        Destructor to clean up FAISS resources.
+        Clean up method for FAISS resources.
         Ensures memory is properly freed when the object is destroyed.
 
         The method handles cleanup by
         explicitly deleting the internal Index and Vectors data
         """
         try:
+            # Cleanup the GPU Index explicitly if the CPU Index was successfully created
+            # OR if the orchestrator fails after creating a GPU Index and
+            # before the CPU Index is created, and replaced GPU Index in the IndexIDMap
             if self.gpu_index:
                 # Delete the internal Index
-                del self.gpu_index
+
+                if self.index_id_map:
+                    self.index_id_map.index = None
+                self.gpu_index.thisown = True
+
+                gpu_index = self.gpu_index
+                self.gpu_index = None
+                gpu_index.__swig_destroy__(gpu_index)
+
+            # A reference to IndexIDMap is maintained until the underlying GPU Index is replaced by a CPU Index
+            # This block runs if the orchestrator fails before the successful replacement
             if self.index_id_map:
-                self.index_id_map.index = None
                 # Delete the vectors, vector ids stored as part of the IndexIDMap
-                self.index_id_map.own_fields = True
-                del self.index_id_map
+                self.index_id_map.thisown = True
+                self.index_id_map.own_fields = False
+
+                self.index_id_map.index = None
+
+                index_id_map = self.index_id_map
+                self.index_id_map = None
+                index_id_map.__swig_destroy__(index_id_map)
         except Exception as e:
             print(f"Error during cleanup of FaissGpuBuildIndexOutput: {e}")

remote_vector_index_builder/core/index_builder/faiss/faiss_index_build_service.py

@@ -107,14 +107,14 @@ def build_index(
             # Clean up GPU Index Response if orchestrator failed after GPU Index Creation
             if faiss_gpu_build_index_output is not None:
                 try:
-                    del faiss_gpu_build_index_output
+                    faiss_gpu_build_index_output.cleanup()
                 except Exception as e:
                     print(f"Warning: Failed to clean up GPU index response: {str(e)}")
 
             # Clean up CPU Index Response if orchestrator failed after CPU Index Creation
             if faiss_cpu_build_index_output is not None:
                 try:
-                    del faiss_cpu_build_index_output
+                    faiss_cpu_build_index_output.cleanup()
                 except Exception as e:
                     print(f"Warning: Failed to clean up CPU index response: {str(e)}")
             raise Exception(