From d78b555c10547a4ea7eaf216208e3939047cf7be Mon Sep 17 00:00:00 2001 From: Rishabh Singh Date: Mon, 14 Aug 2023 11:20:55 -0700 Subject: [PATCH] Add support to select ec2 instance type for data nodes (#3876) Signed-off-by: Rishabh Singh --- src/test_workflow/benchmark_test/benchmark_args.py | 4 ++++ src/test_workflow/benchmark_test/benchmark_test_cluster.py | 1 + .../benchmark_test/test_benchmark_test_cluster.py | 2 ++ 3 files changed, 7 insertions(+) diff --git a/src/test_workflow/benchmark_test/benchmark_args.py b/src/test_workflow/benchmark_test/benchmark_args.py index 8c4cfa965f..bb15069444 100644 --- a/src/test_workflow/benchmark_test/benchmark_args.py +++ b/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 diff --git a/src/test_workflow/benchmark_test/benchmark_test_cluster.py b/src/test_workflow/benchmark_test/benchmark_test_cluster.py index 39e2190c03..74b3545438 100644 --- a/src/test_workflow/benchmark_test/benchmark_test_cluster.py +++ b/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, diff --git a/tests/tests_test_workflow/test_benchmark_workflow/benchmark_test/test_benchmark_test_cluster.py b/tests/tests_test_workflow/test_benchmark_workflow/benchmark_test/test_benchmark_test_cluster.py index bce511d696..54e0493de8 100644 --- a/tests/tests_test_workflow/test_benchmark_workflow/benchmark_test/test_benchmark_test_cluster.py +++ b/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: