From b8478a8b044404407cd1fbba9e521267bba4880f 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?=
 <peter@pmrowla.com>
Date: Thu, 14 Oct 2021 03:40:30 +0000
Subject: [PATCH] machine: re-add instance dir management from TPI (#6748)

* machine: re-add instance dir management from TPI

* update default constructor

* machine: lazy load Repo.machine
---
 dvc/machine/backend/base.py      |  2 +-
 dvc/machine/backend/terraform.py | 49 ++++++++++++++++++++++++++++----
 dvc/repo/__init__.py             | 20 +++++++------
 3 files changed, 56 insertions(+), 15 deletions(-)

diff --git a/dvc/machine/backend/base.py b/dvc/machine/backend/base.py
index 19ca7cdf0..833217b49 100644
--- a/dvc/machine/backend/base.py
+++ b/dvc/machine/backend/base.py
@@ -10,7 +10,7 @@ if TYPE_CHECKING:
 
 class BaseMachineBackend(ABC):
     def __init__(self, tmp_dir: "StrPath", **kwargs):
-        raise NotImplementedError
+        self.tmp_dir = tmp_dir
 
     @abstractmethod
     def create(self, name: Optional[str] = None, **config):
diff --git a/dvc/machine/backend/terraform.py b/dvc/machine/backend/terraform.py
index ce836cdd5..51fb173f0 100644
--- a/dvc/machine/backend/terraform.py
+++ b/dvc/machine/backend/terraform.py
@@ -1,9 +1,11 @@
+import os
 from contextlib import contextmanager
+from functools import partialmethod
 from typing import TYPE_CHECKING, Iterator, Optional
 
-from tpi import TerraformProviderIterative, terraform
-
+from dvc.exceptions import DvcException
 from dvc.fs.ssh import SSHFileSystem
+from dvc.utils.fs import makedirs
 
 from .base import BaseMachineBackend
 
@@ -12,9 +14,42 @@ if TYPE_CHECKING:
     from dvc.types import StrPath
 
 
-class TerraformBackend(terraform.TerraformBackend, BaseMachineBackend):
+class TerraformBackend(BaseMachineBackend):
     def __init__(self, tmp_dir: "StrPath", **kwargs):
-        super().__init__(tmp_dir)
+        super().__init__(tmp_dir, **kwargs)
+        makedirs(tmp_dir, exist_ok=True)
+
+    @contextmanager
+    def make_tpi(self, name: str):
+        from tpi import TPIError
+        from tpi.terraform import TerraformBackend as TPIBackend
+
+        try:
+            working_dir = os.path.join(self.tmp_dir, name)
+            makedirs(working_dir, exist_ok=True)
+            yield TPIBackend(working_dir=working_dir)
+        except TPIError as exc:
+            raise DvcException("TPI operation failed") from exc
+
+    def _tpi_func(self, fname, name: Optional[str] = None, **config):
+        from tpi import TPIError
+
+        assert name
+        with self.make_tpi(name) as tpi:
+            func = getattr(tpi, fname)
+            try:
+                return func(name=name, **config)
+            except TPIError as exc:
+                raise DvcException(f"TPI {fname} failed") from exc
+
+    create = partialmethod(_tpi_func, "create")  # type: ignore[assignment]
+    destroy = partialmethod(_tpi_func, "destroy")  # type: ignore[assignment]
+    instances = partialmethod(
+        _tpi_func, "instances"
+    )  # type: ignore[assignment]
+    run_shell = partialmethod(
+        _tpi_func, "run_shell"
+    )  # type: ignore[assignment]
 
     def get_executor(
         self, name: Optional[str] = None, **config
@@ -25,7 +60,11 @@ class TerraformBackend(terraform.TerraformBackend, BaseMachineBackend):
     def get_sshfs(  # pylint: disable=unused-argument
         self, name: Optional[str] = None, **config
     ) -> Iterator["SSHFileSystem"]:
-        resource = self.default_resource(name)
+        from tpi import TerraformProviderIterative
+
+        assert name
+        with self.make_tpi(name) as tpi:
+            resource = tpi.default_resource(name)
         with TerraformProviderIterative.pemfile(resource) as pem:
             fs = SSHFileSystem(
                 host=resource["instance_ip"],
diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py
index 95d579554..dc12ff853 100644
--- a/dvc/repo/__init__.py
+++ b/dvc/repo/__init__.py
@@ -163,7 +163,6 @@ 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
@@ -224,14 +223,6 @@ 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
@@ -298,6 +289,17 @@ class Repo:
 
         return Experiments(self)
 
+    @cached_property
+    def machine(self):
+        from dvc.machine import MachineManager
+
+        if self.tmp_dir and (
+            self.config["feature"].get("machine", False)
+            or env2bool("DVC_TEST")
+        ):
+            return MachineManager(self)
+        return None
+
     @property
     def fs(self) -> "BaseFileSystem":
         return self._fs