torchx-nightly 2025.3.18__py3-none-any.whl → 2025.3.20__py3-none-any.whl

torchx/specs/file_linter.py

@@ -244,11 +244,18 @@ class TorchFunctionVisitor(ast.NodeVisitor):
 
     """
 
-    def __init__(self, component_function_name: str) -> None:
-        self.validators = [
-            TorchxFunctionArgsValidator(),
-            TorchxReturnValidator(),
-        ]
+    def __init__(
+        self,
+        component_function_name: str,
+        validators: Optional[List[TorchxFunctionValidator]],
+    ) -> None:
+        if validators is None:
+            self.validators: List[TorchxFunctionValidator] = [
+                TorchxFunctionArgsValidator(),
+                TorchxReturnValidator(),
+            ]
+        else:
+            self.validators = validators
         self.linter_errors: List[LinterMessage] = []
         self.component_function_name = component_function_name
         self.visited_function = False
@@ -264,7 +271,11 @@ class TorchFunctionVisitor(ast.NodeVisitor):
             self.linter_errors += validator.validate(node)
 
 
-def validate(path: str, component_function: str) -> List[LinterMessage]:
+def validate(
+    path: str,
+    component_function: str,
+    validators: Optional[List[TorchxFunctionValidator]],
+) -> List[LinterMessage]:
     """
     Validates the function to make sure it complies the component standard.
 
@@ -293,7 +304,7 @@ def validate(path: str, component_function: str) -> List[LinterMessage]:
             severity="error",
         )
         return [linter_message]
-    visitor = TorchFunctionVisitor(component_function)
+    visitor = TorchFunctionVisitor(component_function, validators)
     visitor.visit(module)
     linter_errors = visitor.linter_errors
     if not visitor.visited_function:

torchx/specs/finder.py

@@ -19,7 +19,7 @@ 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 get_fn_docstring, TorchxFunctionValidator, validate
 from torchx.util import entrypoints
 from torchx.util.io import read_conf_file
 from torchx.util.types import none_throws
@@ -59,7 +59,9 @@ class _Component:
 
 class ComponentsFinder(abc.ABC):
     @abc.abstractmethod
-    def find(self) -> List[_Component]:
+    def find(
+        self, validators: Optional[List[TorchxFunctionValidator]]
+    ) -> List[_Component]:
         """
         Retrieves a set of components. A component is defined as a python
         function that conforms to ``torchx.specs.file_linter`` linter.
@@ -203,10 +205,12 @@ class ModuleComponentsFinder(ComponentsFinder):
                 else:
                     yield self._try_import(module_info.name)
 
-    def find(self) -> List[_Component]:
+    def find(
+        self, validators: Optional[List[TorchxFunctionValidator]]
+    ) -> 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:
@@ -221,7 +225,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[TorchxFunctionValidator]]
+    ) -> List[_Component]:
         functions = getmembers(module, isfunction)
         component_defs = []
 
@@ -230,7 +236,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=""
@@ -255,13 +261,20 @@ 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[TorchxFunctionValidator]],
+    ) -> 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[TorchxFunctionValidator]]
+    ) -> 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)
@@ -284,7 +297,9 @@ class CustomComponentsFinder(ComponentsFinder):
         ]
 
 
-def _load_custom_components() -> List[_Component]:
+def _load_custom_components(
+    validators: Optional[List[TorchxFunctionValidator]],
+) -> List[_Component]:
     component_modules = {
         name: load_fn()
         for name, load_fn in
@@ -303,11 +318,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[TorchxFunctionValidator]],
+) -> Dict[str, _Component]:
     """
     Loads either the custom component defs from the entrypoint ``[torchx.components]``
     or the default builtins from ``torchx.components`` module.
@@ -318,19 +335,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[TorchxFunctionValidator]],
+) -> Dict[str, _Component]:
     global _components
     if not _components:
-        _components = _load_components()
+        _components = _load_components(validators)
     return none_throws(_components)
 
 
@@ -338,17 +357,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[TorchxFunctionValidator]]
+) -> 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[TorchxFunctionValidator]] = None,
+) -> Dict[str, _Component]:
     """
     Returns all custom components registered via ``[torchx.components]`` entrypoints
     OR builtin components that ship with TorchX (but not both).
@@ -395,13 +418,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[TorchxFunctionValidator]] = None
+) -> _Component:
     """
     Retrieves components by the provided name.
 
@@ -409,9 +434,9 @@ 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 "
@@ -428,7 +453,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[TorchxFunctionValidator]] = 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
@@ -446,7 +473,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_nightly-2025.3.18.dist-info → torchx_nightly-2025.3.20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: torchx-nightly
-Version: 2025.3.18
+Version: 2025.3.20
 Summary: TorchX SDK and Components
 Home-page: https://github.com/pytorch/torchx
 Author: TorchX Devs

{torchx_nightly-2025.3.18.dist-info → torchx_nightly-2025.3.20.dist-info}/RECORD

@@ -85,8 +85,8 @@ torchx/schedulers/ray/ray_driver.py,sha256=RdaCLfth16ky-5PDVOWRe_RuheWJu9xufWux2
 torchx/specs/__init__.py,sha256=T8xUCz7iVE6OsUL5P4Pzy2B8ZY_YinCVDwUer5Q-XPc,6179
 torchx/specs/api.py,sha256=ISdnfGH1R0-1Ah23UsPOZQT3WH63dIMj9CxA3QRoP0g,38046
 torchx/specs/builders.py,sha256=QDcQrnCO4bdSaiP0216XbCgTsnLutO_1_FW5jDiEIWI,9939
-torchx/specs/file_linter.py,sha256=IeiomB1BgHUlT-ZsvGxar3llY63NOupfLBrOrD_---A,11860
-torchx/specs/finder.py,sha256=MnwxG_UC4a-3X2wQ37ANEQR6D1TvriCLyuVYBh_-wuI,16249
+torchx/specs/file_linter.py,sha256=OOyDnfg9m_9BiyhXc3T0fhE_-ykKCJ32gxCQGbfT-yE,12157
+torchx/specs/finder.py,sha256=IP0dnExXtyuOzI5mCZ0sTe3fU7AiFqfp7kCWXdEVdKs,17206
 torchx/specs/named_resources_aws.py,sha256=ISjHtifRJqB8u7PeAMiyLyO_S0WCaZiK-CFF3qe6JDU,11415
 torchx/specs/named_resources_generic.py,sha256=Sg4tAdqiiWDrDz2Lj_pnfsjzGIXKTou73wPseh6j55w,2646
 torchx/specs/test/components/__init__.py,sha256=J8qjUOysmcMAek2KFN13mViOXZxTYc5vCrF02t3VuFU,223
@@ -115,9 +115,9 @@ torchx/workspace/__init__.py,sha256=FqN8AN4VhR1C_SBY10MggQvNZmyanbbuPuE-JCjkyUY,
 torchx/workspace/api.py,sha256=PtDkGTC5lX03pRoYpuMz2KCmM1ZOycRP1UknqvNb97Y,6341
 torchx/workspace/dir_workspace.py,sha256=npNW_IjUZm_yS5r-8hrRkH46ndDd9a_eApT64m1S1T4,2268
 torchx/workspace/docker_workspace.py,sha256=PFu2KQNVC-0p2aKJ-W_BKA9ZOmXdCY2ABEkCExp3udQ,10269
-torchx_nightly-2025.3.18.dist-info/LICENSE,sha256=WVHfXhFC0Ia8LTKt_nJVYobdqTJVg_4J3Crrfm2A8KQ,1721
-torchx_nightly-2025.3.18.dist-info/METADATA,sha256=DlFiSRfQVyiw6cqp0gRre_MBbUzPofa6Q2vF5L4gXmg,6169
-torchx_nightly-2025.3.18.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-torchx_nightly-2025.3.18.dist-info/entry_points.txt,sha256=T328AMXeKI3JZnnxfkEew2ZcMN1oQDtkXjMz7lkV-P4,169
-torchx_nightly-2025.3.18.dist-info/top_level.txt,sha256=pxew3bc2gsiViS0zADs0jb6kC5v8o_Yy_85fhHj_J1A,7
-torchx_nightly-2025.3.18.dist-info/RECORD,,
+torchx_nightly-2025.3.20.dist-info/LICENSE,sha256=WVHfXhFC0Ia8LTKt_nJVYobdqTJVg_4J3Crrfm2A8KQ,1721
+torchx_nightly-2025.3.20.dist-info/METADATA,sha256=70Gc21VJeFVFPFjzecjVnBT_V8JhacxAQNhXJkZMRag,6169
+torchx_nightly-2025.3.20.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+torchx_nightly-2025.3.20.dist-info/entry_points.txt,sha256=T328AMXeKI3JZnnxfkEew2ZcMN1oQDtkXjMz7lkV-P4,169
+torchx_nightly-2025.3.20.dist-info/top_level.txt,sha256=pxew3bc2gsiViS0zADs0jb6kC5v8o_Yy_85fhHj_J1A,7
+torchx_nightly-2025.3.20.dist-info/RECORD,,