mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-01 23:08:43 +01:00
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 <[email protected]>
This commit is contained in:
Casper da Costa-Luis
co-authored by
Peter Rowlands
Peter Rowlands
parent
2778c05e49
commit
0578643f08
+24
-10
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
@@ -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 = <<EOF
|
||||
{{ startup_script }}
|
||||
EOF
|
||||
{% endif %}
|
||||
}
|
||||
@@ -114,7 +114,7 @@ webdav = ["webdav4>=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="[email protected]",
|
||||
maintainer="Iterative",
|
||||
maintainer_email="[email protected]",
|
||||
download_url="https://github.com/iterative/dvc",
|
||||
license="Apache License 2.0",
|
||||
install_requires=install_requires,
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
Reference in New Issue
Block a user