Added step to update the index thread qty

Signed-off-by: AWS ParallelCluster user <[email protected]>
navneet1v · Nov 16, 2023 · 8b32525 · 8b32525
1 parent 9b313de
commit 8b32525
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 17 deletions.
diff --git a/benchmarks/perf-tool/okpt/test/steps/factory.py b/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, WarmupStep, MultiProcessQueryStep
+    QueryStep, QueryWithFilterStep, GetStatsStep, WarmupStep, MultiProcessQueryStep, UpdateIndexThreadQty
 
 
 def create_step(step_config: StepConfig) -> Step:
@@ -44,5 +44,7 @@ def create_step(step_config: StepConfig) -> Step:
         return WarmupStep(step_config)
     elif step_config.step_name == MultiProcessQueryStep.label:
         return MultiProcessQueryStep(step_config)
+    elif step_config.step_name == UpdateIndexThreadQty.label:
+        return UpdateIndexThreadQty(step_config)
 
     raise ConfigurationError(f'Invalid step {step_config.step_name}')
diff --git a/benchmarks/perf-tool/okpt/test/steps/steps.py b/benchmarks/perf-tool/okpt/test/steps/steps.py
@@ -309,28 +309,33 @@ def _get_measures(self) -> List[str]:
         return ['took']
 
 
-# class UpdateIndexThreadQty(OpenSearchStep):
-#     """See base class."""
+class UpdateIndexThreadQty(OpenSearchStep):
+    """See base class."""
 
-#     label = 'update_index_thread_qty'
+    label = 'update_index_thread_qty'
 
-#     def __init__(self, step_config: StepConfig):
-#         super().__init__(step_config)
+    def __init__(self, step_config: StepConfig):
+        super().__init__(step_config)
 
-#         self.thread_qty = parse_int_param('thread_qty', step_config.config,
-#                                              {}, None)
+        self.thread_qty = parse_int_param('thread_qty', step_config.config,
+                                             {}, None)
 
-#     def _action(self):
-#         """update the index thread qty
+    def _action(self):
+        """update the index thread qty
 
-#         Returns:
-#             An empty dict
-#         """
-#         #self.opensearch
-#         return {}
+        Returns:
+            An empty dict
+        """
+        body = {
+           "persistent" : {
+                "knn.algo_param.index_thread_qty" : self.thread_qty
+            }
+        }
+        self.opensearch.cluster.put_settings(body)
+        return {}
 
-#     def _get_measures(self) -> List[str]:
-#         return ['took']
+    def _get_measures(self) -> List[str]:
+        return []
 
 
 class BaseIngestStep(OpenSearchStep):