diff --git a/dvc/machine/backend/base.py b/dvc/machine/backend/base.py index 22ff27316..19ca7cdf0 100644 --- a/dvc/machine/backend/base.py +++ b/dvc/machine/backend/base.py @@ -1,15 +1,38 @@ -from abc import abstractmethod +from abc import ABC, abstractmethod from contextlib import contextmanager from typing import TYPE_CHECKING, Iterator, Optional -from tpi import base - if TYPE_CHECKING: from dvc.fs.ssh import SSHFileSystem from dvc.repo.experiments.executor.base import BaseExecutor + from dvc.types import StrPath -class BaseMachineBackend(base.BaseMachineBackend): +class BaseMachineBackend(ABC): + def __init__(self, tmp_dir: "StrPath", **kwargs): + raise NotImplementedError + + @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 run_shell(self, name: Optional[str] = None, **config): + """Spawn an interactive SSH shell for the specified machine.""" + @abstractmethod def get_executor( self, name: Optional[str] = None, **config diff --git a/dvc/machine/backend/terraform.py b/dvc/machine/backend/terraform.py index 187b73313..118fb3512 100644 --- a/dvc/machine/backend/terraform.py +++ b/dvc/machine/backend/terraform.py @@ -5,11 +5,17 @@ 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 + from dvc.types import StrPath -class TerraformBackend(terraform.TerraformBackend): +class TerraformBackend(terraform.TerraformBackend, BaseMachineBackend): + def __init__(self, tmp_dir: "StrPath", **kwargs): + super().__init__(tmp_dir) + def get_executor( self, name: Optional[str] = None, **config ) -> "BaseExecutor":