Add free_vectors_space() to finally block of run_tasks handler

Signed-off-by: Rohan Chitale <[email protected]>

Author: rchitale7

.github/workflows/publish_remote_core_image.yml

@@ -23,8 +23,6 @@ jobs:
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
-        with:
-          submodules: 'recursive'
 
       - name: Build Docker Image
         run : |
@@ -55,9 +53,7 @@ jobs:
           docker push opensearchstaging/remote-vector-index-builder:core-1.0.0
           docker push opensearchstaging/remote-vector-index-builder:core-latest
 
-      - name: Runner Cleanups
+      - name: Docker logout
         if: always()
         run: |
-          docker logout
-          docker system prune -a -f
-          rm -rf ${{ github.workspace }}/*
+          docker logout

remote_vector_index_builder/core/common/models/vectors_dataset.py

@@ -31,8 +31,16 @@ class VectorsDataset:
 
     def free_vectors_space(self):
         """Free up memory by deleting the vectors and document IDs arrays."""
-        del self.vectors
-        del self.doc_ids
+
+        try:
+            del self.vectors
+        except AttributeError:
+            pass
+        try:
+            del self.doc_ids
+        except AttributeError:
+            pass
+        return
 
     @staticmethod
     def get_numpy_dtype(dtype: DataType):

remote_vector_index_builder/core/tasks.py

@@ -89,6 +89,8 @@ def run_tasks(
     with tempfile.TemporaryDirectory() as temp_dir, BytesIO() as vector_buffer, BytesIO() as doc_id_buffer:
         if object_store_config is None:
             object_store_config = {}
+
+        vectors_dataset = None
         try:
             logger.info(
                 f"Starting task execution for vector path: {index_build_params.vector_path}"
@@ -137,6 +139,9 @@ def run_tasks(
         except Exception as e:
             logger.error(f"Error running tasks: {e}")
             return TaskResult(error=str(e))
+        finally:
+            if vectors_dataset is not None:
+                vectors_dataset.free_vectors_space()
 
 
 def build_index(

test_remote_vector_index_builder/test_core/common/models/test_vectors_dataset.py

@@ -44,6 +44,16 @@ def test_free_vectors_space(vectors_dataset):
         _ = vectors_dataset.doc_ids
 
 
+def test_free_vectors_space_when_vectors_and_doc_ids_already_deleted(vectors_dataset):
+    vectors_dataset.free_vectors_space()
+    with pytest.raises(AttributeError):
+        _ = vectors_dataset.vectors
+    with pytest.raises(AttributeError):
+        _ = vectors_dataset.doc_ids
+    # test idempotency
+    vectors_dataset.free_vectors_space()
+
+
 @pytest.mark.parametrize(
     "dtype, expected",
     [

test_remote_vector_index_builder/test_core/test_tasks.py

@@ -170,7 +170,7 @@ def test_successful_task_execution(index_build_parameters, mock_vectors_dataset)
         mock_create_dataset.assert_called_once()
         mock_build_index.assert_called_once()
         mock_upload_index.assert_called_once()
-        mock_vectors_dataset.free_vectors_space.assert_called_once()
+        assert mock_vectors_dataset.free_vectors_space.call_count == 2
         mock_os_remove.assert_called_once()
         mock_os_makedirs.assert_called_once()
 
@@ -207,7 +207,7 @@ def test_successful_task_execution_with_object_store_config(
         )
         mock_build_index.assert_called_once()
         mock_upload_index.assert_called_once()
-        mock_vectors_dataset.free_vectors_space.assert_called_once()
+        assert mock_vectors_dataset.free_vectors_space.call_count == 2
         mock_os_remove.assert_called_once()
         mock_os_makedirs.assert_called_once()
 
@@ -222,3 +222,27 @@ def test_create_vectors_dataset_failure(index_build_parameters):
         assert isinstance(result, TaskResult)
         assert result.file_name is None
         assert result.error == "Dataset creation failed"
+
+
+def test_build_index_failure(index_build_parameters, mock_vectors_dataset):
+    with patch("core.tasks.create_vectors_dataset") as mock_create_dataset, patch(
+        "core.tasks.build_index"
+    ) as mock_build_index, patch("os.makedirs") as mock_os_makedirs:
+
+        # Setup mocks
+        mock_create_dataset.return_value = mock_vectors_dataset
+        mock_build_index.side_effect = Exception("Index building failed")
+
+        # Execute function
+        result = run_tasks(index_build_parameters)
+
+        # Verify success
+        assert isinstance(result, TaskResult)
+        assert result.file_name is None
+        assert result.error == "Index building failed"
+
+        # Verify mock calls
+        mock_create_dataset.assert_called_once()
+        mock_build_index.assert_called_once()
+        mock_vectors_dataset.free_vectors_space.assert_called_once()
+        mock_os_makedirs.assert_called_once()