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

torchx/specs/file_linter.py

@@ -11,8 +11,9 @@ import abc
 import argparse
 import ast
 import inspect
+import sys
 from dataclasses import dataclass
-from typing import Callable, cast, Dict, List, Optional, Tuple
+from typing import Callable, Dict, List, Optional, Tuple
 
 from docstring_parser import parse
 from torchx.util.io import read_conf_file
@@ -74,7 +75,7 @@ def get_fn_docstring(fn: Callable[..., object]) -> Tuple[str, Dict[str, str]]:
             if the description
     """
     default_fn_desc = f"""{fn.__name__} TIP: improve this help string by adding a docstring
-to your component (see: https://pytorch.org/torchx/latest/component_best_practices.html)"""
+to your component (see: https://meta-pytorch.org/torchx/latest/component_best_practices.html)"""
     args_description = _get_default_arguments_descriptions(fn)
     func_description = inspect.getdoc(fn)
     if not func_description:
@@ -98,7 +99,7 @@ class LinterMessage:
     severity: str = "error"
 
 
-class TorchxFunctionValidator(abc.ABC):
+class ComponentFunctionValidator(abc.ABC):
     @abc.abstractmethod
     def validate(self, app_specs_func_def: ast.FunctionDef) -> List[LinterMessage]:
         """
@@ -116,7 +117,55 @@ class TorchxFunctionValidator(abc.ABC):
         )
 
 
-class TorchxFunctionArgsValidator(TorchxFunctionValidator):
+def OK() -> list[LinterMessage]:
+    return []  # empty linter error means validation passes
+
+
+def is_primitive(arg: ast.expr) -> bool:
+    # whether the arg is a primitive type (e.g. int, float, str, bool)
+    return isinstance(arg, ast.Name)
+
+
+def get_generic_type(arg: ast.expr) -> ast.expr:
+    # returns the slice expr of a subscripted type
+    # `arg` must be an instance of ast.Subscript (caller checks)
+    # in this validator's context, this is the generic type of a container type
+    # e.g. for Optional[str] returns the expr for str
+
+    assert isinstance(arg, ast.Subscript)  # e.g. arg = C[T]
+
+    if isinstance(arg.slice, ast.Index):  # python>=3.10
+        return arg.slice.value
+    else:  # python-3.9
+        return arg.slice
+
+
+def get_optional_type(arg: ast.expr) -> Optional[ast.expr]:
+    """
+    Returns the type parameter ``T`` of ``Optional[T]`` or ``None`` if `arg``
+    is not an ``Optional``. Handles both:
+        1. ``typing.Optional[T]`` (python<3.10)
+        2.  ``T | None`` or ``None | T`` (python>=3.10 - PEP 604)
+    """
+    # case 1: 'a: Optional[T]'
+    if isinstance(arg, ast.Subscript) and arg.value.id == "Optional":
+        return get_generic_type(arg)
+
+    # case 2: 'a: T | None' or 'a: None | T'
+    if sys.version_info >= (3, 10):  # PEP 604 introduced in python-3.10
+        if isinstance(arg, ast.BinOp) and isinstance(arg.op, ast.BitOr):
+            if isinstance(arg.right, ast.Constant) and arg.right.value is None:
+                return arg.left
+            if isinstance(arg.left, ast.Constant) and arg.left.value is None:
+                return arg.right
+
+    # case 3: is not optional
+    return None
+
+
+class ArgTypeValidator(ComponentFunctionValidator):
+    """Validates component function's argument types."""
+
     def validate(self, app_specs_func_def: ast.FunctionDef) -> List[LinterMessage]:
         linter_errors = []
         for arg_def in app_specs_func_def.args.args:
@@ -133,53 +182,68 @@ class TorchxFunctionArgsValidator(TorchxFunctionValidator):
         return linter_errors
 
     def _validate_arg_def(
-        self, function_name: str, arg_def: ast.arg
+        self, function_name: str, arg: ast.arg
     ) -> List[LinterMessage]:
-        if not arg_def.annotation:
-            return [
-                self._gen_linter_message(
-                    f"Arg {arg_def.arg} missing type annotation", arg_def.lineno
-                )
-            ]
-        if isinstance(arg_def.annotation, ast.Name):
+        arg_type = arg.annotation  # type hint
+
+        def ok() -> list[LinterMessage]:
+            # return value when validation passes (e.g. no linter errors)
             return []
-        complex_type_def = cast(ast.Subscript, none_throws(arg_def.annotation))
-        if complex_type_def.value.id == "Optional":
-            # ast module in python3.9 does not have ast.Index wrapper
-            if isinstance(complex_type_def.slice, ast.Index):
-                complex_type_def = complex_type_def.slice.value
-            else:
-                complex_type_def = complex_type_def.slice
-            # Check if type is Optional[primitive_type]
-            if isinstance(complex_type_def, ast.Name):
-                return []
-        # Check if type is Union[Dict,List]
-        type_name = complex_type_def.value.id
-        if type_name != "Dict" and type_name != "List":
-            desc = (
-                f"`{function_name}` allows only Dict, List as complex types."
-                f"Argument `{arg_def.arg}` has: {type_name}"
-            )
-            return [self._gen_linter_message(desc, arg_def.lineno)]
-        linter_errors = []
-        # ast module in python3.9 does not have objects wrapped in ast.Index
-        if isinstance(complex_type_def.slice, ast.Index):
-            sub_type = complex_type_def.slice.value
+
+        def err(reason: str) -> list[LinterMessage]:
+            msg = f"{reason} for argument {ast.unparse(arg)!r} in function {function_name!r}"
+            return [self._gen_linter_message(msg, arg.lineno)]
+
+        if not arg_type:
+            return err("Missing type annotation")
+
+        # Case 1: optional
+        if T := get_optional_type(arg_type):
+            # NOTE: optional types can be primitives or any of the allowed container types
+            #   so check if arg is an optional, and if so, run the rest of the validation with the unpacked type
+            arg_type = T
+
+        # Case 2: int, float, str, bool
+        if is_primitive(arg_type):
+            return ok()
+        # Case 3: Containers (Dict, List, Tuple)
+        elif isinstance(arg_type, ast.Subscript):
+            container_type = arg_type.value.id
+
+            if container_type in ["Dict", "dict"]:
+                KV = get_generic_type(arg_type)
+
+                assert isinstance(KV, ast.Tuple)  # dict[K,V] has ast.Tuple slice
+
+                K, V = KV.elts
+                if not is_primitive(K):
+                    return err(f"Non-primitive key type {ast.unparse(K)!r}")
+                if not is_primitive(V):
+                    return err(f"Non-primitive value type {ast.unparse(V)!r}")
+                return ok()
+            elif container_type in ["List", "list"]:
+                T = get_generic_type(arg_type)
+                if is_primitive(T):
+                    return ok()
+                else:
+                    return err(f"Non-primitive element type {ast.unparse(T)!r}")
+            elif container_type in ["Tuple", "tuple"]:
+                E_N = get_generic_type(arg_type)
+                assert isinstance(E_N, ast.Tuple)  # tuple[...] has ast.Tuple slice
+
+                for e in E_N.elts:
+                    if not is_primitive(e):
+                        return err(f"Non-primitive element type '{ast.unparse(e)!r}'")
+
+                return ok()
+
+            return err(f"Unsupported container type {container_type!r}")
         else:
-            sub_type = complex_type_def.slice
-        if type_name == "Dict":
-            sub_type_tuple = cast(ast.Tuple, sub_type)
-            for el in sub_type_tuple.elts:
-                if not isinstance(el, ast.Name):
-                    desc = "Dict can only have primitive types"
-                    linter_errors.append(self._gen_linter_message(desc, arg_def.lineno))
-        elif not isinstance(sub_type, ast.Name):
-            desc = "List can only have primitive types"
-            linter_errors.append(self._gen_linter_message(desc, arg_def.lineno))
-        return linter_errors
+            return err(f"Unsupported argument type {ast.unparse(arg_type)!r}")
 
 
-class TorchxReturnValidator(TorchxFunctionValidator):
+class ReturnTypeValidator(ComponentFunctionValidator):
+    """Validates that component functions always return AppDef type"""
 
     def __init__(self, supported_return_type: str) -> None:
         super().__init__()
@@ -231,7 +295,7 @@ class TorchxReturnValidator(TorchxFunctionValidator):
         return linter_errors
 
 
-class TorchFunctionVisitor(ast.NodeVisitor):
+class ComponentFnVisitor(ast.NodeVisitor):
     """
     Visitor that finds the component_function and runs registered validators on it.
     Current registered validators:
@@ -252,12 +316,12 @@ class TorchFunctionVisitor(ast.NodeVisitor):
     def __init__(
         self,
         component_function_name: str,
-        validators: Optional[List[TorchxFunctionValidator]],
+        validators: Optional[List[ComponentFunctionValidator]],
     ) -> None:
         if validators is None:
-            self.validators: List[TorchxFunctionValidator] = [
-                TorchxFunctionArgsValidator(),
-                TorchxReturnValidator("AppDef"),
+            self.validators: List[ComponentFunctionValidator] = [
+                ArgTypeValidator(),
+                ReturnTypeValidator("AppDef"),
             ]
         else:
             self.validators = validators
@@ -279,7 +343,7 @@ class TorchFunctionVisitor(ast.NodeVisitor):
 def validate(
     path: str,
     component_function: str,
-    validators: Optional[List[TorchxFunctionValidator]],
+    validators: Optional[List[ComponentFunctionValidator]] = None,
 ) -> List[LinterMessage]:
     """
     Validates the function to make sure it complies the component standard.
@@ -309,7 +373,7 @@ def validate(
             severity="error",
         )
         return [linter_message]
-    visitor = TorchFunctionVisitor(component_function, validators)
+    visitor = ComponentFnVisitor(component_function, validators)
     visitor.visit(module)
     linter_errors = visitor.linter_errors
     if not visitor.visited_function:

torchx/specs/finder.py

@@ -17,9 +17,15 @@ from dataclasses import dataclass
 from inspect import getmembers, isfunction
 from pathlib import Path
 from types import ModuleType
-from typing import Any, Callable, Dict, Generator, List, Optional, Union
+from typing import Callable, Dict, Generator, List, Optional, Union
 
-from torchx.specs.file_linter import get_fn_docstring, TorchxFunctionValidator, validate
+from torchx.specs import AppDef
+
+from torchx.specs.file_linter import (
+    ComponentFunctionValidator,
+    get_fn_docstring,
+    validate,
+)
 from torchx.util import entrypoints
 from torchx.util.io import read_conf_file
 from torchx.util.types import none_throws
@@ -55,8 +61,7 @@ class _Component:
     description: str
     fn_name: str
 
-    # pyre-ignore[4] TODO temporary until PipelineDef is decoupled and can be exposed as type to OSS
-    fn: Callable[..., Any]
+    fn: Callable[..., AppDef]
 
     validation_errors: List[str]
 
@@ -64,7 +69,7 @@ class _Component:
 class ComponentsFinder(abc.ABC):
     @abc.abstractmethod
     def find(
-        self, validators: Optional[List[TorchxFunctionValidator]]
+        self, validators: Optional[List[ComponentFunctionValidator]]
     ) -> List[_Component]:
         """
         Retrieves a set of components. A component is defined as a python
@@ -210,7 +215,7 @@ class ModuleComponentsFinder(ComponentsFinder):
                     yield self._try_import(module_info.name)
 
     def find(
-        self, validators: Optional[List[TorchxFunctionValidator]]
+        self, validators: Optional[List[ComponentFunctionValidator]]
     ) -> List[_Component]:
         components = []
         for m in self._iter_modules_recursive(self.base_module):
@@ -230,7 +235,7 @@ class ModuleComponentsFinder(ComponentsFinder):
             return module
 
     def _get_components_from_module(
-        self, module: ModuleType, validators: Optional[List[TorchxFunctionValidator]]
+        self, module: ModuleType, validators: Optional[List[ComponentFunctionValidator]]
     ) -> List[_Component]:
         functions = getmembers(module, isfunction)
         component_defs = []
@@ -269,28 +274,17 @@ class CustomComponentsFinder(ComponentsFinder):
         self,
         path: str,
         function_name: str,
-        validators: Optional[List[TorchxFunctionValidator]],
+        validators: Optional[List[ComponentFunctionValidator]],
     ) -> List[str]:
         linter_errors = validate(path, function_name, validators)
         return [linter_error.description for linter_error in linter_errors]
 
-    def _get_path_to_function_decl(
-        self, function: Callable[..., Any]  # pyre-ignore[2]
-    ) -> str:
-        """
-        Attempts to return the path to the file where the function is implemented.
-        This can be different from the path where the function is looked up, for example if we have:
-        my_component defined in some_file.py, imported in other_file.py
-        and the component is invoked as other_file.py:my_component
-        """
-        path_to_function_decl = inspect.getabsfile(function)
-        if path_to_function_decl is None or not os.path.isfile(path_to_function_decl):
-            return self._filepath
-        return path_to_function_decl
-
     def find(
-        self, validators: Optional[List[TorchxFunctionValidator]]
+        self, validators: Optional[List[ComponentFunctionValidator]]
     ) -> List[_Component]:
+        validation_errors = self._get_validation_errors(
+            self._filepath, self._function_name, validators
+        )
 
         file_source = read_conf_file(self._filepath)
         namespace = copy.copy(globals())
@@ -303,12 +297,6 @@ class CustomComponentsFinder(ComponentsFinder):
             )
         app_fn = namespace[self._function_name]
         fn_desc, _ = get_fn_docstring(app_fn)
-
-        func_path = self._get_path_to_function_decl(app_fn)
-        validation_errors = self._get_validation_errors(
-            func_path, self._function_name, validators
-        )
-
         return [
             _Component(
                 name=f"{self._filepath}:{self._function_name}",
@@ -321,7 +309,7 @@ class CustomComponentsFinder(ComponentsFinder):
 
 
 def _load_custom_components(
-    validators: Optional[List[TorchxFunctionValidator]],
+    validators: Optional[List[ComponentFunctionValidator]],
 ) -> List[_Component]:
     component_modules = {
         name: load_fn()
@@ -346,7 +334,7 @@ def _load_custom_components(
 
 
 def _load_components(
-    validators: Optional[List[TorchxFunctionValidator]],
+    validators: Optional[List[ComponentFunctionValidator]],
 ) -> Dict[str, _Component]:
     """
     Loads either the custom component defs from the entrypoint ``[torchx.components]``
@@ -368,7 +356,7 @@ _components: Optional[Dict[str, _Component]] = None
 
 
 def _find_components(
-    validators: Optional[List[TorchxFunctionValidator]],
+    validators: Optional[List[ComponentFunctionValidator]],
 ) -> Dict[str, _Component]:
     global _components
     if not _components:
@@ -381,7 +369,7 @@ def _is_custom_component(component_name: str) -> bool:
 
 
 def _find_custom_components(
-    name: str, validators: Optional[List[TorchxFunctionValidator]]
+    name: str, validators: Optional[List[ComponentFunctionValidator]]
 ) -> Dict[str, _Component]:
     if ":" not in name:
         raise ValueError(
@@ -393,7 +381,7 @@ def _find_custom_components(
 
 
 def get_components(
-    validators: Optional[List[TorchxFunctionValidator]] = None,
+    validators: Optional[List[ComponentFunctionValidator]] = None,
 ) -> Dict[str, _Component]:
     """
     Returns all custom components registered via ``[torchx.components]`` entrypoints
@@ -448,7 +436,7 @@ def get_components(
 
 
 def get_component(
-    name: str, validators: Optional[List[TorchxFunctionValidator]] = None
+    name: str, validators: Optional[List[ComponentFunctionValidator]] = None
 ) -> _Component:
     """
     Retrieves components by the provided name.
@@ -464,7 +452,7 @@ def get_component(
         raise ComponentNotFoundException(
             f"Component `{name}` not found. Please make sure it is one of the "
             "builtins: `torchx builtins`. Or registered via `[torchx.components]` "
-            "entry point (see: https://pytorch.org/torchx/latest/configure.html)"
+            "entry point (see: https://meta-pytorch.org/torchx/latest/configure.html)"
         )
 
     component = components[name]
@@ -477,7 +465,7 @@ def get_component(
 
 
 def get_builtin_source(
-    name: str, validators: Optional[List[TorchxFunctionValidator]] = None
+    name: str, validators: Optional[List[ComponentFunctionValidator]] = None
 ) -> str:
     """
     Returns a string of the the builtin component's function source code

torchx/specs/named_resources_aws.py

@@ -16,7 +16,7 @@ the equvalent resource in mem, cpu and gpu numbers.
 
 .. note::
     These resource definitions may change in future. It is expected for each user to
-    manage their own resources. Follow https://pytorch.org/torchx/latest/specs.html#torchx.specs.get_named_resources
+    manage their own resources. Follow https://meta-pytorch.org/torchx/latest/specs.html#torchx.specs.get_named_resources
     to set up named resources.
 
 Usage:
@@ -47,7 +47,7 @@ NEURON_DEVICE = "aws.amazon.com/neurondevice"
 MEM_TAX = 0.96
 
 # determines instance type for non-honogeneous CEs
-# see https://github.com/pytorch/torchx/issues/780
+# see https://github.com/meta-pytorch/torchx/issues/780
 K8S_ITYPE = "node.kubernetes.io/instance-type"
 GiB: int = int(1024 * MEM_TAX)
 
@@ -120,6 +120,16 @@ def aws_p5_48xlarge() -> Resource:
     )
 
 
+def aws_p5e_48xlarge() -> Resource:
+    return Resource(
+        cpu=192,
+        gpu=8,
+        memMB=2048 * GiB,
+        capabilities={K8S_ITYPE: "p5e.48xlarge"},
+        devices={EFA_DEVICE: 32},
+    )
+
+
 def aws_p5en_48xlarge() -> Resource:
     return Resource(
         cpu=192,
@@ -419,6 +429,7 @@ NAMED_RESOURCES: Mapping[str, Callable[[], Resource]] = {
     "aws_p4d.24xlarge": aws_p4d_24xlarge,
     "aws_p4de.24xlarge": aws_p4de_24xlarge,
     "aws_p5.48xlarge": aws_p5_48xlarge,
+    "aws_p5e.48xlarge": aws_p5e_48xlarge,
     "aws_p5en.48xlarge": aws_p5en_48xlarge,
     "aws_g4dn.xlarge": aws_g4dn_xlarge,
     "aws_g4dn.2xlarge": aws_g4dn_2xlarge,

torchx/tracker/__init__.py

@@ -32,7 +32,7 @@ implementation.
 
 Example usage
 -------------
-Sample `code <https://github.com/pytorch/torchx/blob/main/torchx/examples/apps/tracker/main.py>`__ using tracker API.
+Sample `code <https://github.com/meta-pytorch/torchx/blob/main/torchx/examples/apps/tracker/main.py>`__ using tracker API.
 
 
 Tracker Setup
@@ -111,7 +111,7 @@ Use :py:meth:`~torchx.tracker.app_run_from_env`:
 Reference :py:class:`~torchx.tracker.api.TrackerBase` implementation
 --------------------------------------------------------------------
 :py:class:`~torchx.tracker.backend.fsspec.FsspecTracker` provides reference implementation of a tracker backend.
-GitHub example `directory <https://github.com/pytorch/torchx/blob/main/torchx/examples/apps/tracker/>`__ provides example on how to
+GitHub example `directory <https://github.com/meta-pytorch/torchx/blob/main/torchx/examples/apps/tracker/>`__ provides example on how to
 configure and use it in user application.
 
 

torchx/tracker/api.py

@@ -191,7 +191,7 @@ def build_trackers(
         factory = entrypoint_factories.get(factory_name) or load_module(factory_name)
         if not factory:
             logger.warning(
-                f"No tracker factory `{factory_name}` found in entry_points or modules. See https://pytorch.org/torchx/main/tracker.html#module-torchx.tracker"
+                f"No tracker factory `{factory_name}` found in entry_points or modules. See https://meta-pytorch.org/torchx/main/tracker.html#module-torchx.tracker"
             )
             continue
         if config:

torchx/util/entrypoints.py

@@ -69,9 +69,7 @@ def _defer_load_ep(ep: EntryPoint) -> object:
     return run
 
 
-def load_group(
-    group: str, default: Optional[Dict[str, Any]] = None, skip_defaults: bool = False
-):
+def load_group(group: str, default: Optional[Dict[str, Any]] = None):
     """
     Loads all the entry points specified by ``group`` and returns
     the entry points as a map of ``name (str) -> deferred_load_fn``.
@@ -90,7 +88,6 @@ def load_group(
     1. ``load_group("foo")["bar"]("baz")`` -> equivalent to calling ``this.is.a_fn("baz")``
     1. ``load_group("food")`` -> ``None``
     1. ``load_group("food", default={"hello": this.is.c_fn})["hello"]("world")`` -> equivalent to calling ``this.is.c_fn("world")``
-    1. ``load_group("food", default={"hello": this.is.c_fn}, skip_defaults=True)`` -> ``None``
 
 
     If the entrypoint is a module (versus a function as shown above), then calling the ``deferred_load_fn``
@@ -115,8 +112,6 @@ def load_group(
         entrypoints = metadata.entry_points().get(group, ())
 
     if len(entrypoints) == 0:
-        if skip_defaults:
-            return None
         return default
 
     eps = {}

torchx/util/strings.py

@@ -13,7 +13,7 @@ def normalize_str(data: str) -> str:
     """
     Invokes ``lower`` on thes string and removes all
     characters that do not satisfy ``[a-z0-9\\-]`` pattern.
-    This method is mostly used to make sure kubernetes and gcp_batch scheduler gets
+    This method is mostly used to make sure kubernetes scheduler gets
     the job name that does not violate its restrictions.
     """
     if data.startswith("-"):

torchx/util/types.py

@@ -8,6 +8,7 @@
 
 import inspect
 import re
+from types import UnionType
 from typing import Any, Callable, Optional, Tuple, TypeVar, Union
 
 
@@ -234,10 +235,20 @@ def decode_optional(param_type: Any) -> Any:
         If ``param_type`` is type Optional[INNER_TYPE], method returns INNER_TYPE
         Otherwise returns ``param_type``
     """
+
     if not hasattr(param_type, "__origin__"):
-        return param_type
+        if isinstance(param_type, UnionType):
+            # handle BinOp style Optional (e.g. `T | None`)
+            if len(param_type.__args__) == 2 and param_type.__args__[1] is type(None):
+                return param_type.__args__[0]
+            else:
+                return param_type
+        else:
+            return param_type
+
     if param_type.__origin__ is not Union:
         return param_type
+
     args = param_type.__args__
     if len(args) == 2 and args[1] is type(None):
         return args[0]

torchx/version.py

@@ -1,4 +1,3 @@
-#!/usr/bin/env python3
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -7,6 +6,7 @@
 
 # pyre-strict
 
+from torchx._version import BASE_VERSION
 from torchx.util.entrypoints import load
 
 # Follows PEP-0440 version scheme guidelines
@@ -18,7 +18,7 @@ from torchx.util.entrypoints import load
 # 0.1.0bN  # Beta release
 # 0.1.0rcN  # Release Candidate
 # 0.1.0  # Final release
-__version__ = "0.8.0dev0"
+__version__: str = BASE_VERSION
 
 
 # Use the github container registry images corresponding to the current package