(bugfix)(torchx/components) handle BinOp style optional when decoding component arguments

Differential Revision: D82856462

Pull Request resolved: #1123

Author: kiukchung

torchx/specs/builders.py

@@ -4,7 +4,7 @@
 # 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
+# pyre-unsafe
 
 import argparse
 import inspect
@@ -39,7 +39,7 @@ def _create_args_parser(
 
 
 def _create_args_parser_from_parameters(
-    cmpnt_fn: Callable[..., Any],  # pyre-ignore[2]
+    cmpnt_fn: Callable[..., AppDef],
     parameters: Mapping[str, inspect.Parameter],
     cmpnt_defaults: Optional[Dict[str, str]] = None,
     config: Optional[Dict[str, Any]] = None,
@@ -120,7 +120,7 @@ def _merge_config_values_with_args(
 
 
 def parse_args(
-    cmpnt_fn: Callable[..., Any],  # pyre-ignore[2]
+    cmpnt_fn: Callable[..., AppDef],
     cmpnt_args: List[str],
     cmpnt_defaults: Optional[Dict[str, Any]] = None,
     config: Optional[Dict[str, Any]] = None,
@@ -149,7 +149,7 @@ def parse_args(
 
 
 def component_args_from_str(
-    cmpnt_fn: Callable[..., Any],  # pyre-fixme[2]: Enforce AppDef type
+    cmpnt_fn: Callable[..., AppDef],
     cmpnt_args: list[str],
     cmpnt_args_defaults: Optional[Dict[str, Any]] = None,
     config: Optional[Dict[str, Any]] = None,
@@ -238,7 +238,7 @@ def example_component_fn(foo: str, *args: str, bar: str = "asdf") -> AppDef:
 
 
 def materialize_appdef(
-    cmpnt_fn: Callable[..., Any],  # pyre-ignore[2]
+    cmpnt_fn: Callable[..., AppDef],
     cmpnt_args: List[str],
     cmpnt_defaults: Optional[Dict[str, Any]] = None,
     config: Optional[Dict[str, Any]] = None,

torchx/specs/finder.py

@@ -17,7 +17,9 @@
 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 Callable, Dict, Generator, List, Optional, Union
+
+from torchx.specs import AppDef
 
 from torchx.specs.file_linter import (
     ComponentFunctionValidator,
@@ -59,8 +61,7 @@ class _Component:
     description: str
     fn_name: str
 
-    # pyre-ignore[4] TODO temporary until PipelineDef is decoupled and can be exposed as type to OSS
-    fn: Callable[..., Any]
+    fn: Callable[..., AppDef]
 
     validation_errors: List[str]
 

torchx/specs/test/builders_test.py

@@ -93,6 +93,7 @@ def example_test_complex_fn(
     nnodes: int = 4,
     first_arg: Optional[str] = None,
     nested_arg: Optional[Dict[str, List[str]]] = None,
+    env: dict[str, str] | None = None,
     *roles_args: str,
 ) -> AppDef:
     """Creates complex application, testing all possible complex types
@@ -127,6 +128,7 @@ def example_test_complex_fn(
             args=args,
             resource=Resource(cpu=cpus, gpu=gpus, memMB=1),
             num_replicas=nnodes,
+            env=env or {},
         )
         roles.append(role)
     return AppDef(app_name, roles)
@@ -193,6 +195,7 @@ def _get_expected_app_with_default(self) -> AppDef:
             4,
             None,
             None,
+            None,
             *role_args,
         )
 
@@ -220,6 +223,7 @@ def _get_expected_app_with_all_args(self) -> AppDef:
             8,
             "first_arg",
             None,
+            {"FOO": "BAR", "HELLO": "WORLD"},
             *role_args,
         )
 
@@ -240,6 +244,8 @@ def _get_app_args(self) -> List[str]:
             "8",
             "--first_arg",
             "first_arg",
+            "--env",
+            "FOO=BAR,HELLO=WORLD",
             "--",
             *role_args,
         ]
@@ -256,6 +262,7 @@ def _get_expected_app_with_nested_objects(self) -> AppDef:
             8,
             "first_arg",
             defaults,
+            None,
             *role_args,
         )
 

torchx/util/types.py

@@ -8,6 +8,7 @@
 
 import inspect
 import re
+from types import UnionType
 from typing import Any, Callable, Optional, Tuple, TypeVar, Union
 
 
@@ -234,10 +235,20 @@ def decode_optional(param_type: Any) -> Any:
         If ``param_type`` is type Optional[INNER_TYPE], method returns INNER_TYPE
         Otherwise returns ``param_type``
     """
+
     if not hasattr(param_type, "__origin__"):
-        return param_type
+        if isinstance(param_type, UnionType):
+            # handle BinOp style Optional (e.g. `T | None`)
+            if len(param_type.__args__) == 2 and param_type.__args__[1] is type(None):
+                return param_type.__args__[0]
+            else:
+                return param_type
+        else:
+            return param_type
+
     if param_type.__origin__ is not Union:
         return param_type
+
     args = param_type.__args__
     if len(args) == 2 and args[1] is type(None):
         return args[0]