torchx-nightly 2023.10.21__py3-none-any.whl → 2025.12.24__py3-none-any.whl

torchx/specs/finder.py

@@ -4,7 +4,10 @@
 # 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 copy
 import importlib
 import inspect
 import logging
@@ -17,11 +20,17 @@ from types import ModuleType
 from typing import Callable, Dict, Generator, List, Optional, Union
 
 from torchx.specs import AppDef
-from torchx.specs.file_linter import get_fn_docstring, validate
+
+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
 
+
 logger: logging.Logger = logging.getLogger(__name__)
 
 
@@ -51,13 +60,17 @@ class _Component:
     name: str
     description: str
     fn_name: str
+
     fn: Callable[..., AppDef]
+
     validation_errors: List[str]
 
 
 class ComponentsFinder(abc.ABC):
     @abc.abstractmethod
-    def find(self) -> List[_Component]:
+    def find(
+        self, validators: Optional[List[ComponentFunctionValidator]]
+    ) -> List[_Component]:
         """
         Retrieves a set of components. A component is defined as a python
         function that conforms to ``torchx.specs.file_linter`` linter.
@@ -201,10 +214,12 @@ class ModuleComponentsFinder(ComponentsFinder):
                 else:
                     yield self._try_import(module_info.name)
 
-    def find(self) -> List[_Component]:
+    def find(
+        self, validators: Optional[List[ComponentFunctionValidator]]
+    ) -> List[_Component]:
         components = []
         for m in self._iter_modules_recursive(self.base_module):
-            components += self._get_components_from_module(m)
+            components += self._get_components_from_module(m, validators)
         return components
 
     def _try_import(self, module: Union[str, ModuleType]) -> ModuleType:
@@ -219,7 +234,9 @@ class ModuleComponentsFinder(ComponentsFinder):
         else:
             return module
 
-    def _get_components_from_module(self, module: ModuleType) -> List[_Component]:
+    def _get_components_from_module(
+        self, module: ModuleType, validators: Optional[List[ComponentFunctionValidator]]
+    ) -> List[_Component]:
         functions = getmembers(module, isfunction)
         component_defs = []
 
@@ -228,7 +245,7 @@ class ModuleComponentsFinder(ComponentsFinder):
         module_path = os.path.abspath(module_path)
         rel_module_name = module_relname(module, relative_to=self.base_module)
         for function_name, function in functions:
-            linter_errors = validate(module_path, function_name)
+            linter_errors = validate(module_path, function_name, validators)
             component_desc, _ = get_fn_docstring(function)
 
             # remove empty string to deal with group=""
@@ -253,17 +270,26 @@ class CustomComponentsFinder(ComponentsFinder):
         self._filepath = filepath
         self._function_name = function_name
 
-    def _get_validation_errors(self, path: str, function_name: str) -> List[str]:
-        linter_errors = validate(path, function_name)
+    def _get_validation_errors(
+        self,
+        path: str,
+        function_name: str,
+        validators: Optional[List[ComponentFunctionValidator]],
+    ) -> List[str]:
+        linter_errors = validate(path, function_name, validators)
         return [linter_error.description for linter_error in linter_errors]
 
-    def find(self) -> List[_Component]:
+    def find(
+        self, validators: Optional[List[ComponentFunctionValidator]]
+    ) -> List[_Component]:
         validation_errors = self._get_validation_errors(
-            self._filepath, self._function_name
+            self._filepath, self._function_name, validators
         )
 
         file_source = read_conf_file(self._filepath)
-        namespace = globals()
+        namespace = copy.copy(globals())
+        # so that __file__ used inside the component points to the correct file
+        namespace["__file__"] = os.path.abspath(self._filepath)
         exec(file_source, namespace)  # noqa: P204
         if self._function_name not in namespace:
             raise ComponentNotFoundException(
@@ -282,7 +308,9 @@ class CustomComponentsFinder(ComponentsFinder):
         ]
 
 
-def _load_custom_components() -> List[_Component]:
+def _load_custom_components(
+    validators: Optional[List[ComponentFunctionValidator]],
+) -> List[_Component]:
     component_modules = {
         name: load_fn()
         for name, load_fn in
@@ -301,11 +329,13 @@ def _load_custom_components() -> List[_Component]:
         # _0 = torchx.components.dist
         # _1 = torchx.components.utils
         group = "" if group.startswith("_") else group
-        components += ModuleComponentsFinder(module, group).find()
+        components += ModuleComponentsFinder(module, group).find(validators)
     return components
 
 
-def _load_components() -> Dict[str, _Component]:
+def _load_components(
+    validators: Optional[List[ComponentFunctionValidator]],
+) -> Dict[str, _Component]:
     """
     Loads either the custom component defs from the entrypoint ``[torchx.components]``
     or the default builtins from ``torchx.components`` module.
@@ -316,19 +346,21 @@ def _load_components() -> Dict[str, _Component]:
 
     """
 
-    components = _load_custom_components()
+    components = _load_custom_components(validators)
     if not components:
-        components = ModuleComponentsFinder("torchx.components", "").find()
+        components = ModuleComponentsFinder("torchx.components", "").find(validators)
     return {c.name: c for c in components}
 
 
 _components: Optional[Dict[str, _Component]] = None
 
 
-def _find_components() -> Dict[str, _Component]:
+def _find_components(
+    validators: Optional[List[ComponentFunctionValidator]],
+) -> Dict[str, _Component]:
     global _components
     if not _components:
-        _components = _load_components()
+        _components = _load_components(validators)
     return none_throws(_components)
 
 
@@ -336,17 +368,21 @@ def _is_custom_component(component_name: str) -> bool:
     return ":" in component_name
 
 
-def _find_custom_components(name: str) -> Dict[str, _Component]:
+def _find_custom_components(
+    name: str, validators: Optional[List[ComponentFunctionValidator]]
+) -> Dict[str, _Component]:
     if ":" not in name:
         raise ValueError(
             f"Invalid custom component: {name}, valid template : `FILEPATH`:`FUNCTION_NAME`"
         )
     filepath, component_name = name.split(":")
-    components = CustomComponentsFinder(filepath, component_name).find()
+    components = CustomComponentsFinder(filepath, component_name).find(validators)
     return {component.name: component for component in components}
 
 
-def get_components() -> Dict[str, _Component]:
+def get_components(
+    validators: Optional[List[ComponentFunctionValidator]] = None,
+) -> Dict[str, _Component]:
     """
     Returns all custom components registered via ``[torchx.components]`` entrypoints
     OR builtin components that ship with TorchX (but not both).
@@ -393,13 +429,15 @@ def get_components() -> Dict[str, _Component]:
     """
 
     valid_components: Dict[str, _Component] = {}
-    for component_name, component in _find_components().items():
+    for component_name, component in _find_components(validators).items():
         if len(component.validation_errors) == 0:
             valid_components[component_name] = component
     return valid_components
 
 
-def get_component(name: str) -> _Component:
+def get_component(
+    name: str, validators: Optional[List[ComponentFunctionValidator]] = None
+) -> _Component:
     """
     Retrieves components by the provided name.
 
@@ -407,14 +445,14 @@ def get_component(name: str) -> _Component:
         Component or None if no component with ``name`` exists
     """
     if _is_custom_component(name):
-        components = _find_custom_components(name)
+        components = _find_custom_components(name, validators)
     else:
-        components = _find_components()
+        components = _find_components(validators)
     if name not in components:
         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]
@@ -426,7 +464,9 @@ def get_component(name: str) -> _Component:
     return component
 
 
-def get_builtin_source(name: str) -> str:
+def get_builtin_source(
+    name: str, validators: Optional[List[ComponentFunctionValidator]] = None
+) -> str:
     """
     Returns a string of the the builtin component's function source code
     with all the import statements. Intended to be used to make a copy
@@ -444,7 +484,7 @@ def get_builtin_source(name: str) -> str:
     are optimized and formatting adheres to your organization's standards.
     """
 
-    component = get_component(name)
+    component = get_component(name, validators)
     fn = component.fn
     fn_name = component.name.split(".")[-1]
 

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
@@ -14,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:
@@ -29,22 +31,36 @@ Usage:
 
 """
 
+import warnings
 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
 # otherwise the job will be stuck in the jobqueue forever
-# 97% is based on empirical observation that works well for most instance types
+# 96% is based on empirical observation that works well for most instance types
 # see: https://docs.aws.amazon.com/batch/latest/userguide/memory-management.html
-MEM_TAX = 0.97
+MEM_TAX = 0.96
+
+# determines instance type for non-honogeneous CEs
+# see https://github.com/meta-pytorch/torchx/issues/780
 K8S_ITYPE = "node.kubernetes.io/instance-type"
 GiB: int = int(1024 * MEM_TAX)
 
 
+def instance_type_from_resource(resource: Resource) -> str:
+    instance_type = resource.capabilities.get(K8S_ITYPE)
+    if instance_type is None:
+        warnings.warn(
+            "Cannot determine resource instance type which can cause issues for non-homogeneous CEs and multinode jobs. Consider providing torchx.specs.named_resources_aws:K8S_TYPE resource capability."
+        )
+    return instance_type
+
+
 def aws_p3_2xlarge() -> Resource:
     return Resource(
         cpu=8, gpu=1, memMB=61 * GiB, capabilities={K8S_ITYPE: "p3.2xlarge"}
@@ -94,6 +110,36 @@ 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_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,
+        gpu=8,
+        memMB=2048 * GiB,
+        capabilities={K8S_ITYPE: "p5en.48xlarge"},
+        devices={EFA_DEVICE: 16},
+    )
+
+
 def aws_t3_medium() -> Resource:
     return Resource(cpu=2, gpu=0, memMB=4 * GiB, capabilities={K8S_ITYPE: "t3.medium"})
 
@@ -104,6 +150,16 @@ def aws_m5_2xlarge() -> Resource:
     )
 
 
+def aws_c5_18xlarge() -> Resource:
+    return Resource(
+        # using lower memory size than the spec since MEM_TAX is not enough for adjustment
+        cpu=72,
+        gpu=0,
+        memMB=142 * 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"}
@@ -124,25 +180,41 @@ def aws_g4dn_4xlarge() -> Resource:
 
 def aws_g4dn_8xlarge() -> Resource:
     return Resource(
-        cpu=32, gpu=1, memMB=128 * GiB, capabilities={K8S_ITYPE: "g4dn.8xlarge"}
+        cpu=32,
+        gpu=1,
+        memMB=128 * GiB,
+        capabilities={K8S_ITYPE: "g4dn.8xlarge"},
+        devices={EFA_DEVICE: 1},
     )
 
 
-def aws_g4dn_16xlarge() -> Resource:
+def aws_g4dn_12xlarge() -> Resource:
     return Resource(
-        cpu=64, gpu=1, memMB=256 * GiB, capabilities={K8S_ITYPE: "g4dn.16xlarge"}
+        cpu=48,
+        gpu=4,
+        memMB=192 * GiB,
+        capabilities={K8S_ITYPE: "g4dn.12xlarge"},
+        devices={EFA_DEVICE: 1},
     )
 
 
-def aws_g4dn_12xlarge() -> Resource:
+def aws_g4dn_16xlarge() -> Resource:
     return Resource(
-        cpu=48, gpu=4, memMB=192 * GiB, capabilities={K8S_ITYPE: "g4dn.12xlarge"}
+        cpu=64,
+        gpu=1,
+        memMB=256 * GiB,
+        capabilities={K8S_ITYPE: "g4dn.16xlarge"},
+        devices={EFA_DEVICE: 1},
     )
 
 
 def aws_g4dn_metal() -> Resource:
     return Resource(
-        cpu=96, gpu=8, memMB=384 * GiB, capabilities={K8S_ITYPE: "g4dn.metal"}
+        cpu=96,
+        gpu=8,
+        memMB=384 * GiB,
+        capabilities={K8S_ITYPE: "g4dn.metal"},
+        devices={EFA_DEVICE: 1},
     )
 
 
@@ -164,53 +236,201 @@ def aws_g5_4xlarge() -> Resource:
 
 def aws_g5_8xlarge() -> Resource:
     return Resource(
-        cpu=32, gpu=1, memMB=128 * GiB, capabilities={K8S_ITYPE: "g5.8xlarge"}
+        cpu=32,
+        gpu=1,
+        memMB=128 * GiB,
+        capabilities={K8S_ITYPE: "g5.8xlarge"},
+        devices={EFA_DEVICE: 1},
     )
 
 
-def aws_g5_16xlarge() -> Resource:
+def aws_g5_12xlarge() -> Resource:
     return Resource(
-        cpu=64, gpu=1, memMB=256 * GiB, capabilities={K8S_ITYPE: "g5.16xlarge"}
+        cpu=48,
+        gpu=4,
+        memMB=192 * GiB,
+        capabilities={K8S_ITYPE: "g5.12xlarge"},
+        devices={EFA_DEVICE: 1},
     )
 
 
-def aws_g5_12xlarge() -> Resource:
+def aws_g5_16xlarge() -> Resource:
     return Resource(
-        cpu=48, gpu=4, memMB=192 * GiB, capabilities={K8S_ITYPE: "g5.12xlarge"}
+        cpu=64,
+        gpu=1,
+        memMB=256 * GiB,
+        capabilities={K8S_ITYPE: "g5.16xlarge"},
+        devices={EFA_DEVICE: 1},
     )
 
 
 def aws_g5_24xlarge() -> Resource:
     return Resource(
-        cpu=96, gpu=4, memMB=384 * GiB, capabilities={K8S_ITYPE: "g5.24xlarge"}
+        cpu=96,
+        gpu=4,
+        memMB=384 * GiB,
+        capabilities={K8S_ITYPE: "g5.24xlarge"},
+        devices={EFA_DEVICE: 1},
     )
 
 
 def aws_g5_48xlarge() -> Resource:
     return Resource(
-        cpu=192, gpu=8, memMB=768 * GiB, capabilities={K8S_ITYPE: "g5.48xlarge"}
+        cpu=192,
+        gpu=8,
+        memMB=768 * GiB,
+        capabilities={K8S_ITYPE: "g5.48xlarge"},
+        devices={EFA_DEVICE: 1},
+    )
+
+
+def aws_g6e_xlarge() -> Resource:
+    return Resource(
+        cpu=4,
+        gpu=1,
+        memMB=32 * GiB,
+        capabilities={K8S_ITYPE: "g6e.xlarge"},
     )
 
 
-def aws_trn1_2xl() -> Resource:
-    return Resource(cpu=8, gpu=0, memMB=32 * GiB, capabilities={K8S_ITYPE: "trn1.2xl"})
+def aws_g6e_2xlarge() -> Resource:
+    return Resource(
+        cpu=8,
+        gpu=1,
+        memMB=64 * GiB,
+        capabilities={K8S_ITYPE: "g6e.2xlarge"},
+    )
 
 
-def aws_trn1_32xl() -> Resource:
+def aws_g6e_4xlarge() -> Resource:
     return Resource(
-        cpu=128, gpu=0, memMB=512 * GiB, capabilities={K8S_ITYPE: "trn1.32xl"}
+        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"},
+        devices={NEURON_DEVICE: 1},
+    )
+
+
+def aws_trn1_32xlarge() -> Resource:
+    return Resource(
+        cpu=128,
+        gpu=0,
+        memMB=512 * GiB,
+        capabilities={K8S_ITYPE: "trn1.32xlarge"},
+        devices={EFA_DEVICE: 8, NEURON_DEVICE: 16},
+    )
+
+
+def aws_inf2_xlarge() -> Resource:
+    return Resource(
+        cpu=4,
+        gpu=0,
+        memMB=16 * GiB,
+        capabilities={K8S_ITYPE: "inf2.xlarge"},
+        devices={NEURON_DEVICE: 1},
+    )
+
+
+def aws_inf2_8xlarge() -> Resource:
+    return Resource(
+        cpu=32,
+        gpu=0,
+        memMB=128 * GiB,
+        capabilities={K8S_ITYPE: "inf2.8xlarge"},
+        devices={NEURON_DEVICE: 1},
+    )
+
+
+def aws_inf2_24xlarge() -> Resource:
+    return Resource(
+        cpu=96,
+        gpu=0,
+        memMB=384 * GiB,
+        capabilities={K8S_ITYPE: "inf2.24xlarge"},
+        devices={NEURON_DEVICE: 6},
+    )
+
+
+def aws_inf2_48xlarge() -> Resource:
+    return Resource(
+        cpu=192,
+        gpu=0,
+        memMB=768 * GiB,
+        capabilities={K8S_ITYPE: "inf2.48xlarge"},
+        devices={NEURON_DEVICE: 12},
     )
 
 
 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_p5e.48xlarge": aws_p5e_48xlarge,
+    "aws_p5en.48xlarge": aws_p5en_48xlarge,
     "aws_g4dn.xlarge": aws_g4dn_xlarge,
     "aws_g4dn.2xlarge": aws_g4dn_2xlarge,
     "aws_g4dn.4xlarge": aws_g4dn_4xlarge,
@@ -226,6 +446,18 @@ 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_trn1.2xl": aws_trn1_2xl,
-    "aws_trn1.32xl": aws_trn1_32xl,
+    "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,
+    "aws_inf2.xlarge": aws_inf2_xlarge,
+    "aws_inf2.8xlarge": aws_inf2_8xlarge,
+    "aws_inf2.24xlarge": aws_inf2_24xlarge,
+    "aws_inf2.48xlarge": aws_inf2_48xlarge,
 }

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