Standalone Workspace build for TorchX (#965)

Summary:

Option to build workspace independently of scheduler with async option.

Differential Revision: D63145971

torchx/runner/api.py

@@ -14,7 +14,7 @@
 import warnings
 from datetime import datetime
 from types import TracebackType
-from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple, Type
+from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple, Type, TypeVar
 
 from torchx.runner.events import log_event
 from torchx.schedulers import get_scheduler_factories, SchedulerFactory
@@ -41,14 +41,16 @@
 )
 
 from torchx.util.types import none_throws
-from torchx.workspace.api import WorkspaceMixin
+from torchx.workspace.api import PkgInfo, WorkspaceBuilder, WorkspaceMixin
 
 from .config import get_config, get_configs
 
 logger: logging.Logger = logging.getLogger(__name__)
 
 
 NONE: str = "<NONE>"
+S = TypeVar("S")
+T = TypeVar("T")
 
 
 def get_configured_trackers() -> Dict[str, Optional[str]]:
@@ -147,6 +149,18 @@ def close(self) -> None:
         for scheduler in self._scheduler_instances.values():
             scheduler.close()
 
+    def build_standalone_workspace(
+        self,
+        workspace_builder: WorkspaceBuilder[S, T],
+        sync: bool = True,
+    ) -> PkgInfo[S]:
+        """
+        Build a standalone workspace for the given role.
+        This method is used to build a workspace for a role independent of the scheduler and
+        also enables asynchronous workspace building using the Role overrides.
+        """
+        return workspace_builder.build_workspace(sync)
+
     def run_component(
         self,
         component: str,

torchx/specs/api.py

@@ -11,6 +11,7 @@
 import copy
 import inspect
 import json
+import logging as logger
 import re
 import typing
 from dataclasses import asdict, dataclass, field
@@ -189,7 +190,16 @@ def apply(self, role: "Role") -> "Role":
             apply applies the values to a copy the specified role and returns it.
             """
 
+            # Overrides might contain future values which can't be serialized so taken out for the copy
+            overrides = role.overrides
+            if len(overrides) > 0:
+                logger.warning(
+                    "Role overrides are not supported for macros. Overrides will not be copied"
+                )
+                role.overrides = {}
             role = copy.deepcopy(role)
+            role.overrides = overrides
+
             role.args = [self.substitute(arg) for arg in role.args]
             role.env = {key: self.substitute(arg) for key, arg in role.env.items()}
             role.metadata = self._apply_nested(role.metadata)

torchx/workspace/api.py

@@ -9,7 +9,8 @@
 import abc
 import fnmatch
 import posixpath
-from typing import Generic, Iterable, Mapping, Tuple, TYPE_CHECKING, TypeVar
+from dataclasses import dataclass
+from typing import Any, Dict, Generic, Iterable, Mapping, Tuple, TYPE_CHECKING, TypeVar
 
 from torchx.specs import AppDef, CfgVal, Role, runopts
 
@@ -20,6 +21,36 @@
 
 T = TypeVar("T")
 
+PackageType = TypeVar("PackageType")
+WorkspaceConfigType = TypeVar("WorkspaceConfigType")
+
+
+@dataclass
+class PkgInfo(Generic[PackageType]):
+    """
+    Convenience class used to specify information regarding the built workspace
+    """
+
+    img: str
+    lazy_overrides: Dict[str, Any]
+    metadata: PackageType
+
+
+@dataclass
+class WorkspaceBuilder(Generic[PackageType, WorkspaceConfigType]):
+    cfg: WorkspaceConfigType
+
+    @abc.abstractmethod
+    def build_workspace(self, sync: bool = True) -> PkgInfo[PackageType]:
+        """
+        Builds the specified ``workspace`` with respect to ``img``.
+        In the simplest case, this method builds a new image.
+        Certain (more efficient) implementations build
+        incremental diff patches that overlay on top of the role's image.
+
+        """
+        pass
+
 
 class WorkspaceMixin(abc.ABC, Generic[T]):
     """