Componenets to get fn returning Any (#1018)

Summary:

Allow `Any` so that we can have different function signatures for pipeline definition

Reviewed By: lgarg26

Differential Revision: D70936712

Author: bobyangyf

torchx/specs/builders.py

@@ -101,7 +101,7 @@ def _merge_config_values_with_args(
 
 
 def parse_args(
-    cmpnt_fn: Callable[..., AppDef],
+    cmpnt_fn: Callable[..., Any],  # pyre-ignore[2]
     cmpnt_args: List[str],
     cmpnt_defaults: Optional[Dict[str, Any]] = None,
     config: Optional[Dict[str, Any]] = None,
@@ -130,7 +130,7 @@ def parse_args(
 
 
 def materialize_appdef(
-    cmpnt_fn: Callable[..., AppDef],
+    cmpnt_fn: Callable[..., Any],  # pyre-ignore[2]
     cmpnt_args: List[str],
     cmpnt_defaults: Optional[Dict[str, Any]] = None,
     config: Optional[Dict[str, Any]] = None,
@@ -187,6 +187,10 @@ def materialize_appdef(
         var_arg = var_arg[1:]
 
     appdef = cmpnt_fn(*function_args, *var_arg, **kwargs)
+    if not isinstance(appdef, AppDef):
+        raise TypeError(
+            f"Expected a component that returns `AppDef`, but got `{type(appdef)}`"
+        )
 
     return appdef
 

torchx/specs/finder.py

@@ -16,14 +16,15 @@
 from inspect import getmembers, isfunction
 from pathlib import Path
 from types import ModuleType
-from typing import Callable, Dict, Generator, List, Optional, Union
+from typing import Any, Callable, Dict, Generator, List, Optional, Union
 
 from torchx.specs 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
 from torchx.util.types import none_throws
 
+
 logger: logging.Logger = logging.getLogger(__name__)
 
 
@@ -53,7 +54,10 @@ class _Component:
     name: str
     description: str
     fn_name: str
-    fn: Callable[..., AppDef]
+
+    # pyre-ignore[4] TODO temporary until PipelineDef is decoupled and can be exposed as type to OSS
+    fn: Callable[..., Any]
+
     validation_errors: List[str]