exp: add initial ssh executor (#6833)

* executor: move local specific git initialization out of base

* machine: return kwargs for executor init instead of executor object

* executor: add initial ssh executor

* ssh-executor: disable strict host checking for dvc-machine instances

* initial tests for ssh executor

* drop GIT_SSH_COMMAND usage

* push: accept destination odb as a parameter

* separate cache initialization into executor.init_cache

* add test for ssh executor cache init

* remove dead code

* machine: fixture-ize default test scenarios

* test SSHExecutor.from_machine

* tpi: return default SSH port for constructing executors

dvc/machine/__init__.py

@@ -211,3 +211,7 @@ class MachineManager:
         """move configuration to a new location if the machine is running."""
         config, backend = self.get_config_and_backend(name)
         return backend.rename(new=new, **config)
+
+    def get_executor_kwargs(self, name: Optional[str]):
+        config, backend = self.get_config_and_backend(name)
+        return backend.get_executor_kwargs(**config)

dvc/machine/backend/base.py

@@ -4,7 +4,6 @@ from typing import TYPE_CHECKING, Iterator, Optional
 
 if TYPE_CHECKING:
     from dvc.fs.ssh import SSHFileSystem
-    from dvc.repo.experiments.executor.base import BaseExecutor
     from dvc.types import StrPath
 
 
@@ -34,10 +33,8 @@ class BaseMachineBackend(ABC):
         """Spawn an interactive SSH shell for the specified machine."""
 
     @abstractmethod
-    def get_executor(
-        self, name: Optional[str] = None, **config
-    ) -> "BaseExecutor":
-        """Return an executor instance which can be used for DVC
+    def get_executor_kwargs(self, name: str, **config) -> dict:
+        """Return SSHExecutor kwargs which can be used for DVC
         experiment/pipeline execution on the specified machine.
         """
 

dvc/machine/backend/terraform.py

@@ -1,6 +1,6 @@
 import os
 from contextlib import contextmanager
-from functools import partialmethod
+from functools import partial, partialmethod
 from typing import TYPE_CHECKING, Iterator, Optional
 
 from dvc.exceptions import DvcException
@@ -10,10 +10,22 @@ from dvc.utils.fs import makedirs
 from .base import BaseMachineBackend
 
 if TYPE_CHECKING:
-    from dvc.repo.experiments.executor.base import BaseExecutor
     from dvc.types import StrPath
 
 
+@contextmanager
+def _sshfs(resource: dict):
+    from tpi import TerraformProviderIterative
+
+    with TerraformProviderIterative.pemfile(resource) as pem:
+        fs = SSHFileSystem(
+            host=resource["instance_ip"],
+            user="ubuntu",
+            keyfile=pem,
+        )
+        yield fs
+
+
 class TerraformBackend(BaseMachineBackend):
     def __init__(self, tmp_dir: "StrPath", **kwargs):
         super().__init__(tmp_dir, **kwargs)
@@ -51,26 +63,24 @@ class TerraformBackend(BaseMachineBackend):
         _tpi_func, "run_shell"
     )  # type: ignore[assignment]
 
-    def get_executor(
-        self, name: Optional[str] = None, **config
-    ) -> "BaseExecutor":
-        raise NotImplementedError
+    def get_executor_kwargs(self, name: str, **config) -> dict:
+        with self.make_tpi(name) as tpi:
+            resource = tpi.default_resource(name)
+        return {
+            "host": resource["instance_ip"],
+            "port": SSHFileSystem.DEFAULT_PORT,
+            "username": "ubuntu",
+            "fs_factory": partial(_sshfs, dict(resource)),
+        }
 
     @contextmanager
     def get_sshfs(  # pylint: disable=unused-argument
         self, name: Optional[str] = None, **config
     ) -> Iterator["SSHFileSystem"]:
-        from tpi import TerraformProviderIterative
-
         assert name
         with self.make_tpi(name) as tpi:
             resource = tpi.default_resource(name)
-        with TerraformProviderIterative.pemfile(resource) as pem:
-            fs = SSHFileSystem(
-                host=resource["instance_ip"],
-                user="ubuntu",
-                keyfile=pem,
-            )
+        with _sshfs(resource) as fs:
             yield fs
 
     def rename(self, name: str, new: str, **config):

dvc/repo/experiments/__init__.py

@@ -674,8 +674,8 @@ class Experiments:
                 name=item.name,
                 branch=item.branch,
                 tmp_dir=base_tmp_dir,
-                cache_dir=self.repo.odb.local.cache_dir,
             )
+            executor.init_cache(self.repo, stash_rev)
             executors[stash_rev] = executor
 
         for ref in (EXEC_HEAD, EXEC_MERGE, EXEC_BASELINE):

dvc/repo/experiments/executor/base.py

@@ -14,8 +14,6 @@ from typing import (
     Union,
 )
 
-from funcy import cached_property
-
 from dvc.env import DVC_EXP_AUTO_PUSH, DVC_EXP_GIT_REMOTE
 from dvc.exceptions import DvcException
 from dvc.path_info import PathInfo
@@ -24,8 +22,6 @@ from dvc.repo.experiments.base import (
     EXEC_BASELINE,
     EXEC_BRANCH,
     EXEC_CHECKPOINT,
-    EXEC_HEAD,
-    EXEC_MERGE,
     EXEC_NAMESPACE,
     EXPS_NAMESPACE,
     EXPS_STASH,
@@ -34,7 +30,6 @@ from dvc.repo.experiments.base import (
     ExpRefInfo,
     UnchangedExperimentError,
 )
-from dvc.scm import SCM
 from dvc.stage import PipelineStage
 from dvc.stage.monitor import CheckpointKilledError
 from dvc.stage.serialize import to_lockfile
@@ -117,44 +112,21 @@ class BaseExecutor(ABC):
         self._init_git(src, branch)
         self.name = name
 
-    def _init_git(self, scm: "Git", branch: Optional[str] = None):
-        """Init git repo and collect executor refs from the specified SCM."""
-        from dulwich.repo import Repo as DulwichRepo
-
-        DulwichRepo.init(os.fspath(self.root_dir))
-
-        cwd = os.getcwd()
-        os.chdir(self.root_dir)
-        try:
-            refspec = f"{EXEC_NAMESPACE}/"
-            scm.push_refspec(self.git_url, refspec, refspec)
-            if branch:
-                scm.push_refspec(self.git_url, branch, branch)
-                self.scm.set_ref(EXEC_BRANCH, branch, symbolic=True)
-            elif self.scm.get_ref(EXEC_BRANCH):
-                self.scm.remove_ref(EXEC_BRANCH)
-
-            if self.scm.get_ref(EXEC_CHECKPOINT):
-                self.scm.remove_ref(EXEC_CHECKPOINT)
-
-            # checkout EXEC_HEAD and apply EXEC_MERGE on top of it without
-            # committing
-            head = EXEC_BRANCH if branch else EXEC_HEAD
-            self.scm.checkout(head, detach=True)
-            merge_rev = self.scm.get_ref(EXEC_MERGE)
-            self.scm.merge(merge_rev, squash=True, commit=False)
-        finally:
-            os.chdir(cwd)
-
-    @cached_property
-    def scm(self):
-        return SCM(self.root_dir)
+    @abstractmethod
+    def _init_git(self, scm: "Git", branch: Optional[str] = None, **kwargs):
+        """Init git repo and populate it using exp refs from the specified
+        SCM instance.
+        """
 
     @property
     @abstractmethod
     def git_url(self) -> str:
         pass
 
+    @abstractmethod
+    def init_cache(self, dvc: "Repo", rev: str, run_cache: bool = True):
+        """Initialize DVC (cache)."""
+
     @property
     def dvc_dir(self) -> str:
         return os.path.join(self.root_dir, self._dvc_dir)
@@ -168,8 +140,7 @@ class BaseExecutor(ABC):
         return dict_sha256(exp_data)
 
     def cleanup(self):
-        self.scm.close()
-        del self.scm
+        pass
 
     # TODO: come up with better way to stash repro arguments
     @staticmethod

dvc/repo/experiments/executor/local.py

@@ -1,10 +1,25 @@
 import logging
 import os
 from tempfile import TemporaryDirectory
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
+
+from funcy import cached_property
+
+from dvc.repo.experiments.base import (
+    EXEC_BRANCH,
+    EXEC_CHECKPOINT,
+    EXEC_HEAD,
+    EXEC_MERGE,
+    EXEC_NAMESPACE,
+)
+from dvc.scm import SCM
 
 from .base import BaseExecutor
 
+if TYPE_CHECKING:
+    from dvc.repo import Repo
+    from dvc.scm.git import Git
+
 logger = logging.getLogger(__name__)
 
 
@@ -36,6 +51,15 @@ class BaseLocalExecutor(BaseExecutor):
             root_dir = root_dir.replace(os.sep, "/")
         return f"file://{root_dir}"
 
+    @cached_property
+    def scm(self):
+        return SCM(self.root_dir)
+
+    def cleanup(self):
+        super().cleanup()
+        self.scm.close()
+        del self.scm
+
 
 class TempDirExecutor(BaseLocalExecutor):
     """Temp directory experiment executor."""
@@ -51,22 +75,46 @@ class TempDirExecutor(BaseLocalExecutor):
         self,
         *args,
         tmp_dir: Optional[str] = None,
-        cache_dir: Optional[str] = None,
         **kwargs,
     ):
         self._tmp_dir = ExpTemporaryDirectory(dir=tmp_dir)
         kwargs["root_dir"] = self._tmp_dir.name
         super().__init__(*args, **kwargs)
-        if cache_dir:
-            self._config(cache_dir)
         logger.debug("Init temp dir executor in dir '%s'", self._tmp_dir)
 
+    def _init_git(self, scm: "Git", branch: Optional[str] = None, **kwargs):
+        from dulwich.repo import Repo as DulwichRepo
+
+        DulwichRepo.init(os.fspath(self.root_dir))
+
+        refspec = f"{EXEC_NAMESPACE}/"
+        scm.push_refspec(self.git_url, refspec, refspec, **kwargs)
+        if branch:
+            scm.push_refspec(self.git_url, branch, branch, **kwargs)
+            self.scm.set_ref(EXEC_BRANCH, branch, symbolic=True)
+        elif self.scm.get_ref(EXEC_BRANCH):
+            self.scm.remove_ref(EXEC_BRANCH)
+
+        if self.scm.get_ref(EXEC_CHECKPOINT):
+            self.scm.remove_ref(EXEC_CHECKPOINT)
+
+        # checkout EXEC_HEAD and apply EXEC_MERGE on top of it without
+        # committing
+        head = EXEC_BRANCH if branch else EXEC_HEAD
+        self.scm.checkout(head, detach=True)
+        merge_rev = self.scm.get_ref(EXEC_MERGE)
+        self.scm.merge(merge_rev, squash=True, commit=False)
+
     def _config(self, cache_dir):
         local_config = os.path.join(self.dvc_dir, "config.local")
         logger.debug("Writing experiments local config '%s'", local_config)
         with open(local_config, "w", encoding="utf-8") as fobj:
             fobj.write(f"[cache]\n    dir = {cache_dir}")
 
+    def init_cache(self, dvc: "Repo", rev: str, run_cache: bool = True):
+        """Initialize DVC (cache)."""
+        self._config(dvc.odb.local.cache_dir)
+
     def cleanup(self):
         super().cleanup()
         logger.debug("Removing tmpdir '%s'", self._tmp_dir)

dvc/repo/experiments/executor/ssh.py

@@ -0,0 +1,159 @@
+import logging
+import posixpath
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Callable, Optional
+
+from funcy import first
+
+from dvc.fs.ssh import SSHFileSystem
+from dvc.repo.experiments.base import (
+    EXEC_BRANCH,
+    EXEC_CHECKPOINT,
+    EXEC_HEAD,
+    EXEC_MERGE,
+    EXEC_NAMESPACE,
+)
+
+from .base import BaseExecutor
+
+if TYPE_CHECKING:
+    from dvc.machine import MachineManager
+    from dvc.repo import Repo
+    from dvc.scm.git import Git
+
+logger = logging.getLogger(__name__)
+
+
+@contextmanager
+def _sshfs(fs_factory, **kwargs):
+    if fs_factory:
+        with fs_factory() as fs:
+            yield fs
+        return
+    yield SSHFileSystem(**kwargs)
+
+
+class SSHExecutor(BaseExecutor):
+    """SSH experiment executor."""
+
+    WARN_UNTRACKED = True
+    QUIET = True
+
+    def __init__(
+        self,
+        *args,
+        host: Optional[str] = None,
+        port: Optional[int] = None,
+        username: Optional[str] = None,
+        fs_factory: Optional[Callable] = None,
+        **kwargs,
+    ):
+        assert host
+
+        self.host: str = host
+        self.port = port
+        self.username = username
+        self._fs_factory = fs_factory
+
+        super().__init__(*args, **kwargs)
+        logger.debug("Init SSH executor for host '%s'", self.host)
+
+    @classmethod
+    def gen_dirname(cls, name: Optional[str] = None):
+        from shortuuid import uuid
+
+        return "-".join([name or "dvc-exp", "executor", uuid()])
+
+    @classmethod
+    def from_machine(
+        cls,
+        manager: "MachineManager",
+        machine_name: Optional[str],
+        *args,
+        **kwargs,
+    ):
+        kwargs["root_dir"] = cls.gen_dirname(kwargs.get("name"))
+        kwargs.update(manager.get_executor_kwargs(machine_name))
+        return cls(*args, **kwargs)
+
+    def sshfs(self):
+        return _sshfs(self._fs_factory, host=self.host, port=self.port)
+
+    @property
+    def git_url(self) -> str:
+        user = f"{self.username}@" if self.username else ""
+        port = f":{self.port}" if self.port is not None else ""
+        path = f"{self.root_dir}" if self.root_dir else ""
+        if path and not posixpath.isabs(path):
+            path = f"/~/{path}"
+        return f"ssh://{user}{self.host}{port}{path}"
+
+    @property
+    def abs_url(self) -> str:
+        assert self._repo_abspath
+        user = f"{self.username}@" if self.username else ""
+        port = f":{self.port}" if self.port is not None else ""
+        return f"ssh://{user}{self.host}{port}{self._repo_abspath}"
+
+    @staticmethod
+    def _git_client_args(fs):
+        kwargs = {
+            "password": fs.fs_args.get("password"),
+            "key_filename": first(fs.fs_args.get("client_keys", [])),
+        }
+        return kwargs
+
+    def _init_git(self, scm: "Git", branch: Optional[str] = None, **kwargs):
+        with self.sshfs() as fs:
+            fs.makedirs(self.root_dir)
+            self._ssh_cmd(fs, "git init .")
+            self._ssh_cmd(fs, "git config user.name dvc-exp")
+            self._ssh_cmd(
+                fs, "git config user.email [email protected]"
+            )
+
+            result = self._ssh_cmd(fs, "pwd")
+            path = result.stdout.strip()
+            self._repo_abspath = path
+
+            # TODO: support multiple client key retries in git backends
+            # (see https://github.com/iterative/dvc/issues/6508)
+            kwargs.update(self._git_client_args(fs))
+            refspec = f"{EXEC_NAMESPACE}/"
+            scm.push_refspec(self.git_url, refspec, refspec, **kwargs)
+            if branch:
+                scm.push_refspec(self.git_url, branch, branch, **kwargs)
+                self._ssh_cmd(fs, f"git symbolic-ref {EXEC_BRANCH} {branch}")
+            else:
+                self._ssh_cmd(
+                    fs, f"git symbolic-ref -d {EXEC_BRANCH}", check=False
+                )
+            self._ssh_cmd(
+                fs, f"git update-ref -d {EXEC_CHECKPOINT}", check=False
+            )
+
+            # checkout EXEC_HEAD and apply EXEC_MERGE on top of it without
+            # committing
+            head = EXEC_BRANCH if branch else EXEC_HEAD
+            self._ssh_cmd(fs, f"git checkout {head}")
+            merge_rev = scm.get_ref(EXEC_MERGE)
+            self._ssh_cmd(fs, f"git merge --squash --no-commit {merge_rev}")
+
+    def _ssh_cmd(self, sshfs, cmd, chdir=None, **kwargs):
+        working_dir = chdir or self.root_dir
+        return sshfs.fs.execute(f"cd {working_dir};{cmd}", **kwargs)
+
+    def init_cache(self, dvc: "Repo", rev: str, run_cache: bool = True):
+        from dvc.objects.db import ODBManager, get_odb
+        from dvc.repo import Repo
+        from dvc.repo.push import push
+
+        cache_url = posixpath.join(
+            self.abs_url,
+            self._dvc_dir,
+            Repo.DVC_DIR,
+            ODBManager.CACHE_DIR,
+        )
+        with self.sshfs() as fs:
+            odb = get_odb(fs, fs.PATH_CLS(cache_url), **fs.config)
+            push(dvc, revs=[rev], run_cache=run_cache, odb=odb)

dvc/repo/push.py

@@ -1,8 +1,13 @@
+from typing import TYPE_CHECKING, Optional
+
 from dvc.exceptions import FileTransferError, UploadError
 
 from ..utils import glob_targets
 from . import locked
 
+if TYPE_CHECKING:
+    from dvc.objects.db.base import ObjectDB
+
 
 @locked
 def push(
@@ -18,8 +23,11 @@ def push(
     run_cache=False,
     revs=None,
     glob=False,
+    odb: Optional["ObjectDB"] = None,
 ):
-    used_run_cache = self.stage_cache.push(remote) if run_cache else []
+    used_run_cache = (
+        self.stage_cache.push(remote, odb=odb) if run_cache else []
+    )
 
     if isinstance(targets, str):
         targets = [targets]
@@ -41,11 +49,13 @@ def push(
     )
 
     pushed = len(used_run_cache)
-    for odb, obj_ids in used.items():
-        if odb and odb.read_only:
+    for dest_odb, obj_ids in used.items():
+        if dest_odb and dest_odb.read_only:
             continue
         try:
-            pushed += self.cloud.push(obj_ids, jobs, remote=remote, odb=odb)
+            pushed += self.cloud.push(
+                obj_ids, jobs, remote=remote, odb=odb or dest_odb
+            )
         except FileTransferError as exc:
             raise UploadError(exc.amount)
     return pushed

dvc/stage/cache.py

@@ -2,7 +2,7 @@ import logging
 import os
 import tempfile
 from contextlib import contextmanager
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 from funcy import cached_property, first
 
@@ -10,6 +10,9 @@ from dvc.exceptions import DvcException
 from dvc.path_info import PathInfo
 from dvc.utils import dict_sha256, relpath
 
+if TYPE_CHECKING:
+    from dvc.objects.db.base import ObjectDB
+
 logger = logging.getLogger(__name__)
 
 
@@ -229,10 +232,11 @@ class StageCache:
 
         return ret
 
-    def push(self, remote: Optional[str]):
+    def push(self, remote: Optional[str], odb: Optional["ObjectDB"] = None):
         from dvc.objects.transfer import _log_exceptions
 
-        odb = self.repo.cloud.get_remote_odb(remote)
+        if odb is None:
+            odb = self.repo.cloud.get_remote_odb(remote)
         return self._transfer(
             _log_exceptions(odb.fs.upload),
             self.repo.odb.local,

tests/func/experiments/executor/test_ssh.py

@@ -0,0 +1,87 @@
+import posixpath
+from contextlib import contextmanager
+from functools import partial
+
+import pytest
+
+from dvc.fs.ssh import SSHFileSystem
+from dvc.path_info import URLInfo
+from dvc.repo.experiments.base import EXEC_HEAD, EXEC_MERGE
+from dvc.repo.experiments.executor.ssh import SSHExecutor
+from tests.func.machine.conftest import *  # noqa, pylint: disable=wildcard-import
+from tests.remotes.ssh import TEST_SSH_KEY_PATH, TEST_SSH_USER
+
+
+@contextmanager
+def _ssh_factory(cloud):
+    yield SSHFileSystem(
+        host=cloud.host,
+        port=cloud.port,
+        user=TEST_SSH_USER,
+        keyfile=TEST_SSH_KEY_PATH,
+    )
+
+
+def test_from_machine(tmp_dir, scm, dvc, machine_instance, mocker):
+    mocker.patch.object(SSHExecutor, "_init_git")
+    executor = SSHExecutor.from_machine(dvc.machine, "foo", scm, "")
+    assert executor.host == machine_instance["instance_ip"]
+
+
+@pytest.mark.needs_internet
+@pytest.mark.parametrize("cloud", [pytest.lazy_fixture("git_ssh")])
+def test_init_git(tmp_dir, scm, cloud):
+    tmp_dir.scm_gen({"foo": "foo", "dir": {"bar": "bar"}}, commit="init")
+    rev = scm.get_rev()
+    scm.set_ref(EXEC_HEAD, rev)
+    tmp_dir.gen("foo", "stashed")
+    scm.gitpython.git.stash()
+    rev = scm.resolve_rev("stash@{0}")
+    scm.set_ref(EXEC_MERGE, rev)
+
+    root_url = URLInfo(str(cloud)) / SSHExecutor.gen_dirname()
+
+    executor = SSHExecutor(
+        scm,
+        ".",
+        root_dir=root_url.path,
+        host=root_url.host,
+        port=root_url.port,
+        username=TEST_SSH_USER,
+        fs_factory=partial(_ssh_factory, cloud),
+    )
+    assert root_url.path == executor._repo_abspath
+
+    fs = cloud._ssh
+    assert fs.exists(posixpath.join(executor._repo_abspath, "foo"))
+    assert fs.exists(posixpath.join(executor._repo_abspath, "dir"))
+    assert fs.exists(posixpath.join(executor._repo_abspath, "dir", "bar"))
+
+
+@pytest.mark.needs_internet
+@pytest.mark.parametrize("cloud", [pytest.lazy_fixture("git_ssh")])
+def test_init_cache(tmp_dir, dvc, scm, cloud):
+    foo = tmp_dir.dvc_gen("foo", "foo", commit="init")[0].outs[0]
+    rev = scm.get_rev()
+    scm.set_ref(EXEC_HEAD, rev)
+    scm.set_ref(EXEC_MERGE, rev)
+    root_url = URLInfo(str(cloud)) / SSHExecutor.gen_dirname()
+
+    executor = SSHExecutor(
+        scm,
+        ".",
+        root_dir=root_url.path,
+        host=root_url.host,
+        port=root_url.port,
+        username=TEST_SSH_USER,
+        fs_factory=partial(_ssh_factory, cloud),
+    )
+    executor.init_cache(dvc, rev)
+
+    fs = cloud._ssh
+    foo_hash = foo.hash_info.value
+    assert fs.exists(
+        posixpath.join(
+            executor._repo_abspath, ".dvc", "cache", foo_hash[:2], foo_hash[2:]
+        )
+    )

tests/func/machine/__init__.py

@@ -1,10 +0,0 @@
-import textwrap
-
-CONFIG_TEXT = textwrap.dedent(
-    """\
-        [feature]
-            machine = true
-        ['machine "foo"']
-            cloud = aws
-    """
-)

tests/func/machine/conftest.py

@@ -0,0 +1,53 @@
+import textwrap
+
+import pytest
+
+BASIC_CONFIG = textwrap.dedent(
+    """\
+        [feature]
+            machine = true
+        ['machine "foo"']
+            cloud = aws
+    """
+)
+
+TEST_INSTANCE = {
+    "aws_security_group": None,
+    "cloud": "aws",
+    "id": "iterative-2jyhw8j9ieov6",
+    "image": "ubuntu-bionic-18.04-amd64-server-20210818",
+    "instance_gpu": None,
+    "instance_hdd_size": 35,
+    "instance_ip": "123.123.123.123",
+    "instance_launch_time": "2021-08-25T07:13:03Z",
+    "instance_type": "m",
+    "name": "test-resource",
+    "region": "us-west",
+    "spot": False,
+    "spot_price": -1,
+    "ssh_name": None,
+    "ssh_private": "-----BEGIN RSA PRIVATE KEY-----\\n",
+    "startup_script": "IyEvYmluL2Jhc2g=",
+    "timeouts": None,
+}
+
+
+@pytest.fixture
+def machine_config(tmp_dir):
+    (tmp_dir / ".dvc" / "config").write_text(BASIC_CONFIG)
+    yield BASIC_CONFIG
+
+
+@pytest.fixture
+def machine_instance(tmp_dir, dvc, mocker):
+    from tpi.terraform import TerraformBackend
+
+    with dvc.config.edit() as conf:
+        conf["machine"]["foo"] = {"cloud": "aws"}
+    mocker.patch.object(
+        TerraformBackend,
+        "instances",
+        autospec=True,
+        return_value=[TEST_INSTANCE],
+    )
+    yield TEST_INSTANCE

tests/func/machine/test_machine_config.py

@@ -6,7 +6,7 @@ import tpi
 
 from dvc.main import main
 
-from . import CONFIG_TEXT
+from .conftest import BASIC_CONFIG
 
 
 @pytest.mark.parametrize(
@@ -24,25 +24,23 @@ from . import CONFIG_TEXT
         ("ssh_private", "secret"),
     ],
 )
-def test_machine_modify_susccess(tmp_dir, dvc, slot, value):
-    (tmp_dir / ".dvc" / "config").write_text(CONFIG_TEXT)
+def test_machine_modify_susccess(tmp_dir, dvc, machine_config, slot, value):
     assert main(["machine", "modify", "foo", slot, value]) == 0
     assert (
         tmp_dir / ".dvc" / "config"
-    ).read_text() == CONFIG_TEXT + f"    {slot} = {value}\n"
+    ).read_text() == machine_config + f"    {slot} = {value}\n"
     assert main(["machine", "modify", "--unset", "foo", slot]) == 0
-    assert (tmp_dir / ".dvc" / "config").read_text() == CONFIG_TEXT
+    assert (tmp_dir / ".dvc" / "config").read_text() == machine_config
 
 
-def test_machine_modify_startup_script(tmp_dir, dvc):
+def test_machine_modify_startup_script(tmp_dir, dvc, machine_config):
     slot, value = "startup_script", "start.sh"
-    (tmp_dir / ".dvc" / "config").write_text(CONFIG_TEXT)
     assert main(["machine", "modify", "foo", slot, value]) == 0
     assert (
         tmp_dir / ".dvc" / "config"
-    ).read_text() == CONFIG_TEXT + f"    {slot} = ../{value}\n"
+    ).read_text() == machine_config + f"    {slot} = ../{value}\n"
     assert main(["machine", "modify", "--unset", "foo", slot]) == 0
-    assert (tmp_dir / ".dvc" / "config").read_text() == CONFIG_TEXT
+    assert (tmp_dir / ".dvc" / "config").read_text() == machine_config
 
 
 @pytest.mark.parametrize(
@@ -57,11 +55,11 @@ def test_machine_modify_startup_script(tmp_dir, dvc):
         ("instance_hdd_size", "BIG", "expected int"),
     ],
 )
-def test_machine_modify_fail(tmp_dir, dvc, caplog, slot, value, msg):
-    (tmp_dir / ".dvc" / "config").write_text(CONFIG_TEXT)
-
+def test_machine_modify_fail(
+    tmp_dir, dvc, machine_config, caplog, slot, value, msg
+):
     assert main(["machine", "modify", "foo", slot, value]) == 251
-    assert (tmp_dir / ".dvc" / "config").read_text() == CONFIG_TEXT
+    assert (tmp_dir / ".dvc" / "config").read_text() == machine_config
     assert msg in caplog.text
 
 
@@ -117,9 +115,10 @@ def test_machine_list(tmp_dir, dvc, capsys):
     )
 
 
-def test_machine_rename_success(tmp_dir, scm, dvc, capsys, mocker):
+def test_machine_rename_success(
+    tmp_dir, scm, dvc, machine_config, capsys, mocker
+):
     config_file = tmp_dir / ".dvc" / "config"
-    config_file.write_text(CONFIG_TEXT)
 
     mocker.patch.object(
         tpi.terraform.TerraformBackend,
@@ -133,7 +132,7 @@ def test_machine_rename_success(tmp_dir, scm, dvc, capsys, mocker):
     assert main(["machine", "rename", "foo", "bar"]) == 0
     cap = capsys.readouterr()
     assert "Rename machine 'foo' to 'bar'." in cap.out
-    assert config_file.read_text() == CONFIG_TEXT.replace("foo", "bar")
+    assert config_file.read_text() == machine_config.replace("foo", "bar")
     assert not (
         tmp_dir / ".dvc" / "tmp" / "machine" / "terraform" / "foo"
     ).exists()
@@ -143,7 +142,7 @@ def test_machine_rename_success(tmp_dir, scm, dvc, capsys, mocker):
 
 
 def test_machine_rename_none_exist(tmp_dir, scm, dvc, caplog):
-    config_alice = CONFIG_TEXT.replace("foo", "alice")
+    config_alice = BASIC_CONFIG.replace("foo", "alice")
     config_file = tmp_dir / ".dvc" / "config"
     config_file.write_text(config_alice)
     assert main(["machine", "rename", "foo", "bar"]) == 251
@@ -152,7 +151,7 @@ def test_machine_rename_none_exist(tmp_dir, scm, dvc, caplog):
 
 
 def test_machine_rename_exist(tmp_dir, scm, dvc, caplog):
-    config_bar = CONFIG_TEXT + "['machine \"bar\"']\n    cloud = aws"
+    config_bar = BASIC_CONFIG + "['machine \"bar\"']\n    cloud = aws"
     config_file = tmp_dir / ".dvc" / "config"
     config_file.write_text(config_bar)
     assert main(["machine", "rename", "foo", "bar"]) == 251
@@ -160,10 +159,10 @@ def test_machine_rename_exist(tmp_dir, scm, dvc, caplog):
     assert "Machine 'bar' already exists." in caplog.text
 
 
-def test_machine_rename_error(tmp_dir, scm, dvc, caplog, mocker):
+def test_machine_rename_error(
+    tmp_dir, scm, dvc, machine_config, caplog, mocker
+):
     config_file = tmp_dir / ".dvc" / "config"
-    config_file.write_text(CONFIG_TEXT)
-
     os.makedirs((tmp_dir / ".dvc" / "tmp" / "machine" / "terraform" / "foo"))
 
     def cmd_error(self, source, destination, **kwargs):
@@ -172,5 +171,5 @@ def test_machine_rename_error(tmp_dir, scm, dvc, caplog, mocker):
     mocker.patch.object(tpi.terraform.TerraformBackend, "state_mv", cmd_error)
 
     assert main(["machine", "rename", "foo", "bar"]) == 251
-    assert config_file.read_text() == CONFIG_TEXT
+    assert config_file.read_text() == machine_config
     assert "rename failed" in caplog.text

tests/func/machine/test_machine_status.py

@@ -1,41 +1,11 @@
-import json
-
-from tpi.terraform import TerraformBackend
-
 from dvc.command.machine import CmdMachineStatus
 from dvc.main import main
 
-from . import CONFIG_TEXT
-
-STATUS = """{
-    "aws_security_group": null,
-    "cloud": "aws",
-    "id": "iterative-2jyhw8j9ieov6",
-    "image": "ubuntu-bionic-18.04-amd64-server-20210818",
-    "instance_gpu": null,
-    "instance_hdd_size": 35,
-    "instance_ip": "123.123.123.123",
-    "instance_launch_time": "2021-08-25T07:13:03Z",
-    "instance_type": "m",
-    "name": "test-resource",
-    "region": "us-west",
-    "spot": false,
-    "spot_price": -1,
-    "ssh_name": null,
-    "ssh_private": "-----BEGIN RSA PRIVATE KEY-----\\n",
-    "startup_script": "IyEvYmluL2Jhc2g=",
-    "timeouts": null
-}"""
-
-
-def test_status(tmp_dir, scm, dvc, mocker, capsys):
-    (tmp_dir / ".dvc" / "config").write_text(CONFIG_TEXT)
-    status = json.loads(STATUS)
-
-    mocker.patch.object(
-        TerraformBackend, "instances", autospec=True, return_value=[status]
-    )
 
+def test_status(
+    tmp_dir, scm, dvc, machine_config, machine_instance, mocker, capsys
+):
+    status = machine_instance
     assert main(["machine", "status", "foo"]) == 0
     cap = capsys.readouterr()
     assert "instance_num_1:" in cap.out