torchx-nightly 2025.8.5__py3-none-any.whl → 2026.1.11__py3-none-any.whl

torchx/specs/api.py

@@ -1,4 +1,3 @@
-#!/usr/bin/env python3
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -12,11 +11,15 @@ import copy
 import inspect
 import json
 import logging as logger
+import os
+import pathlib
 import re
+import shutil
 import typing
+import warnings
 from dataclasses import asdict, dataclass, field
 from datetime import datetime
-from enum import Enum
+from enum import Enum, IntEnum
 from json import JSONDecodeError
 from string import Template
 from typing import (
@@ -67,6 +70,32 @@ YELLOW_BOLD = "\033[1;33m"
 RESET = "\033[0m"
 
 
+def TORCHX_HOME(*subdir_paths: str) -> pathlib.Path:
+    """
+    Path to the "dot-directory" for torchx.
+    Defaults to `~/.torchx` and is overridable via the `TORCHX_HOME` environment variable.
+
+    Usage:
+
+    .. doc-test::
+
+        from pathlib import Path
+        from torchx.specs import TORCHX_HOME
+
+        assert TORCHX_HOME() == Path.home() / ".torchx"
+        assert TORCHX_HOME("conda-pack-out") ==  Path.home() / ".torchx" / "conda-pack-out"
+    ```
+    """
+
+    default_dir = str(pathlib.Path.home() / ".torchx")
+    torchx_home = pathlib.Path(os.getenv("TORCHX_HOME", default_dir))
+
+    torchx_home = torchx_home / os.path.sep.join(subdir_paths)
+    torchx_home.mkdir(parents=True, exist_ok=True)
+
+    return torchx_home
+
+
 # ========================================
 # ==== Distributed AppDef API =======
 # ========================================
@@ -83,6 +112,8 @@ class Resource:
         memMB: MB of ram
         capabilities: additional hardware specs (interpreted by scheduler)
         devices: a list of named devices with their quantities
+        tags: metadata tags for the resource (not interpreted by schedulers)
+          used to add non-functional information about resources (e.g. whether it is an alias of another resource)
 
     Note: you should prefer to use named_resources instead of specifying the raw
     resource requirement directly.
@@ -93,6 +124,7 @@ class Resource:
     memMB: int
     capabilities: Dict[str, Any] = field(default_factory=dict)
     devices: Dict[str, int] = field(default_factory=dict)
+    tags: Dict[str, object] = field(default_factory=dict)
 
     @staticmethod
     def copy(original: "Resource", **capabilities: Any) -> "Resource":
@@ -101,6 +133,7 @@ class Resource:
         are present in the original resource and as parameter, the one from parameter
         will be used.
         """
+
         res_capabilities = dict(original.capabilities)
         res_capabilities.update(capabilities)
         return Resource(
@@ -220,7 +253,9 @@ class macros:
                         current_dict[k] = self.substitute(v)
                     elif isinstance(v, list):
                         for i in range(len(v)):
-                            if isinstance(v[i], str):
+                            if isinstance(v[i], dict):
+                                stack.append(v[i])
+                            elif isinstance(v[i], str):
                                 v[i] = self.substitute(v[i])
             return d
 
@@ -319,6 +354,121 @@ class DeviceMount:
     permissions: str = "rwm"
 
 
+@dataclass
+class Workspace:
+    """
+    Specifies a local "workspace" (a set of directories). Workspaces are ad-hoc built
+    into an (usually ephemeral) image. This effectively mirrors the local code changes
+    at job submission time.
+
+    For example:
+
+      1. ``projects={"~/github/torch": "torch"}`` copies ``~/github/torch/**`` into ``$REMOTE_WORKSPACE_ROOT/torch/**``
+      2. ``projects={"~/github/torch": ""}`` copies ``~/github/torch/**`` into ``$REMOTE_WORKSPACE_ROOT/**``
+
+    The exact location of ``$REMOTE_WORKSPACE_ROOT`` is implementation dependent and varies between
+    different implementations of :py:class:`~torchx.workspace.api.WorkspaceMixin`.
+    Check the scheduler documentation for details on which workspace it supports.
+
+    Note: ``projects`` maps the location of the local project to a sub-directory in the remote workspace root directory.
+    Typically the local project location is a directory path (e.g. ``/home/foo/github/torch``).
+
+
+    Attributes:
+        projects: mapping of local project to the sub-dir in the remote workspace dir.
+    """
+
+    projects: dict[str, str]
+
+    def __bool__(self) -> bool:
+        """False if no projects mapping. Lets us use workspace object in an if-statement"""
+        return bool(self.projects)
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Workspace):
+            return False
+        return self.projects == other.projects
+
+    def __hash__(self) -> int:
+        # makes it possible to use Workspace as the key in the workspace build cache
+        # see WorkspaceMixin.caching_build_workspace_and_update_role
+        return hash(frozenset(self.projects.items()))
+
+    def is_unmapped_single_project(self) -> bool:
+        """
+        Returns ``True`` if this workspace only has 1 project
+        and its target mapping is an empty string.
+        """
+        return len(self.projects) == 1 and not next(iter(self.projects.values()))
+
+    def merge_into(self, outdir: str | pathlib.Path) -> None:
+        """
+        Copies each project dir of this workspace into the specified ``outdir``.
+        Each project dir is copied into ``{outdir}/{target}`` where ``target`` is
+        the target mapping of the project dir.
+
+        For example:
+
+        .. code-block:: python
+            from os.path import expanduser
+
+            workspace = Workspace(
+                projects={
+                    expanduser("~/workspace/torch"): "torch",
+                    expanduser("~/workspace/my_project": "")
+                }
+            )
+            workspace.merge_into(expanduser("~/tmp"))
+
+        Copies:
+
+            * ``~/workspace/torch/**`` into ``~/tmp/torch/**``
+            * ``~/workspace/my_project/**`` into ``~/tmp/**``
+
+        """
+
+        for src, dst in self.projects.items():
+            dst_path = pathlib.Path(outdir) / dst
+            if pathlib.Path(src).is_file():
+                shutil.copy2(src, dst_path)
+            else:  # src is dir
+                shutil.copytree(src, dst_path, dirs_exist_ok=True)
+
+    @staticmethod
+    def from_str(workspace: str | None) -> "Workspace":
+        import yaml
+
+        if not workspace:
+            return Workspace({})
+
+        projects = yaml.safe_load(workspace)
+        if isinstance(projects, str):  # single project workspace
+            projects = {projects: ""}
+        else:  # multi-project workspace
+            # Replace None mappings with "" (empty string)
+            projects = {k: ("" if v is None else v) for k, v in projects.items()}
+
+        return Workspace(projects)
+
+    def __str__(self) -> str:
+        """
+        Returns a string representation of the Workspace by concatenating
+        the project mappings using ';' as a delimiter and ':' between key and value.
+        If the single-project workspace with no target mapping, then simply
+        returns the src (local project dir)
+
+        NOTE: meant to be used for logging purposes not serde.
+          Therefore not symmetric with :py:func:`Workspace.from_str`.
+
+        """
+        if self.is_unmapped_single_project():
+            return next(iter(self.projects))
+        else:
+            return ";".join(
+                k if not v else f"{k}:{v}" for k, v in self.projects.items()
+            )
+
+
 @dataclass
 class Role:
     """
@@ -371,12 +521,15 @@ class Role:
             metadata: Free form information that is associated with the role, for example
                 scheduler specific data. The key should follow the pattern: ``$scheduler.$key``
             mounts: a list of mounts on the machine
+            workspace: local project directories to be mirrored on the remote job.
+              NOTE: The workspace argument provided to the :py:class:`~torchx.runner.api.Runner` APIs
+              only takes effect on ``appdef.role[0]`` and overrides this attribute.
+
     """
 
     name: str
     image: str
     min_replicas: Optional[int] = None
-    base_image: Optional[str] = None  # DEPRECATED DO NOT SET, WILL BE REMOVED SOON
     entrypoint: str = MISSING
     args: List[str] = field(default_factory=list)
     env: Dict[str, str] = field(default_factory=dict)
@@ -386,9 +539,10 @@ class Role:
     resource: Resource = field(default_factory=_null_resource)
     port_map: Dict[str, int] = field(default_factory=dict)
     metadata: Dict[str, Any] = field(default_factory=dict)
-    mounts: List[Union[BindMount, VolumeMount, DeviceMount]] = field(
-        default_factory=list
-    )
+    mounts: List[BindMount | VolumeMount | DeviceMount] = field(default_factory=list)
+    workspace: Workspace | None = None
+
+    # DEPRECATED DO NOT SET, WILL BE REMOVED SOON
     overrides: Dict[str, Any] = field(default_factory=dict)
 
     # pyre-ignore
@@ -788,6 +942,8 @@ class runopt:
     opt_type: Type[CfgVal]
     is_required: bool
     help: str
+    aliases: list[str] | None = None
+    deprecated_aliases: list[str] | None = None
 
     @property
     def is_type_list_of_str(self) -> bool:
@@ -823,7 +979,7 @@ class runopt:
 
         NOTE: dict parsing uses ":" as the kv separator (rather than the standard "=") because "=" is used
         at the top-level cfg to parse runopts (notice the plural) from the CLI. Originally torchx only supported
-        primitives and list[str] as CfgVal but dict[str,str] was added in https://github.com/pytorch/torchx/pull/855
+        primitives and list[str] as CfgVal but dict[str,str] was added in https://github.com/meta-pytorch/torchx/pull/855
         """
 
         if self.opt_type is None:
@@ -879,6 +1035,7 @@ class runopts:
 
     def __init__(self) -> None:
         self._opts: Dict[str, runopt] = {}
+        self._alias_to_key: dict[str, str] = {}
 
     def __iter__(self) -> Iterator[Tuple[str, runopt]]:
         return self._opts.items().__iter__()
@@ -906,9 +1063,16 @@ class runopts:
 
     def get(self, name: str) -> Optional[runopt]:
         """
-        Returns option if any was registered, or None otherwise
+        Returns option if any was registered, or None otherwise.
+        First searches for the option by ``name``, then falls-back to matching ``name`` with any
+        registered aliases.
+
         """
-        return self._opts.get(name, None)
+        if name in self._opts:
+            return self._opts[name]
+        if name in self._alias_to_key:
+            return self._opts[self._alias_to_key[name]]
+        return None
 
     def resolve(self, cfg: Mapping[str, CfgVal]) -> Dict[str, CfgVal]:
         """
@@ -923,6 +1087,36 @@ class runopts:
 
         for cfg_key, runopt in self._opts.items():
             val = resolved_cfg.get(cfg_key)
+            resolved_name = None
+            aliases = runopt.aliases or []
+            deprecated_aliases = runopt.deprecated_aliases or []
+            if val is None:
+                for alias in aliases:
+                    val = resolved_cfg.get(alias)
+                    if alias in cfg or val is not None:
+                        resolved_name = alias
+                        break
+                for alias in deprecated_aliases:
+                    val = resolved_cfg.get(alias)
+                    if val is not None:
+                        resolved_name = alias
+                        use_instead = self._alias_to_key.get(alias)
+                        warnings.warn(
+                            f"Run option `{alias}` is deprecated, use `{use_instead}` instead",
+                            UserWarning,
+                            stacklevel=2,
+                        )
+                        break
+            else:
+                resolved_name = cfg_key
+                for alias in aliases:
+                    duplicate_val = resolved_cfg.get(alias)
+                    if alias in cfg or duplicate_val is not None:
+                        raise InvalidRunConfigException(
+                            f"Duplicate opt name. runopt: `{resolved_name}``, is an alias of runopt: `{alias}`",
+                            resolved_name,
+                            cfg,
+                        )
 
             # check required opt
             if runopt.is_required and val is None:
@@ -942,7 +1136,7 @@ class runopts:
                 )
 
             # not required and not set, set to default
-            if val is None:
+            if val is None and resolved_name is None:
                 resolved_cfg[cfg_key] = runopt.default
         return resolved_cfg
 
@@ -1042,12 +1236,16 @@ class runopts:
         help: str,
         default: CfgVal = None,
         required: bool = False,
+        aliases: Optional[list[str]] = None,
+        deprecated_aliases: Optional[list[str]] = None,
     ) -> None:
         """
         Adds the ``config`` option with the given help string and ``default``
         value (if any). If the ``default`` is not specified then this option
         is a required option.
         """
+        aliases = aliases or []
+        deprecated_aliases = deprecated_aliases or []
         if required and default is not None:
             raise ValueError(
                 f"Required option: {cfg_key} must not specify default value. Given: {default}"
@@ -1059,7 +1257,19 @@ class runopts:
                     f" Given: {default} ({type(default).__name__})"
                 )
 
-        self._opts[cfg_key] = runopt(default, type_, required, help)
+        opt = runopt(
+            default,
+            type_,
+            required,
+            help,
+            list(set(aliases)),
+            list(set(deprecated_aliases)),
+        )
+        for alias in aliases:
+            self._alias_to_key[alias] = cfg_key
+        for deprecated_alias in deprecated_aliases:
+            self._alias_to_key[deprecated_alias] = cfg_key
+        self._opts[cfg_key] = opt
 
     def update(self, other: "runopts") -> None:
         self._opts.update(other._opts)

torchx/specs/builders.py

@@ -4,13 +4,13 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-# pyre-strict
+# pyre-unsafe
 
 import argparse
 import inspect
 import os
 from argparse import Namespace
-from typing import Any, Callable, Dict, List, Mapping, Optional, Union
+from typing import Any, Callable, Dict, List, Mapping, NamedTuple, Optional, Union
 
 from torchx.specs.api import BindMount, MountType, VolumeMount
 from torchx.specs.file_linter import get_fn_docstring, TorchXArgumentHelpFormatter
@@ -19,6 +19,14 @@ from torchx.util.types import decode, decode_optional, get_argparse_param_type,
 from .api import AppDef, DeviceMount
 
 
+class ComponentArgs(NamedTuple):
+    """Parsed component function arguments"""
+
+    positional_args: dict[str, Any]
+    var_args: list[str]
+    kwargs: dict[str, Any]
+
+
 def _create_args_parser(
     cmpnt_fn: Callable[..., AppDef],
     cmpnt_defaults: Optional[Dict[str, str]] = None,
@@ -31,7 +39,7 @@ def _create_args_parser(
 
 
 def _create_args_parser_from_parameters(
-    cmpnt_fn: Callable[..., Any],  # pyre-ignore[2]
+    cmpnt_fn: Callable[..., AppDef],
     parameters: Mapping[str, inspect.Parameter],
     cmpnt_defaults: Optional[Dict[str, str]] = None,
     config: Optional[Dict[str, Any]] = None,
@@ -112,7 +120,7 @@ def _merge_config_values_with_args(
 
 
 def parse_args(
-    cmpnt_fn: Callable[..., Any],  # pyre-ignore[2]
+    cmpnt_fn: Callable[..., AppDef],
     cmpnt_args: List[str],
     cmpnt_defaults: Optional[Dict[str, Any]] = None,
     config: Optional[Dict[str, Any]] = None,
@@ -140,8 +148,97 @@ def parse_args(
     return parsed_args
 
 
+def component_args_from_str(
+    cmpnt_fn: Callable[..., AppDef],
+    cmpnt_args: list[str],
+    cmpnt_args_defaults: Optional[Dict[str, Any]] = None,
+    config: Optional[Dict[str, Any]] = None,
+) -> ComponentArgs:
+    """
+    Parses and decodes command-line arguments for a component function.
+
+    This function takes a component function and its arguments, parses them using argparse,
+    and decodes the arguments into their expected types based on the function's signature.
+    It separates positional arguments, variable positional arguments (*args), and keyword-only arguments.
+
+    Args:
+        cmpnt_fn: The component function whose arguments are to be parsed and decoded.
+        cmpnt_args: List of command-line arguments to be parsed. Supports both space separated and '=' separated arguments.
+        cmpnt_args_defaults: Optional dictionary of default values for the component function's parameters.
+        config: Optional dictionary containing additional configuration values.
+
+    Returns:
+        ComponentArgs representing the input args to a component function containing:
+            - positional_args: Dictionary of positional and positional-or-keyword arguments.
+            - var_args: List of variable positional arguments (*args).
+            - kwargs: Dictionary of keyword-only arguments.
+
+        Usage:
+
+        .. doctest::
+            from torchx.specs.api import AppDef
+            from torchx.specs.builders import component_args_from_str
+
+            def example_component_fn(foo: str, *args: str, bar: str = "asdf") -> AppDef:
+                return AppDef(name="example")
+
+            # Supports space separated arguments
+            args = ["--foo", "fooval", "--bar", "barval", "arg1", "arg2"]
+            parsed_args = component_args_from_str(example_component_fn, args)
+
+            assert parsed_args.positional_args == {"foo": "fooval"}
+            assert parsed_args.var_args == ["arg1", "arg2"]
+            assert parsed_args.kwargs == {"bar": "barval"}
+
+            # Supports '=' separated arguments
+            args = ["--foo=fooval", "--bar=barval", "arg1", "arg2"]
+            parsed_args = component_args_from_str(example_component_fn, args)
+
+            assert parsed_args.positional_args == {"foo": "fooval"}
+            assert parsed_args.var_args == ["arg1", "arg2"]
+            assert parsed_args.kwargs == {"bar": "barval"}
+
+
+    """
+    parsed_args: Namespace = parse_args(
+        cmpnt_fn, cmpnt_args, cmpnt_args_defaults, config
+    )
+
+    positional_args = {}
+    var_args = []
+    kwargs = {}
+
+    parameters = inspect.signature(cmpnt_fn).parameters
+    for param_name, parameter in parameters.items():
+        arg_value = getattr(parsed_args, param_name)
+        parameter_type = parameter.annotation
+        parameter_type = decode_optional(parameter_type)
+        if (
+            parameter_type != arg_value.__class__
+            and parameter.kind != inspect.Parameter.VAR_POSITIONAL
+        ):
+            arg_value = decode(arg_value, parameter_type)
+        if parameter.kind == inspect.Parameter.VAR_POSITIONAL:
+            var_args = arg_value
+        elif parameter.kind == inspect.Parameter.KEYWORD_ONLY:
+            kwargs[param_name] = arg_value
+        elif parameter.kind == inspect.Parameter.VAR_KEYWORD:
+            raise TypeError(
+                f"component fn param `{param_name}` is a '**kwargs' which is not supported; consider changing the "
+                f"type to a dict or explicitly declare the params"
+            )
+        else:
+            # POSITIONAL or POSITIONAL_OR_KEYWORD
+            positional_args[param_name] = arg_value
+
+    if len(var_args) > 0 and var_args[0] == "--":
+        var_args = var_args[1:]
+
+    return ComponentArgs(positional_args, var_args, kwargs)
+
+
 def materialize_appdef(
-    cmpnt_fn: Callable[..., Any],  # pyre-ignore[2]
+    cmpnt_fn: Callable[..., AppDef],
     cmpnt_args: List[str],
     cmpnt_defaults: Optional[Dict[str, Any]] = None,
     config: Optional[Dict[str, Any]] = None,
@@ -174,30 +271,14 @@ def materialize_appdef(
         An application spec
     """
 
-    function_args = []
-    var_arg = []
-    kwargs = {}
-
-    parsed_args = parse_args(cmpnt_fn, cmpnt_args, cmpnt_defaults, config)
-
-    parameters = inspect.signature(cmpnt_fn).parameters
-    for param_name, parameter in parameters.items():
-        arg_value = getattr(parsed_args, param_name)
-        parameter_type = parameter.annotation
-        parameter_type = decode_optional(parameter_type)
-        arg_value = decode(arg_value, parameter_type)
-        if parameter.kind == inspect.Parameter.VAR_POSITIONAL:
-            var_arg = arg_value
-        elif parameter.kind == inspect.Parameter.KEYWORD_ONLY:
-            kwargs[param_name] = arg_value
-        elif parameter.kind == inspect.Parameter.VAR_KEYWORD:
-            raise TypeError("**kwargs are not supported for component definitions")
-        else:
-            function_args.append(arg_value)
-    if len(var_arg) > 0 and var_arg[0] == "--":
-        var_arg = var_arg[1:]
+    component_args: ComponentArgs = component_args_from_str(
+        cmpnt_fn, cmpnt_args, cmpnt_defaults, config
+    )
+    positional_arg_values = list(component_args.positional_args.values())
+    appdef = cmpnt_fn(
+        *positional_arg_values, *component_args.var_args, **component_args.kwargs
+    )
 
-    appdef = cmpnt_fn(*function_args, *var_arg, **kwargs)
     if not isinstance(appdef, AppDef):
         raise TypeError(
             f"Expected a component that returns `AppDef`, but got `{type(appdef)}`"