Add support for instance based storge in nightly benchmark runs

Signed-off-by: Rishabh Singh <[email protected]>

src/test_workflow/benchmark_test/benchmark_args.py

@@ -57,6 +57,8 @@ class BenchmarkArgs:
     results_file: str
     show_in_results: str
     command: str
+    enable_instance_storage: bool
+
 
     def __init__(self) -> None:
         parser = argparse.ArgumentParser(description="Test an OpenSearch Bundle or compare two tests")
@@ -143,6 +145,8 @@ def __init__(self) -> None:
                                          help="Show more verbose output.")
         execute_test_parser.add_argument("--ml-node-storage", dest="ml_node_storage",
                                          help="User provided ml-node ebs block storage size defaults to 100Gb")
+        execute_test_parser.add_argument("--enable-instance-storage", dest="enable_instance_storage", action="store_true",
+                                         help="Use instance based SSD storage instead of EBS for data nodes", default=False)
 
         # command to run comparison
         compare_parser = subparsers.add_parser("compare", parents=[parent_parser],
@@ -184,6 +188,7 @@ def __init__(self) -> None:
             self.jvm_sys_props = args.jvm_sys_props if args.jvm_sys_props else None
             self.data_node_storage = args.data_node_storage if args.data_node_storage else None
             self.ml_node_storage = args.ml_node_storage if args.ml_node_storage else None
+            self.enable_instance_storage = args.enable_instance_storage
             self.enable_remote_store = args.enable_remote_store
             self.data_instance_type = args.data_instance_type if args.data_instance_type else None
             self.workload = args.workload

src/test_workflow/benchmark_test/benchmark_create_cluster.py

@@ -134,6 +134,7 @@ def setup_cdk_params(self, config: dict) -> dict:
             "mlNodeCount": self.args.ml_node_count,
             "dataNodeStorage": self.args.data_node_storage,
             "mlNodeStorage": self.args.ml_node_storage,
+            "useInstanceBasedStorage": str(self.args.enable_instance_storage).lower(),
             "jvmSysProps": self.args.jvm_sys_props,
             "use50PercentHeap": str(self.args.use_50_percent_heap).lower(),
             "isInternal": config["Constants"]["isInternal"],

tests/tests_test_workflow/test_benchmark_args.py

@@ -38,6 +38,7 @@ def test_benchmark_with_default_parameters(self) -> None:
         self.assertFalse(test_args.single_node)
         self.assertFalse(test_args.min_distribution)
         self.assertFalse(test_args.enable_remote_store)
+        self.assertFalse(test_args.enable_instance_storage)
 
     @patch("argparse._sys.argv",
            [ARGS_PY, "execute-test", "--bundle-manifest", TEST_DIST_MANIFEST_PATH, "--config", TEST_CONFIG_PATH, "--workload", "test",
@@ -53,11 +54,12 @@ def test_benchmark_with_optional_node_count_parameters(self) -> None:
 
     @patch("argparse._sys.argv",
            [ARGS_PY, "execute-test", "--bundle-manifest", TEST_DIST_MANIFEST_PATH, "--config", TEST_CONFIG_PATH, "--workload", "test",
-            "--data-node-storage", "200", "--ml-node-storage", "100"])
+            "--data-node-storage", "200", "--ml-node-storage", "100", "--enable-instance-storage"])
     def test_benchmark_with_optional_node_storage_parameters(self) -> None:
         test_args = BenchmarkArgs()
         self.assertEqual(test_args.data_node_storage, "200")
         self.assertEqual(test_args.ml_node_storage, "100")
+        self.assertTrue(test_args.enable_instance_storage)
 
     @patch("argparse._sys.argv", [ARGS_PY, "execute-test", "--bundle-manifest", TEST_DIST_MANIFEST_PATH, "--config", TEST_CONFIG_PATH, "--workload", "test",
                                   "--additional-config", 'opensearch.experimental.feature.replication_type.enabled:true','key:value',  # noqa: E231

tests/tests_test_workflow/test_benchmark_workflow/benchmark_test/test_benchmark_create_cluster.py

@@ -28,6 +28,7 @@ def setUp(self, args: Optional[Mock] = None, use_manifest: bool = True) -> None:
             self.args.insecure = False
             self.args.single_node = True
             self.args.min_distribution = False
+            self.args.enable_instance_storage = False
         self.manifest = BundleManifest.from_path(self.BUNDLE_MANIFEST) if use_manifest else None
         self.stack_name = "stack"
         self.security = True
@@ -51,6 +52,7 @@ def test_create_single_node_secure(self, mock_wait_for_processing: Optional[Mock
         self.assertTrue("adminPassword=admin" in self.benchmark_create_cluster.params)
         self.assertTrue("singleNodeCluster=true" in self.benchmark_create_cluster.params)
         self.assertTrue("isInternal=true" in self.benchmark_create_cluster.params)
+        self.assertTrue("useInstanceBasedStorage=false" in self.benchmark_create_cluster.params)
         self.assertTrue("distributionUrl=https://artifacts.opensearch.org/bundles/1.0.0/41d5ae25183d4e699e92debfbe3f83bd/opensearch-1.0.0-linux-x64.tar.gz" in self.benchmark_create_cluster.params)
         self.assertTrue(isinstance(self.manifest, BundleManifest))
         with patch("subprocess.check_call") as mock_check_call:
@@ -67,6 +69,8 @@ def test_port(self) -> None:
     def test_create_single_node_insecure(self, mock_wait_for_processing: Optional[Mock]) -> None:
         self.args.insecure = True
         self.args.data_instance_type = 'r5.4xlarge'
+        self.args.enable_instance_storage = True
+
         TestBenchmarkCreateCluster.setUp(self, self.args)
         mock_file = MagicMock(side_effect=[{"opensearch-infra-stack-test-suffix-007-x64": {"loadbalancerurl": "www.example.com"}}])
         with patch("subprocess.check_call") as mock_check_call:
@@ -80,6 +84,7 @@ def test_create_single_node_insecure(self, mock_wait_for_processing: Optional[Mo
         self.assertTrue("securityDisabled=true" in self.benchmark_create_cluster.params)
         self.assertTrue("dataInstanceType=r5.4xlarge" in self.benchmark_create_cluster.params)
         self.assertTrue("customRoleArn=arn:aws:iam::12344567890:role/customRole" in self.benchmark_create_cluster.params)
+        self.assertTrue("useInstanceBasedStorage=true" in self.benchmark_create_cluster.params)
 
     @patch("test_workflow.benchmark_test.benchmark_create_cluster.BenchmarkCreateCluster.wait_for_processing")
     def test_create_multi_node(self, mock_wait_for_processing: Optional[Mock]) -> None: