Add RemoteEnv tests using SshService testing infrastructure

.vscode/launch.json

@@ -11,7 +11,11 @@
             "justMyCode": false,
             "env": {
                 // When debugging tests, only run a single pytest worker instance.
-                "PYTEST_ADDOPTS": "-n1 --no-cov --dist=no --log-level=DEBUG"
+                "PYTEST_ADDOPTS": "-n1 --no-cov --dist=no --log-level=DEBUG",
+                // For debugging the ssh containers in particular, let them
+                // live a little longer before automatically stopping them.
+                // See Also: mlos_bench/tests/services/remote/ssh/
+                "TIMEOUT": "1800"
             }
         }
     ]

Makefile

@@ -170,7 +170,7 @@ PYTEST_OPTIONS :=
 SKIP_COVERAGE := $(shell echo $${SKIP_COVERAGE:-} | grep -i -x -e 1 -e true)
 
 ifeq ($(SKIP_COVERAGE),)
-    PYTEST_OPTIONS += --cov=. --cov-append --cov-fail-under=0.91 --cov-report=xml --cov-report=html --junitxml=junit/test-results.xml --local-badge-output-dir=doc/source/badges/
+    PYTEST_OPTIONS += --cov=. --cov-append --cov-fail-under=0.92 --cov-report=xml --cov-report=html --junitxml=junit/test-results.xml --local-badge-output-dir=doc/source/badges/
 endif
 
 # Run the pytest target on only the modules that have changed recently, but

mlos_bench/mlos_bench/config/services/local/local_exec_service.jsonc

@@ -0,0 +1,3 @@
+{
+    "class": "mlos_bench.services.local.local_exec.LocalExecService"
+}

mlos_bench/mlos_bench/config/services/remote/ssh/ssh_fileshare_service.jsonc

@@ -0,0 +1,7 @@
+{
+    "class": "mlos_bench.services.remote.ssh.ssh_fileshare.SshFileShareService",
+    "config": {
+        // don't check host keys
+        "ssh_known_hosts_path": null
+    }
+}

mlos_bench/mlos_bench/config/services/remote/ssh/ssh_host_service.jsonc

@@ -0,0 +1,7 @@
+{
+    "class": "mlos_bench.services.remote.ssh.ssh_host_service.SshHostService",
+    "config": {
+        // don't check host keys
+        "ssh_known_hosts_path": null
+    }
+}

mlos_bench/mlos_bench/environments/local/local_fileshare_env.py

@@ -121,7 +121,7 @@ def setup(self, tunables: TunableGroups, global_config: Optional[dict] = None) -
             params = self._get_env_params()
             params["PWD"] = self._temp_dir
             for (path_from, path_to) in self._expand(self._upload, params):
-                self._file_share_service.upload(self.config, self._config_loader_service.resolve_path(
+                self._file_share_service.upload(self._params, self._config_loader_service.resolve_path(
                     path_from, extra_paths=[self._temp_dir]), path_to)
         return self._is_ready
 
@@ -140,13 +140,16 @@ def _download_files(self, ignore_missing: bool = False) -> None:
         params["PWD"] = self._temp_dir
         for (path_from, path_to) in self._expand(self._download, params):
             try:
-                self._file_share_service.download(self.config,
+                self._file_share_service.download(self._params,
                                                   path_from, self._config_loader_service.resolve_path(
                                                       path_to, extra_paths=[self._temp_dir]))
             except FileNotFoundError as ex:
                 _LOG.warning("Cannot download: %s", path_from)
                 if not ignore_missing:
                     raise ex
+            except Exception as ex:
+                _LOG.exception("Cannot download %s to %s", path_from, path_to)
+                raise ex
 
     def run(self) -> Tuple[Status, Optional[Dict[str, TunableValue]]]:
         """

mlos_bench/mlos_bench/environments/remote/remote_env.py

@@ -67,9 +67,10 @@ def __init__(self,
             "RemoteEnv requires a service that supports remote execution operations"
         self._remote_exec_service: SupportsRemoteExec = self._service
 
-        assert self._service is not None and isinstance(self._service, SupportsHostOps), \
-            "RemoteEnv requires a service that supports host operations"
-        self._host_service: SupportsHostOps = self._service
+        if self._wait_boot:
+            assert self._service is not None and isinstance(self._service, SupportsHostOps), \
+                "RemoteEnv requires a service that supports host operations"
+            self._host_service: SupportsHostOps = self._service
 
     def setup(self, tunables: TunableGroups, global_config: Optional[dict] = None) -> bool:
         """

mlos_bench/mlos_bench/environments/script_env.py

@@ -94,6 +94,7 @@ def _get_env_params(self) -> Dict[str, str]:
                 # Only rename specified - use it.
                 rename = self._shell_env_params_rename.copy()
             else:
+                # FIXME: We should not be exposing all params by default.
                 # Neither `shell_env_params` nor rename are specified - use all params.
                 rename = {self._RE_INVALID.sub("_", key): key for key in self._params}
         else:

mlos_bench/mlos_bench/services/remote/ssh/ssh_fileshare.py

@@ -11,7 +11,7 @@
 
 import logging
 
-from asyncssh import scp, SFTPError, SSHClientConnection
+from asyncssh import scp, SFTPError, SFTPNoSuchFile, SFTPFailure, SSHClientConnection
 
 from mlos_bench.services.base_fileshare import FileShareService
 from mlos_bench.services.remote.ssh.ssh_service import SshService
@@ -58,6 +58,8 @@ async def _start_file_copy(self, params: dict, mode: CopyMode,
             If the local OS returns an error.
         SFTPError
             If the remote OS returns an error.
+        FileNotFoundError
+            If the remote file does not exist, the SFTPError is converted to a FileNotFoundError.
         """
         connection, _ = await self._get_client_connection(params)
         srcpaths: Union[str, Tuple[SSHClientConnection, str]]
@@ -87,6 +89,12 @@ def download(self, params: dict, remote_path: str, local_path: str, recursive: b
             file_copy_future.result()
         except (OSError, SFTPError) as ex:
             _LOG.error("Failed to download %s to %s from %s: %s", remote_path, local_path, params, ex)
+            if isinstance(ex, SFTPNoSuchFile) or (
+                isinstance(ex, SFTPFailure) and ex.code == 4
+                    and any(msg.lower() in ex.reason.lower() for msg in ("File not found", "No such file or directory"))
+            ):
+                _LOG.warning("File %s does not exist on %s", remote_path, params)
+                raise FileNotFoundError(f"File {remote_path} does not exist on {params}") from ex
             raise ex
 
     def upload(self, params: dict, local_path: str, remote_path: str, recursive: bool = True) -> None:

mlos_bench/mlos_bench/services/remote/ssh/ssh_host_service.py

@@ -88,7 +88,9 @@ async def _run_cmd(self, params: dict, script: Iterable[str], env_params: dict)
         env_script_lines = [f"export {name}='{value}'" for (name, value) in env_params.items()]
         script_lines = env_script_lines + [line_split for line in script for line_split in line.splitlines()]
         # Note: connection.run() uses "exec" with a shell by default.
-        return await connection.run('\n'.join(script_lines),
+        script_str = '\n'.join(script_lines)
+        _LOG.debug("Running script on %s:\n%s", connection, script_str)
+        return await connection.run(script_str,
                                     check=False,
                                     timeout=self._request_timeout,
                                     env=env_params)

mlos_bench/mlos_bench/tests/config/environments/remote/test_ssh_env.jsonc

@@ -0,0 +1,81 @@
+// Test config for test_ssh_env.py
+{
+    "class": "mlos_bench.environments.composite_env.CompositeEnv",
+    "description": "Test Environment for RemoteEnv via local SSH test infrastructure",
+    "name": "composite_remote_ssh_env",
+
+    "include_services": [
+        // from the built in configs
+        "services/local/local_exec_service.jsonc",
+        "services/remote/ssh/ssh_host_service.jsonc",
+        "services/remote/ssh/ssh_fileshare_service.jsonc"
+    ],
+
+    "config": {
+        "children": [
+            {
+                "class": "mlos_bench.environments.remote.remote_env.RemoteEnv",
+                "name": "remote_ssh_env",
+                "config": {
+                    "required_args": [
+                        "ssh_hostname",
+                        "ssh_port",
+                        "ssh_username",
+                        "ssh_priv_key_path"
+                    ],
+                    "shell_env_params": [
+                        "ssh_username"
+                    ],
+                    "setup": [
+                        "echo 'Preparing output.csv' >&2",
+                        "echo 'metric,value' | tee output.csv"
+                    ],
+                    "run": [
+                        // write a basic results csv file we can download and compare with
+                        "echo 'Generating test metrics output ...' >&2",
+                        // Check that we connected to the test server.
+                        "echo \"hostname,`hostname`\" | tee -a output.csv",
+                        // Test environment variables
+                        "echo \"username,${ssh_username}\" | tee -a output.csv",
+                        "echo score,0.9 | tee -a output.csv",
+                        "echo 'Done' >&2",
+                        // Wait a moment (mostly for debugging)
+                        "echo 'Sleeping ...' >&2",
+                        //"sleep 10",
+                        "echo 'Done' >&2"
+                    ],
+                    "teardown": [
+                        "rm output.csv",
+                        "echo 'teardown,complete' | tee teardown.txt"
+                    ]
+                }
+            },
+            {
+                "class": "mlos_bench.environments.local.local_fileshare_env.LocalFileShareEnv",
+                "name": "local_ssh_env",
+                "config": {
+                    // these are needed by the SshFileShareService
+                    // FIXME: Can we only include these in the CompositeEnv and have them propagate down?
+                    "required_args": [
+                        "ssh_hostname",
+                        "ssh_port",
+                        "ssh_username",
+                        "ssh_priv_key_path"
+                    ],
+                    // Download happens *prior* to teardown, so this should succeed.
+                    "download": [
+                        {
+                            "from": "output.csv",
+                            // Download to the temporary directory, not the current working directory.
+                            "to": "$PWD/output-downloaded.csv"
+                        }
+                    ],
+                    "read_results_file": "output-downloaded.csv",
+                    "teardown": [
+                        "test ! -e output-downloaded.csv"
+                    ]
+                }
+            }
+        ]
+    }
+}

mlos_bench/mlos_bench/tests/environments/remote/__init__.py

@@ -0,0 +1,38 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""
+Helpers for RemoteEnv tests.
+"""
+
+from typing import Any, Dict, List, Tuple
+
+from mlos_bench.environments.remote.remote_env import RemoteEnv
+from mlos_bench.services.config_persistence import ConfigPersistenceService
+from mlos_bench.services.remote.ssh.ssh_host_service import SshHostService
+from mlos_bench.services.remote.ssh.ssh_fileshare import SshFileShareService
+from mlos_bench.tunables.tunable_groups import TunableGroups
+
+
+def create_remote_ssh_env(tunable_groups: TunableGroups, config: Dict[str, Any]) -> RemoteEnv:
+    """
+    Create a RemoteEnv with the given configuration that uses SshServices.
+
+    Parameters
+    ----------
+    tunable_groups : TunableGroups
+        Tunable parameters (usually come from a fixture).
+    config : Dict[str, Any]
+        Environment configuration.
+
+    Returns
+    -------
+    env : RemoteEnv
+        A new instance of the local environment.
+    """
+    service = ConfigPersistenceService()
+    service.register(SshHostService().export())
+    service.register(SshFileShareService().export())
+
+    return RemoteEnv(name="TestRemoveEnv", config=config, tunables=tunable_groups, service=service)

mlos_bench/mlos_bench/tests/environments/remote/conftest.py

@@ -0,0 +1,16 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""
+Fixtures for the RemoteEnv tests using SSH Services.
+"""
+
+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
+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

mlos_bench/mlos_bench/tests/environments/remote/test_ssh_env.py

@@ -0,0 +1,61 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""
+Unit tests for RemoveEnv benchmark environment via local SSH test services.
+"""
+
+from typing import Dict
+
+import os
+import sys
+
+import pytest
+
+from mlos_bench.services.config_persistence import ConfigPersistenceService
+from mlos_bench.tunables.tunable import TunableValue
+from mlos_bench.tunables.tunable_groups import TunableGroups
+
+from mlos_bench.tests import requires_docker
+from mlos_bench.tests.environments import check_env_success
+from mlos_bench.tests.services.remote.ssh import SshTestServerInfo
+
+if sys.version_info < (3, 10):
+    from importlib_resources import files
+else:
+    from importlib.resources import files
+
+
+@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.
+    """
+    global_config: Dict[str, TunableValue] = {
+        "ssh_hostname": ssh_test_server.hostname,
+        "ssh_port": ssh_test_server.get_port(),
+        "ssh_username": ssh_test_server.username,
+        "ssh_priv_key_path": ssh_test_server.id_rsa_path,
+    }
+
+    service = ConfigPersistenceService(config={"config_path": [str(files("mlos_bench.tests.config"))]})
+    config_path = service.resolve_path("environments/remote/test_ssh_env.jsonc")
+    env = service.load_environment(config_path, tunable_groups, global_config=global_config, service=service)
+
+    check_env_success(
+        env, tunable_groups,
+        expected_results={
+            "hostname": ssh_test_server.service_name,
+            "username": ssh_test_server.username,
+            "score": 0.9,
+        },
+        expected_telemetry=[],
+    )
+    assert not os.path.exists(os.path.join(os.getcwd(), "output-downloaded.csv")), \
+        "output-downloaded.csv should have been cleaned up by temp_dir context"
+
+
+if __name__ == "__main__":
+    pytest.main(["-n1", __file__])

mlos_bench/mlos_bench/tests/services/remote/ssh/test_ssh_fileshare.py

@@ -17,8 +17,6 @@
 
 import pytest
 
-from asyncssh import SFTPError
-
 from mlos_bench.services.remote.ssh.ssh_host_service import SshHostService
 from mlos_bench.services.remote.ssh.ssh_fileshare import SshFileShareService
 from mlos_bench.util import path_join
@@ -167,7 +165,7 @@ def test_ssh_fileshare_download_file_dne(ssh_test_server: SshTestServerInfo,
             temp_file.flush()
             temp_file.close()
 
-            with pytest.raises(SFTPError):
+            with pytest.raises(FileNotFoundError):
                 ssh_fileshare_service.download(
                     params=config,
                     remote_path="/tmp/file-dne.txt",