Support retries in MultiClusterWrapper

Makefile

@@ -21,7 +21,7 @@ PYENV_REGEX = .pyenv/shims
 PY_BIN = python3
 # https://github.com/pypa/pip/issues/5599
 PIP_WRAPPER = $(PY_BIN) -m pip
-export PY38 = "3.8.13"
+export PY312 = "3.12.2"
 VIRTUAL_ENV ?= .venv
 VENV_ACTIVATE_FILE = $(VIRTUAL_ENV)/bin/activate
 VENV_ACTIVATE = . $(VENV_ACTIVATE_FILE)
@@ -31,8 +31,8 @@ PYENV_PREREQ_HELP = "\033[0;31mIMPORTANT\033[0m: please type \033[0;31mpyenv ini
 VE_MISSING_HELP = "\033[0;31mIMPORTANT\033[0m: Couldn't find $(PWD)/$(VIRTUAL_ENV); have you executed make venv-create?\033[0m\n"
 
 prereq:
-	pyenv install --skip-existing $(PY38)
-	pyenv local $(PY38)
+	pyenv install --skip-existing $(PY312)
+	pyenv local $(PY312)
 
 venv-create:
 	@if [[ ! -x $$(command -v pyenv) ]]; then \

elastic/shared/runners/remote_cluster.py

@@ -215,10 +215,16 @@ async def __call__(self, multi_es, params):
             if cluster_name in params.get("ignore-clusters", []):
                 self.logger.info(f"Multi cluster wrapped runner [{base_runner}] ignoring cluster [{cluster_name}].")
                 continue
-            runner_for_op = unwrap(runner_for(base_runner))
+            runner_for_op = runner_for(base_runner)
             self.logger.info(f"Multi cluster wrapped runner [{base_runner}] executing on cluster [{cluster_name}].")
-            # just call base runner op, don't mess with 'return' values
-            coroutines.append(runner_for_op(cluster_client, params))
+            # determine type of runner and call accordingly assuming the presence of MultiClientWrapper (!)
+            # see https://github.com/elastic/rally/pull/488 and https://github.com/elastic/rally/pull/1563
+            #
+            # don't mess with 'return' values
+            if not hasattr(unwrap(runner_for_op), "multi_cluster"):
+                coroutines.append(runner_for_op({"default": cluster_client}, params))
+            else:
+                coroutines.append(runner_for_op(multi_es, params))
         await asyncio.gather(*coroutines)
 
     def __repr__(self, *args, **kwargs):

elastic/tests/runners/remote_cluster_test.py

@@ -19,6 +19,7 @@
 from unittest import mock
 
 import pytest
+from esrally.driver.runner import MultiClientRunner
 from shared.runners.remote_cluster import (
     ConfigureCrossClusterReplication,
     ConfigureRemoteClusters,
@@ -468,7 +469,7 @@ def setup_params(self):
 
     @pytest.mark.asyncio
     @mock.patch("shared.runners.remote_cluster.runner_for")
-    async def test_wraps_correctly(self, mocked_runner_for, setup_es, setup_params):
+    async def test_wraps_single_cluster_runner_correctly(self, mocked_runner_for, setup_es, setup_params):
         class UnitTestSingleClusterRunner:
             async def __call__(self, es, params):
                 es.test_method(params["base-runner-param"])
@@ -477,8 +478,37 @@ async def __call__(self, es, params):
             def __str__(self):
                 return "UnitTestSingleClusterRunner"
 
-        base_runner = UnitTestSingleClusterRunner()
-        mocked_runner_for.return_value = base_runner
+        scr = UnitTestSingleClusterRunner()
+        # mimicks Rally runner internals, ugly but necessary
+        mocked_runner_for.return_value = MultiClientRunner(scr, str(scr), lambda es: es["default"])
+
+        mcw = MultiClusterWrapper()
+        r = await mcw(setup_es, setup_params)
+
+        for cluster_name, _ in setup_es.items():
+            # skipped clusters
+            if cluster_name not in ["cluster_0", "cluster_1"]:
+                setup_es[cluster_name].test_method.assert_has_calls([mock.call(setup_params["base-runner-param"])])
+
+    @pytest.mark.asyncio
+    @mock.patch("shared.runners.remote_cluster.runner_for")
+    async def test_wraps_multi_cluster_runner_correctly(self, mocked_runner_for, setup_es, setup_params):
+        # in this test multi-cluster wrapper is used together with multi-cluster runner creating a double loop
+        # the external loop is executed by the wrapper, while the internal by the runner
+        class UnitTestMultiClusterRunner:
+            multi_cluster = True
+
+            async def __call__(self, es, params):
+                for _, cluster_client in es.items():
+                    cluster_client.test_method(params["base-runner-param"])
+                return {"weight": len(es), "unit": "ops", "test": "value"}
+
+            def __str__(self):
+                return "UnitTestMultiClusterRunner"
+
+        mcr = UnitTestMultiClusterRunner()
+        # mimicks Rally runner internals, ugly but necessary
+        mocked_runner_for.return_value = MultiClientRunner(mcr, str(mcr), lambda es: es)
 
         mcw = MultiClusterWrapper()
         r = await mcw(setup_es, setup_params)

pyproject.toml

@@ -1,15 +1,15 @@
 [build-system]
-requires = ["hatchling<=1.18.0", "hatch-vcs"]
+requires = ["hatchling", "hatch-vcs"]
 build-backend = "hatchling.build"
 
 [project.optional-dependencies]
 develop = [
-    "hatch==1.7.0",
-    "hatchling==1.18.0",
-    "black==23.3.0",
+    "hatch==1.14.0",
+    "hatchling==1.26.3",
+    "black==24.10.0",
     "isort==5.12.0",
     "pre-commit==3.3.3",
-    "pip==22.2",
+    "pip==24.0",
 ]
 
 [tool.hatch.metadata]
@@ -22,15 +22,19 @@ source = "vcs"
 name = "rally-tracks"
 readme = "README.md"
 dynamic = ["version"]
-requires-python = ">=3.8"
+requires-python = ">=3.9"
 
 [tool.hatch.build.targets.sdist]
+include = [ "*/**" ]
 exclude = [
   "/.buildkite",
   "/.github",
   "/.ci",
 ]
 
+[tool.hatch.build.targets.wheel]
+include = [ "pyproject.toml" ]
+
 [tool.hatch.envs.default]
 dependencies = [
     "esrally[develop] @ git+https://github.com/elastic/rally.git@master",
@@ -63,7 +67,7 @@ asyncio_mode = "strict"
 
 [tool.black]
 line-length = 140
-target-version = ['py38']
+target-version = ['py39', 'py310', 'py311', 'py312']
 
 [tool.isort]
 profile = 'black'