Adding new branch matching machanisum to switch to the correct new br…

…anches with type names

Signed-off-by: beaioun <mshi@ucsd.edu>

osbenchmark/builder/__init__.py

@@ -24,4 +24,4 @@
 
 # expose only the minimum API
 from .builder import StartEngine, EngineStarted, StopEngine, EngineStopped, ResetRelativeTime, BuilderActor, \
-    cluster_distribution_version, download, install, start, stop
+    cluster_distribution_version, cluster_distribution_type, download, install, start, stop

osbenchmark/builder/builder.py

@@ -268,6 +268,28 @@ def cluster_distribution_version(cfg, client_factory=client.OsClientFactory):
     return opensearch.info()["version"]["number"]
 
 
+def cluster_distribution_type(cfg, client_factory=client.OsClientFactory):
+    """
+    Attempt to get the cluster's distribution type even before it is actually started (which makes only sense for externally
+    provisioned clusters).
+
+    :param cfg: The current config object.
+    :param client_factory: Factory class that creates the OpenSearch client.
+    :return: The distribution type.
+    """
+    hosts = cfg.opts("client", "hosts").default
+    client_options = cfg.opts("client", "options").default
+    opensearch = client_factory(hosts, client_options).create()
+    # unconditionally wait for the REST layer - if it's not up by then, we'll intentionally raise the original error
+    client.wait_for_rest_layer(opensearch)
+    try:
+        distribution_type = opensearch.info()["version"]["distribution"]
+    except Exception:
+        console.warn("Could not determine distribution type from endpoint, use --distribution-version to specify")
+        distribution_type = None
+    return distribution_type
+
+
 def to_ip_port(hosts):
     ip_port_pairs = []
     for host in hosts:

osbenchmark/test_execution_orchestrator.py

@@ -184,8 +184,10 @@ def setup(self, sources=False):
         # but there are rare cases (external pipeline and user did not specify the distribution version) where we need
         # to derive it ourselves. For source builds we always assume "master"
         if not sources and not self.cfg.exists("builder", "distribution.version"):
+            distribution_type = builder.cluster_distribution_type(self.cfg)
             distribution_version = builder.cluster_distribution_version(self.cfg)
             self.logger.info("Automatically derived distribution version [%s]", distribution_version)
+            self.cfg.add(config.Scope.benchmark, "builder", "distribution.type", distribution_type)
             self.cfg.add(config.Scope.benchmark, "builder", "distribution.version", distribution_version)
             min_os_version = versions.Version.from_string(version.minimum_os_version())
             specified_version = versions.Version.from_string(distribution_version)

osbenchmark/utils/repo.py

@@ -62,10 +62,10 @@ def __init__(self, default_directory, root_dir, repo_name, resource_name, offlin
                     raise exceptions.SystemSetupError("[{src}] must be a git repository.\n\nPlease run:\ngit -C {src} init"
                                                       .format(src=self.repo_dir))
 
-    def update(self, distribution_version):
+    def update(self, distribution_version, distribution_type=None):
         try:
             if self.remote:
-                branch = versions.best_match(git.branches(self.repo_dir, remote=self.remote), distribution_version)
+                branch = versions.best_matching_branch(git.branches(self.repo_dir, remote=self.remote), distribution_version, distribution_type)
                 if branch:
                     # Allow uncommitted changes iff we do not have to change the branch
                     self.logger.info(
@@ -85,7 +85,7 @@ def update(self, distribution_version):
                     msg = "Could not find %s remotely for distribution version [%s]. Trying to find %s locally." % \
                           (self.resource_name, distribution_version, self.resource_name)
                     self.logger.warning(msg)
-            branch = versions.best_match(git.branches(self.repo_dir, remote=False), distribution_version)
+            branch = versions.best_matching_branch(git.branches(self.repo_dir, remote=False), distribution_version, distribution_type)
             if branch:
                 if git.current_branch(self.repo_dir) != branch:
                     self.logger.info("Checking out [%s] in [%s] for distribution version [%s].",

osbenchmark/utils/versions.py

@@ -31,6 +31,10 @@
 
 VERSIONS_OPTIONAL = re.compile(r"^(\d+)(?:\.(\d+))?(?:\.(\d+))?(?:-(.+))?$")
 
+OS_VERSIONS = ["1", "2", "3"]
+
+ES_VERSIONS = ["6", "7"]
+
 
 def _versions_pattern(strict):
     return VERSIONS if strict else VERSIONS_OPTIONAL
@@ -151,9 +155,9 @@ def all_versions(self):
         return versions
 
 
-def best_match(available_alternatives, distribution_version):
+def best_matching_version(available_alternatives, distribution_version):
     """
-    Finds the most specific branch for a given distribution version assuming that versions have the pattern:
+    Finds the most specific branch version number for a given distribution version assuming that versions have the pattern:
 
         major.minor.patch-suffix
 
@@ -188,6 +192,38 @@ def best_match(available_alternatives, distribution_version):
     return None
 
 
+def best_matching_branch(available_alternatives, distribution_version, distribution_type=None):
+    """
+    Finds the most specific branch for a given distribution version by calling best_matching_version().
+    See best_matching_version() for more explainations.
+
+    :param available_alternatives: A list of possible distribution versions (or shortened versions).
+    :param distribution_version: An OpenSearch distribution version.
+    :param distribution_type: A type of cluster engine.
+    :return: The most specific alternative that is available or None.
+    """
+    version = best_matching_version(available_alternatives, distribution_version)
+    if distribution_type == "":
+        distribution_type = "elasticsearch" 
+    if version and version != "main":
+        version_list = [OS_VERSIONS, ES_VERSIONS]
+        matched_type = next((i for i, lst in enumerate(version_list) if version in lst), "")
+        matched_type = "opensearch" if matched_type == 0 else "elasticsearch"
+        if distribution_type and distribution_type != matched_type:
+            raise exceptions.BuildError("mismatched expected ('%s') and actual distribution type ('%s')" % (matched_type, distribution_type))
+        elif distribution_type == "opensearch" or matched_type == "opensearch":
+            prefix = "OS-"
+        elif distribution_type == "elasticsearch" or matched_type == "elasticsearch":
+            prefix = "ES-"
+        else:
+            raise exceptions.InvalidSyntax("unlisted distribution type '%s'" % (distribution_type))
+        return prefix + version
+    else:
+        #the returing version variable at this point should be either None or 'main'
+        return version
+
+
+
 def _latest_major(alternatives):
     max_major = -1
     for a in alternatives:

osbenchmark/workload/loader.py

@@ -315,6 +315,7 @@ def __init__(self, cfg, fetch, update, repo_class=repo.BenchmarkRepository):
         # current workload name (if any)
         self.workload_name = cfg.opts("workload", "workload.name", mandatory=False)
         distribution_version = cfg.opts("builder", "distribution.version", mandatory=False)
+        distribution_type = cfg.opts("builder", "distribution.type", mandatory=False)
         repo_name = cfg.opts("workload", "repository.name")
         repo_revision = cfg.opts("workload", "repository.revision", mandatory=False)
         offline = cfg.opts("system", "offline.mode")
@@ -328,7 +329,7 @@ def __init__(self, cfg, fetch, update, repo_class=repo.BenchmarkRepository):
             if repo_revision:
                 self.repo.checkout(repo_revision)
             else:
-                self.repo.update(distribution_version)
+                self.repo.update(distribution_version, distribution_type)
                 cfg.add(config.Scope.applicationOverride, "workload", "repository.revision", self.repo.revision)
 
     @property

tests/utils/repo_test.py

@@ -158,11 +158,11 @@ def test_updates_from_remote(self, rebase, checkout, branches, fetch, is_working
             resource_name="unittest-resources",
             offline=random.choice([True, False]))
 
-        r.update(distribution_version="1.7.3")
+        r.update(distribution_version="1.7.3", distribution_type="opensearch")
 
         branches.assert_called_with("/benchmark-resources/unit-test", remote=True)
-        rebase.assert_called_with("/benchmark-resources/unit-test", branch="1")
-        checkout.assert_called_with("/benchmark-resources/unit-test", branch="1")
+        rebase.assert_called_with("/benchmark-resources/unit-test", branch="OS-1")
+        checkout.assert_called_with("/benchmark-resources/unit-test", branch="OS-1")
 
     @mock.patch("osbenchmark.utils.git.head_revision")
     @mock.patch("osbenchmark.utils.git.is_working_copy", autospec=True)
@@ -184,7 +184,7 @@ def test_updates_locally(self, curr_branch, rebase, checkout, branches, fetch, i
             resource_name="unittest-resources",
             offline=False)
 
-        r.update(distribution_version="6.0.0")
+        r.update(distribution_version="6.0.0", distribution_type="")
 
         branches.assert_called_with("/benchmark-resources/unit-test", remote=False)
         self.assertEqual(0, rebase.call_count)
@@ -212,7 +212,7 @@ def test_fallback_to_tags(self, curr_branch, rebase, checkout, branches, tags, f
             resource_name="unittest-resources",
             offline=False)
 
-        r.update(distribution_version="1.7.4")
+        r.update(distribution_version="1.7.4", distribution_type="opensearch")
 
         branches.assert_called_with("/benchmark-resources/unit-test", remote=False)
         self.assertEqual(0, rebase.call_count)
@@ -280,7 +280,7 @@ def test_does_not_update_unknown_branch_remotely_local_fallback(self, curr_branc
             resource_name="unittest-resources",
             offline=False)
 
-        r.update(distribution_version="1.7.3")
+        r.update(distribution_version="1.7.3", distribution_type="opensearch")
 
         calls = [
             # first try to find it remotely...
@@ -291,7 +291,7 @@ def test_does_not_update_unknown_branch_remotely_local_fallback(self, curr_branc
 
         branches.assert_has_calls(calls)
         self.assertEqual(0, tags.call_count)
-        checkout.assert_called_with("/benchmark-resources/unit-test", branch="1")
+        checkout.assert_called_with("/benchmark-resources/unit-test", branch="OS-1")
         self.assertEqual(0, rebase.call_count)
 
     @mock.patch("osbenchmark.utils.git.is_working_copy", autospec=True)

tests/utils/versions_test.py

@@ -110,40 +110,40 @@ def test_versions_rejects_invalid_version_strings(self):
             versions.VersionVariants("5.0.0a-SNAPSHOT")
 
     def test_find_best_match(self):
-        assert versions.best_match(["1.7", "2", "5.0.0-alpha1", "5", "main"], "6.0.0-alpha1") == "main",\
+        assert versions.best_matching_version(["1.7", "2", "5.0.0-alpha1", "5", "main"], "6.0.0-alpha1") == "main",\
             "Assume main for versions newer than latest alternative available"
 
-        assert versions.best_match(["1.7", "2", "5.0.0-alpha1", "5", "main"], "5.1.0-SNAPSHOT") == "5",\
+        assert versions.best_matching_version(["1.7", "2", "5.0.0-alpha1", "5", "main"], "5.1.0-SNAPSHOT") == "5",\
             "Best match for specific version"
 
-        assert versions.best_match(["1.7", "2", "5.0.0-alpha1", "5", "main"], None) == "main",\
+        assert versions.best_matching_version(["1.7", "2", "5.0.0-alpha1", "5", "main"], None) == "main",\
             "Assume main on unknown version"
 
-        assert versions.best_match(["1.7", "2", "5.0.0-alpha1", "5", "main"], "0.4") is None,\
+        assert versions.best_matching_version(["1.7", "2", "5.0.0-alpha1", "5", "main"], "0.4") is None,\
             "Reject versions that are too old"
 
-        assert versions.best_match(["7", "7.10.2", "7.11", "7.2", "5", "6", "main"], "7.10.2") == "7.10.2", \
+        assert versions.best_matching_version(["7", "7.10.2", "7.11", "7.2", "5", "6", "main"], "7.10.2") == "7.10.2", \
             "Exact match"
 
-        assert versions.best_match(["7", "7.10", "main"], "7.1.0") == "7", \
+        assert versions.best_matching_version(["7", "7.10", "main"], "7.1.0") == "7", \
             "Best match is major version"
 
-        assert versions.best_match(["7", "7.11", "7.2", "5", "6", "main"], "7.11.0") == "7.11",\
+        assert versions.best_matching_version(["7", "7.11", "7.2", "5", "6", "main"], "7.11.0") == "7.11",\
             "Best match for specific minor version"
 
-        assert versions.best_match(["7", "7.11", "7.2", "5", "6", "main"], "7.12.0") == "7.11",\
+        assert versions.best_matching_version(["7", "7.11", "7.2", "5", "6", "main"], "7.12.0") == "7.11",\
             "If no exact match, best match is the nearest prior minor"
 
-        assert versions.best_match(["7", "7.11", "7.2", "5", "6", "main"], "7.3.0") == "7.2",\
+        assert versions.best_matching_version(["7", "7.11", "7.2", "5", "6", "main"], "7.3.0") == "7.2",\
             "If no exact match, best match is the nearest prior minor"
 
-        assert versions.best_match(["7", "7.11", "7.2", "5", "6", "main"], "7.10.0") == "7.2", \
+        assert versions.best_matching_version(["7", "7.11", "7.2", "5", "6", "main"], "7.10.0") == "7.2", \
             "If no exact match, best match is the nearest prior minor"
 
-        assert versions.best_match(["7", "7.1", "7.11.1", "7.11.0", "7.2", "5", "6", "main"], "7.12.0") == "7.2",\
+        assert versions.best_matching_version(["7", "7.1", "7.11.1", "7.11.0", "7.2", "5", "6", "main"], "7.12.0") == "7.2",\
             "Patch or patch-suffix branches are not supported and ignored, best match is nearest prior minor"
 
-        assert versions.best_match(["7", "7.11", "7.2", "5", "6", "main"], "7.1.0") == "7",\
+        assert versions.best_matching_version(["7", "7.11", "7.2", "5", "6", "main"], "7.1.0") == "7",\
             "If no exact match and no minor match, next best match is major version"
 
     def test_version_comparison(self):