diff --git a/dvc/machine/__init__.py b/dvc/machine/__init__.py index 211d72373..ccf767f46 100644 --- a/dvc/machine/__init__.py +++ b/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) diff --git a/dvc/machine/backend/base.py b/dvc/machine/backend/base.py index 9a0c8aa28..f780acff9 100644 --- a/dvc/machine/backend/base.py +++ b/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. """ diff --git a/dvc/machine/backend/terraform.py b/dvc/machine/backend/terraform.py index b4a9189f9..9a37732bc 100644 --- a/dvc/machine/backend/terraform.py +++ b/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): diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 43d1f1db1..d25619ed0 100644 --- a/dvc/repo/experiments/__init__.py +++ b/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): diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py index d7d977ca7..bdb149a0e 100644 --- a/dvc/repo/experiments/executor/base.py +++ b/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 diff --git a/dvc/repo/experiments/executor/local.py b/dvc/repo/experiments/executor/local.py index a90831db2..bd869e9b7 100644 --- a/dvc/repo/experiments/executor/local.py +++ b/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) diff --git a/dvc/repo/experiments/executor/ssh.py b/dvc/repo/experiments/executor/ssh.py new file mode 100644 index 000000000..efbcbb4db --- /dev/null +++ b/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 dvc-exp@noreply.localhost" + ) + + 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) diff --git a/dvc/repo/push.py b/dvc/repo/push.py index 88f9adc6a..7f5320ab3 100644 --- a/dvc/repo/push.py +++ b/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 diff --git a/dvc/stage/cache.py b/dvc/stage/cache.py index 1c96c30f4..76771a59f 100644 --- a/dvc/stage/cache.py +++ b/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, diff --git a/tests/func/experiments/executor/__init__.py b/tests/func/experiments/executor/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/func/experiments/executor/test_ssh.py b/tests/func/experiments/executor/test_ssh.py new file mode 100644 index 000000000..9d84edc6b --- /dev/null +++ b/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:] + ) + ) diff --git a/tests/func/machine/__init__.py b/tests/func/machine/__init__.py index 12b6782f5..e69de29bb 100644 --- a/tests/func/machine/__init__.py +++ b/tests/func/machine/__init__.py @@ -1,10 +0,0 @@ -import textwrap - -CONFIG_TEXT = textwrap.dedent( - """\ - [feature] - machine = true - ['machine "foo"'] - cloud = aws - """ -) diff --git a/tests/func/machine/conftest.py b/tests/func/machine/conftest.py new file mode 100644 index 000000000..2795a62b5 --- /dev/null +++ b/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 diff --git a/tests/func/machine/test_machine_config.py b/tests/func/machine/test_machine_config.py index 701c69d33..a788088d6 100644 --- a/tests/func/machine/test_machine_config.py +++ b/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 diff --git a/tests/func/machine/test_machine_status.py b/tests/func/machine/test_machine_status.py index 5c5081e28..9bea45e55 100644 --- a/tests/func/machine/test_machine_status.py +++ b/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