mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-02 07:18:25 +01:00
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
This commit is contained in:
Peter Rowlands (변기호)
@@ -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
|
||||
|
||||
@@ -30,6 +30,7 @@ from .command import (
|
||||
install,
|
||||
live,
|
||||
ls,
|
||||
machine,
|
||||
metrics,
|
||||
move,
|
||||
params,
|
||||
@@ -87,6 +88,7 @@ COMMANDS = [
|
||||
experiments,
|
||||
check_ignore,
|
||||
live,
|
||||
machine,
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
+25
-11
@@ -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
|
||||
|
||||
|
||||
@@ -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},
|
||||
}
|
||||
|
||||
@@ -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 <name>\n"
|
||||
)
|
||||
else:
|
||||
error_msg = (
|
||||
"no machine specified. Create a default machine with\n"
|
||||
" dvc machine add -d <name> <cloud>"
|
||||
)
|
||||
|
||||
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)
|
||||
@@ -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."""
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -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 = <<EOF
|
||||
{{ startup_script }}
|
||||
EOF
|
||||
{% endif %}
|
||||
}
|
||||
@@ -112,6 +112,8 @@ 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"]
|
||||
|
||||
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"
|
||||
|
||||
@@ -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")
|
||||
@@ -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()
|
||||
@@ -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]
|
||||
Reference in New Issue
Block a user