Generalized to update settings and added max slice param

Signed-off-by: Martin Gaievski <[email protected]>
opensearch-project · Jul 23, 2024 · 034def4 · 034def4
1 parent d25b145
commit 034def4
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 8 deletions.
diff --git a/osbenchmark/worker_coordinator/runner.py b/osbenchmark/worker_coordinator/runner.py
@@ -107,7 +107,7 @@ def register_default_runners():
     register_runner(workload.OperationType.DeployMlModel, Retry(DeployMlModel()), async_runner=True)
     register_runner(workload.OperationType.TrainKnnModel, Retry(TrainKnnModel()), async_runner=True)
     register_runner(workload.OperationType.DeleteKnnModel, Retry(DeleteKnnModel()), async_runner=True)
-    register_runner(workload.OperationType.EnableConcurrentSegmentSearch, Retry(EnableConcurrentSegmentSearch()), async_runner=True)
+    register_runner(workload.OperationType.UpdateConcurrentSegmentSearchSettings, Retry(UpdateConcurrentSegmentSearchSettings()), async_runner=True)
 
 def runner_for(operation_type):
     try:
@@ -2698,16 +2698,19 @@ async def __call__(self, opensearch, params):
     def __repr__(self, *args, **kwargs):
         return "deploy-ml-model"
 
-class EnableConcurrentSegmentSearch(Runner):
+class UpdateConcurrentSegmentSearchSettings(Runner):
     @time_func
     async def __call__(self, opensearch, params):
         enable_setting = params.get("enable", "false")
+        max_slice_count = params.get("max_slice_count", None)
         body = {
             "persistent": {
                 "search.concurrent_segment_search.enabled": enable_setting
             }
         }
+        if max_slice_count is not None:
+            body["persistent"]["search.concurrent_segment_search.max_slice_count"] = max_slice_count
         await opensearch.cluster.put_settings(body=body)
 
     def __repr__(self, *args, **kwargs):
-        return "enable-concurrent-segment-search"
+        return "update-concurrent-segment-search-settings"
diff --git a/osbenchmark/workload/workload.py b/osbenchmark/workload/workload.py
@@ -635,7 +635,7 @@ class OperationType(Enum):
     DeleteMlModel = 1041
     RegisterMlModel = 1042
     DeployMlModel = 1043
-    EnableConcurrentSegmentSearch = 1044
+    UpdateConcurrentSegmentSearchSettings = 1044
 
     @property
     def admin_op(self):
@@ -753,8 +753,8 @@ def from_hyphenated_string(cls, v):
             return OperationType.TrainKnnModel
         elif v == "delete-knn-model":
             return OperationType.DeleteKnnModel
-        elif v == "enable-concurrent-segment-search":
-            return OperationType.EnableConcurrentSegmentSearch
+        elif v == "update-concurrent-segment-search-settings":
+            return OperationType.UpdateConcurrentSegmentSearchSettings
         else:
             raise KeyError(f"No enum value for [{v}]")
 

diff --git a/tests/worker_coordinator/runner_test.py b/tests/worker_coordinator/runner_test.py
@@ -6786,7 +6786,7 @@ async def test_param_id_mandatory(self, opensearch, on_client_request_start, on_
 
         self.assertEqual(0, opensearch.transport.perform_request.call_count)
 
-class EnableConcurrentSegmentSearchTests(TestCase):
+class UpdateConcurrentSegmentSearchSettingsTests(TestCase):
     @mock.patch('osbenchmark.client.RequestContextHolder.on_client_request_end')
     @mock.patch('osbenchmark.client.RequestContextHolder.on_client_request_start')
     @mock.patch("opensearchpy.OpenSearch")
@@ -6797,11 +6797,52 @@ async def test_enable_concurrent_segment_search(self, opensearch, on_client_requ
             "enable": "true"
         }
 
-        r = runner.EnableConcurrentSegmentSearch()
+        r = runner.UpdateConcurrentSegmentSearchSettings()
         await r(opensearch, params)
 
         opensearch.cluster.put_settings.assert_called_once_with(body={
             "persistent": {
                 "search.concurrent_segment_search.enabled": "true"
             }
         })
+
+    @mock.patch('osbenchmark.client.RequestContextHolder.on_client_request_end')
+    @mock.patch('osbenchmark.client.RequestContextHolder.on_client_request_start')
+    @mock.patch("opensearchpy.OpenSearch")
+    @run_async
+    async def test_max_slice_count(self, opensearch, on_client_request_start, on_client_request_end):
+        opensearch.cluster.put_settings.return_value = as_future()
+        params = {
+            "max_slice_count": 2
+        }
+
+        r = runner.UpdateConcurrentSegmentSearchSettings()
+        await r(opensearch, params)
+
+        opensearch.cluster.put_settings.assert_called_once_with(body={
+            "persistent": {
+                "search.concurrent_segment_search.enabled": "false",
+                "search.concurrent_segment_search.max_slice_count": 2
+            }
+        })
+
+    @mock.patch('osbenchmark.client.RequestContextHolder.on_client_request_end')
+    @mock.patch('osbenchmark.client.RequestContextHolder.on_client_request_start')
+    @mock.patch("opensearchpy.OpenSearch")
+    @run_async
+    async def test_concurrent_segment_search_settings(self, opensearch, on_client_request_start, on_client_request_end):
+        opensearch.cluster.put_settings.return_value = as_future()
+        params = {
+            "enable": "true",
+            "max_slice_count": 2
+        }
+
+        r = runner.UpdateConcurrentSegmentSearchSettings()
+        await r(opensearch, params)
+
+        opensearch.cluster.put_settings.assert_called_once_with(body={
+            "persistent": {
+                "search.concurrent_segment_search.enabled": "true",
+                "search.concurrent_segment_search.max_slice_count": 2
+            }
+        })