From 4ee98299ee4c2931ec0fd682c1d009ea36ab7d0e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Peter=20Rowlands=20=28=EB=B3=80=EA=B8=B0=ED=98=B8=29?= Date: Fri, 27 Aug 2021 09:46:24 +0000 Subject: [PATCH] `dvc machine`: implement initial terraform support (#6462) * tpi: add naive terraform-provider-iterative wrapper * executor: add naive backend/terraform wrapper * setup: add dvc[terraform] install extras * config: add feature.executor feature flag * fix named executor config sections * executor: add poc add/modify/remove/init/destroy commands * relax python-terraform requirement * ci: install terraform extras * tpi: support .tf.json format * s/executor/machine * terraform: use separate dirs per machine * s/init/create/ * reserve local machine name * add initial ssh skeleton * add initial machine unit tests * add initial machine command tests * add tests for add/remove * enable feature by default in test env * close temp pemfile object so it can be used on windows * cleanup pemfile syntax --- .github/workflows/tests.yaml | 4 +- dvc/cli.py | 2 + dvc/command/machine.py | 274 +++++++++++++++++++++++++++ dvc/config.py | 36 ++-- dvc/config_schema.py | 9 +- dvc/machine/__init__.py | 166 ++++++++++++++++ dvc/machine/backend/__init__.py | 0 dvc/machine/backend/base.py | 52 +++++ dvc/machine/backend/terraform.py | 125 ++++++++++++ dvc/repo/__init__.py | 10 + dvc/tpi/__init__.py | 49 +++++ dvc/tpi/templates/main.tf | 26 +++ setup.py | 3 + tests/unit/command/test_machine.py | 66 +++++++ tests/unit/machine/__init__.py | 0 tests/unit/machine/test_machine.py | 36 ++++ tests/unit/machine/test_terraform.py | 129 +++++++++++++ 17 files changed, 973 insertions(+), 14 deletions(-) create mode 100644 dvc/command/machine.py create mode 100644 dvc/machine/__init__.py create mode 100644 dvc/machine/backend/__init__.py create mode 100644 dvc/machine/backend/base.py create mode 100644 dvc/machine/backend/terraform.py create mode 100644 dvc/tpi/__init__.py create mode 100644 dvc/tpi/templates/main.tf create mode 100644 tests/unit/command/test_machine.py create mode 100644 tests/unit/machine/__init__.py create mode 100644 tests/unit/machine/test_machine.py create mode 100644 tests/unit/machine/test_terraform.py diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index e9060a17d..862de5784 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -27,7 +27,7 @@ jobs: - name: Install requirements run: | pip install wheel - pip install ".[all,tests]" Pygments collective.checkdocs pre-commit + pip install ".[all,tests,terraform]" Pygments collective.checkdocs pre-commit - name: install temporary dependencies run: | pip install git+https://github.com/isidentical/gdrivefs@service-account @@ -77,7 +77,7 @@ jobs: if: steps.cache.pip-cache-dir.cache-hit != 'true' run: | pip install --upgrade pip setuptools wheel - pip install -e ".[all,tests]" + pip install -e ".[all,tests,terraform]" - name: install temporary dependencies run: | pip install git+https://github.com/isidentical/gdrivefs@service-account diff --git a/dvc/cli.py b/dvc/cli.py index 4306cd784..38845f079 100644 --- a/dvc/cli.py +++ b/dvc/cli.py @@ -30,6 +30,7 @@ from .command import ( install, live, ls, + machine, metrics, move, params, @@ -87,6 +88,7 @@ COMMANDS = [ experiments, check_ignore, live, + machine, ] diff --git a/dvc/command/machine.py b/dvc/command/machine.py new file mode 100644 index 000000000..c5a51f3da --- /dev/null +++ b/dvc/command/machine.py @@ -0,0 +1,274 @@ +import argparse + +from dvc.command.base import CmdBase, append_doc_link, fix_subparsers +from dvc.command.config import CmdConfig +from dvc.config import ConfigError +from dvc.ui import ui +from dvc.utils import format_link + + +class MachineDisabledError(ConfigError): + def __init__(self): + super().__init__("Machine feature is disabled") + + +class CmdMachineConfig(CmdConfig): + def __init__(self, args): + + super().__init__(args) + if not self.config["feature"].get("machine", False): + raise MachineDisabledError + + if getattr(self.args, "name", None): + self.args.name = self.args.name.lower() + + def _check_exists(self, conf): + if self.args.name not in conf["machine"]: + raise ConfigError(f"machine '{self.args.name}' doesn't exist.") + + +class CmdMachineAdd(CmdMachineConfig): + def run(self): + from dvc.machine import validate_name + + validate_name(self.args.name) + + if self.args.default: + ui.write(f"Setting '{self.args.name}' as a default machine.") + + with self.config.edit(self.args.level) as conf: + if self.args.name in conf["machine"] and not self.args.force: + raise ConfigError( + "machine '{}' already exists. Use `-f|--force` to " + "overwrite it.".format(self.args.name) + ) + + conf["machine"][self.args.name] = {"cloud": self.args.cloud} + if self.args.default: + conf["core"]["machine"] = self.args.name + + return 0 + + +class CmdMachineRemove(CmdMachineConfig): + def run(self): + with self.config.edit(self.args.level) as conf: + self._check_exists(conf) + del conf["machine"][self.args.name] + + up_to_level = self.args.level or "repo" + # Remove core.machine refs to this machine in any shadowing configs + for level in reversed(self.config.LEVELS): + with self.config.edit(level) as conf: + if conf["core"].get("machine") == self.args.name: + del conf["core"]["machine"] + + if level == up_to_level: + break + + return 0 + + +class CmdMachineModify(CmdMachineConfig): + def run(self): + from dvc.config import merge + + with self.config.edit(self.args.level) as conf: + merged = self.config.load_config_to_level(self.args.level) + merge(merged, conf) + self._check_exists(merged) + + if self.args.name not in conf["machine"]: + conf["machine"][self.args.name] = {} + section = conf["machine"][self.args.name] + if self.args.unset: + section.pop(self.args.option, None) + else: + section[self.args.option] = self.args.value + return 0 + + +class CmdMachineDefault(CmdMachineConfig): + def run(self): + if self.args.name is None and not self.args.unset: + conf = self.config.read(self.args.level) + try: + print(conf["core"]["machine"]) + except KeyError: + ui.write("No default machine set") + return 1 + else: + with self.config.edit(self.args.level) as conf: + if self.args.unset: + conf["core"].pop("machine", None) + else: + merged_conf = self.config.load_config_to_level( + self.args.level + ) + if ( + self.args.name in conf["machine"] + or self.args.name in merged_conf["machine"] + ): + conf["core"]["machine"] = self.args.name + else: + raise ConfigError( + "default machine must be present in machine " + "list." + ) + return 0 + + +class CmdMachineCreate(CmdBase): + def run(self): + if self.repo.machine is None: + raise MachineDisabledError + + self.repo.machine.create(self.args.name) + return 0 + + +class CmdMachineDestroy(CmdBase): + def run(self): + if self.repo.machine is None: + raise MachineDisabledError + + self.repo.machine.destroy(self.args.name) + return 0 + + +def add_parser(subparsers, parent_parser): + from dvc.command.config import parent_config_parser + + machine_HELP = "Set up and manage cloud machines." + machine_parser = subparsers.add_parser( + "machine", + parents=[parent_parser], + description=append_doc_link(machine_HELP, "machine"), + # NOTE: suppress help during development to hide command + # help=machine_HELP, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + + machine_subparsers = machine_parser.add_subparsers( + dest="cmd", + help="Use `dvc machine CMD --help` for " "command-specific help.", + ) + + fix_subparsers(machine_subparsers) + + machine_ADD_HELP = "Add a new data machine." + machine_add_parser = machine_subparsers.add_parser( + "add", + parents=[parent_config_parser, parent_parser], + description=append_doc_link(machine_ADD_HELP, "machine/add"), + help=machine_ADD_HELP, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + machine_add_parser.add_argument("name", help="Name of the machine") + machine_add_parser.add_argument( + "cloud", + help="Machine cloud. See full list of supported clouds at {}".format( + format_link( + "https://github.com/iterative/" + "terraform-provider-iterative#machine" + ) + ), + ) + machine_add_parser.add_argument( + "-d", + "--default", + action="store_true", + default=False, + help="Set as default machine.", + ) + machine_add_parser.add_argument( + "-f", + "--force", + action="store_true", + default=False, + help="Force overwriting existing configs", + ) + machine_add_parser.set_defaults(func=CmdMachineAdd) + + machine_DEFAULT_HELP = "Set/unset the default machine." + machine_default_parser = machine_subparsers.add_parser( + "default", + parents=[parent_config_parser, parent_parser], + description=append_doc_link(machine_DEFAULT_HELP, "machine/default"), + help=machine_DEFAULT_HELP, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + machine_default_parser.add_argument( + "name", nargs="?", help="Name of the machine" + ) + machine_default_parser.add_argument( + "-u", + "--unset", + action="store_true", + default=False, + help="Unset default machine.", + ) + machine_default_parser.set_defaults(func=CmdMachineDefault) + + machine_MODIFY_HELP = "Modify the configuration of an machine." + machine_modify_parser = machine_subparsers.add_parser( + "modify", + parents=[parent_config_parser, parent_parser], + description=append_doc_link(machine_MODIFY_HELP, "machine/modify"), + help=machine_MODIFY_HELP, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + machine_modify_parser.add_argument("name", help="Name of the machine") + machine_modify_parser.add_argument( + "option", help="Name of the option to modify." + ) + machine_modify_parser.add_argument( + "value", nargs="?", help="(optional) Value of the option." + ) + machine_modify_parser.add_argument( + "-u", + "--unset", + default=False, + action="store_true", + help="Unset option.", + ) + machine_modify_parser.set_defaults(func=CmdMachineModify) + + machine_REMOVE_HELP = "Remove an machine." + machine_remove_parser = machine_subparsers.add_parser( + "remove", + parents=[parent_config_parser, parent_parser], + description=append_doc_link(machine_REMOVE_HELP, "machine/remove"), + help=machine_REMOVE_HELP, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + machine_remove_parser.add_argument( + "name", help="Name of the machine to remove." + ) + machine_remove_parser.set_defaults(func=CmdMachineRemove) + + machine_CREATE_HELP = "Create and start a machine instance." + machine_create_parser = machine_subparsers.add_parser( + "create", + parents=[parent_config_parser, parent_parser], + description=append_doc_link(machine_CREATE_HELP, "machine/create"), + help=machine_CREATE_HELP, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + machine_create_parser.add_argument( + "name", help="Name of the machine to create." + ) + machine_create_parser.set_defaults(func=CmdMachineCreate) + + machine_DESTROY_HELP = "Destroy an machine instance." + machine_destroy_parser = machine_subparsers.add_parser( + "destroy", + parents=[parent_config_parser, parent_parser], + description=append_doc_link(machine_DESTROY_HELP, "machine/destroy"), + help=machine_DESTROY_HELP, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + machine_destroy_parser.add_argument( + "name", help="Name of the machine instance to destroy." + ) + machine_destroy_parser.set_defaults(func=CmdMachineDestroy) diff --git a/dvc/config.py b/dvc/config.py index 98fe3086f..ac927c5c6 100644 --- a/dvc/config.py +++ b/dvc/config.py @@ -32,6 +32,18 @@ class RemoteNotFoundError(RemoteConfigError): pass +class MachineConfigError(ConfigError): + pass + + +class NoMachineError(MachineConfigError): + pass + + +class MachineNotFoundError(MachineConfigError): + pass + + @memoize def get_compiled_schema(): from voluptuous import Schema @@ -167,7 +179,7 @@ class Config(dict): conf_obj = ConfigObj(fobj) else: conf_obj = ConfigObj() - return _parse_remotes(_lower_keys(conf_obj.dict())) + return _parse_named(_lower_keys(conf_obj.dict())) def _save_config(self, level, conf_dict): from configobj import ConfigObj @@ -179,7 +191,7 @@ class Config(dict): fs.makedirs(os.path.dirname(filename)) - config = ConfigObj(_pack_remotes(conf_dict)) + config = ConfigObj(_pack_named(conf_dict)) with fs.open(filename, "wb") as fobj: config.write(fobj) config.filename = filename @@ -300,27 +312,29 @@ class Config(dict): raise ConfigError(str(exc)) from None -def _parse_remotes(conf): - result = {"remote": {}} +def _parse_named(conf): + result = {"remote": {}, "machine": {}} for section, val in conf.items(): - name = re_find(r'^\s*remote\s*"(.*)"\s*$', section) - if name: - result["remote"][name] = val + match = re_find(r'^\s*(remote|machine)\s*"(.*)"\s*$', section) + if match: + key, name = match + result[key][name] = val else: result[section] = val return result -def _pack_remotes(conf): +def _pack_named(conf): # Drop empty sections result = compact(conf) # Transform remote.name -> 'remote "name"' - for name, val in conf["remote"].items(): - result[f'remote "{name}"'] = val - result.pop("remote", None) + for key in ("remote", "machine"): + for name, val in conf[key].items(): + result[f'{key} "{name}"'] = val + result.pop(key, None) return result diff --git a/dvc/config_schema.py b/dvc/config_schema.py index 0ce6d0916..06ae3245f 100644 --- a/dvc/config_schema.py +++ b/dvc/config_schema.py @@ -118,6 +118,7 @@ SCHEMA = { Optional("autostage", default=False): Bool, Optional("experiments"): Bool, # obsoleted Optional("check_update", default=True): Bool, + "machine": Lower, }, "cache": { "local": str, @@ -224,10 +225,16 @@ SCHEMA = { "row_limit": All(Coerce(int), Range(1)), # obsoleted "row_cleanup_quota": All(Coerce(int), Range(0, 100)), # obsoleted }, + "machine": { + str: { + "cloud": All(Lower, Choices("aws", "azure")), + }, + }, # section for experimental features "feature": { + Optional("machine", default=False): Bool, # enabled by default. It's of no use, kept for backward compatibility. - Optional("parametrization", default=True): Bool + Optional("parametrization", default=True): Bool, }, "plots": {"html_template": str}, } diff --git a/dvc/machine/__init__.py b/dvc/machine/__init__.py new file mode 100644 index 000000000..164864257 --- /dev/null +++ b/dvc/machine/__init__.py @@ -0,0 +1,166 @@ +import logging +import os +from typing import ( + TYPE_CHECKING, + Dict, + Iterable, + Mapping, + Optional, + Tuple, + Type, +) + +from dvc.types import StrPath + +from .backend.base import BaseMachineBackend +from .backend.terraform import TerraformBackend + +if TYPE_CHECKING: + from dvc.repo import Repo + +logger = logging.getLogger(__name__) + +BackendCls = Type[BaseMachineBackend] + + +RESERVED_NAMES = {"local", "localhost"} + + +def validate_name(name: str): + from dvc.exceptions import InvalidArgumentError + + name = name.lower() + if name in RESERVED_NAMES: + raise InvalidArgumentError( + f"Machine name '{name}' is reserved for internal DVC use." + ) + + +class MachineBackends(Mapping): + DEFAULT: Dict[str, BackendCls] = { + "terraform": TerraformBackend, + } + + def __getitem__(self, key: str) -> BaseMachineBackend: + """Lazily initialize backends and cache it afterwards""" + initialized = self.initialized.get(key) + if not initialized: + backend = self.backends[key] + initialized = backend( + os.path.join(self.tmp_dir, key), **self.kwargs + ) + self.initialized[key] = initialized + return initialized + + def __init__( + self, + selected: Optional[Iterable[str]], + tmp_dir: StrPath, + **kwargs, + ) -> None: + selected = selected or list(self.DEFAULT) + self.backends = {key: self.DEFAULT[key] for key in selected} + + self.initialized: Dict[str, BaseMachineBackend] = {} + + self.tmp_dir = tmp_dir + self.kwargs = kwargs + + def __iter__(self): + return iter(self.backends) + + def __len__(self) -> int: + return len(self.backends) + + def close_initialized(self) -> None: + for backend in self.initialized.values(): + backend.close() + + +class MachineManager: + """Class that manages dvc cloud machines. + + Args: + repo (dvc.repo.Repo): repo instance that belongs to the repo that + we are working on. + + Raises: + config.ConfigError: thrown when config has invalid format. + """ + + CLOUD_BACKENDS = { + "aws": "terraform", + "azure": "terraform", + } + + def __init__( + self, repo: "Repo", backends: Optional[Iterable[str]] = None, **kwargs + ): + self.repo = repo + tmp_dir = os.path.join(self.repo.tmp_dir, "machine") + self.backends = MachineBackends(backends, tmp_dir=tmp_dir, **kwargs) + + def get_config_and_backend( + self, + name: Optional[str] = None, + ) -> Tuple[dict, "BaseMachineBackend"]: + from dvc.config import NoMachineError + + if not name: + name = self.repo.config["core"].get("machine") + + if name: + config = self._get_config(name=name) + backend = self._get_backend(config["cloud"]) + return config, backend + + if bool(self.repo.config["machine"]): + error_msg = ( + "no machine specified. Setup default machine with\n" + " dvc machine default \n" + ) + else: + error_msg = ( + "no machine specified. Create a default machine with\n" + " dvc machine add -d " + ) + + raise NoMachineError(error_msg) + + def _get_config(self, **kwargs): + config = self.repo.config + name = kwargs.get("name") + if name: + try: + conf = config["machine"][name.lower()] + conf["name"] = name + except KeyError: + from dvc.config import MachineNotFoundError + + raise MachineNotFoundError(f"machine '{name}' doesn't exist") + else: + conf = kwargs + return conf + + def _get_backend(self, cloud: str) -> BaseMachineBackend: + from dvc.config import NoMachineError + + try: + backend = self.CLOUD_BACKENDS[cloud] + return self.backends[backend] + except KeyError: + raise NoMachineError(f"Machine platform '{cloud}' unsupported") + + def create(self, name: Optional[str]): + """Create and start the specified machine instance.""" + config, backend = self.get_config_and_backend(name) + return backend.create(**config) + + def destroy(self, name: Optional[str]): + """Destroy the specified machine instance.""" + config, backend = self.get_config_and_backend(name) + return backend.destroy(**config) + + def get_sshfs(self, name: Optional[str]): + config, backend = self.get_config_and_backend(name) + return backend.get_sshfs(**config) diff --git a/dvc/machine/backend/__init__.py b/dvc/machine/backend/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/dvc/machine/backend/base.py b/dvc/machine/backend/base.py new file mode 100644 index 000000000..4946823bd --- /dev/null +++ b/dvc/machine/backend/base.py @@ -0,0 +1,52 @@ +import logging +from abc import ABC, abstractmethod +from contextlib import contextmanager +from typing import TYPE_CHECKING, Iterator, Optional + +from dvc.types import StrPath +from dvc.utils.fs import makedirs + +if TYPE_CHECKING: + from dvc.fs.ssh import SSHFileSystem + from dvc.repo.experiments.executor.base import BaseExecutor + +logger = logging.getLogger(__name__) + + +class BaseMachineBackend(ABC): + def __init__(self, tmp_dir: StrPath, **kwargs): + self.tmp_dir = tmp_dir + makedirs(self.tmp_dir, exist_ok=True) + + @abstractmethod + def create(self, name: Optional[str] = None, **config): + """Create and start an instance of the specified machine.""" + + @abstractmethod + def destroy(self, name: Optional[str] = None, **config): + """Stop and destroy all instances of the specified machine.""" + + @abstractmethod + def instances( + self, name: Optional[str] = None, **config + ) -> Iterator[dict]: + """Iterate over status of all instances of the specified machine.""" + + def close(self): + pass + + @abstractmethod + def get_executor( + self, name: Optional[str] = None, **config + ) -> "BaseExecutor": + """Return an executor instance which can be used for DVC + experiment/pipeline execution on the specified machine. + """ + + @abstractmethod + @contextmanager + def get_sshfs( + self, name: Optional[str] = None, **config + ) -> Iterator["SSHFileSystem"]: + """Return an sshfs instance for the default directory on the + specified machine.""" diff --git a/dvc/machine/backend/terraform.py b/dvc/machine/backend/terraform.py new file mode 100644 index 000000000..848a97e74 --- /dev/null +++ b/dvc/machine/backend/terraform.py @@ -0,0 +1,125 @@ +import logging +import os +import tempfile +from contextlib import contextmanager +from typing import TYPE_CHECKING, Iterator, Optional + +from funcy import first +from python_terraform import Terraform + +from dvc.exceptions import DvcException +from dvc.fs.ssh import SSHFileSystem + +from .base import BaseMachineBackend + +if TYPE_CHECKING: + from dvc.repo.experiments.executor.base import BaseExecutor + +logger = logging.getLogger(__name__) + + +class TerraformError(DvcException): + pass + + +class DvcTerraform(Terraform): + def cmd(self, *args, **kwargs): + logger.debug(" ".join(self.generate_cmd_string(*args, **kwargs))) + kwargs["capture_output"] = False + ret, _stdout, _stderr = super().cmd(*args, **kwargs) + if ret != 0: + raise TerraformError("Cmd 'terraform {cmd}' failed") + + def iter_instances(self, name: str): + """Iterate over active iterative_machine instances.""" + self.read_state_file() + resources = getattr(self.tfstate, "resources", []) + for resource in resources: + if ( + resource.get("type") == "iterative_machine" + and resource.get("name") == name + ): + yield from ( + instance.get("attributes", {}) + for instance in resource.get("instances", []) + ) + + @staticmethod + @contextmanager + def pemfile(resource: dict): + """Generate a temporary SSH key (PEM) file for the specified resource. + + The path to the file will be yielded, and the file will be removed + upon exiting the yielded context. + """ + with tempfile.NamedTemporaryFile(delete=False) as fobj: + fobj.write(resource["ssh_private"].encode("ascii")) + try: + yield fobj.name + finally: + os.unlink(fobj.name) + + +class TerraformBackend(BaseMachineBackend): + @contextmanager + def make_tf(self, name: str): + from dvc.utils.fs import makedirs + + try: + working_dir = os.path.join(self.tmp_dir, name) + makedirs(working_dir, exist_ok=True) + yield DvcTerraform(working_dir=working_dir) + except TerraformError: + raise + except Exception as exc: + raise TerraformError("terraform failed") from exc + + def create(self, name: Optional[str] = None, **config): + from python_terraform import IsFlagged + + from dvc.tpi import render_json + + assert name and "cloud" in config + with self.make_tf(name) as tf: + tf_file = os.path.join(tf.working_dir, "main.tf.json") + with open(tf_file, "w", encoding="utf-8") as fobj: + fobj.write(render_json(name=name, **config, indent=2)) + tf.cmd("init") + tf.cmd("apply", auto_approve=IsFlagged) + + def destroy(self, name: Optional[str] = None, **config): + from python_terraform import IsFlagged + + assert name + + with self.make_tf(name) as tf: + if first(tf.iter_instances(name)): + tf.cmd("destroy", auto_approve=IsFlagged) + + def instances( + self, name: Optional[str] = None, **config + ) -> Iterator[dict]: + assert name + + with self.make_tf(name) as tf: + yield from tf.iter_instances(name) + + def get_executor( + self, name: Optional[str] = None, **config + ) -> "BaseExecutor": + raise NotImplementedError + + @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") + with DvcTerraform.pemfile(resource) as pem: + fs = SSHFileSystem( + host=resource["instance_ip"], + user="ubuntu", + keyfile=pem, + ) + yield fs diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index e853be219..49c2afe16 100644 --- a/dvc/repo/__init__.py +++ b/dvc/repo/__init__.py @@ -12,6 +12,7 @@ from dvc.exceptions import IsADirectoryError as DvcIsADirectoryError from dvc.exceptions import NotDvcRepoError, OutputNotFoundError from dvc.ignore import DvcIgnoreFilter from dvc.path_info import PathInfo +from dvc.utils import env2bool from dvc.utils.fs import path_isin if TYPE_CHECKING: @@ -134,6 +135,7 @@ class Repo: from dvc.data_cloud import DataCloud from dvc.fs.local import LocalFileSystem from dvc.lock import LockNoop, make_lock + from dvc.machine import MachineManager from dvc.objects.db import ODBManager from dvc.repo.live import Live from dvc.repo.metrics import Metrics @@ -196,6 +198,14 @@ class Repo: self.params = Params(self) self.live = Live(self) + if self.tmp_dir and ( + self.config["feature"].get("machine", False) + or env2bool("DVC_TEST") + ): + self.machine = MachineManager(self) + else: + self.machine = None + self.stage_collection_error_handler: Optional[ Callable[[str, Exception], None] ] = None diff --git a/dvc/tpi/__init__.py b/dvc/tpi/__init__.py new file mode 100644 index 000000000..de137038c --- /dev/null +++ b/dvc/tpi/__init__.py @@ -0,0 +1,49 @@ +import json +import logging +import os +from collections import OrderedDict + +from jinja2 import Environment, FileSystemLoader, select_autoescape + +logger = logging.getLogger(__name__) + +TEMPLATE_DIR = os.path.join(os.path.dirname(__file__), "templates") +TEMPLATE_MAIN = "main.tf" + +BASE_CONFIG = { + "//": "This file auto-generated by dvc-tpi, do not edit manually.", + "terraform": { + "required_providers": { + "iterative": { + "source": "iterative/iterative", + }, + }, + }, + "provider": { + "iterative": {}, + }, +} + +_jinja_env = Environment( + loader=FileSystemLoader(TEMPLATE_DIR), + autoescape=select_autoescape(), +) + + +def render_config(**config) -> str: + """Render TPI iterative_machine config as terraform HCL.""" + assert "name" in config and "cloud" in config + template = _jinja_env.get_template(TEMPLATE_MAIN) + return template.render(**config) + + +def render_json(indent=None, **config) -> str: + """Render TPI iterative_machine config as terraform JSON.""" + assert "name" in config and "cloud" in config + tf_config = OrderedDict(BASE_CONFIG) + name = config["name"] + resource = { + name: {key: value for key, value in config.items()}, + } + tf_config["resource"] = {"iterative_machine": resource} + return json.dumps(tf_config, indent=indent) diff --git a/dvc/tpi/templates/main.tf b/dvc/tpi/templates/main.tf new file mode 100644 index 000000000..efec829c9 --- /dev/null +++ b/dvc/tpi/templates/main.tf @@ -0,0 +1,26 @@ +terraform { + required_providers { + iterative = { + source = "iterative/iterative" + } + } +} + +provider "iterative" {} + +resource "iterative_machine" "executor" { + {% if name %}name = "{{name}}"{% endif %} + {% if cloud %}cloud = "{{cloud}}"{% endif %} + {% if region %}region = "{{region}}"{% endif %} + {% if image %}image = "{{image}}"{% endif %} + {% if instance_type %}instance_type = "{{instance_type}}"{% endif %} + {% if instance_gpu %}instance_gpu = "{{instance_gpu}}"{% endif %} + {% if instance_hdd_size %}instance_hdd_size = {{instance_hdd_size}}{% endif %} + {% if ssh_private %}ssh_private = "{{ssh_private}}"{% endif %} + {% if spot %}spot = {{spot}}{% endif %} + {% if spot_price %}spot_price = {{spot_price}}{% endif %} + {% if startup_script %}startup_script = <=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"] + tests_requirements = ( Path("test_requirements.txt").read_text().strip().splitlines() ) @@ -138,6 +140,7 @@ setup( "hdfs": hdfs, "webhdfs": webhdfs, "webdav": webdav, + "terraform": terraform, "tests": tests_requirements, }, keywords="data-science data-version-control machine-learning git" diff --git a/tests/unit/command/test_machine.py b/tests/unit/command/test_machine.py new file mode 100644 index 000000000..2e4458a0a --- /dev/null +++ b/tests/unit/command/test_machine.py @@ -0,0 +1,66 @@ +import configobj + +from dvc.cli import parse_args +from dvc.command.machine import ( + CmdMachineAdd, + CmdMachineCreate, + CmdMachineDestroy, + CmdMachineRemove, +) + + +def test_add(tmp_dir): + tmp_dir.gen({".dvc": {"config": "[feature]\n machine = true"}}) + cli_args = parse_args(["machine", "add", "foo", "aws"]) + assert cli_args.func == CmdMachineAdd + cmd = cli_args.func(cli_args) + assert cmd.run() == 0 + config = configobj.ConfigObj(str(tmp_dir / ".dvc" / "config")) + assert config['machine "foo"']["cloud"] == "aws" + + +def test_remove(tmp_dir): + tmp_dir.gen( + { + ".dvc": { + "config": ( + "[feature]\n" + " machine = true\n" + "['machine \"foo\"']\n" + " cloud = aws" + ) + } + } + ) + cli_args = parse_args(["machine", "remove", "foo"]) + assert cli_args.func == CmdMachineRemove + cmd = cli_args.func(cli_args) + assert cmd.run() == 0 + config = configobj.ConfigObj(str(tmp_dir / ".dvc" / "config")) + assert list(config.keys()) == ["feature"] + + +def test_create(tmp_dir, dvc, mocker): + cli_args = parse_args(["machine", "create", "foo"]) + assert cli_args.func == CmdMachineCreate + + cmd = cli_args.func(cli_args) + m = mocker.patch.object( + cmd.repo.machine, "create", autospec=True, return_value=0 + ) + + assert cmd.run() == 0 + m.assert_called_once_with("foo") + + +def test_destroy(tmp_dir, dvc, mocker): + cli_args = parse_args(["machine", "destroy", "foo"]) + assert cli_args.func == CmdMachineDestroy + + cmd = cli_args.func(cli_args) + m = mocker.patch.object( + cmd.repo.machine, "destroy", autospec=True, return_value=0 + ) + + assert cmd.run() == 0 + m.assert_called_once_with("foo") diff --git a/tests/unit/machine/__init__.py b/tests/unit/machine/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/machine/test_machine.py b/tests/unit/machine/test_machine.py new file mode 100644 index 000000000..69e0754b5 --- /dev/null +++ b/tests/unit/machine/test_machine.py @@ -0,0 +1,36 @@ +import pytest + + +def test_validate_name(): + from dvc.exceptions import InvalidArgumentError + from dvc.machine import RESERVED_NAMES, validate_name + + for name in RESERVED_NAMES: + with pytest.raises(InvalidArgumentError): + validate_name(name) + + +@pytest.mark.parametrize("cloud", ["aws", "azure"]) +def test_get_config_and_backend(tmp_dir, dvc, cloud): + from dvc.machine.backend.terraform import TerraformBackend + + name = "foo" + with dvc.config.edit() as conf: + conf["machine"][name] = {"cloud": cloud} + config, backend = dvc.machine.get_config_and_backend(name) + assert config == {"cloud": cloud, "name": name} + assert isinstance(backend, TerraformBackend) + + +def test_get_config_and_backend_nonexistent(tmp_dir, dvc): + from dvc.config import MachineNotFoundError + + with pytest.raises(MachineNotFoundError): + dvc.machine.get_config_and_backend("foo") + + +def test_get_config_and_backend_default(tmp_dir, dvc): + from dvc.config import NoMachineError + + with pytest.raises(NoMachineError): + dvc.machine.get_config_and_backend() diff --git a/tests/unit/machine/test_terraform.py b/tests/unit/machine/test_terraform.py new file mode 100644 index 000000000..cfa81d4e3 --- /dev/null +++ b/tests/unit/machine/test_terraform.py @@ -0,0 +1,129 @@ +import json +import os + +import pytest +from python_terraform import IsFlagged + +from dvc.machine.backend.terraform import DvcTerraform, TerraformError + +TEST_RESOURCE_STATE = """{ + "version": 4, + "terraform_version": "1.0.4", + "serial": 12, + "lineage": "7f93cd8b-3abf-b4a4-d939-60b55699c3ed", + "outputs": {}, + "resources": [ + { + "mode": "managed", + "type": "iterative_machine", + "name": "test-resource", + "provider": "provider[\\"registry.terraform.io/iterative/iterative\\"]", + "instances": [ + { + "schema_version": 0, + "attributes": { + "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 + }, + "sensitive_attributes": [], + "private": "" + } + ] + } + ] +} +""" + + +@pytest.fixture +def terraform(tmp_dir, mocker): + from python_terraform import Terraform + + from dvc.machine.backend.terraform import TerraformBackend + + cmd_mock = mocker.patch.object( + Terraform, "cmd", return_value=(0, None, None) + ) + tf = TerraformBackend(tmp_dir) + tf.cmd_mock = cmd_mock + yield tf + + +@pytest.fixture +def resource(tmp_dir): + name = "test-resource" + (tmp_dir / name).gen("terraform.tfstate", TEST_RESOURCE_STATE) + yield name + + +def test_run_cmd(mocker): + from python_terraform import Terraform + + tf = DvcTerraform() + cmd_mock = mocker.patch.object( + Terraform, "cmd", return_value=(0, None, None) + ) + tf.cmd("init") + cmd_mock.assert_called_once_with("init", capture_output=False) + + cmd_mock = mocker.patch.object( + Terraform, "cmd", return_value=(-1, None, None) + ) + with pytest.raises(TerraformError): + tf.cmd("init") + + +def test_make_pemfile(): + content = "foo\nbar\n" + with DvcTerraform.pemfile({"ssh_private": content}) as path: + with open(path, "r") as fobj: + assert fobj.read() == content + + +def test_make_tf(terraform): + name = "foo" + with terraform.make_tf(name) as tf: + assert tf.working_dir == os.path.join(terraform.tmp_dir, name) + + +def test_create(tmp_dir, terraform): + name = "foo" + terraform.create(name=name, cloud="aws") + terraform.cmd_mock.assert_any_call("init", capture_output=False) + terraform.cmd_mock.assert_any_call( + "apply", capture_output=False, auto_approve=IsFlagged + ) + assert (tmp_dir / name / "main.tf.json").exists() + with open(tmp_dir / name / "main.tf.json") as fobj: + data = json.load(fobj) + assert data["resource"]["iterative_machine"] == { + name: {"name": name, "cloud": "aws"}, + } + + +def test_destroy(tmp_dir, terraform, resource): + terraform.destroy(name=resource, cloud="aws") + terraform.cmd_mock.assert_called_with( + "destroy", capture_output=False, auto_approve=IsFlagged + ) + + +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]