handle components whose implementation lives in a different file

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

@@ -19,6 +19,7 @@
 from types import ModuleType
 from typing import Any, Callable, Dict, Generator, List, Optional, Union
 
+from torchx.specs.api import AppDef
 from torchx.specs.file_linter import get_fn_docstring, TorchxFunctionValidator, validate
 from torchx.util import entrypoints
 from torchx.util.io import read_conf_file
@@ -274,12 +275,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[..., AppDef]) -> 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 = function.__code__.co_filename  # pyre-ignore[16]
+        if 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 +302,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}",