Increase the minimum doc count to 5 (#52)

Signed-off-by: Rohan Chitale <[email protected]>
Co-authored-by: Navneet Verma <[email protected]>

Author: rchitale7

API.md

@@ -15,7 +15,7 @@ Request Parameters
     "doc_id_path" : "DocIdPath", // File path all doc Ids are written to
     "tenant_id": "UniqueClusterID", // Unique identifier for the cluster making the request
     "dimension": 768,
-    "doc_count": 1000000,
+    "doc_count": 1000000, // Must be greater than 4
     "data_type": "float",
     "engine": "faiss",
     "index_parameters": {
@@ -41,6 +41,7 @@ Request Response:
 * The top two parameters are derived by the KNN plugin from the customer's repo setting. `container_name` is used to specifically refer to the "bucket" (or non-S3 equivalent), rather than the name of the repository itself.
 * Tenant ID is included to be used for billing, authorization, etc
 * Dimension, doc count, and data type are specifically placed on the first JSON level so the build service can quickly use them first to calculate workload size
+* Doc count must be greater than 4, since doc counts of 1, 3, and 4 have been observed to consistently fail during index builds
 * Engine: If in the future we have different workers hosting different engines, like Lucene etc. this parameter will act as an extension point. This is not a required parameter.
 * Space type is inside `index_parameters` since it is only used in index creation
 * By including an `algorithm` field in index parameters, we leave the door open for IVF, future algorithms

remote_vector_index_builder/app/models/workflow.py

@@ -6,7 +6,7 @@
 # compatible open source license.
 
 from pydantic import BaseModel
-from core.common.models.index_build_parameters import IndexBuildParameters
+from core.common.models import IndexBuildParameters
 
 
 class BuildWorkflow(BaseModel):

remote_vector_index_builder/core/common/models/__init__.py

@@ -6,28 +6,8 @@
 # compatible open source license.
 
 from .index_build_parameters import SpaceType
-
-from .index_builder.cagra_graph_build_algo import CagraGraphBuildAlgo
-
-
 from .index_build_parameters import IndexBuildParameters
 from .vectors_dataset import VectorsDataset
-from .index_builder.response.faiss_gpu_build_index_output import (
-    FaissGpuBuildIndexOutput,
-)
-from .index_builder.response.faiss_cpu_build_index_output import (
-    FaissCpuBuildIndexOutput,
-)
-from .index_builder.faiss_gpu_index_builder import FaissGPUIndexBuilder
-from .index_builder.faiss_cpu_index_builder import FaissCPUIndexBuilder
 
-__all__ = [
-    "SpaceType",
-    "CagraGraphBuildAlgo",
-    "IndexBuildParameters",
-    "VectorsDataset",
-    "FaissGpuBuildIndexOutput",
-    "FaissCpuBuildIndexOutput",
-    "FaissGPUIndexBuilder",
-    "FaissCPUIndexBuilder",
-]
+
+__all__ = ["SpaceType", "IndexBuildParameters", "VectorsDataset"]

remote_vector_index_builder/core/common/models/index_build_parameters.py

@@ -144,7 +144,7 @@ class IndexBuildParameters(BaseModel):
     doc_id_path: str
     tenant_id: str = ""
     dimension: int = Field(gt=0)
-    doc_count: int = Field(gt=1)
+    doc_count: int = Field(gt=4)
     data_type: DataType = DataType.FLOAT
     engine: Engine = Engine.FAISS
     index_parameters: IndexParameters = Field(default_factory=IndexParameters)

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

@@ -0,0 +1,24 @@
+# Copyright OpenSearch Contributors
+# SPDX-License-Identifier: Apache-2.0
+#
+# The OpenSearch Contributors require contributions made to
+# this file be licensed under the Apache-2.0 license or a
+# compatible open source license.
+
+from .cagra_graph_build_algo import CagraGraphBuildAlgo
+from .response.faiss_gpu_build_index_output import (
+    FaissGpuBuildIndexOutput,
+)
+from .response.faiss_cpu_build_index_output import (
+    FaissCpuBuildIndexOutput,
+)
+from .faiss_gpu_index_builder import FaissGPUIndexBuilder
+from .faiss_cpu_index_builder import FaissCPUIndexBuilder
+
+__all__ = [
+    "CagraGraphBuildAlgo",
+    "FaissGpuBuildIndexOutput",
+    "FaissCpuBuildIndexOutput",
+    "FaissGPUIndexBuilder",
+    "FaissCPUIndexBuilder",
+]

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

@@ -10,12 +10,15 @@
 from dataclasses import field, dataclass
 
 from core.common.models import (
-    CagraGraphBuildAlgo,
     VectorsDataset,
     SpaceType,
+)
+from core.common.models.index_builder import (
+    CagraGraphBuildAlgo,
     FaissGpuBuildIndexOutput,
     FaissGPUIndexBuilder,
 )
+
 from core.index_builder.index_builder_utils import configure_metric
 from .ivf_pq_build_cagra_config import IVFPQBuildCagraConfig
 from .ivf_pq_search_cagra_config import IVFPQSearchCagraConfig

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

@@ -8,7 +8,7 @@
 import faiss
 from dataclasses import dataclass
 from typing import Dict, Any
-from core.common.models import (
+from core.common.models.index_builder import (
     FaissCpuBuildIndexOutput,
     FaissGpuBuildIndexOutput,
     FaissCPUIndexBuilder,

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

@@ -6,7 +6,7 @@
 # compatible open source license.
 
 from abc import ABC, abstractmethod
-from core.common.models import (
+from core.common.models.index_builder import (
     FaissCpuBuildIndexOutput,
     FaissGpuBuildIndexOutput,
 )

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

@@ -10,6 +10,9 @@
 from core.common.models import (
     SpaceType,
     VectorsDataset,
+)
+
+from core.common.models.index_builder import (
     FaissGpuBuildIndexOutput,
 )
 

remote_vector_index_builder/core/index_builder/index_builder_utils.py

@@ -10,7 +10,7 @@
 import faiss
 
 
-from core.common.models.index_build_parameters import (
+from core.common.models import (
     SpaceType,
 )
 

remote_vector_index_builder/core/object_store/object_store_factory.py

@@ -8,7 +8,7 @@
 from typing import Any, Dict
 
 from core.common.exceptions import UnsupportedObjectStoreTypeError
-from core.common.models.index_build_parameters import IndexBuildParameters
+from core.common.models import IndexBuildParameters
 from core.object_store.object_store import ObjectStore
 from core.object_store.s3.s3_object_store import S3ObjectStore
 from core.object_store.types import ObjectStoreType

remote_vector_index_builder/core/object_store/s3/s3_object_store.py

@@ -19,7 +19,7 @@
 from botocore.config import Config
 from botocore.exceptions import ClientError
 from core.common.exceptions import BlobError
-from core.common.models.index_build_parameters import IndexBuildParameters
+from core.common.models import IndexBuildParameters
 from core.object_store.object_store import ObjectStore
 
 logger = logging.getLogger(__name__)

remote_vector_index_builder/core/tasks.py

@@ -33,8 +33,8 @@
 from io import BytesIO
 from typing import Any, Dict, Optional
 
-from core.common.models.index_build_parameters import IndexBuildParameters
-from core.common.models.vectors_dataset import VectorsDataset
+from core.common.models import IndexBuildParameters
+from core.common.models import VectorsDataset
 from core.index_builder.faiss.faiss_index_build_service import FaissIndexBuildService
 from core.object_store.object_store_factory import ObjectStoreFactory
 

test_remote_vector_index_builder/conftest.py

@@ -5,187 +5,29 @@
 # this file be licensed under the Apache-2.0 license or a
 # compatible open source license.
 
-import sys
-from types import ModuleType
-from unittest.mock import Mock
-
 import pytest
-
-
-class DeletionTracker:
-    """Helper class to track object deletions"""
-
-    def __init__(self):
-        self.deleted_objects = set()
-
-    def mark_deleted(self, obj_id):
-        self.deleted_objects.add(obj_id)
-
-    def is_deleted(self, obj_id):
-        return obj_id in self.deleted_objects
-
-    def reset(self):
-        self.deleted_objects.clear()
-
-
-# Create global deletion tracker
-_deletion_tracker = DeletionTracker()
+from core.common.models.index_build_parameters import (
+    AlgorithmParameters,
+    IndexBuildParameters,
+    IndexParameters,
+    SpaceType,
+)
 
 
 @pytest.fixture
-def deletion_tracker():
-    """Fixture to provide access to deletion tracker"""
-    _deletion_tracker.reset()  # Reset before each test
-    return _deletion_tracker
-
-
-@pytest.fixture(autouse=True)
-def reset_deletion_tracker():
-    """Reset deletion tracker before each test"""
-    _deletion_tracker.reset()
-    yield
-
-
-class MockGpuIndexCagra:
-    """Mock for faiss.GpuIndexCagra with deletion tracking"""
-
-    def __init__(self, *args, **kwargs):
-        self.id = id(self)
-        self.thisown = False
-        self.args = args
-        self.kwargs = kwargs
-
-    def __del__(self):
-        print("deleting MockGpuIndexCagra:", self.id)
-        _deletion_tracker.mark_deleted(self.id)
-
-    @property
-    def is_deleted(self):
-        return _deletion_tracker.is_deleted(self.id)
-
-    def copyTo(self, cpu_index):
-        """Mock implementation of copyTo method"""
-        if not isinstance(cpu_index, MockIndexHNSWCagra):
-            raise TypeError("Target must be IndexHNSWCagra")
-        # Simulate copying data to CPU index
-        return True
-
-
-class MockIndexIDMap:
-    """Mock for faiss.IndexIDMap with deletion tracking"""
-
-    def __init__(self, *args, **kwargs):
-        self.id = id(self)
-        self.own_fields = False
-        self.index = None
-        self.args = args
-        self.kwargs = kwargs
-
-    def __del__(self):
-        print("deleting MockIndexIDMap:", self.id)
-        _deletion_tracker.mark_deleted(self.id)
-
-    @property
-    def is_deleted(self):
-        return _deletion_tracker.is_deleted(self.id)
-
-    def add_with_ids(self, vectors, ids):
-        pass
-
-
-class MockIndexHNSWCagra(Mock):
-    """Mock for faiss.IndexHNSWCagra"""
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.hnsw = Mock()
-        self.base_level_only = True
-
-    def __del__(self):
-        _deletion_tracker.mark_deleted(self.id)
-
-    @property
-    def is_deleted(self):
-        return _deletion_tracker.is_deleted(self.id)
-
-
-class MockIVFPQBuildCagraConfig:
-    """Mock class for faiss.IVFPQBuildCagraConfig"""
-
-    def __init__(self):
-        self.n_lists = 1024
-        self.kmeans_n_iters = 20
-        self.kmeans_trainset_fraction = 0.5
-        self.pq_bits = 8
-        self.pq_dim = 0
-        self.conservative_memory_allocation = True
-
-
-class MockIVFPQSearchCagraConfig:
-    """Mock class for faiss.IVFPQSearchCagraConfig"""
-
-    def __init__(self):
-        self.n_probes = 20
-
-
-class MockGpuIndexCagraConfig:
-    """Mock class for faiss.GpuIndexCagraConfig"""
-
-    def __init__(self):
-        self.intermediate_graph_degree = 64
-        self.graph_degree = 32
-        self.store_dataset = False
-        self.device = 0
-        self.refine_rate = 2.0
-        self.build_algo = None
-        self.ivf_pq_build_config = None
-        self.ivf_pq_search_config = None
-
-
-class FaissMock(ModuleType):
-    """Complete mock for faiss module"""
-
-    def __init__(self):
-        super().__init__("faiss")
-        # Classes
-        self.StandardGpuResources = Mock()
-        self.GpuIndexCagra = MockGpuIndexCagra
-        self.IndexIDMap = MockIndexIDMap
-        self.IndexHNSWCagra = MockIndexHNSWCagra
-        self.IVFPQBuildCagraConfig = MockIVFPQBuildCagraConfig
-        self.IVFPQSearchCagraConfig = MockIVFPQSearchCagraConfig
-        self.GpuIndexCagraConfig = MockGpuIndexCagraConfig
-
-        # Enums
-        self.graph_build_algo_IVF_PQ = 1
-
-        self.METRIC_L2 = 0
-        self.METRIC_INNER_PRODUCT = 1
-
-        self._num_threads = None
-        self.omp_set_num_threads = self._omp_set_num_threads
-        self.omp_get_num_threads = self._omp_get_num_threads
-
-        self.write_index = self._write_index
-
-    def _omp_set_num_threads(self, num_threads: int) -> None:
-        self._num_threads = num_threads
-
-    def _omp_get_num_threads(self) -> int:
-        return self._num_threads
-
-    def _write_index(self, index, filepath):
-        if not isinstance(filepath, str):
-            raise TypeError("Filepath must be a string")
-        if not index:
-            raise ValueError("Index cannot be None")
-        try:
-            with open(filepath, "wb") as f:
-                f.write(b"MOCK_INDEX")
-        except IOError as e:
-            raise IOError(f"Failed to write to {filepath}: {str(e)}")
-
-
-# Create the mock and patch faiss
-faiss_mock = FaissMock()
-sys.modules["faiss"] = faiss_mock
+def index_build_parameters():
+    """Create sample IndexBuildParameters for testing"""
+    return IndexBuildParameters(
+        container_name="testbucket",
+        vector_path="vec.knnvec",
+        doc_id_path="doc.knndid",
+        dimension=3,
+        doc_count=5,
+        index_parameters=IndexParameters(
+            space_type=SpaceType.INNERPRODUCT,
+            algorithm_parameters=AlgorithmParameters(
+                ef_construction=200, ef_search=200
+            ),
+        ),
+        repository_type="s3",
+    )