Introduce index serialization mode option to support serializing graph in memory (#124)

Signed-off-by: Rohan Chitale <rchital@amazon.com>

Author: rchitale7

DEVELOPER_GUIDE.md

@@ -177,7 +177,7 @@ Follow the steps below to use run the API image locally. Note that s3 is current
     ```
 7. Run the docker image: 
     ```
-    docker run --gpus all -p 80:1025 opensearchstaging/remote-vector-index-builder:api-latest
+    docker run -e AWS_DEFAULT_REGION=us-east-1 --gpus all -p 80:1025 opensearchstaging/remote-vector-index-builder:api-latest
     ```
 8. In a separate terminal, issue a build request:
     ```

benchmarking/utils/common_utils.py

@@ -5,8 +5,7 @@
 import sys
 
 from benchmarking.dataset.dataset import HDF5DataSet
-import config
-
+import benchmarking.config as config
 
 def recall_at_r(results, neighbor_dataset: HDF5DataSet, r, k, query_count):
     """

e2e/api/docker-compose.yml

@@ -25,6 +25,7 @@ services:
       - AWS_ACCESS_KEY_ID=test
       - AWS_SECRET_ACCESS_KEY=test
       - AWS_DEFAULT_REGION=us-east-1
+      - LOG_LEVEL=DEBUG
     ports:
       - "8080:1025"
 

remote_vector_index_builder/app/services/index_builder.py

@@ -10,6 +10,7 @@
 from typing import Optional, Tuple
 from app.models.workflow import BuildWorkflow
 from core.object_store.s3.s3_object_store_config import S3ClientConfig
+from core.common.models import IndexSerializationMode
 from core.tasks import run_tasks
 
 logger = logging.getLogger(__name__)
@@ -36,6 +37,9 @@ def build_index(
                 - Error message if failed, None otherwise
         """
         s3_endpoint_url = os.environ.get("S3_ENDPOINT_URL", None)
+        index_serialization_mode = os.environ.get(
+            "INDEX_SERIALIZATION_MODE", IndexSerializationMode.DISK
+        )
         result = run_tasks(
             workflow.index_build_parameters,
             {
@@ -44,6 +48,7 @@ def build_index(
                     endpoint_url=s3_endpoint_url,
                 ),
             },
+            index_serialization_mode,
         )
         if not result.file_name:
             return False, None, result.error

remote_vector_index_builder/core/common/models/__init__.py

@@ -7,7 +7,13 @@
 
 from .index_build_parameters import SpaceType
 from .index_build_parameters import IndexBuildParameters
+from .index_build_parameters import IndexSerializationMode
 from .vectors_dataset import VectorsDataset
 
 
-__all__ = ["SpaceType", "IndexBuildParameters", "VectorsDataset"]
+__all__ = [
+    "SpaceType",
+    "IndexBuildParameters",
+    "VectorsDataset",
+    "IndexSerializationMode",
+]

remote_vector_index_builder/core/common/models/index_build_parameters.py

@@ -121,6 +121,18 @@ class IndexParameters(BaseModel):
     )
 
 
+class IndexSerializationMode(str, Enum):
+    """Serialization mode for the constructed vector index
+
+    Attributes:
+        MEMORY: Store index in memory after processing
+        DISK: Store index on disk after processing
+    """
+
+    MEMORY = "memory"
+    DISK = "disk"
+
+
 class IndexBuildParameters(BaseModel):
     """Parameters required for building a vector index.
 

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

@@ -149,41 +149,3 @@ def convert_gpu_to_cpu_index(
             return self._do_convert_gpu_to_cpu_index(faiss_gpu_build_index_output)
 
         return self._do_convert_gpu_to_cpu_binary_index(faiss_gpu_build_index_output)
-
-    def write_cpu_index(
-        self,
-        cpu_build_index_output: FaissCpuBuildIndexOutput,
-        cpu_index_output_file_path: str,
-    ) -> None:
-        """
-        Method to write the CPU index and vector dataset id mapping to persistent local file path
-        for uploading later to remote object store.
-        Uses faiss write_index library method to achieve this
-
-        Args:
-        cpu_build_index_output (FaissCpuBuildIndexOutput): A datamodel containing the created GPU Faiss Index
-        and dataset Vector Ids components
-        cpu_index_output_file_path (str): File path to persist Index-Vector IDs map to
-        """
-        try:
-
-            # TODO: Investigate what issues may arise while writing index to local file
-            # Write the final cpu index - vectors id mapping to disk
-            if self.vector_dtype != DataType.BINARY:
-                faiss.write_index(
-                    cpu_build_index_output.index_id_map, cpu_index_output_file_path
-                )
-            else:
-                faiss.write_index_binary(
-                    cpu_build_index_output.index_id_map, cpu_index_output_file_path
-                )
-            # Free memory taken by CPU Index
-            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)}"
-            ) from io_error
-        except Exception as e:
-            raise Exception(
-                f"Unexpected error while writing index to file: {str(e)}"
-            ) from e

remote_vector_index_builder/core/common/models/index_builder/faiss_cpu_index_builder.py

@@ -16,7 +16,6 @@ class FaissCPUIndexBuilder(ABC):
     """
     Base class for CPU Index Configuration
     Also exposes methods to convert gpu index to cpu index from the configuration
-    and writing cpu index to file
     """
 
     @abstractmethod
@@ -37,19 +36,3 @@ def convert_gpu_to_cpu_index(
         """
 
         pass
-
-    @abstractmethod
-    def write_cpu_index(
-        self,
-        cpu_build_index_output: FaissCpuBuildIndexOutput,
-        cpu_index_output_file_path: str,
-    ) -> None:
-        """
-        Implement this abstract method to write the cpu index to specified output file path
-
-        Args:
-        cpu_build_index_output (FaissCpuBuildIndexOutput): A datamodel containing the created GPU Faiss Index
-        and dataset Vector Ids components
-        cpu_index_output_file_path (str): File path to persist Index-Vector IDs map to
-        """
-        pass

remote_vector_index_builder/core/index_builder/faiss/faiss_index_build_service.py

@@ -9,6 +9,7 @@
 from core.common.models import (
     IndexBuildParameters,
     VectorsDataset,
+    IndexSerializationMode,
 )
 
 from core.common.models.index_builder.faiss import (
@@ -20,10 +21,15 @@
     get_omp_num_threads,
 )
 from core.common.models.index_build_parameters import DataType
-from core.common.models.index_builder import CagraGraphBuildAlgo
+from core.common.models.index_builder import (
+    CagraGraphBuildAlgo,
+    FaissCpuBuildIndexOutput,
+)
 from core.index_builder.interface import IndexBuildService
 from timeit import default_timer as timer
 
+from typing import Union, Any
+from io import BytesIO
 import logging
 
 logger = logging.getLogger(__name__)
@@ -42,19 +48,19 @@ def build_index(
         self,
         index_build_parameters: IndexBuildParameters,
         vectors_dataset: VectorsDataset,
-        cpu_index_output_file_path: str,
-    ) -> None:
+    ) -> Any:
         """
         Orchestrates the workflow of
         - creating a GPU Index for the specified vectors dataset,
         - converting into CPU compatible Index
-        - and writing the CPU Index to disc
         Uses the faiss library methods to achieve this.
 
         Args:
             vectors_dataset: The set of vectors to index
             index_build_parameters: The API Index Build parameters
-            cpu_index_output_file_path: The complete file path on disc to write the cpuIndex to.
+
+        Returns:
+            The CPU compatible Index
         """
         faiss_gpu_index_cagra_builder = None
         faiss_index_hnsw_cagra_builder = None
@@ -139,17 +145,7 @@ def build_index(
                 f"{index_conversion_time:.2f} seconds"
             )
 
-            # Step 3: Write CPU Index to persistent storage
-            t1 = timer()
-            faiss_index_hnsw_cagra_builder.write_cpu_index(
-                faiss_cpu_build_index_output, cpu_index_output_file_path
-            )
-            t2 = timer()
-            index_write_time = t2 - t1
-            logger.debug(
-                f"Index write time for vector path {index_build_parameters.vector_path}: "
-                f"{index_write_time:.2f} seconds"
-            )
+            return faiss_cpu_build_index_output
 
         except Exception as exception:
             # Clean up GPU Index Response if orchestrator failed after GPU Index Creation
@@ -161,17 +157,70 @@ def build_index(
                         f"Failed to clean up GPU index response for vector path "
                         f"{index_build_parameters.vector_path}: {e}"
                     )
+            raise Exception(
+                f"Faiss Index Build Service build_index workflow failed: {exception}"
+            ) from exception
 
-            # Clean up CPU Index Response if orchestrator failed after CPU Index Creation
-            if faiss_cpu_build_index_output is not None:
-                try:
-                    faiss_cpu_build_index_output.cleanup()
-                except Exception as e:
-                    logger.error(
-                        f"Failed to clean up CPU index response for vector path "
-                        f"{index_build_parameters.vector_path}: {e}"
-                    )
+    def write_cpu_index(
+        self,
+        cpu_build_index_output: FaissCpuBuildIndexOutput,
+        index_build_parameters: IndexBuildParameters,
+        index_serialization_mode: IndexSerializationMode,
+        output_destination: Union[str, BytesIO],
+    ) -> None:
+        """
+        Method to write the CPU index and vector dataset id mapping to storage destination,
+        for uploading later to remote object store.
+        Uses faiss write_index library method to achieve this
+
+        Args:
+        cpu_build_index_output (FaissCpuBuildIndexOutput): A datamodel containing the created GPU Faiss Index
+            and dataset Vector Ids components
+        index_build_parameters: The API Index Build parameters
+        index_serialization_mode: The serialization mode for the index
+        output_destination (Union[str, BytesIO]): Output destination for the index.
+            - str: File path for disk writing
+            - BytesIO: Existing buffer to write to
+        """
+        try:
+            t1 = timer()
+            writer = None
+            if (
+                index_serialization_mode == IndexSerializationMode.MEMORY
+                and isinstance(
+                    output_destination, BytesIO
+                )  # for resolving mypy errors, we add this check
+            ):
+                # Use a faiss callback to serialize directly to the buffer
+                # We use a faiss callback instead of faiss.serialize_index
+                # to avoid the extra memory overhead of the c++ vector data structure
+                # and numpy arrays created in serialize_index method
+                writer = faiss.PyCallbackIOWriter(output_destination.write)
+            else:
+                # Otherwise, treat the output destination like a file path
+                writer = output_destination
+            if index_build_parameters.data_type != DataType.BINARY:
+                faiss.write_index(cpu_build_index_output.index_id_map, writer)
+            else:
+                faiss.write_index_binary(cpu_build_index_output.index_id_map, writer)
+            # Free memory taken by CPU Index
+            cpu_build_index_output.cleanup()
+            t2 = timer()
+            index_write_time = t2 - t1
+            logger.debug(
+                f"Index write time for vector path {index_build_parameters.vector_path}: "
+                f"{index_write_time:.2f} seconds"
+            )
+        except Exception as exception:
+            # Clean up CPU Index Response if write failed
+            try:
+                cpu_build_index_output.cleanup()
+            except Exception as e:
+                logger.error(
+                    f"Failed to clean up CPU index response for vector path "
+                    f"{index_build_parameters.vector_path}: {e}"
+                )
 
             raise Exception(
-                f"Faiss Index Build Service build_index workflow failed: {exception}"
+                f"Faiss Index Build Service write_index workflow failed: {exception}"
             ) from exception

remote_vector_index_builder/core/index_builder/interface/index_build_service.py

@@ -6,23 +6,22 @@
 # compatible open source license.
 
 from abc import ABC, abstractmethod
-
+from typing import Any
 from core.common.models import VectorsDataset, IndexBuildParameters
 
 
 class IndexBuildService(ABC):
     """
     The Index Build Service orchestrates the workflow of building a vector search index
-    New engines extending this class must call the necessary worlflow steps in the build_index method
+    New engines extending this class must call the necessary workflow steps in the build_index method
     """
 
     @abstractmethod
     def build_index(
         self,
         index_build_parameters: IndexBuildParameters,
         vectors_dataset: VectorsDataset,
-        cpu_index_output_file_path: str,
-    ) -> None:
+    ) -> Any:
         """
         Implement this abstract method orchestrating an index build for the specified vectors dataset
         and input index build parameters