Add SSH Support

.cspell.json

@@ -4,6 +4,7 @@
         "abstractmethods",
         "arange",
         "astype",
+        "asyncssh",
         "autotune",
         "autotuning",
         "bincount",
@@ -12,11 +13,13 @@
         "commandline",
         "conda",
         "configspace",
+        "coro",
         "dataframe",
         "devcontainer",
         "discretization",
         "discretize",
         "drivername",
+        "dstpath",
         "dtype",
         "duckdb",
         "emukit",
@@ -31,6 +34,7 @@
         "iterrows",
         "jsonschema",
         "jupyterlab",
+        "keepalive",
         "kwargs",
         "libmamba",
         "linalg",
@@ -47,10 +51,12 @@
         "obvs",
         "perc",
         "pinv",
+        "poweroff",
         "pylint",
         "pyplot",
         "pytest",
         "Quickstart",
+        "refcnt",
         "rexec",
         "rootfs",
         "runhistory",
@@ -63,11 +69,13 @@
         "skopt",
         "smac",
         "sqlalchemy",
+        "srcpaths",
         "subcmd",
         "subschema",
         "subschemas",
         "tolist",
         "tunables",
+        "xdist",
         "xlabel",
         "ylabel"
     ]

.devcontainer/devcontainer.json

@@ -17,6 +17,7 @@
     "postCreateCommand": "/opt/conda/bin/conda env update --solver=libmamba -v -n mlos -f ${containerWorkspaceFolder}/conda-envs/mlos.yml",
     // Various mounting, run, post-create, and user settings
     "containerEnv": {
+        "LOCAL_USER_NAME": "${localEnv:USER}${localEnv:USERNAME}",
         "LOCAL_WORKSPACE_FOLDER": "${localWorkspaceFolder}",
         "CONTAINER_WORKSPACE_FOLDER": "${containerWorkspaceFolder}",
         "SSH_AUTH_SOCK": "${localEnv:SSH_AUTH_SOCK}",

.devcontainer/docker-compose.yml

@@ -24,6 +24,9 @@ services:
 
     env_file: ../.env
 
+    extra_hosts:
+      - "host.docker.internal:host-gateway"
+
   docs-www:
     build:
       context: ./tmp

.devcontainer/scripts/run-devcontainer.sh

@@ -28,6 +28,7 @@ else
     docker_gid=$(stat -c%g /var/run/docker.sock)
 fi
 
+set -x
 mkdir -p "/tmp/$container_name/dc/shellhistory"
 docker run -it --rm \
     --name "$container_name" \

.github/workflows/devcontainer.yml

@@ -90,6 +90,7 @@ jobs:
           --env LOCAL_WORKSPACE_FOLDER=$(pwd) \
           --env PYTEST_EXTRA_OPTIONS=$PYTEST_EXTRA_OPTIONS \
           --workdir /workspaces/MLOS \
+          --add-host host.docker.internal:host-gateway \
           --name mlos-devcontainer mlos-devcontainer sleep 1800
     - name: Fixup vscode uid/gid in the running container
       timeout-minutes: 3

.vscode/settings.json

@@ -134,6 +134,7 @@
     },
     "python.testing.pytestArgs": [
         "-n1",  // don't run tests in parallel inside vscode - makes attaching the debugger more cumbersome
+        "--dist=no",
         "--log-level=DEBUG",
         "."
     ],

conda-envs/mlos-3.10.yml

@@ -32,6 +32,7 @@ dependencies:
     - types-colorama
     - types-jsonschema
     - types-pygments
+    - types-pytest-lazy-fixture
     - types-requests
     - types-setuptools
     - "--editable ../mlos_core[full-tests]"

conda-envs/mlos-3.11.yml

@@ -32,6 +32,7 @@ dependencies:
     - types-colorama
     - types-jsonschema
     - types-pygments
+    - types-pytest-lazy-fixture
     - types-requests
     - types-setuptools
     - "--editable ../mlos_core[full-tests]"

conda-envs/mlos-3.8.yml

@@ -32,6 +32,7 @@ dependencies:
     - types-colorama
     - types-jsonschema
     - types-pygments
+    - types-pytest-lazy-fixture
     - types-requests
     - types-setuptools
     - "--editable ../mlos_core[full-tests]"

conda-envs/mlos-3.9.yml

@@ -32,6 +32,7 @@ dependencies:
     - types-colorama
     - types-jsonschema
     - types-pygments
+    - types-pytest-lazy-fixture
     - types-requests
     - types-setuptools
     - "--editable ../mlos_core[full-tests]"

conda-envs/mlos-windows.yml

@@ -36,6 +36,7 @@ dependencies:
     - types-colorama
     - types-jsonschema
     - types-pygments
+    - types-pytest-lazy-fixture
     - types-requests
     - types-setuptools
     - "--editable ../mlos_core[full-tests]"

conda-envs/mlos.yml

@@ -30,6 +30,7 @@ dependencies:
     - types-colorama
     - types-jsonschema
     - types-pygments
+    - types-pytest-lazy-fixture
     - types-requests
     - types-setuptools
     - "--editable ../mlos_core[full-tests]"

mlos_bench/mlos_bench/config/schemas/config_schemas.py

@@ -137,5 +137,5 @@ def validate(self, config: dict) -> None:
         else:
             jsonschema.Draft202012Validator(
                 schema=self.schema,
-                registry=SCHEMA_STORE.registry,     # type: ignore[call-arg]
+                registry=SCHEMA_STORE.registry,
             ).validate(config)

mlos_bench/mlos_bench/config/schemas/services/remote/ssh/common-defs-subschemas.json

@@ -0,0 +1,47 @@
+{
+    "$schema": "https://json-schema.org/draft/2020-12/schema",
+    "$id": "https://raw.githubusercontent.com/microsoft/MLOS/main/mlos_bench/mlos_bench/config/schemas/services/remote/ssh/common-defs-subschemas.json",
+    "title": "mlos_bench SSH Service common defs config",
+    "description": "mlos_bench SSH Service common defs config",
+    "$defs": {
+        "ssh_service_config": {
+            "description": "SSH Service config.",
+            "type": "object",
+            "properties": {
+                "ssh_request_timeout": {
+                    "description": "Request timeout in seconds, or null to disable timeout.",
+                    "type": ["null", "number"],
+                    "minimum": 1
+                },
+                "ssh_keepalive_interval": {
+                    "description": "Whether to send keep alive packets to the remote machine(s).",
+                    "type": ["null", "number"],
+                    "minimum": 1
+                },
+                "ssh_port": {
+                    "description": "Default port to use when connecting to the remote machine(s).",
+                    "type": "integer",
+                    "minimum": 1,
+                    "maximum": 65535,
+                    "examples": [22]
+                },
+                "ssh_username": {
+                    "description": "Default username to use when connecting to the remote machine(s). If null or unspecified, will default to the value in ~/.ssh/config or the current user if not provided in another config.",
+                    "type": ["null", "string"],
+                    "examples": ["root"]
+                },
+                "ssh_priv_key_path": {
+                    "$comment": "TODO: add support for multiple private keys or dynamically fetched private keys.",
+                    "description": "Optional path to the private key to use when connecting to the remote machine(s). Keys are automatically searched for if not specified. Note: these should be passwordless or else an ssh-agent should be available via the SSH_AUTH_SOCK environment variable.",
+                    "type": "string",
+                    "examples": ["~/.ssh/id_rsa", "~/.ssh/id_ecdsa"]
+                },
+                "ssh_known_hosts_path": {
+                    "description": "Path to known_hosts file. Set to null to disable host key checking.",
+                    "type": ["null", "string"],
+                    "examples": [null, "~/.ssh/known_hosts"]
+                }
+            }
+        }
+    }
+}

mlos_bench/mlos_bench/config/schemas/services/remote/ssh/ssh-fileshare-service-subschema.json

@@ -0,0 +1,21 @@
+{
+    "$schema": "https://json-schema.org/draft/2020-12/schema",
+    "$id": "https://raw.githubusercontent.com/microsoft/MLOS/main/mlos_bench/mlos_bench/config/schemas/services/remote/ssh/ssh-fileshare-service-subschema.json",
+    "title": "mlos_bench SSH Fileshare Service config",
+    "description": "config for an mlos_bench SSH Fileshare Service",
+    "type": "object",
+    "properties": {
+        "class": {
+            "enum": [
+                "mlos_bench.services.remote.ssh.SshFileShareService",
+                "mlos_bench.services.remote.ssh.ssh_fileshare.SshFileShareService"
+            ]
+        },
+        "config": {
+            "$ref": "./common-defs-subschemas.json#/$defs/ssh_service_config",
+            "minProperties": 1,
+            "unevaluatedProperties": false
+        }
+    },
+    "required": ["class"]
+}

mlos_bench/mlos_bench/config/schemas/services/remote/ssh/ssh-host-service-subschema.json

@@ -0,0 +1,37 @@
+{
+    "$schema": "https://json-schema.org/draft/2020-12/schema",
+    "$id": "https://raw.githubusercontent.com/microsoft/MLOS/main/mlos_bench/mlos_bench/config/schemas/services/remote/ssh/ssh-host-service-subschema.json",
+    "title": "mlos_bench SSH Host Service config",
+    "description": "config for an mlos_bench SSH Host Service",
+    "type": "object",
+    "properties": {
+        "class": {
+            "enum": [
+                "mlos_bench.services.remote.ssh.SshHostService",
+                "mlos_bench.services.remote.ssh.ssh_host_service.SshHostService"
+            ]
+        },
+        "config": {
+            "allOf": [
+                {
+                    "$ref": "./common-defs-subschemas.json#/$defs/ssh_service_config"
+                },
+                {
+                    "type": "object",
+                    "properties": {
+                        "ssh_shell": {
+                            "type": "string",
+                            "examples": [
+                                "/bin/bash",
+                                "/bin/sh"
+                            ]
+                        }
+                    }
+                }
+            ],
+            "minProperties": 1,
+            "unevaluatedProperties": false
+        }
+    },
+    "required": ["class"]
+}

mlos_bench/mlos_bench/config/schemas/services/service-schema.json

@@ -34,6 +34,12 @@
                         {
                             "$ref": "./remote/mock/mock-auth-service-subschema.json"
                         },
+                        {
+                            "$ref": "./remote/ssh/ssh-host-service-subschema.json"
+                        },
+                        {
+                            "$ref": "./remote/ssh/ssh-fileshare-service-subschema.json"
+                        },
                         {
                             "$ref": "./remote/azure/azure-auth-service-subschema.json"
                         },

mlos_bench/mlos_bench/services/remote/ssh/__init__.py

@@ -4,4 +4,10 @@
 #
 """SSH remote service."""
 
-# TODO
+from mlos_bench.services.remote.ssh.ssh_host_service import SshHostService
+from mlos_bench.services.remote.ssh.ssh_fileshare import SshFileShareService
+
+__all__ = [
+    "SshHostService",
+    "SshFileShareService",
+]

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

@@ -0,0 +1,107 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""
+A collection functions for interacting with SSH servers as file shares.
+"""
+
+from enum import Enum
+from typing import Tuple, Union
+
+import logging
+
+from asyncssh import scp, SFTPError, SSHClientConnection
+
+from mlos_bench.services.base_fileshare import FileShareService
+from mlos_bench.services.remote.ssh.ssh_service import SshService
+from mlos_bench.util import merge_parameters
+
+_LOG = logging.getLogger(__name__)
+
+
+class CopyMode(Enum):
+    """
+    Copy mode enum.
+    """
+
+    DOWNLOAD = 1
+    UPLOAD = 2
+
+
+class SshFileShareService(FileShareService, SshService):
+    """A collection of functions for interacting with SSH servers as file shares."""
+
+    async def _start_file_copy(self, params: dict, mode: CopyMode,
+                               local_path: str, remote_path: str,
+                               recursive: bool = True) -> None:
+        # pylint: disable=too-many-arguments
+        """
+        Starts a file copy operation
+
+        Parameters
+        ----------
+        params : dict
+            Flat dictionary of (key, value) pairs of parameters (used for establishing the connection).
+        mode : CopyMode
+            Whether to download or upload the file.
+        local_path : str
+            Local path to the file/dir.
+        remote_path : str
+            Remote path to the file/dir.
+        recursive : bool, optional
+            _description_, by default True
+
+        Raises
+        ------
+        OSError
+            If the local OS returns an error.
+        SFTPError
+            If the remote OS returns an error.
+        """
+        connection, _ = await self._get_client_connection(params)
+        srcpaths: Union[str, Tuple[SSHClientConnection, str]]
+        dstpath: Union[str, Tuple[SSHClientConnection, str]]
+        if mode == CopyMode.DOWNLOAD:
+            srcpaths = (connection, remote_path)
+            dstpath = local_path
+        elif mode == CopyMode.UPLOAD:
+            srcpaths = local_path
+            dstpath = (connection, remote_path)
+        else:
+            raise ValueError(f"Unknown copy mode: {mode}")
+        return await scp(srcpaths=srcpaths, dstpath=dstpath, recurse=recursive, preserve=True)
+
+    def download(self, params: dict, remote_path: str, local_path: str, recursive: bool = True) -> None:
+        params = merge_parameters(
+            dest=self.config.copy(),
+            source=params,
+            required_keys=[
+                "ssh_hostname",
+            ]
+        )
+        super().download(params, remote_path, local_path, recursive)
+        file_copy_future = self._run_coroutine(
+            self._start_file_copy(params, CopyMode.DOWNLOAD, local_path, remote_path, recursive))
+        try:
+            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)
+            raise ex
+
+    def upload(self, params: dict, local_path: str, remote_path: str, recursive: bool = True) -> None:
+        params = merge_parameters(
+            dest=self.config.copy(),
+            source=params,
+            required_keys=[
+                "ssh_hostname",
+            ]
+        )
+        super().upload(params, local_path, remote_path, recursive)
+        file_copy_future = self._run_coroutine(
+            self._start_file_copy(params, CopyMode.UPLOAD, local_path, remote_path, recursive))
+        try:
+            file_copy_future.result()
+        except (OSError, SFTPError) as ex:
+            _LOG.error("Failed to upload %s to %s on %s: %s", local_path, remote_path, params, ex)
+            raise ex