From 0578643f0826515f1ebe08cbd6ae04e0d6fce707 Mon Sep 17 00:00:00 2001 From: Casper da Costa-Luis Date: Fri, 3 Sep 2021 06:37:24 +0100 Subject: [PATCH] machine: move to `tpi[ssh]` (#6528) * machine: move to tpi[ssh] - depends on https://github.com/iterative/tpi * lazy import tpi * fix lint warning Co-authored-by: Peter Rowlands --- dvc/machine/__init__.py | 34 +++++-- dvc/machine/backend/base.py | 67 +------------ dvc/machine/backend/terraform.py | 80 +-------------- dvc/tpi/__init__.py | 96 ------------------ dvc/tpi/templates/main.tf | 26 ----- setup.py | 4 +- tests/unit/machine/test_terraform.py | 143 --------------------------- 7 files changed, 34 insertions(+), 416 deletions(-) delete mode 100644 dvc/tpi/__init__.py delete mode 100644 dvc/tpi/templates/main.tf delete mode 100644 tests/unit/machine/test_terraform.py diff --git a/dvc/machine/__init__.py b/dvc/machine/__init__.py index 20148d5e3..8d87b1c39 100644 --- a/dvc/machine/__init__.py +++ b/dvc/machine/__init__.py @@ -10,17 +10,17 @@ from typing import ( Type, ) +from dvc.exceptions import DvcException 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__) + from .backend.base import BaseMachineBackend -BackendCls = Type[BaseMachineBackend] + BackendCls = Type[BaseMachineBackend] + +logger = logging.getLogger(__name__) RESERVED_NAMES = {"local", "localhost"} @@ -37,11 +37,16 @@ def validate_name(name: str): class MachineBackends(Mapping): - DEFAULT: Dict[str, BackendCls] = { + try: + from .backend.terraform import TerraformBackend + except ImportError: + TerraformBackend = None # type: ignore[assignment, misc] + + DEFAULT: Dict[str, Optional["BackendCls"]] = { "terraform": TerraformBackend, } - def __getitem__(self, key: str) -> BaseMachineBackend: + def __getitem__(self, key: str) -> "BaseMachineBackend": """Lazily initialize backends and cache it afterwards""" initialized = self.initialized.get(key) if not initialized: @@ -59,9 +64,18 @@ class MachineBackends(Mapping): **kwargs, ) -> None: selected = selected or list(self.DEFAULT) - self.backends = {key: self.DEFAULT[key] for key in selected} + self.backends: Dict[str, "BackendCls"] = {} + for key in selected: + cls = self.DEFAULT.get(key) + if cls is None: + raise DvcException( + f"'dvc machine' backend '{key}' is missing required " + "dependencies. Install them with:\n" + f"\tpip install dvc[{key}]" + ) + self.backends[key] = cls - self.initialized: Dict[str, BaseMachineBackend] = {} + self.initialized: Dict[str, "BaseMachineBackend"] = {} self.tmp_dir = tmp_dir self.kwargs = kwargs @@ -142,7 +156,7 @@ class MachineManager: conf = kwargs return conf - def _get_backend(self, cloud: str) -> BaseMachineBackend: + def _get_backend(self, cloud: str) -> "BaseMachineBackend": from dvc.config import NoMachineError try: diff --git a/dvc/machine/backend/base.py b/dvc/machine/backend/base.py index 5a90353c6..22ff27316 100644 --- a/dvc/machine/backend/base.py +++ b/dvc/machine/backend/base.py @@ -1,44 +1,15 @@ -import asyncio -import logging -import os -import sys -from abc import ABC, abstractmethod +from abc import abstractmethod from contextlib import contextmanager from typing import TYPE_CHECKING, Iterator, Optional -from dvc.exceptions import DvcException -from dvc.types import StrPath -from dvc.utils.fs import makedirs +from tpi import base 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 +class BaseMachineBackend(base.BaseMachineBackend): @abstractmethod def get_executor( self, name: Optional[str] = None, **config @@ -54,35 +25,3 @@ class BaseMachineBackend(ABC): ) -> Iterator["SSHFileSystem"]: """Return an sshfs instance for the default directory on the specified machine.""" - - @abstractmethod - def run_shell(self, name: Optional[str] = None, **config): - """Spawn an interactive SSH shell for the specified machine.""" - - def _shell(self, *args, **kwargs): - """Sync wrapper for an asyncssh shell session. - - Args will be passed into asyncssh.connect(). - """ - import asyncssh - - loop = asyncio.new_event_loop() - try: - asyncio.set_event_loop(loop) - loop.run_until_complete(self._shell_async(*args, **kwargs)) - except (OSError, asyncssh.Error) as exc: - raise DvcException("SSH connection failed") from exc - finally: - asyncio.set_event_loop(None) - loop.close() - - async def _shell_async(self, *args, **kwargs): - import asyncssh - - async with asyncssh.connect(*args, **kwargs) as conn: - await conn.run( - term_type=os.environ.get("TERM", "xterm"), - stdin=sys.stdin, - stdout=sys.stdout, - stderr=sys.stderr, - ) diff --git a/dvc/machine/backend/terraform.py b/dvc/machine/backend/terraform.py index 43b89ddf9..187b73313 100644 --- a/dvc/machine/backend/terraform.py +++ b/dvc/machine/backend/terraform.py @@ -1,101 +1,29 @@ -import logging -import os from contextlib import contextmanager from typing import TYPE_CHECKING, Iterator, Optional -from funcy import first +from tpi import TerraformProviderIterative, terraform 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 TerraformBackend(BaseMachineBackend): - @contextmanager - def make_tf(self, name: str): - from dvc.tpi import DvcTerraform, TerraformError - 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) +class TerraformBackend(terraform.TerraformBackend): def get_executor( self, name: Optional[str] = None, **config ) -> "BaseExecutor": raise NotImplementedError - def _default_resource(self, name): - from dvc.tpi import TerraformError - - resource = first(self.instances(name)) - if not resource: - raise TerraformError(f"No active '{name}' instances") - return resource - @contextmanager - def get_sshfs( + def get_sshfs( # pylint: disable=unused-argument self, name: Optional[str] = None, **config ) -> Iterator["SSHFileSystem"]: - from dvc.tpi import DvcTerraform - resource = self._default_resource(name) - with DvcTerraform.pemfile(resource) as pem: + with TerraformProviderIterative.pemfile(resource) as pem: fs = SSHFileSystem( host=resource["instance_ip"], user="ubuntu", keyfile=pem, ) yield fs - - def run_shell(self, name: Optional[str] = None, **config): - from dvc.tpi import DvcTerraform - - resource = self._default_resource(name) - with DvcTerraform.pemfile(resource) as pem: - self._shell( - host=resource["instance_ip"], - username="ubuntu", - client_keys=pem, - known_hosts=None, - ) diff --git a/dvc/tpi/__init__.py b/dvc/tpi/__init__.py deleted file mode 100644 index f39efd63f..000000000 --- a/dvc/tpi/__init__.py +++ /dev/null @@ -1,96 +0,0 @@ -import json -import logging -import os -import tempfile -from collections import OrderedDict -from contextlib import contextmanager - -from jinja2 import Environment, FileSystemLoader, select_autoescape -from python_terraform import Terraform - -from dvc.exceptions import DvcException - -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(), -) - - -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) - - -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 deleted file mode 100644 index efec829c9..000000000 --- a/dvc/tpi/templates/main.tf +++ /dev/null @@ -1,26 +0,0 @@ -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 = ssh + ["python-terraform>=0.10.1", "jinja2>=2.0.0"] +terraform = ["tpi[ssh]>=0.0.0"] tests_requirements = ( Path("test_requirements.txt").read_text().strip().splitlines() @@ -127,6 +127,8 @@ setup( long_description=open("README.rst", encoding="UTF-8").read(), author="Dmitry Petrov", author_email="dmitry@dvc.org", + maintainer="Iterative", + maintainer_email="support@dvc.org", download_url="https://github.com/iterative/dvc", license="Apache License 2.0", install_requires=install_requires, diff --git a/tests/unit/machine/test_terraform.py b/tests/unit/machine/test_terraform.py deleted file mode 100644 index 8dc5bbd76..000000000 --- a/tests/unit/machine/test_terraform.py +++ /dev/null @@ -1,143 +0,0 @@ -import json -import os - -import pytest -from mock import ANY, MagicMock -from python_terraform import IsFlagged - -from dvc.tpi 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] - - -def test_shell(tmp_dir, terraform, resource, mocker): - data = json.loads(TEST_RESOURCE_STATE) - expected = data["resources"][0]["instances"][0]["attributes"] - mock_connect = mocker.patch("asyncssh.connect", return_value=MagicMock()) - terraform.run_shell(name=resource) - mock_connect.assert_called_once_with( - host=expected["instance_ip"], - username="ubuntu", - client_keys=ANY, - known_hosts=None, - )