Add Warmup Operation in Perf-tool Benchmarking Tests (#1137)

Signed-off-by: Naveen Tatikonda <[email protected]>

benchmarks/perf-tool/okpt/io/config/parsers/test.py

@@ -24,6 +24,7 @@ class TestConfig:
     test_id: str
     endpoint: str
     port: int
+    timeout: int
     num_runs: int
     show_runs: bool
     setup: List[Step]
@@ -67,6 +68,7 @@ def parse(self, file_obj: TextIOWrapper) -> TestConfig:
         test_config = TestConfig(
             endpoint=config_obj['endpoint'],
             port=config_obj['port'],
+            timeout=config_obj['timeout'],
             test_name=config_obj['test_name'],
             test_id=config_obj['test_id'],
             num_runs=config_obj['num_runs'],

benchmarks/perf-tool/okpt/io/config/schemas/test.yml

@@ -12,6 +12,9 @@ endpoint:
 port:
   type: integer
   default: 9200
+timeout:
+  type: integer
+  default: 60  
 test_name:
   type: string
 test_id:

benchmarks/perf-tool/okpt/test/steps/factory.py

@@ -10,7 +10,7 @@
 
 from okpt.test.steps.steps import CreateIndexStep, DisableRefreshStep, RefreshIndexStep, DeleteIndexStep, \
     TrainModelStep, DeleteModelStep, ForceMergeStep, ClearCacheStep, IngestStep, IngestMultiFieldStep, \
-    QueryStep, QueryWithFilterStep, GetStatsStep
+    QueryStep, QueryWithFilterStep, GetStatsStep, WarmupStep
 
 
 def create_step(step_config: StepConfig) -> Step:
@@ -40,5 +40,7 @@ def create_step(step_config: StepConfig) -> Step:
         return ClearCacheStep(step_config)
     elif step_config.step_name == GetStatsStep.label:
         return GetStatsStep(step_config)
+    elif step_config.step_name == WarmupStep.label:
+        return WarmupStep(step_config)
 
     raise ConfigurationError(f'Invalid step {step_config.step_name}')

benchmarks/perf-tool/okpt/test/steps/steps.py

@@ -38,8 +38,9 @@ def __init__(self, step_config: StepConfig):
         default_port = 9200 if self.endpoint == 'localhost' else 80
         self.port = parse_int_param('port', step_config.config,
                                     step_config.implicit_config, default_port)
+        self.timeout = parse_int_param('timeout', step_config.config, {}, 60)
         self.opensearch = get_opensearch_client(str(self.endpoint),
-                                                int(self.port))
+                                                int(self.port), int(self.timeout))
 
 
 class CreateIndexStep(OpenSearchStep):
@@ -163,6 +164,25 @@ def _get_measures(self) -> List[str]:
         return ['took']
 
 
+class WarmupStep(OpenSearchStep):
+    """See base class."""
+
+    label = 'warmup_operation'
+
+    def __init__(self, step_config: StepConfig):
+        super().__init__(step_config)
+        self.index_name = parse_string_param('index_name', step_config.config, {},
+                                           None)
+
+    def _action(self):
+        """Performs warmup operation on an index."""
+        warmup_operation(self.endpoint, self.port, self.index_name)
+        return {}
+
+    def _get_measures(self) -> List[str]:
+        return ['took']
+
+
 class TrainModelStep(OpenSearchStep):
     """See base class."""
 
@@ -690,6 +710,23 @@ def delete_model(endpoint, port, model_id):
     return response.json()
 
 
+def warmup_operation(endpoint, port, index):
+    """
+    Performs warmup operation on index to load native library files
+    of that index to reduce query latencies.
+    Args:
+        endpoint: Endpoint OpenSearch is running on
+        port: Port OpenSearch is running on
+        index: index name
+    Returns:
+        number of shards the plugin succeeded and failed to warm up.
+    """
+    response = requests.get('http://' + endpoint + ':' + str(port) +
+                               '/_plugins/_knn/warmup/' + index,
+                               headers={'content-type': 'application/json'})
+    return response.json()
+
+
 def get_opensearch_client(endpoint: str, port: int, timeout=60):
     """
     Get an opensearch client from an endpoint and port

benchmarks/perf-tool/release-configs/faiss-hnsw/filtering/relaxed-filter/relaxed-filter-test.yml

@@ -19,6 +19,8 @@ steps:
     attribute_spec: [ { name: 'color', type: 'str' }, { name: 'taste', type: 'str' }, { name: 'age', type: 'int' } ]
   - name: refresh_index
     index_name: target_index
+  - name: warmup_operation
+    index_name: target_index
   - name: query_with_filter
     k: 100
     r: 1

benchmarks/perf-tool/release-configs/faiss-hnsw/filtering/restrictive-filter/restrictive-filter-test.yml

@@ -22,6 +22,8 @@ steps:
   - name: force_merge
     index_name: target_index
     max_num_segments: 1
+  - name: warmup_operation
+    index_name: target_index
   - name: query_with_filter
     k: 100
     r: 1

benchmarks/perf-tool/release-configs/faiss-hnsw/test.yml

@@ -20,6 +20,8 @@ steps:
   - name: force_merge
     index_name: target_index
     max_num_segments: 1
+  - name: warmup_operation
+    index_name: target_index
   - name: query
     k: 100
     r: 1

benchmarks/perf-tool/release-configs/faiss-hnswpq/test.yml

@@ -41,6 +41,8 @@ steps:
     dataset_path: /home/ec2-user/data/sift-128-euclidean.hdf5
   - name: refresh_index
     index_name: target_index
+  - name: warmup_operation
+    index_name: target_index  
   - name: query
     k: 100
     r: 1

benchmarks/perf-tool/release-configs/faiss-ivf/test.yml

@@ -41,6 +41,8 @@ steps:
     dataset_path: /home/ec2-user/data/sift-128-euclidean.hdf5
   - name: refresh_index
     index_name: target_index
+  - name: warmup_operation
+    index_name: target_index  
   - name: query
     k: 100
     r: 1

benchmarks/perf-tool/release-configs/faiss-ivfpq/test.yml

@@ -41,6 +41,8 @@ steps:
     dataset_path: /home/ec2-user/data/sift-128-euclidean.hdf5
   - name: refresh_index
     index_name: target_index
+  - name: warmup_operation
+    index_name: target_index  
   - name: query
     k: 100
     r: 1

benchmarks/perf-tool/release-configs/nmslib-hnsw/test.yml

@@ -17,6 +17,8 @@ steps:
     dataset_path: /home/ec2-user/data/sift-128-euclidean.hdf5
   - name: refresh_index
     index_name: target_index
+  - name: warmup_operation
+    index_name: target_index  
   - name: query
     k: 100
     r: 1

benchmarks/perf-tool/sample-configs/faiss-sift-ivf/test.yml

@@ -43,6 +43,8 @@ steps:
   - name: force_merge
     index_name: target_index
     max_num_segments: 10
+  - name: warmup_operation
+    index_name: target_index  
   - name: query
     k: 100
     r: 1

benchmarks/perf-tool/sample-configs/nmslib-sift-hnsw/test.yml

@@ -21,6 +21,8 @@ steps:
   - name: force_merge
     index_name: target_index
     max_num_segments: 10
+  - name: warmup_operation
+    index_name: target_index
   - name: query
     k: 100
     r: 1