handle components whose implementation lives in a different file (meta-pytorch#1075)

Summary:

Add support for cases like:
```lang=python
# some_file.py
# ====================
def my_component(...) -> specs.AppDef: ...

# other_file.py
# ====================
from some_file import my_component
```

where the component is invoked with
`torchx run ... other_file.py:my_component`

This was currently failing with a validation error because in the step where we inspect the AST of the component, we assume that the file where the component is being looked up is the same as the file where it is implemented.

Differential Revision: D75496839

Author: Daniel Ohayon

torchx/specs/finder.py

@@ -17,7 +17,7 @@
 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 Any, Callable, Dict, Generator, List, Optional, TypeVar, Union
 
 from torchx.specs.file_linter import get_fn_docstring, TorchxFunctionValidator, validate
 from torchx.util import entrypoints
@@ -36,6 +36,9 @@ class ComponentNotFoundException(Exception):
     pass
 
 
+T = TypeVar("T")
+
+
 @dataclass
 class _Component:
     """
@@ -274,12 +277,21 @@ def _get_validation_errors(
         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[..., T]) -> 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]]
     ) -> 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())
@@ -292,6 +304,12 @@ def find(
             )
         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}",

torchx/specs/test/finder_test.py

@@ -29,6 +29,7 @@
     get_components,
     ModuleComponentsFinder,
 )
+from torchx.specs.test.components.a import comp_a
 from torchx.util.test.entrypoints_test import EntryPoint_from_text
 from torchx.util.types import none_throws
 
@@ -238,6 +239,10 @@ def test_get_component_invalid(self) -> None:
         with self.assertRaises(ComponentValidationException):
             get_component(f"{current_file_path()}:invalid_component")
 
+    def test_get_component_imported_from_other_file(self) -> None:
+        component = get_component(f"{current_file_path()}:comp_a")
+        self.assertListEqual([], component.validation_errors)
+
 
 class GetBuiltinSourceTest(unittest.TestCase):
     def setUp(self) -> None: