Add support to select ec2 instance type for data nodes (#3876)

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

src/test_workflow/benchmark_test/benchmark_args.py

@@ -33,6 +33,7 @@ class BenchmarkArgs:
     ml_node_storage: int
     jvm_sys_props: str
     additional_config: str
+    data_instance_type: str
     use_50_percent_heap: bool
     enable_remote_store: bool
     workload: str
@@ -80,6 +81,8 @@ def __init__(self) -> None:
                             help="User provided data-node ebs block storage size, defaults to 100Gb")
         parser.add_argument("--enable-remote-store", dest="enable_remote_store", action="store_true",
                             help="Enable Remote Store feature in OpenSearch")
+        parser.add_argument("--data-instance-type", dest="data_instance_type",
+                            help="EC2 instance type for data node, defaults to r5.xlarge.")
         parser.add_argument("--workload", dest="workload", required=True,
                             help="Name of the workload that OpenSearch Benchmark should run")
         parser.add_argument("--benchmark-config", dest="benchmark_config",
@@ -112,6 +115,7 @@ def __init__(self) -> 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_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
         self.workload_params = args.workload_params if args.workload_params else None
         self.benchmark_config = args.benchmark_config if args.benchmark_config else None

src/test_workflow/benchmark_test/benchmark_test_cluster.py

@@ -137,6 +137,7 @@ def setup_cdk_params(self, config: dict) -> dict:
             "serverAccessType": config["Constants"]["serverAccessType"],
             "restrictServerAccessTo": config["Constants"]["restrictServerAccessTo"],
             "additionalConfig": self.args.additional_config,
+            "dataInstanceType": self.args.data_instance_type,
             "managerNodeCount": self.args.manager_node_count,
             "dataNodeCount": self.args.data_node_count,
             "clientNodeCount": self.args.client_node_count,

tests/tests_test_workflow/test_benchmark_workflow/benchmark_test/test_benchmark_test_cluster.py

@@ -65,6 +65,7 @@ def test_port(self) -> None:
     @patch("test_workflow.benchmark_test.benchmark_test_cluster.BenchmarkTestCluster.wait_for_processing")
     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'
         TestBenchmarkTestCluster.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:
@@ -76,6 +77,7 @@ def test_create_single_node_insecure(self, mock_wait_for_processing: Optional[Mo
         self.assertEqual(self.benchmark_test_cluster.endpoint_with_port, 'www.example.com:80')
         self.assertEqual(self.benchmark_test_cluster.port, 80)
         self.assertTrue("securityDisabled=true" in self.benchmark_test_cluster.params)
+        self.assertTrue("dataInstanceType=r5.4xlarge" in self.benchmark_test_cluster.params)
 
     @patch("test_workflow.benchmark_test.benchmark_test_cluster.BenchmarkTestCluster.wait_for_processing")
     def test_create_multi_node(self, mock_wait_for_processing: Optional[Mock]) -> None: