Add support to run benchmarks using OpenSearch distribution url (#4098)

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

src/run_benchmark_test.py

@@ -22,9 +22,13 @@ def main() -> int:
     """
     benchmark_args = BenchmarkArgs()
     console.configure(level=benchmark_args.logging_level)
-    manifest: Union[BundleManifest, BuildManifest] = BundleManifest.from_file(benchmark_args.bundle_manifest) if not benchmark_args.min_distribution else \
-        BuildManifest.from_file(benchmark_args.bundle_manifest)
-    BenchmarkTestRunners.from_args(benchmark_args, manifest).run()
+    if benchmark_args.bundle_manifest:
+        manifest: Union[BundleManifest, BuildManifest] = BundleManifest.from_file(benchmark_args.bundle_manifest) if not benchmark_args.min_distribution else \
+            BuildManifest.from_file(benchmark_args.bundle_manifest)
+        BenchmarkTestRunners.from_args(benchmark_args, manifest).run()
+    else:
+        BenchmarkTestRunners.from_args(benchmark_args).run()
+
     return 0
 
 

src/test_workflow/benchmark_test/benchmark_args.py

@@ -18,6 +18,8 @@
 # Contains the arguments required to run a perf test.
 class BenchmarkArgs:
     bundle_manifest: IO
+    distribution_url: str
+    distribution_version: str
     stack_suffix: str
     config: IO
     keep: bool
@@ -47,8 +49,10 @@ class BenchmarkArgs:
 
     def __init__(self) -> None:
         parser = argparse.ArgumentParser(description="Test an OpenSearch Bundle")
-        parser.add_argument("--bundle-manifest", type=argparse.FileType("r"), help="Bundle Manifest file.",
-                            required=True)
+        parser.add_argument("--bundle-manifest", type=argparse.FileType("r"), help="Bundle Manifest file.")
+        parser.add_argument("--distribution-url", dest="distribution_url", help="Link to a downloadable OpenSearch tarball.")
+        parser.add_argument("--distribution-version", dest="distribution_version",
+                            help="provide OpenSearch version if using distribution-url param.")
         parser.add_argument("--suffix", dest="suffix", help="Suffix to be added to stack name for performance test")
         parser.add_argument("--component", dest="component", default="OpenSearch",
                             help="Component name that needs to be performance tested")
@@ -104,7 +108,9 @@ def __init__(self) -> None:
                             const=logging.DEBUG, dest="logging_level")
 
         args = parser.parse_args()
-        self.bundle_manifest = args.bundle_manifest
+        self.bundle_manifest = args.bundle_manifest if args.bundle_manifest else None
+        self.distribution_url = args.distribution_url if args.distribution_url else None
+        self.distribution_version = args.distribution_version if args.distribution_version else None
         self.stack_suffix = args.suffix if args.suffix else None
         self.config = args.config
         self.keep = args.keep
@@ -131,3 +137,8 @@ def __init__(self) -> None:
         self.telemetry = args.telemetry
         self.telemetry_params = args.telemetry_params if args.telemetry_params else None
         self.logging_level = args.logging_level
+
+        if self.bundle_manifest is None and self.distribution_url is None:
+            raise Exception('Please provide either --bundle-manifest or --distribution-url to run the performance test.')
+        elif self.distribution_url and self.distribution_version is None:
+            raise Exception("--distribution-version is required parameter while using --distribution-url param.")

src/test_workflow/benchmark_test/benchmark_test_cluster.py

@@ -71,7 +71,9 @@ def __init__(
         self.is_endpoint_public = False
         self.cluster_endpoint = None
         self.cluster_endpoint_with_port = None
-        self.stack_name = f"opensearch-infra-stack-{self.args.stack_suffix}-{self.manifest.build.id}-{self.manifest.build.architecture}"
+        self.stack_name = f"opensearch-infra-stack-{self.args.stack_suffix}"
+        if self.manifest:
+            self.stack_name += f"-{self.manifest.build.id}-{self.manifest.build.architecture}"
 
     def start(self) -> None:
         command = f"npm install && cdk deploy \"*\" {self.params} --outputs-file {self.output_file}"
@@ -112,27 +114,37 @@ def wait_for_processing(self, tries: int = 3, delay: int = 15, backoff: int = 2)
         logging.info(f"Waiting for domain at {self.endpoint} to be up")
         protocol = "http://" if self.args.insecure else "https://"
         url = "".join([protocol, self.endpoint, "/_cluster/health"])
-        request_args = {"url": url} if self.args.insecure else {"url": url, "auth": HTTPBasicAuth("admin", "admin"), "verify": False}  # type: ignore
+        request_args = {"url": url} if self.args.insecure else {"url": url, "auth": HTTPBasicAuth("admin", "admin"),  # type: ignore
+                                                                "verify": False}  # type: ignore
         retry_call(requests.get, fkwargs=request_args,
                    tries=tries, delay=delay, backoff=backoff)
 
     def setup_cdk_params(self, config: dict) -> dict:
-        if self.args.stack_suffix:
+        suffix = ''
+        if self.args.stack_suffix and self.manifest:
             suffix = self.args.stack_suffix + '-' + self.manifest.build.id + '-' + self.manifest.build.architecture
-        else:
+        elif self.manifest:
             suffix = self.manifest.build.id + '-' + self.manifest.build.architecture
+        elif self.args.stack_suffix:
+            suffix = self.args.stack_suffix
+
+        if self.manifest:
+            artifact_url = self.manifest.build.location if isinstance(self.manifest, BundleManifest) else \
+                f"https://artifacts.opensearch.org/snapshots/core/opensearch/{self.manifest.build.version}/opensearch-min-" \
+                f"{self.manifest.build.version}-linux-{self.manifest.build.architecture}-latest.tar.gz"
+        else:
+            artifact_url = self.args.distribution_url.strip()
+
         return {
-            "distributionUrl": self.manifest.build.location if isinstance(self.manifest, BundleManifest) else
-            f"https://artifacts.opensearch.org/snapshots/core/opensearch/{self.manifest.build.version}/opensearch-min-"
-            f"{self.manifest.build.version}-linux-{self.manifest.build.architecture}-latest.tar.gz",
+            "distributionUrl": artifact_url,
             "vpcId": config["Constants"]["VpcId"],
             "account": config["Constants"]["AccountId"],
             "region": config["Constants"]["Region"],
             "suffix": suffix,
             "securityDisabled": str(self.args.insecure).lower(),
-            "cpuArch": self.manifest.build.architecture,
+            "cpuArch": self.manifest.build.architecture if self.manifest else 'x64',
             "singleNodeCluster": str(self.args.single_node).lower(),
-            "distVersion": self.manifest.build.version,
+            "distVersion": self.manifest.build.version if self.manifest else self.args.distribution_version,
             "minDistribution": str(self.args.min_distribution).lower(),
             "serverAccessType": config["Constants"]["serverAccessType"],
             "restrictServerAccessTo": config["Constants"]["restrictServerAccessTo"],

src/test_workflow/benchmark_test/benchmark_test_runner.py

@@ -24,7 +24,11 @@ def __init__(self, args: BenchmarkArgs, test_manifest: Union[BundleManifest, Bui
         self.args = args
         self.test_manifest = test_manifest
 
-        self.security = "security" in self.test_manifest.components and not self.args.insecure
+        if self.test_manifest:
+            self.security = "security" in self.test_manifest.components and not self.args.insecure
+        else:
+            self.security = False
+
         self.tests_dir = os.path.join(os.getcwd(), "test-results", "benchmark-test", f"{'with' if self.security else 'without'}-security")
         os.makedirs(self.tests_dir, exist_ok=True)
 

src/test_workflow/benchmark_test/benchmark_test_runners.py

@@ -21,5 +21,5 @@ class BenchmarkTestRunners:
     }
 
     @classmethod
-    def from_args(cls, args: BenchmarkArgs, test_manifest: Union[BundleManifest, BuildManifest]) -> BenchmarkTestRunner:
+    def from_args(cls, args: BenchmarkArgs, test_manifest: Union[BundleManifest, BuildManifest] = None) -> BenchmarkTestRunner:
         return cls.RUNNERS.get(args.component, BenchmarkTestRunnerOpenSearchPlugins)(args, test_manifest)

tests/test_run_benchmark_test.py

@@ -15,7 +15,7 @@
 from run_benchmark_test import main
 
 
-class TestRunPerfTest(unittest.TestCase):
+class TestRunBenchmarkTest(unittest.TestCase):
     @pytest.fixture(autouse=True)
     def _capfd(self, capfd: Any) -> None:
         self.capfd = capfd
@@ -66,3 +66,10 @@ def test_run_benchmark_test_plugins(self, os_mock_runner: Mock, plugin_mock_runn
         main()
         self.assertEqual(0, os_mock_runner.call_count)
         self.assertEqual(1, plugin_mock_runner.call_count)
+
+    @patch("argparse._sys.argv", ["run_benchmark_test.py", "--distribution-url", "test.url", "--distribution-version", "2.10.0",
+                                  "--config", BENCHMARK_TEST_CONFIG, "--workload", "test", "--suffix", "test"])
+    @patch("run_benchmark_test.BenchmarkTestRunners.from_args")
+    def test_default_execute_benchmark_test_without_manifest(self, mock_runner: Mock) -> None:
+        main()
+        self.assertEqual(1, mock_runner.call_count)

tests/tests_test_workflow/test_benchmark_args.py

@@ -65,3 +65,23 @@ def test_benchmark_with_optional_config_parameters(self) -> None:
         self.assertEqual(test_args.additional_config,
                          '{"opensearch.experimental.feature.replication_type.enabled": "true", "key": "value"}')
         self.assertEqual(test_args.jvm_sys_props, "key1=value1,key2=value2")
+
+    @patch("argparse._sys.argv", [ARGS_PY, "--distribution-url", "https://artifacts.opensearch.org/2.10.0/opensearch-2.10.0-linux-x64.tar.gz",
+                                  "--distribution-version", "2.10.0", "--config", TEST_CONFIG_PATH, "--workload", "test"])
+    def test_benchmark_with_distribution_url_and_version(self) -> None:
+        test_args = BenchmarkArgs()
+        self.assertEqual(test_args.distribution_url, "https://artifacts.opensearch.org/2.10.0/opensearch-2.10.0-linux-x64.tar.gz")
+        self.assertEqual(test_args.distribution_version, "2.10.0")
+
+    @patch("argparse._sys.argv", [ARGS_PY, "--distribution-url", "https://artifacts.opensearch.org/2.10.0/opensearch-2.10.0-linux-x64.tar.gz",
+                                  "--distribution-version", None, "--config", TEST_CONFIG_PATH, "--workload", "test"])
+    def test_benchmark_with_distribution_url_and_without_version(self) -> None:
+        with self.assertRaises(Exception) as context:
+            BenchmarkArgs()
+        self.assertEqual(str(context.exception), "--distribution-version is required parameter while using --distribution-url param.")
+
+    @patch("argparse._sys.argv", [ARGS_PY, "--config", TEST_CONFIG_PATH, "--workload", "test"])
+    def test_benchmark_without_distribution_url_and_without_manifest(self) -> None:
+        with self.assertRaises(Exception) as context:
+            BenchmarkArgs()
+        self.assertEqual(str(context.exception), "Please provide either --bundle-manifest or --distribution-url to run the performance test.")

tests/tests_test_workflow/test_benchmark_workflow/benchmark_test/test_benchmark_test_cluster.py

@@ -18,7 +18,7 @@ class TestBenchmarkTestCluster(unittest.TestCase):
     DATA = os.path.join(os.path.dirname(__file__), "data")
     BUNDLE_MANIFEST = os.path.join(DATA, "bundle_manifest.yml")
 
-    def setUp(self, args: Optional[Mock] = None) -> None:
+    def setUp(self, args: Optional[Mock] = None, use_manifest: bool = True) -> None:
         self.args = Mock()
         if args:
             self.args = args
@@ -28,7 +28,7 @@ def setUp(self, args: Optional[Mock] = None) -> None:
             self.args.insecure = False
             self.args.single_node = True
             self.args.min_distribution = False
-        self.manifest = BundleManifest.from_path(self.BUNDLE_MANIFEST)
+        self.manifest = BundleManifest.from_path(self.BUNDLE_MANIFEST) if use_manifest else None
         self.stack_name = "stack"
         self.security = True
         self.config = {"Constants": {"SecurityGroupId": "sg-00000000", "VpcId": "vpc-12345", "AccountId": "12345678",
@@ -95,3 +95,19 @@ def test_create_multi_node(self, mock_wait_for_processing: Optional[Mock]) -> No
         self.assertTrue("singleNodeCluster=false" in self.benchmark_test_cluster.params)
         self.assertTrue("use50PercentHeap=true" in self.benchmark_test_cluster.params)
         self.assertTrue("enableRemoteStore=true" in self.benchmark_test_cluster.params)
+
+    @patch("test_workflow.benchmark_test.benchmark_test_cluster.BenchmarkTestCluster.wait_for_processing")
+    def test_create_multi_node_without_manifest(self, mock_wait_for_processing: Optional[Mock]) -> None:
+        self.args.distribution_url = "https://artifacts.opensearch.org/2.10.0/opensearch.tar.gz"
+        self.args.distribution_version = '2.10.0'
+        TestBenchmarkTestCluster.setUp(self, self.args, False)
+        mock_file = MagicMock(side_effect=[{"opensearch-infra-stack-test-suffix": {"loadbalancerurl": "www.example.com"}}])
+        with patch("subprocess.check_call") as mock_check_call:
+            with patch("builtins.open", MagicMock()):
+                with patch("json.load", mock_file):
+                    self.benchmark_test_cluster.start()
+                    self.assertEqual(mock_check_call.call_count, 1)
+        self.assertTrue("opensearch-infra-stack-test-suffix" in self.benchmark_test_cluster.stack_name)
+        self.assertTrue("cpuArch=x64" in self.benchmark_test_cluster.params)
+        self.assertTrue("distVersion=2.10.0" in self.benchmark_test_cluster.params)
+        self.assertTrue("distributionUrl=https://artifacts.opensearch.org/2.10.0/opensearch.tar.gz" in self.benchmark_test_cluster.params)

tests/tests_test_workflow/test_benchmark_workflow/benchmark_test/test_benchmark_test_runner_opensearch.py

@@ -15,7 +15,7 @@
 from test_workflow.benchmark_test.benchmark_test_runners import BenchmarkTestRunners
 
 
-class TestPerfTestRunnerOpenSearch(unittest.TestCase):
+class TestBenchmarkTestRunnerOpenSearch(unittest.TestCase):
 
     @patch("argparse._sys.argv", ["run_benchmark_test.py",
                                   "--bundle-manifest",
@@ -44,3 +44,33 @@ def test_run(self, mock_suite: Mock, mock_cluster: Mock, mock_git: Mock, mock_te
         self.assertEqual(mock_cluster.call_count, 1)
         self.assertEqual(mock_git.call_count, 1)
         self.assertEqual(mock_temp_directory.call_count, 1)
+
+    @patch("argparse._sys.argv", ["run_benchmark_test.py",
+                                  "--distribution-url",
+                                  "https://artifacts.opensearch.org/2.10.0/opensearch.tar.gz",
+                                  "--distribution-version",
+                                  "2.10.0",
+                                  "--config", os.path.join(os.path.dirname(__file__), "data", "test-config.yml"),
+                                  "--workload", "test",
+                                  "--suffix", "test"])
+    @patch("os.chdir")
+    @patch("test_workflow.benchmark_test.benchmark_test_runner_opensearch.TemporaryDirectory")
+    @patch("test_workflow.benchmark_test.benchmark_test_runner_opensearch.GitRepository")
+    @patch("test_workflow.benchmark_test.benchmark_test_runner_opensearch.BenchmarkTestCluster.create")
+    @patch("test_workflow.benchmark_test.benchmark_test_runner_opensearch.BenchmarkTestSuite")
+    def test_run_with_dist_url_and_version(self, mock_suite: Mock, mock_cluster: Mock, mock_git: Mock,
+                                           mock_temp_directory: Mock,
+                                           *mocks: Any) -> None:
+        mock_temp_directory.return_value.__enter__.return_value.name = tempfile.gettempdir()
+        mock_cluster.return_value.__enter__.return_value = mock_cluster
+
+        benchmark_args = BenchmarkArgs()
+        runner = BenchmarkTestRunners.from_args(benchmark_args)
+        runner.run()
+
+        mock_git.assert_called_with("https://github.com/opensearch-project/opensearch-cluster-cdk.git", "main",
+                                    os.path.join(tempfile.gettempdir(), "opensearch-cluster-cdk"))
+        self.assertEqual(mock_suite.call_count, 1)
+        self.assertEqual(mock_cluster.call_count, 1)
+        self.assertEqual(mock_git.call_count, 1)
+        self.assertEqual(mock_temp_directory.call_count, 1)