machine: add dvc machine ssh (#6501)

* machine: add naive asyncssh shell client

* machine: add `dvc machine ssh` command

* add ssh requirements to terraform install extras

* update asyncssh comments

* fix py <3.7

* fix async mocking in py<3.8

* use mock instead of unittest.mock

dvc/command/machine.py

@@ -136,6 +136,15 @@ class CmdMachineDestroy(CmdBase):
         return 0
 
 
+class CmdMachineSsh(CmdBase):
+    def run(self):
+        if self.repo.machine is None:
+            raise MachineDisabledError
+
+        self.repo.machine.run_shell(self.args.name)
+        return 0
+
+
 def add_parser(subparsers, parent_parser):
     from dvc.command.config import parent_config_parser
 
@@ -272,3 +281,16 @@ def add_parser(subparsers, parent_parser):
         "name", help="Name of the machine instance to destroy."
     )
     machine_destroy_parser.set_defaults(func=CmdMachineDestroy)
+
+    machine_SSH_HELP = "Connect to a machine via SSH."
+    machine_ssh_parser = machine_subparsers.add_parser(
+        "ssh",
+        parents=[parent_config_parser, parent_parser],
+        description=append_doc_link(machine_SSH_HELP, "machine/ssh"),
+        help=machine_SSH_HELP,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    machine_ssh_parser.add_argument(
+        "name", help="Name of the machine instance to connect to."
+    )
+    machine_ssh_parser.set_defaults(func=CmdMachineSsh)

dvc/machine/__init__.py

@@ -164,3 +164,7 @@ class MachineManager:
     def get_sshfs(self, name: Optional[str]):
         config, backend = self.get_config_and_backend(name)
         return backend.get_sshfs(**config)
+
+    def run_shell(self, name: Optional[str]):
+        config, backend = self.get_config_and_backend(name)
+        return backend.run_shell(**config)

dvc/machine/backend/base.py

@@ -1,8 +1,12 @@
+import asyncio
 import logging
+import os
+import sys
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Iterator, Optional
 
+from dvc.exceptions import DvcException
 from dvc.types import StrPath
 from dvc.utils.fs import makedirs
 
@@ -50,3 +54,35 @@ class BaseMachineBackend(ABC):
     ) -> Iterator["SSHFileSystem"]:
         """Return an sshfs instance for the default directory on the
         specified machine."""
+
+    @abstractmethod
+    def run_shell(self, name: Optional[str] = None, **config):
+        """Spawn an interactive SSH shell for the specified machine."""
+
+    def _shell(self, *args, **kwargs):
+        """Sync wrapper for an asyncssh shell session.
+
+        Args will be passed into asyncssh.connect().
+        """
+        import asyncssh
+
+        loop = asyncio.new_event_loop()
+        try:
+            asyncio.set_event_loop(loop)
+            loop.run_until_complete(self._shell_async(*args, **kwargs))
+        except (OSError, asyncssh.Error) as exc:
+            raise DvcException("SSH connection failed") from exc
+        finally:
+            asyncio.set_event_loop(None)
+            loop.close()
+
+    async def _shell_async(self, *args, **kwargs):
+        import asyncssh
+
+        async with asyncssh.connect(*args, **kwargs) as conn:
+            await conn.run(
+                term_type=os.environ.get("TERM", "xterm"),
+                stdin=sys.stdin,
+                stdout=sys.stdout,
+                stderr=sys.stderr,
+            )

dvc/machine/backend/terraform.py

@@ -109,13 +109,17 @@ class TerraformBackend(BaseMachineBackend):
     ) -> "BaseExecutor":
         raise NotImplementedError
 
+    def _default_resource(self, name):
+        resource = first(self.instances(name))
+        if not resource:
+            raise TerraformError(f"No active '{name}' instances")
+        return resource
+
     @contextmanager
     def get_sshfs(
         self, name: Optional[str] = None, **config
     ) -> Iterator["SSHFileSystem"]:
-        resource = first(self.instances(name))
-        if not resource:
-            raise TerraformError(f"No active '{name}' instances")
+        resource = self._default_resource(name)
         with DvcTerraform.pemfile(resource) as pem:
             fs = SSHFileSystem(
                 host=resource["instance_ip"],
@@ -123,3 +127,13 @@ class TerraformBackend(BaseMachineBackend):
                 keyfile=pem,
             )
             yield fs
+
+    def run_shell(self, name: Optional[str] = None, **config):
+        resource = self._default_resource(name)
+        with DvcTerraform.pemfile(resource) as pem:
+            self._shell(
+                host=resource["instance_ip"],
+                username="ubuntu",
+                client_keys=pem,
+                known_hosts=None,
+            )

setup.py

@@ -114,7 +114,7 @@ webdav = ["webdav4>=0.9.0"]
 ssh_gssapi = ["sshfs[gssapi]>=2021.7.1"]
 all_remotes = gs + s3 + azure + ssh + oss + gdrive + hdfs + webhdfs + webdav
 
-terraform = ["python-terraform>=0.10.1", "jinja2>=2.0.0"]
+terraform = ssh + ["python-terraform>=0.10.1", "jinja2>=2.0.0"]
 
 tests_requirements = (
     Path("test_requirements.txt").read_text().strip().splitlines()

tests/unit/command/test_machine.py

@@ -6,6 +6,7 @@ from dvc.command.machine import (
     CmdMachineCreate,
     CmdMachineDestroy,
     CmdMachineRemove,
+    CmdMachineSsh,
 )
 
 
@@ -64,3 +65,16 @@ def test_destroy(tmp_dir, dvc, mocker):
 
     assert cmd.run() == 0
     m.assert_called_once_with("foo")
+
+
+def test_ssh(tmp_dir, dvc, mocker):
+    cli_args = parse_args(["machine", "ssh", "foo"])
+    assert cli_args.func == CmdMachineSsh
+
+    cmd = cli_args.func(cli_args)
+    m = mocker.patch.object(
+        cmd.repo.machine, "run_shell", autospec=True, return_value=0
+    )
+
+    assert cmd.run() == 0
+    m.assert_called_once_with("foo")

tests/unit/machine/test_terraform.py

@@ -2,6 +2,7 @@ import json
 import os
 
 import pytest
+from mock import ANY, MagicMock
 from python_terraform import IsFlagged
 
 from dvc.machine.backend.terraform import DvcTerraform, TerraformError
@@ -127,3 +128,16 @@ def test_instances(tmp_dir, terraform, resource):
     data = json.loads(TEST_RESOURCE_STATE)
     expected = data["resources"][0]["instances"][0]["attributes"]
     assert list(terraform.instances(resource)) == [expected]
+
+
+def test_shell(tmp_dir, terraform, resource, mocker):
+    data = json.loads(TEST_RESOURCE_STATE)
+    expected = data["resources"][0]["instances"][0]["attributes"]
+    mock_connect = mocker.patch("asyncssh.connect", return_value=MagicMock())
+    terraform.run_shell(name=resource)
+    mock_connect.assert_called_once_with(
+        host=expected["instance_ip"],
+        username="ubuntu",
+        client_keys=ANY,
+        known_hosts=None,
+    )