diff --git a/torchx/runner/api.py b/torchx/runner/api.py index c33e70fc7..39efc0555 100644 --- a/torchx/runner/api.py +++ b/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,7 +41,7 @@ ) 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 @@ -49,6 +49,8 @@ NONE: str = "" +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, diff --git a/torchx/specs/api.py b/torchx/specs/api.py index 489e3498a..835dea53d 100644 --- a/torchx/specs/api.py +++ b/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) diff --git a/torchx/workspace/api.py b/torchx/workspace/api.py index 2eeb764fe..4805a14af 100644 --- a/torchx/workspace/api.py +++ b/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]): """