Repository: microsoft/MLOS

Issue:
Investigate alternative to `--dist=loadgroup`
See if we can use something besides `--dist=loadgroup` for ensuring reliable ssh test runs and still get parallel execution and avoid pytest in vscode plugin parsing issues.


Discussion:

Resolving patch:
diff --git a/conftest.py b/conftest.py
--- a/conftest.py
+++ b/conftest.py
@@ -9,13 +9,25 @@
 # Note: This file is named conftest.py so that pytest picks it up automatically
 # without the need to adjust PYTHONPATH or sys.path as much.
 
-import os
 from warnings import warn
+from tempfile import mkdtemp
+
+import os
+import shutil
 
 import pytest
+from xdist.workermanage import WorkerController
+
+
+def is_master(config: pytest.Config) -> bool:
+    """
+    True if the code running the given pytest.config object is running in a
+    xdist master node or not running xdist at all.
+    """
+    return not hasattr(config, "workerinput")
 
 
-def pytest_configure(config: pytest.Config) -> None:   # pylint: disable=unused-argument
+def pytest_configure(config: pytest.Config) -> None:
     """
     Add some additional (global) configuration steps for pytest.
     """
@@ -23,5 +35,51 @@ def pytest_configure(config: pytest.Config) -> None:   # pylint: disable=unused-
     if os.environ.get('DISPLAY', None):
         import matplotlib   # pylint: disable=import-outside-toplevel
         matplotlib.rcParams['backend'] = 'agg'
-        warn(UserWarning('DISPLAY environment variable is set, which can cause problems in some setups (e.g. WSL). '
-                         + f'Adjusting matplotlib backend to "{matplotlib.rcParams["backend"]}" to compensate.'))
+        if is_master(config) or dict(getattr(config, 'workerinput', {}))['workerid'] == 'gw0':
+            # Only warn once.
+            warn(UserWarning('DISPLAY environment variable is set, which can cause problems in some setups (e.g. WSL). '
+                             + f'Adjusting matplotlib backend to "{matplotlib.rcParams["backend"]}" to compensate.'))
+
+    # Create a temporary directory for sharing files between master and worker nodes.
+    if is_master(config):
+        # Add it to the config so that it can passed to the worker nodes.
+        setattr(config, "shared_temp_dir", mkdtemp())
+
+
+def pytest_configure_node(node: WorkerController) -> None:
+    """
+    Xdist hook used to inform workers of the location of the shared temp dir.
+    """
+    workerinput: dict = getattr(node, "workerinput")
+    workerinput["shared_temp_dir"] = getattr(node.config, "shared_temp_dir")
+
+
+@pytest.fixture(scope="session")
+def shared_temp_dir(request: pytest.FixtureRequest) -> str:
+    """
+    Returns a unique and temporary directory which can be shared by
+    master or worker nodes in xdist runs.
+    """
+    if is_master(request.config):
+        return str(getattr(request.config, "shared_temp_dir"))
+    else:
+        workerinput: dict = getattr(request.config, "workerinput")
+        return str(workerinput["shared_temp_dir"])
+
+
+def pytest_unconfigure(config: pytest.Config) -> None:
+    """
+    Called after all tests have completed.
+    """
+    if is_master(config):
+        shared_tmp_dir = str(getattr(config, "shared_temp_dir"))
+        shutil.rmtree(shared_tmp_dir)
+
+
+@pytest.fixture(scope="session")
+def short_testrun_uid(testrun_uid: str) -> str:
+    """
+    Shorten the unique test run id that xdist provides so we can use it with
+    other systems (e.g., docker).
+    """
+    return testrun_uid[0:8]
diff --git a/mlos_bench/mlos_bench/services/remote/ssh/ssh_host_service.py b/mlos_bench/mlos_bench/services/remote/ssh/ssh_host_service.py
--- a/mlos_bench/mlos_bench/services/remote/ssh/ssh_host_service.py
+++ b/mlos_bench/mlos_bench/services/remote/ssh/ssh_host_service.py
@@ -196,8 +196,9 @@ def _exec_os_op(self, cmd_opts_list: List[str], params: dict) -> Tuple[Status, d
                 sudo=${{sudo:+$sudo -n}}
             fi
 
+            set -x
             for cmd in {cmd_opts}; do
-                $sudo $cmd && exit 0
+                $sudo /bin/bash -c "$cmd" && exit 0
             done
 
             echo 'ERROR: Failed to shutdown/reboot the system.'
@@ -252,7 +253,7 @@ def reboot(self, params: dict, force: bool = False) -> Tuple[Status, dict]:
             'reboot',
             'halt --reboot',
             'systemctl reboot',
-            'kill -9 1' if force else 'kill 1',
+            'kill -KILL 1; kill -KILL -1' if force else 'kill -TERM 1; kill -TERM -1',
         ]
         return self._exec_os_op(cmd_opts_list=cmd_opts_list, params=params)
 
diff --git a/mlos_bench/setup.py b/mlos_bench/setup.py
--- a/mlos_bench/setup.py
+++ b/mlos_bench/setup.py
@@ -50,6 +50,7 @@
     'pytest-local-badge',
     'pytest-lazy-fixture',
     'pytest-docker',
+    'fasteners',
 ]
 
 # pylint: disable=duplicate-code


Test patch:
diff --git a/mlos_bench/mlos_bench/tests/__init__.py b/mlos_bench/mlos_bench/tests/__init__.py
--- a/mlos_bench/mlos_bench/tests/__init__.py
+++ b/mlos_bench/mlos_bench/tests/__init__.py
@@ -25,7 +25,7 @@
 # Use with @requires_docker above a test_...() function.
 DOCKER = shutil.which('docker')
 if DOCKER:
-    cmd = run("docker builder inspect default", shell=True, check=False, capture_output=True)
+    cmd = run("docker builder inspect default || docker buildx inspect default", shell=True, check=False, capture_output=True)
     stdout = cmd.stdout.decode()
     if cmd.returncode != 0 or not any(line for line in stdout.splitlines() if 'Platform' in line and 'linux' in line):
         debug("Docker is available but missing support for targeting linux platform.")
diff --git a/mlos_bench/mlos_bench/tests/conftest.py b/mlos_bench/mlos_bench/tests/conftest.py
--- a/mlos_bench/mlos_bench/tests/conftest.py
+++ b/mlos_bench/mlos_bench/tests/conftest.py
@@ -6,10 +6,13 @@
 Common fixtures for mock TunableGroups and Environment objects.
 """
 
-from typing import Any, Dict, List
+from typing import Any, Dict, Generator, List
 
 import os
 
+from fasteners import InterProcessLock, InterProcessReaderWriterLock
+from pytest_docker.plugin import get_docker_services, Services as DockerServices
+
 import json5 as json
 import pytest
 
@@ -168,7 +171,7 @@ def docker_compose_file(pytestconfig: pytest.Config) -> List[str]:
 
 
 @pytest.fixture(scope="session")
-def docker_compose_project_name() -> str:
+def docker_compose_project_name(short_testrun_uid: str) -> str:
     """
     Returns the name of the docker-compose project.
 
@@ -177,4 +180,84 @@ def docker_compose_project_name() -> str:
     str
         Name of the docker-compose project.
     """
-    return f"mlos_bench-test-{os.getpid()}"
+    # Use the xdist testrun UID to ensure that the docker-compose project name
+    # is unique across sessions, but shared amongst workers.
+    return f"mlos_bench-test-{short_testrun_uid}"
+
+
+@pytest.fixture(scope="session")
+def docker_services_lock(shared_temp_dir: str, short_testrun_uid: str) -> InterProcessReaderWriterLock:
+    """
+    Gets a pytest session lock for xdist workers to mark when they're using the
+    docker services.
+
+    Yields
+    ------
+        A lock to ensure that setup/teardown operations don't happen while a
+        worker is using the docker services.
+    """
+    return InterProcessReaderWriterLock(f"{shared_temp_dir}/pytest_docker_services-{short_testrun_uid}.lock")
+
+
+@pytest.fixture(scope="session")
+def docker_setup_teardown_lock(shared_temp_dir: str, short_testrun_uid: str) -> InterProcessLock:
+    """
+    Gets a pytest session lock between xdist workers for the docker
+    setup/teardown operations.
+
+    Yields
+    ------
+        A lock to ensure that only one worker is doing setup/teardown at a time.
+    """
+    return InterProcessLock(f"{shared_temp_dir}/pytest_docker_services-setup-teardown-{short_testrun_uid}.lock")
+
+
+@pytest.fixture(scope="session")
+def locked_docker_services(
+    docker_compose_command: Any,
+    docker_compose_file: Any,
+    docker_compose_project_name: Any,
+    docker_setup: Any,
+    docker_cleanup: Any,
+    docker_setup_teardown_lock: InterProcessLock,
+    docker_services_lock: InterProcessReaderWriterLock,
+) -> Generator[DockerServices, Any, None]:
+    """
+    A locked version of the docker_services fixture to implement xdist single instance locking.
+    """
+    # pylint: disable=too-many-arguments
+    # Mark the services as in use with the reader lock.
+    docker_services_lock.acquire_read_lock()
+    # Acquire the setup lock to prevent multiple setup operations at once.
+    docker_setup_teardown_lock.acquire()
+    # This "with get_docker_services(...)"" pattern is in the default fixture.
+    # We call it instead of docker_services() to avoid pytest complaints about
+    # calling fixtures directly.
+    with get_docker_services(
+        docker_compose_command,
+        docker_compose_file,
+        docker_compose_project_name,
+        docker_setup,
+        docker_cleanup,
+    ) as docker_services:
+        # Release the setup/tear down lock in order to let the setup operation
+        # continue for other workers (should be a no-op at this point).
+        docker_setup_teardown_lock.release()
+        # Yield the services so that tests within this worker can use them.
+        yield docker_services
+        # Now tests that use those services get to run on this worker...
+        # Once the tests are done, release the read lock that marks the services as in use.
+        docker_services_lock.release_read_lock()
+        # Now as we prepare to execute the cleanup code on context exit we need
+        # to acquire the setup/teardown lock again.
+        # First we attempt to get the write lock so that we wait for other
+        # readers to finish and guard against a lock inversion possibility.
+        docker_services_lock.acquire_write_lock()
+        # Next, acquire the setup/teardown lock
+        # First one here is the one to do actual work, everyone else is basically a no-op.
+        # Upon context exit, we should execute the docker_cleanup code.
+        # And try to get the setup/tear down lock again.
+        docker_setup_teardown_lock.acquire()
+    # Finally, after the docker_cleanup code has finished, remove both locks.
+    docker_setup_teardown_lock.release()
+    docker_services_lock.release_write_lock()
diff --git a/mlos_bench/mlos_bench/tests/environments/remote/conftest.py b/mlos_bench/mlos_bench/tests/environments/remote/conftest.py
--- a/mlos_bench/mlos_bench/tests/environments/remote/conftest.py
+++ b/mlos_bench/mlos_bench/tests/environments/remote/conftest.py
@@ -8,9 +8,6 @@
 
 import mlos_bench.tests.services.remote.ssh.fixtures as ssh_fixtures
 
-# Expose those as local names.
-ssh_test_server_hostname = ssh_fixtures.ssh_test_server_hostname
+# Expose some of those as local names so they can be picked up as fixtures by pytest.
 ssh_test_server = ssh_fixtures.ssh_test_server
-alt_test_server = ssh_fixtures.alt_test_server
-ssh_host_service = ssh_fixtures.ssh_host_service
-ssh_fileshare_service = ssh_fixtures.ssh_fileshare_service
+ssh_test_server_hostname = ssh_fixtures.ssh_test_server_hostname
diff --git a/mlos_bench/mlos_bench/tests/environments/remote/test_ssh_env.py b/mlos_bench/mlos_bench/tests/environments/remote/test_ssh_env.py
--- a/mlos_bench/mlos_bench/tests/environments/remote/test_ssh_env.py
+++ b/mlos_bench/mlos_bench/tests/environments/remote/test_ssh_env.py
@@ -28,7 +28,6 @@
 
 
 @requires_docker
-@pytest.mark.xdist_group("ssh_test_server")
 def test_remote_ssh_env(tunable_groups: TunableGroups, ssh_test_server: SshTestServerInfo) -> None:
     """
     Produce benchmark and telemetry data in a local script and read it.
diff --git a/mlos_bench/mlos_bench/tests/event_loop_context_test.py b/mlos_bench/mlos_bench/tests/event_loop_context_test.py
--- a/mlos_bench/mlos_bench/tests/event_loop_context_test.py
+++ b/mlos_bench/mlos_bench/tests/event_loop_context_test.py
@@ -48,6 +48,7 @@ def __exit__(self, ex_type: Optional[Type[BaseException]],
         return False
 
 
+@pytest.mark.filterwarnings("ignore:.*(coroutine 'sleep' was never awaited).*:RuntimeWarning:.*event_loop_context_test.*:0")
 def test_event_loop_context() -> None:
     """Test event loop context background thread setup/cleanup handling."""
     # pylint: disable=protected-access
@@ -111,8 +112,9 @@ def test_event_loop_context() -> None:
     assert EventLoopContextCaller.EVENT_LOOP_CONTEXT._event_loop_thread_refcnt == 0
     assert EventLoopContextCaller.EVENT_LOOP_CONTEXT._event_loop is None
 
-    with pytest.raises(AssertionError):
-        event_loop_caller_instance_1.EVENT_LOOP_CONTEXT.run_coroutine(asyncio.sleep(0.1))
+    with pytest.raises(AssertionError), pytest.warns(RuntimeWarning, match=r".*coroutine 'sleep' was never awaited"):
+        future = event_loop_caller_instance_1.EVENT_LOOP_CONTEXT.run_coroutine(asyncio.sleep(0.1, result='foo'))
+        raise ValueError(f"Future should not have been available to wait on {future.result()}")
 
 
 if __name__ == '__main__':
diff --git a/mlos_bench/mlos_bench/tests/services/remote/ssh/Dockerfile b/mlos_bench/mlos_bench/tests/services/remote/ssh/Dockerfile
--- a/mlos_bench/mlos_bench/tests/services/remote/ssh/Dockerfile
+++ b/mlos_bench/mlos_bench/tests/services/remote/ssh/Dockerfile
@@ -13,4 +13,5 @@ RUN echo "Port ${PORT}" > /etc/ssh/sshd_config.d/local.conf \
     && cat /root/.ssh/id_rsa.pub > /root/.ssh/authorized_keys
 
 ENV TIMEOUT=180
-CMD ["/bin/bash", "-eux", "-c", "trap : TERM INT; service ssh start && sleep ${TIMEOUT:-180} & wait"]
+# NOTE: The "-n" is important in the call to "wait" so that an exit failure can be detected.
+CMD ["/bin/bash", "-eux", "-c", "trap : TERM INT; service ssh start && sleep ${TIMEOUT:-180} & wait -n"]
diff --git a/mlos_bench/mlos_bench/tests/services/remote/ssh/README.md b/mlos_bench/mlos_bench/tests/services/remote/ssh/README.md
--- a/mlos_bench/mlos_bench/tests/services/remote/ssh/README.md
+++ b/mlos_bench/mlos_bench/tests/services/remote/ssh/README.md
@@ -20,8 +20,9 @@ These are brought up as session fixtures under a unique (PID based) compose proj
 
 > For manual testing, to bring up/down the test infrastructure the [`up.sh`](./up.sh) and [`down.sh`](./down.sh) scripts can be used, which assigns a known project name.
 
-In the case of `pytest`, since the `SshService` base class implements a shared connection cache that we wish to test, and testing "rebooting" of servers (containers) is also necessary, tests are run serially across a single worker by using the `pytest-xdist` plugin's `--dist loadgroup` feature and the `@pytest.mark.xdist_group("ssh_test_server")` decorator.
-In some cases we explicitly call the python garbage collector via `gc.collect()` to make sure that the shared cache cleanup handler is operating as expected.
+In the case of `pytest`, since the `SshService` base class implements a shared connection cache that we wish to test, and testing "rebooting" of servers (containers) is also necessary, but we want to avoid single threaded execution for tests, we start a third container only for testing reboots.
+
+Additionally, since `scope="session"` fixtures are executed once per worker, which is excessive in our case, we use lockfiles (one of the only portal synchronization methods) to ensure that the usual `docker_services` fixture which starts and stops the containers is only executed once per test run and uses a shared compose instance.
 
 ## See Also
 
diff --git a/mlos_bench/mlos_bench/tests/services/remote/ssh/__init__.py b/mlos_bench/mlos_bench/tests/services/remote/ssh/__init__.py
--- a/mlos_bench/mlos_bench/tests/services/remote/ssh/__init__.py
+++ b/mlos_bench/mlos_bench/tests/services/remote/ssh/__init__.py
@@ -20,6 +20,7 @@
 SSH_TEST_SERVER_PORT = 2254
 SSH_TEST_SERVER_NAME = 'ssh-server'
 ALT_TEST_SERVER_NAME = 'alt-server'
+REBOOT_TEST_SERVER_NAME = 'reboot-server'
 
 
 @dataclass
diff --git a/mlos_bench/mlos_bench/tests/services/remote/ssh/conftest.py b/mlos_bench/mlos_bench/tests/services/remote/ssh/conftest.py
--- a/mlos_bench/mlos_bench/tests/services/remote/ssh/conftest.py
+++ b/mlos_bench/mlos_bench/tests/services/remote/ssh/conftest.py
@@ -8,9 +8,10 @@
 
 import mlos_bench.tests.services.remote.ssh.fixtures as ssh_fixtures
 
-# Expose those as local names.
+# Expose some of those as local names so they can be picked up as fixtures by pytest.
 ssh_test_server_hostname = ssh_fixtures.ssh_test_server_hostname
 ssh_test_server = ssh_fixtures.ssh_test_server
 alt_test_server = ssh_fixtures.alt_test_server
+reboot_test_server = ssh_fixtures.reboot_test_server
 ssh_host_service = ssh_fixtures.ssh_host_service
 ssh_fileshare_service = ssh_fixtures.ssh_fileshare_service
diff --git a/mlos_bench/mlos_bench/tests/services/remote/ssh/docker-compose.yml b/mlos_bench/mlos_bench/tests/services/remote/ssh/docker-compose.yml
--- a/mlos_bench/mlos_bench/tests/services/remote/ssh/docker-compose.yml
+++ b/mlos_bench/mlos_bench/tests/services/remote/ssh/docker-compose.yml
@@ -28,7 +28,29 @@ services:
       - ssh-server
     hostname: alt-server
     attach: false
-    restart: always
+    image: mlos_bench-test-ssh-server:latest
+    ports:
+      - ${PORT:-2254}
+    extra_hosts:
+      - host.docker.internal:host-gateway
+    environment:
+      # Let the environment variable TIMEOUT override the default.
+      - TIMEOUT=${TIMEOUT:-180}
+  # Add a third for testing only reboot operations.
+  reboot-server:
+    depends_on:
+      - ssh-server
+    hostname: reboot-server
+    attach: false
+    deploy:
+      # Only restart a few times so that we eventually cleanup after
+      # ourselves in the case of an interrupted pytest failure.
+      # Note: this also relies on "wait -n" in the CMD to exit properly on
+      # various signals.
+      restart_policy:
+        condition: on-failure
+        delay: 0s
+        max_attempts: 3
     image: mlos_bench-test-ssh-server:latest
     ports:
       - ${PORT:-2254}
diff --git a/mlos_bench/mlos_bench/tests/services/remote/ssh/fixtures.py b/mlos_bench/mlos_bench/tests/services/remote/ssh/fixtures.py
--- a/mlos_bench/mlos_bench/tests/services/remote/ssh/fixtures.py
+++ b/mlos_bench/mlos_bench/tests/services/remote/ssh/fixtures.py
@@ -24,6 +24,7 @@
 from mlos_bench.tests import resolve_host_name
 from mlos_bench.tests.services.remote.ssh import (SshTestServerInfo,
                                                   ALT_TEST_SERVER_NAME,
+                                                  REBOOT_TEST_SERVER_NAME,
                                                   SSH_TEST_SERVER_NAME,
                                                   wait_docker_service_socket)
 
@@ -48,7 +49,7 @@ def ssh_test_server_hostname() -> str:
 @pytest.fixture(scope="session")
 def ssh_test_server(ssh_test_server_hostname: str,
                     docker_compose_project_name: str,
-                    docker_services: DockerServices) -> Generator[SshTestServerInfo, None, None]:
+                    locked_docker_services: DockerServices) -> Generator[SshTestServerInfo, None, None]:
     """
     Fixture for getting the ssh test server services setup via docker-compose
     using pytest-docker.
@@ -67,7 +68,7 @@ def ssh_test_server(ssh_test_server_hostname: str,
             hostname=ssh_test_server_hostname,
             username='root',
             id_rsa_path=id_rsa_file.name)
-        wait_docker_service_socket(docker_services, ssh_test_server_info.hostname, ssh_test_server_info.get_port())
+        wait_docker_service_socket(locked_docker_services, ssh_test_server_info.hostname, ssh_test_server_info.get_port())
         id_rsa_src = f"/{ssh_test_server_info.username}/.ssh/id_rsa"
         docker_cp_cmd = f"docker compose -p {docker_compose_project_name} cp {SSH_TEST_SERVER_NAME}:{id_rsa_src} {id_rsa_file.name}"
         cmd = run(docker_cp_cmd.split(), check=True, cwd=os.path.dirname(__file__), capture_output=True, text=True)
@@ -81,7 +82,7 @@ def ssh_test_server(ssh_test_server_hostname: str,
 
 @pytest.fixture(scope="session")
 def alt_test_server(ssh_test_server: SshTestServerInfo,
-                    docker_services: DockerServices) -> SshTestServerInfo:
+                    locked_docker_services: DockerServices) -> SshTestServerInfo:
     """
     Fixture for getting the second ssh test server info from the docker-compose.yml.
     See additional notes in the ssh_test_server fixture above.
@@ -95,10 +96,30 @@ def alt_test_server(ssh_test_server: SshTestServerInfo,
         hostname=ssh_test_server.hostname,
         username=ssh_test_server.username,
         id_rsa_path=ssh_test_server.id_rsa_path)
-    wait_docker_service_socket(docker_services, alt_test_server_info.hostname, alt_test_server_info.get_port())
+    wait_docker_service_socket(locked_docker_services, alt_test_server_info.hostname, alt_test_server_info.get_port())
     return alt_test_server_info
 
 
+@pytest.fixture(scope="session")
+def reboot_test_server(ssh_test_server: SshTestServerInfo,
+                       locked_docker_services: DockerServices) -> SshTestServerInfo:
+    """
+    Fixture for getting the third ssh test server info from the docker-compose.yml.
+    See additional notes in the ssh_test_server fixture above.
+    """
+    # Note: The reboot-server uses the same image as the ssh-server container, so
+    # the id_rsa key and username should all match.
+    # Only the host port it is allocate is different.
+    reboot_test_server_info = SshTestServerInfo(
+        compose_project_name=ssh_test_server.compose_project_name,
+        service_name=REBOOT_TEST_SERVER_NAME,
+        hostname=ssh_test_server.hostname,
+        username=ssh_test_server.username,
+        id_rsa_path=ssh_test_server.id_rsa_path)
+    wait_docker_service_socket(locked_docker_services, reboot_test_server_info.hostname, reboot_test_server_info.get_port())
+    return reboot_test_server_info
+
+
 @pytest.fixture
 def ssh_host_service(ssh_test_server: SshTestServerInfo) -> SshHostService:
     """Generic SshHostService fixture."""
diff --git a/mlos_bench/mlos_bench/tests/services/remote/ssh/test_ssh_fileshare.py b/mlos_bench/mlos_bench/tests/services/remote/ssh/test_ssh_fileshare.py
--- a/mlos_bench/mlos_bench/tests/services/remote/ssh/test_ssh_fileshare.py
+++ b/mlos_bench/mlos_bench/tests/services/remote/ssh/test_ssh_fileshare.py
@@ -53,7 +53,6 @@ def closeable_temp_file(**kwargs: Any) -> Generator[_TemporaryFileWrapper, None,
             os.unlink(fname)
 
 
-@pytest.mark.xdist_group("ssh_test_server")
 @requires_docker
 def test_ssh_fileshare_single_file(ssh_test_server: SshTestServerInfo,
                                    ssh_fileshare_service: SshFileShareService) -> None:
@@ -94,7 +93,6 @@ def test_ssh_fileshare_single_file(ssh_test_server: SshTestServerInfo,
                 assert read_lines == lines
 
 
-@pytest.mark.xdist_group("ssh_test_server")
 @requires_docker
 def test_ssh_fileshare_recursive(ssh_test_server: SshTestServerInfo,
                                  ssh_fileshare_service: SshFileShareService) -> None:
@@ -150,7 +148,6 @@ def test_ssh_fileshare_recursive(ssh_test_server: SshTestServerInfo,
             assert are_dir_trees_equal(tempdir1, path_join(tempdir2, basename(remote_file_path)))
 
 
-@pytest.mark.xdist_group("ssh_test_server")
 @requires_docker
 def test_ssh_fileshare_download_file_dne(ssh_test_server: SshTestServerInfo,
                                          ssh_fileshare_service: SshFileShareService) -> None:
@@ -176,7 +173,6 @@ def test_ssh_fileshare_download_file_dne(ssh_test_server: SshTestServerInfo,
             assert read_lines == [canary_str]
 
 
-@pytest.mark.xdist_group("ssh_test_server")
 @requires_docker
 def test_ssh_fileshare_upload_file_dne(ssh_test_server: SshTestServerInfo,
                                        ssh_host_service: SshHostService,
diff --git a/mlos_bench/mlos_bench/tests/services/remote/ssh/test_ssh_host_service.py b/mlos_bench/mlos_bench/tests/services/remote/ssh/test_ssh_host_service.py
--- a/mlos_bench/mlos_bench/tests/services/remote/ssh/test_ssh_host_service.py
+++ b/mlos_bench/mlos_bench/tests/services/remote/ssh/test_ssh_host_service.py
@@ -6,9 +6,11 @@
 Tests for mlos_bench.services.remote.ssh.ssh_host_service
 """
 
+from subprocess import CalledProcessError, run
+
+import logging
 import time
 
-import pytest
 from pytest_docker.plugin import Services as DockerServices
 
 from mlos_bench.services.remote.ssh.ssh_host_service import SshHostService
@@ -17,12 +19,14 @@
 from mlos_bench.tests import requires_docker
 from mlos_bench.tests.services.remote.ssh import (SshTestServerInfo,
                                                   ALT_TEST_SERVER_NAME,
+                                                  REBOOT_TEST_SERVER_NAME,
                                                   SSH_TEST_SERVER_NAME,
                                                   wait_docker_service_socket)
 
+_LOG = logging.getLogger(__name__)
+
 
 @requires_docker
-@pytest.mark.xdist_group("ssh_test_server")
 def test_ssh_service_remote_exec(ssh_test_server: SshTestServerInfo,
                                  alt_test_server: SshTestServerInfo,
                                  ssh_host_service: SshHostService) -> None:
@@ -36,8 +40,6 @@ def test_ssh_service_remote_exec(ssh_test_server: SshTestServerInfo,
     with ssh_host_service:
         config = ssh_test_server.to_ssh_service_config()
 
-        # Note: this may interact with state from other tests, so we group them
-        # into the same xdist_group, run with --dist loadgroup.
         connection_id = SshClient.id_from_params(ssh_test_server.to_connect_params())
         assert ssh_host_service._EVENT_LOOP_THREAD_SSH_CLIENT_CACHE is not None
         connection_client = ssh_host_service._EVENT_LOOP_THREAD_SSH_CLIENT_CACHE._cache.get(connection_id)
@@ -132,75 +134,86 @@ def test_ssh_service_remote_exec(ssh_test_server: SshTestServerInfo,
     assert len(SshHostService._EVENT_LOOP_THREAD_SSH_CLIENT_CACHE) == 0
 
 
-@requires_docker
-@pytest.mark.parametrize("graceful", [True, False])
-@pytest.mark.xdist_group("ssh_test_server")
-def test_ssh_service_reboot(docker_services: DockerServices,
-                            alt_test_server: SshTestServerInfo,
-                            ssh_host_service: SshHostService,
-                            graceful: bool) -> None:
+def check_ssh_service_reboot(docker_services: DockerServices,
+                             reboot_test_server: SshTestServerInfo,
+                             ssh_host_service: SshHostService,
+                             graceful: bool) -> None:
     """
-    Test the SshHostService reboot.
+    Check the SshHostService reboot operation.
     """
     # Note: rebooting changes the port number unfortunately, but makes it
     # easier to check for success.
     # Also, it may cause issues with other parallel unit tests, so we run it as
     # a part of the same unit test for now.
     with ssh_host_service:
-        alt_test_server_ssh_service_config = alt_test_server.to_ssh_service_config()
+        reboot_test_srv_ssh_svc_conf = reboot_test_server.to_ssh_service_config(uncached=True)
         (status, results_info) = ssh_host_service.remote_exec(
             script=[
                 'echo "sleeping..."',
                 'sleep 30',
                 'echo "should not reach this point"'
             ],
-            config=alt_test_server_ssh_service_config,
+            config=reboot_test_srv_ssh_svc_conf,
             env_params={},
         )
         assert status.is_pending()
         # Wait a moment for that to start in the background thread.
-        time.sleep(0.5)
-
-        # Now try to restart the server (gracefully).
-        # TODO: Test graceful vs. forceful.
-        if graceful:
-            (status, reboot_results_info) = ssh_host_service.reboot(params=alt_test_server_ssh_service_config)
-            assert status.is_pending()
-
-            (status, reboot_results_info) = ssh_host_service.wait_os_operation(reboot_results_info)
-            # NOTE: reboot/shutdown ops mostly return FAILED, even though the reboot succeeds.
-            print(f"reboot status: {status} {reboot_results_info}")
-        else:
-            (status, kill_results_info) = ssh_host_service.remote_exec(
-                script=["kill -9 1; kill -9 -1"],
-                config=alt_test_server_ssh_service_config,
-                env_params={},
-            )
-            (status, kill_results_info) = ssh_host_service.get_remote_exec_results(kill_results_info)
-            print(f"kill status: {status} {kill_results_info}")
-
-        # TODO: Check for decent error handling on disconnects.
+        time.sleep(1)
+
+        # Now try to restart the server.
+        (status, reboot_results_info) = ssh_host_service.reboot(params=reboot_test_srv_ssh_svc_conf,
+                                                                force=not graceful)
+        assert status.is_pending()
+
+        (status, reboot_results_info) = ssh_host_service.wait_os_operation(reboot_results_info)
+        # NOTE: reboot/shutdown ops mostly return FAILED, even though the reboot succeeds.
+        _LOG.debug("reboot status: %s: %s", status, reboot_results_info)
+
+        # Check for decent error handling on disconnects.
         status, results = ssh_host_service.get_remote_exec_results(results_info)
         assert status.is_failed()
         stdout = str(results["stdout"])
         assert "sleeping" in stdout
         assert "should not reach this point" not in stdout
 
-        # Give docker some time to restart the service after the "reboot".
-        # Note: this relies on having `restart: always` in the docker-compose.yml file.
-        time.sleep(1)
-
-        # try to reconnect and see if the port changed
-        alt_test_server_ssh_service_config_new = alt_test_server.to_ssh_service_config(uncached=True)
-        assert alt_test_server_ssh_service_config_new["ssh_port"] != alt_test_server_ssh_service_config["ssh_port"]
-
-        wait_docker_service_socket(docker_services, alt_test_server.hostname, alt_test_server_ssh_service_config_new["ssh_port"])
+        reboot_test_srv_ssh_svc_conf_new: dict = {}
+        for _ in range(0, 3):
+            # Give docker some time to restart the service after the "reboot".
+            # Note: this relies on having a `restart_policy` in the docker-compose.yml file.
+            time.sleep(1)
+            # try to reconnect and see if the port changed
+            try:
+                run_res = run("docker ps | grep mlos_bench-test- | grep reboot", shell=True, capture_output=True, check=False)
+                print(run_res.stdout.decode())
+                print(run_res.stderr.decode())
+                reboot_test_srv_ssh_svc_conf_new = reboot_test_server.to_ssh_service_config(uncached=True)
+                if reboot_test_srv_ssh_svc_conf_new["ssh_port"] != reboot_test_srv_ssh_svc_conf["ssh_port"]:
+                    break
+            except CalledProcessError as ex:
+                _LOG.info("Failed to check port for reboot test server: %s", ex)
+        assert reboot_test_srv_ssh_svc_conf_new["ssh_port"] != reboot_test_srv_ssh_svc_conf["ssh_port"]
+
+        wait_docker_service_socket(docker_services,
+                                   reboot_test_server.hostname,
+                                   reboot_test_srv_ssh_svc_conf_new["ssh_port"])
 
         (status, results_info) = ssh_host_service.remote_exec(
             script=["hostname"],
-            config=alt_test_server_ssh_service_config_new,
+            config=reboot_test_srv_ssh_svc_conf_new,
             env_params={},
         )
         status, results = ssh_host_service.get_remote_exec_results(results_info)
         assert status.is_succeeded()
-        assert results["stdout"].strip() == ALT_TEST_SERVER_NAME
+        assert results["stdout"].strip() == REBOOT_TEST_SERVER_NAME
+
+
+@requires_docker
+def test_ssh_service_reboot(locked_docker_services: DockerServices,
+                            reboot_test_server: SshTestServerInfo,
+                            ssh_host_service: SshHostService) -> None:
+    """
+    Test the SshHostService reboot operation.
+    """
+    # Grouped together to avoid parallel runner interactions.
+    check_ssh_service_reboot(locked_docker_services, reboot_test_server, ssh_host_service, graceful=True)
+    check_ssh_service_reboot(locked_docker_services, reboot_test_server, ssh_host_service, graceful=False)
diff --git a/mlos_bench/mlos_bench/tests/services/remote/ssh/test_ssh_service.py b/mlos_bench/mlos_bench/tests/services/remote/ssh/test_ssh_service.py
--- a/mlos_bench/mlos_bench/tests/services/remote/ssh/test_ssh_service.py
+++ b/mlos_bench/mlos_bench/tests/services/remote/ssh/test_ssh_service.py
@@ -25,7 +25,6 @@
 
 @requires_docker
 @requires_ssh
-@pytest.mark.xdist_group("ssh_test_server")
 @pytest.mark.parametrize(["ssh_test_server_info", "server_name"], [
     (lazy_fixture("ssh_test_server"), SSH_TEST_SERVER_NAME),
     (lazy_fixture("alt_test_server"), ALT_TEST_SERVER_NAME),
@@ -50,7 +49,7 @@ def test_ssh_service_test_infra(ssh_test_server_info: SshTestServerInfo,
     assert cmd.stdout.strip() == server_name
 
 
-@pytest.mark.xdist_group("ssh_test_server")
+@pytest.mark.filterwarnings("ignore:.*(coroutine 'sleep' was never awaited).*:RuntimeWarning:.*event_loop_context_test.*:0")
 def test_ssh_service_context_handler() -> None:
     """
     Test the SSH service context manager handling.
@@ -92,8 +91,9 @@ def test_ssh_service_context_handler() -> None:
 
         assert not ssh_fileshare_service._in_context
         # And that instance should be unusable after we are outside the context.
-        with pytest.raises(AssertionError):
-            ssh_fileshare_service._run_coroutine(asyncio.sleep(0.1))
+        with pytest.raises(AssertionError), pytest.warns(RuntimeWarning, match=r".*coroutine 'sleep' was never awaited"):
+            future = ssh_fileshare_service._run_coroutine(asyncio.sleep(0.1, result='foo'))
+            raise ValueError(f"Future should not have been available to wait on {future.result()}")
 
         # The background thread should remain running since we have another context still open.
         assert isinstance(SshService._EVENT_LOOP_CONTEXT._event_loop_thread, Thread)