torchx-nightly 2024.2.12__py3-none-any.whl → 2025.1.14__py3-none-any.whl

torchx/specs/api.py

@@ -5,12 +5,19 @@
 # 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
+
+import asyncio
 import copy
+import inspect
 import json
+import logging as logger
 import re
+import typing
 from dataclasses import asdict, dataclass, field
 from datetime import datetime
 from enum import Enum
+from json import JSONDecodeError
 from string import Template
 from typing import (
     Any,
@@ -183,11 +190,37 @@ class macros:
             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)
+
             return role
 
+        def _apply_nested(self, d: typing.Dict[str, Any]) -> typing.Dict[str, Any]:
+            stack = [d]
+            while stack:
+                current_dict = stack.pop()
+                for k, v in current_dict.items():
+                    if isinstance(v, dict):
+                        stack.append(v)
+                    elif isinstance(v, str):
+                        current_dict[k] = self.substitute(v)
+                    elif isinstance(v, list):
+                        for i in range(len(v)):
+                            if isinstance(v[i], str):
+                                v[i] = self.substitute(v[i])
+            return d
+
         def substitute(self, arg: str) -> str:
             """
             substitute applies the values to the template arg.
@@ -216,11 +249,13 @@ class RetryPolicy(str, Enum):
                 application to deal with failed replica departures and
                 replacement replica admittance.
     2. APPLICATION: Restarts the entire application.
-
+    3. ROLE: Restarts the role when any error occurs in that role. This does not
+             restart the whole job.
     """
 
     REPLICA = "REPLICA"
     APPLICATION = "APPLICATION"
+    ROLE = "ROLE"
 
 
 class MountType(str, Enum):
@@ -347,6 +382,24 @@ class Role:
     mounts: List[Union[BindMount, VolumeMount, DeviceMount]] = field(
         default_factory=list
     )
+    overrides: Dict[str, Any] = field(default_factory=dict)
+
+    # pyre-ignore
+    def __getattribute__(self, attrname: str) -> Any:
+        if attrname == "overrides":
+            return super().__getattribute__(attrname)
+        try:
+            ov = super().__getattribute__("overrides")
+        except AttributeError:
+            ov = {}
+        if attrname in ov:
+            if inspect.isawaitable(ov[attrname]):
+                result = asyncio.get_event_loop().run_until_complete(ov[attrname])
+            else:
+                result = ov[attrname]()
+            setattr(self, attrname, result)
+            del ov[attrname]
+        return super().__getattribute__(attrname)
 
     def pre_proc(
         self,
@@ -552,7 +605,10 @@ class AppStatus:
 
     def _format_replica_status(self, replica_status: ReplicaStatus) -> str:
         if replica_status.structured_error_msg != NONE:
-            error_data = json.loads(replica_status.structured_error_msg)
+            try:
+                error_data = json.loads(replica_status.structured_error_msg)
+            except JSONDecodeError:
+                return replica_status.structured_error_msg
             error_message = self._format_error_message(
                 msg=error_data["message"]["message"], header="    error_msg: "
             )
@@ -637,11 +693,11 @@ class AppStatusError(Exception):
         self.status = status
 
 
-# valid run cfg values; only support primitives (str, int, float, bool, List[str])
+# valid run cfg values; only support primitives (str, int, float, bool, List[str], Dict[str, str])
 # TODO(wilsonhong): python 3.9+ supports list[T] in typing, which can be used directly
 # in isinstance(). Should replace with that.
 # see: https://docs.python.org/3/library/stdtypes.html#generic-alias-type
-CfgVal = Union[str, int, float, bool, List[str], None]
+CfgVal = Union[str, int, float, bool, List[str], Dict[str, str], None]
 
 
 T = TypeVar("T")
@@ -755,6 +811,10 @@ class runopts:
         except TypeError:
             if isinstance(obj, list):
                 return all(isinstance(e, str) for e in obj)
+            elif isinstance(obj, dict):
+                return all(
+                    isinstance(k, str) and isinstance(v, str) for k, v in obj.items()
+                )
             else:
                 return False
 
@@ -863,8 +923,13 @@ class runopts:
                 # lists may be ; or , delimited
                 # also deal with trailing "," by removing empty strings
                 return [v for v in value.replace(";", ",").split(",") if v]
+            elif opt_type == Dict[str, str]:
+                return {
+                    s.split(":", 1)[0]: s.split(":", 1)[1]
+                    for s in value.replace(";", ",").split(",")
+                }
             else:
-                # pyre-ignore[19]
+                # pyre-ignore[19, 6] type won't be dict here as we handled it above
                 return opt_type(value)
 
         cfg: Dict[str, CfgVal] = {}
@@ -874,6 +939,23 @@ class runopts:
                 cfg[key] = _cast_to_type(val, runopt_.opt_type)
         return cfg
 
+    def cfg_from_json_repr(self, json_repr: str) -> Dict[str, CfgVal]:
+        """
+        Converts the given dict to a valid cfg for this ``runopts`` object.
+        """
+        cfg: Dict[str, CfgVal] = {}
+        cfg_dict = json.loads(json_repr)
+        for key, val in cfg_dict.items():
+            runopt_ = self.get(key)
+            if runopt_:
+                if runopt_.opt_type == List[str]:
+                    cfg[key] = [str(v) for v in val]
+                elif runopt_.opt_type == Dict[str, str]:
+                    cfg[key] = {str(k): str(v) for k, v in val.items()}
+                else:
+                    cfg[key] = val
+        return cfg
+
     def add(
         self,
         cfg_key: str,

torchx/specs/builders.py

@@ -4,25 +4,25 @@
 # 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
+
 import argparse
 import inspect
+import os
+from argparse import Namespace
 from typing import Any, Callable, Dict, List, Mapping, Optional, Union
 
 from torchx.specs.api import BindMount, MountType, VolumeMount
 from torchx.specs.file_linter import get_fn_docstring, TorchXArgumentHelpFormatter
-from torchx.util.types import (
-    decode_from_string,
-    decode_optional,
-    get_argparse_param_type,
-    is_bool,
-    is_primitive,
-)
+from torchx.util.types import decode, decode_optional, get_argparse_param_type, is_bool
 
 from .api import AppDef, DeviceMount
 
 
 def _create_args_parser(
-    cmpnt_fn: Callable[..., AppDef], cmpnt_defaults: Optional[Dict[str, str]] = None
+    cmpnt_fn: Callable[..., AppDef],
+    cmpnt_defaults: Optional[Dict[str, str]] = None,
+    config: Optional[Dict[str, Any]] = None,
 ) -> argparse.ArgumentParser:
     parameters = inspect.signature(cmpnt_fn).parameters
     function_desc, args_desc = get_fn_docstring(cmpnt_fn)
@@ -85,15 +85,55 @@ def _create_args_parser(
             if len(param_name) == 1:
                 arg_names = [f"-{param_name}"] + arg_names
             if "default" not in args:
-                args["required"] = True
+                if (config and param_name not in config) or not config:
+                    args["required"] = True
+
             script_parser.add_argument(*arg_names, **args)
     return script_parser
 
 
+def _merge_config_values_with_args(
+    parsed_args: argparse.Namespace, config: Dict[str, Any]
+) -> None:
+    for key, val in config.items():
+        if key in parsed_args:
+            setattr(parsed_args, key, val)
+
+
+def parse_args(
+    cmpnt_fn: Callable[..., AppDef],
+    cmpnt_args: List[str],
+    cmpnt_defaults: Optional[Dict[str, Any]] = None,
+    config: Optional[Dict[str, Any]] = None,
+) -> Namespace:
+    """
+    Parse passed arguments, defaults, and config values into a namespace for
+    a component function.
+
+    Args:
+    cmpnt_fn: Component function
+    cmpnt_args: Function args
+    cmpnt_defaults: Additional default values for parameters of ``app_fn``
+                        (overrides the defaults set on the fn declaration)
+    config: Optional dict containing additional configuration for the component from a passed config file
+
+    Returns:
+    A Namespace object with the args, defaults, and config values incorporated.
+    """
+
+    script_parser = _create_args_parser(cmpnt_fn, cmpnt_defaults, config)
+    parsed_args = script_parser.parse_args(cmpnt_args)
+    if config:
+        _merge_config_values_with_args(parsed_args, config)
+
+    return parsed_args
+
+
 def materialize_appdef(
     cmpnt_fn: Callable[..., AppDef],
     cmpnt_args: List[str],
-    cmpnt_defaults: Optional[Dict[str, str]] = None,
+    cmpnt_defaults: Optional[Dict[str, Any]] = None,
+    config: Optional[Dict[str, Any]] = None,
 ) -> AppDef:
     """
     Creates an application by running user defined ``app_fn``.
@@ -118,26 +158,23 @@ def materialize_appdef(
         cmpnt_args: Function args
         cmpnt_defaults: Additional default values for parameters of ``app_fn``
                           (overrides the defaults set on the fn declaration)
+        config: Optional dict containing additional configuration for the component from a passed config file
     Returns:
         An application spec
     """
 
-    script_parser = _create_args_parser(cmpnt_fn, cmpnt_defaults)
-    parsed_args = script_parser.parse_args(cmpnt_args)
-
     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)
-        if is_bool(parameter_type):
-            arg_value = arg_value and arg_value.lower() == "true"
-        elif not is_primitive(parameter_type):
-            arg_value = decode_from_string(arg_value, 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:
@@ -149,7 +186,9 @@ def materialize_appdef(
     if len(var_arg) > 0 and var_arg[0] == "--":
         var_arg = var_arg[1:]
 
-    return cmpnt_fn(*function_args, *var_arg, **kwargs)
+    appdef = cmpnt_fn(*function_args, *var_arg, **kwargs)
+
+    return appdef
 
 
 def make_app_handle(scheduler_backend: str, session_name: str, app_id: str) -> str:
@@ -205,9 +244,12 @@ def parse_mounts(opts: List[str]) -> List[Union[BindMount, VolumeMount, DeviceMo
     for opts in mount_opts:
         typ = opts.get("type")
         if typ == MountType.BIND:
+            src_path = opts["src"]
+            if src_path.startswith("~"):
+                src_path = os.path.expanduser(src_path)
             mounts.append(
                 BindMount(
-                    src_path=opts["src"],
+                    src_path=src_path,
                     dst_path=opts["dst"],
                     read_only="readonly" in opts,
                 )

torchx/specs/file_linter.py

@@ -5,6 +5,8 @@
 # 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
+
 import abc
 import argparse
 import ast
@@ -29,7 +31,11 @@ def _get_default_arguments_descriptions(fn: Callable[..., object]) -> Dict[str,
     return args_decs
 
 
-class TorchXArgumentHelpFormatter(argparse.HelpFormatter):
+class TorchXArgumentHelpFormatter(
+    argparse.RawDescriptionHelpFormatter,
+    argparse.ArgumentDefaultsHelpFormatter,
+    argparse.MetavarTypeHelpFormatter,
+):
     """Help message formatter which adds default values and required to argument help.
 
     If the argument is required, the class appends `(required)` at the end of the help message.
@@ -79,7 +85,7 @@ to your component (see: https://pytorch.org/torchx/latest/component_best_practic
             args_description[param.arg_name] = param.description
     short_func_description = docstring.short_description or default_fn_desc
     if docstring.long_description:
-        short_func_description += " ..."
+        short_func_description += "\n" + docstring.long_description
     return (short_func_description or default_fn_desc, args_description)
 
 

torchx/specs/finder.py

@@ -4,6 +4,8 @@
 # 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
+
 import abc
 import importlib
 import inspect

torchx/specs/named_resources_aws.py

@@ -4,6 +4,8 @@
 # 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
+
 r"""
 `torchx.specs.named_resources_aws` contains resource definitions that represent corresponding AWS instance types
 taken from https://aws.amazon.com/ec2/instance-types/. The resources are exposed
@@ -35,6 +37,7 @@ from typing import Callable, Mapping
 from torchx.specs.api import Resource
 
 EFA_DEVICE = "vpc.amazonaws.com/efa"
+NEURON_DEVICE = "aws.amazon.com/neurondevice"
 
 # ecs and ec2 have memtax and currently AWS Batch uses hard memory limits
 # so we have to account for mem tax when registering these resources for AWS
@@ -107,6 +110,16 @@ def aws_p4de_24xlarge() -> Resource:
     )
 
 
+def aws_p5_48xlarge() -> Resource:
+    return Resource(
+        cpu=192,
+        gpu=8,
+        memMB=2048 * GiB,
+        capabilities={K8S_ITYPE: "p5.48xlarge"},
+        devices={EFA_DEVICE: 32},
+    )
+
+
 def aws_t3_medium() -> Resource:
     return Resource(cpu=2, gpu=0, memMB=4 * GiB, capabilities={K8S_ITYPE: "t3.medium"})
 
@@ -117,6 +130,12 @@ def aws_m5_2xlarge() -> Resource:
     )
 
 
+def aws_c5_18xlarge() -> Resource:
+    return Resource(
+        cpu=72, gpu=0, memMB=144 * GiB, capabilities={K8S_ITYPE: "c5.18xlarge"}
+    )
+
+
 def aws_g4dn_xlarge() -> Resource:
     return Resource(
         cpu=4, gpu=1, memMB=16 * GiB, capabilities={K8S_ITYPE: "g4dn.xlarge"}
@@ -241,9 +260,87 @@ def aws_g5_48xlarge() -> Resource:
     )
 
 
+def aws_g6e_xlarge() -> Resource:
+    return Resource(
+        cpu=4,
+        gpu=1,
+        memMB=32 * GiB,
+        capabilities={K8S_ITYPE: "g6e.xlarge"},
+    )
+
+
+def aws_g6e_2xlarge() -> Resource:
+    return Resource(
+        cpu=8,
+        gpu=1,
+        memMB=64 * GiB,
+        capabilities={K8S_ITYPE: "g6e.2xlarge"},
+    )
+
+
+def aws_g6e_4xlarge() -> Resource:
+    return Resource(
+        cpu=16,
+        gpu=1,
+        memMB=128 * GiB,
+        capabilities={K8S_ITYPE: "g6e.4xlarge"},
+    )
+
+
+def aws_g6e_8xlarge() -> Resource:
+    return Resource(
+        cpu=32,
+        gpu=1,
+        memMB=256 * GiB,
+        capabilities={K8S_ITYPE: "g6e.8xlarge"},
+    )
+
+
+def aws_g6e_16xlarge() -> Resource:
+    return Resource(
+        cpu=64,
+        gpu=1,
+        memMB=512 * GiB,
+        capabilities={K8S_ITYPE: "g6e.16xlarge"},
+    )
+
+
+def aws_g6e_12xlarge() -> Resource:
+    return Resource(
+        cpu=48,
+        gpu=4,
+        memMB=384 * GiB,
+        capabilities={K8S_ITYPE: "g6e.12xlarge"},
+    )
+
+
+def aws_g6e_24xlarge() -> Resource:
+    return Resource(
+        cpu=96,
+        gpu=4,
+        memMB=768 * GiB,
+        capabilities={K8S_ITYPE: "g6e.24xlarge"},
+        devices={EFA_DEVICE: 2},
+    )
+
+
+def aws_g6e_48xlarge() -> Resource:
+    return Resource(
+        cpu=192,
+        gpu=8,
+        memMB=1536 * GiB,
+        capabilities={K8S_ITYPE: "g6e.48xlarge"},
+        devices={EFA_DEVICE: 4},
+    )
+
+
 def aws_trn1_2xlarge() -> Resource:
     return Resource(
-        cpu=8, gpu=0, memMB=32 * GiB, capabilities={K8S_ITYPE: "trn1.2xlarge"}
+        cpu=8,
+        gpu=0,
+        memMB=32 * GiB,
+        capabilities={K8S_ITYPE: "trn1.2xlarge"},
+        devices={NEURON_DEVICE: 1},
     )
 
 
@@ -253,19 +350,21 @@ def aws_trn1_32xlarge() -> Resource:
         gpu=0,
         memMB=512 * GiB,
         capabilities={K8S_ITYPE: "trn1.32xlarge"},
-        devices={EFA_DEVICE: 8},
+        devices={EFA_DEVICE: 8, NEURON_DEVICE: 16},
     )
 
 
 NAMED_RESOURCES: Mapping[str, Callable[[], Resource]] = {
     "aws_t3.medium": aws_t3_medium,
     "aws_m5.2xlarge": aws_m5_2xlarge,
+    "aws_c5.18xlarge": aws_c5_18xlarge,
     "aws_p3.2xlarge": aws_p3_2xlarge,
     "aws_p3.8xlarge": aws_p3_8xlarge,
     "aws_p3.16xlarge": aws_p3_16xlarge,
     "aws_p3dn.24xlarge": aws_p3dn_24xlarge,
     "aws_p4d.24xlarge": aws_p4d_24xlarge,
     "aws_p4de.24xlarge": aws_p4de_24xlarge,
+    "aws_p5.48xlarge": aws_p5_48xlarge,
     "aws_g4dn.xlarge": aws_g4dn_xlarge,
     "aws_g4dn.2xlarge": aws_g4dn_2xlarge,
     "aws_g4dn.4xlarge": aws_g4dn_4xlarge,
@@ -281,6 +380,14 @@ NAMED_RESOURCES: Mapping[str, Callable[[], Resource]] = {
     "aws_g5.12xlarge": aws_g5_12xlarge,
     "aws_g5.24xlarge": aws_g5_24xlarge,
     "aws_g5.48xlarge": aws_g5_48xlarge,
+    "aws_g6e.xlarge": aws_g6e_xlarge,
+    "aws_g6e.2xlarge": aws_g6e_2xlarge,
+    "aws_g6e.4xlarge": aws_g6e_4xlarge,
+    "aws_g6e.8xlarge": aws_g6e_8xlarge,
+    "aws_g6e.16xlarge": aws_g6e_16xlarge,
+    "aws_g6e.12xlarge": aws_g6e_12xlarge,
+    "aws_g6e.24xlarge": aws_g6e_24xlarge,
+    "aws_g6e.48xlarge": aws_g6e_48xlarge,
     "aws_trn1.2xlarge": aws_trn1_2xlarge,
     "aws_trn1.32xlarge": aws_trn1_32xlarge,
 }

torchx/specs/named_resources_generic.py

@@ -4,6 +4,8 @@
 # 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
+
 """
 Defines generic named resources that are not specific to any cloud provider's
 instance types. These generic named resources are meant to be used as

torchx/specs/test/components/__init__.py

@@ -3,3 +3,5 @@
 #
 # 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

torchx/specs/test/components/a/__init__.py

@@ -3,6 +3,8 @@
 #
 # 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
 import torchx
 from torchx import specs
 

torchx/specs/test/components/a/b/__init__.py

@@ -3,3 +3,5 @@
 #
 # 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

torchx/specs/test/components/a/b/c.py

@@ -3,6 +3,8 @@
 #
 # 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
 import torchx
 from torchx import specs
 

torchx/specs/test/components/c/__init__.py

@@ -4,3 +4,5 @@
 #
 # 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

torchx/specs/test/components/c/d.py

@@ -3,6 +3,8 @@
 #
 # 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
 import torchx
 from torchx import specs
 

torchx/tracker/__init__.py

@@ -4,6 +4,8 @@
 # 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
+
 """
 .. note:: PROTOTYPE, USE AT YOUR OWN RISK, APIs SUBJECT TO CHANGE
 

torchx/tracker/api.py

@@ -4,6 +4,8 @@
 # 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
+
 from __future__ import annotations
 
 import logging
@@ -67,8 +69,7 @@ class AppRunTrackableSource:
     artifact_name: Optional[str]
 
 
-class Lineage:
-    ...
+class Lineage: ...
 
 
 class TrackerBase(ABC):
@@ -332,5 +333,4 @@ class AppRun:
 
         return model_run_sources
 
-    def children(self) -> Iterable[AppRun]:
-        ...
+    def children(self) -> Iterable[AppRun]: ...

torchx/tracker/backend/fsspec.py

@@ -4,6 +4,8 @@
 # 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
+
 from __future__ import annotations
 
 import json

torchx/util/cuda.py

@@ -4,6 +4,8 @@
 # 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
+
 import torch
 
 

torchx/util/datetime.py

@@ -4,6 +4,8 @@
 # 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
+
 from datetime import datetime, timedelta