torchx-nightly 2024.2.12__py3-none-any.whl → 2024.4.10__py3-none-any.whl

torchx/schedulers/gcp_batch_scheduler.py

@@ -5,6 +5,8 @@
 # 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
+
 """
 
 This contains the TorchX GCP Batch scheduler which can be used to run TorchX
@@ -205,14 +207,12 @@ class GCPBatchScheduler(Scheduler[GCPBatchOpts]):
             if cpu <= 0:
                 cpu = 1
             MILLI = 1000
-            # pyre-fixme[8]: Attribute has type `Field`; used as `int`.
             res.cpu_milli = cpu * MILLI
             memMB = resource.memMB
             if memMB < 0:
                 raise ValueError(
                     f"memMB should to be set to a positive value, got {memMB}"
                 )
-            # pyre-fixme[8]: Attribute has type `Field`; used as `int`.
             res.memory_mib = memMB
 
             # TODO support named resources
@@ -360,13 +360,11 @@ class GCPBatchScheduler(Scheduler[GCPBatchOpts]):
             return None
 
         gpu = 0
-        # pyre-fixme[16]: `Field` has no attribute `instances`.
         if len(job.allocation_policy.instances) != 0:
             gpu_type = job.allocation_policy.instances[0].policy.machine_type
             gpu = GPU_TYPE_TO_COUNT[gpu_type]
 
         roles = {}
-        # pyre-fixme[16]: `RepeatedField` has no attribute `__iter__`.
         for tg in job.task_groups:
             env = tg.task_spec.environment.variables
             role = env["TORCHX_ROLE_NAME"]
@@ -390,7 +388,6 @@ class GCPBatchScheduler(Scheduler[GCPBatchOpts]):
         # TODO map role/replica status
         desc = DescribeAppResponse(
             app_id=app_id,
-            # pyre-fixme[16]: `Field` has no attribute `state`.
             state=JOB_STATE[job.status.state.name],
             roles=list(roles.values()),
         )
@@ -415,8 +412,10 @@ class GCPBatchScheduler(Scheduler[GCPBatchOpts]):
             raise ValueError(f"app not found: {app_id}")
 
         job_uid = job.uid
-        filters = [f"labels.job_uid={job_uid}"]
-        filters.append(f"resource.labels.task_id:task/{job_uid}-group0-{k}")
+        filters = [
+            f"labels.job_uid={job_uid}",
+            f"labels.task_id:{job_uid}-group0-{k}",
+        ]
 
         if since is not None:
             filters.append(f'timestamp>="{str(since.isoformat())}"')
@@ -437,7 +436,7 @@ class GCPBatchScheduler(Scheduler[GCPBatchOpts]):
 
         logger = logging.Client().logger(BATCH_LOGGER_NAME)
         for entry in logger.list_entries(filter_=filter):
-            yield entry.payload
+            yield entry.payload + "\n"
 
     def _job_full_name_to_app_id(self, job_full_name: str) -> str:
         """

torchx/schedulers/ids.py

@@ -5,6 +5,8 @@
 # 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
+
 import os
 import random
 import struct

torchx/schedulers/kubernetes_scheduler.py

@@ -5,6 +5,8 @@
 # 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
+
 """
 
 This contains the TorchX Kubernetes scheduler which can be used to run TorchX
@@ -23,7 +25,7 @@ Install Volcano:
     kubectl apply -f https://raw.githubusercontent.com/volcano-sh/volcano/v1.6.0/installer/volcano-development.yaml
 
 See the
-`Volcano Quickstart <https://github.com/volcano-sh/volcano#quick-start-guide>`_
+`Volcano Quickstart <https://github.com/volcano-sh/volcano>`_
 for more information.
 """
 

torchx/schedulers/local_scheduler.py

@@ -5,6 +5,8 @@
 # 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
+
 """
 This contains the TorchX local scheduler which can be used to run TorchX
 components locally via subprocesses.
@@ -35,6 +37,7 @@ from typing import (
     Iterable,
     List,
     Optional,
+    Protocol,
     TextIO,
     Tuple,
 )
@@ -262,6 +265,26 @@ AppName = str
 RoleName = str
 
 
+class PopenProtocol(Protocol):
+    """
+    Protocol wrapper around python's ``subprocess.Popen``. Keeps track of
+    the a list of interface methods that the process scheduled by the `LocalScheduler`
+    must implement.
+    """
+
+    @property
+    def pid(self) -> int: ...
+
+    @property
+    def returncode(self) -> int: ...
+
+    def wait(self, timeout: Optional[float] = None) -> int: ...
+
+    def poll(self) -> Optional[int]: ...
+
+    def kill(self) -> None: ...
+
+
 @dataclass
 class _LocalReplica:
     """
@@ -270,8 +293,7 @@ class _LocalReplica:
 
     role_name: RoleName
     replica_id: int
-    # pyre-fixme[24]: Generic type `subprocess.Popen` expects 1 type parameter.
-    proc: subprocess.Popen
+    proc: PopenProtocol
 
     # IO streams:
     # None means no log_dir (out to console)

torchx/schedulers/lsf_scheduler.py

@@ -5,6 +5,8 @@
 # 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
+
 """
 This contains the TorchX LSF scheduler which can be used to run TorchX
 components on a LSF cluster.

torchx/schedulers/ray/ray_driver.py

@@ -148,12 +148,12 @@ class RayDriver:
         else:
             self.min_replicas = replicas[0].min_replicas  # pyre-ignore[8]
 
-        self.placement_groups: List[
-            PlacementGroup
-        ] = []  # all the placement groups, shall never change
-        self.actor_info_of_id: Dict[
-            str, ActorInfo
-        ] = {}  # store the info used to recover an actor
+        self.placement_groups: List[PlacementGroup] = (
+            []
+        )  # all the placement groups, shall never change
+        self.actor_info_of_id: Dict[str, ActorInfo] = (
+            {}
+        )  # store the info used to recover an actor
         self.active_tasks: List["ray.ObjectRef"] = []  # list of active tasks
 
         self.terminating: bool = False  # if the job has finished and being terminated

torchx/schedulers/slurm_scheduler.py

@@ -5,6 +5,8 @@
 # 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
+
 """
 This contains the TorchX Slurm scheduler which can be used to run TorchX
 components on a Slurm cluster.

torchx/schedulers/streams.py

@@ -5,6 +5,8 @@
 # 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
+
 import io
 import os
 import threading

torchx/specs/__init__.py

@@ -5,6 +5,8 @@
 # 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
+
 """
 This contains the TorchX AppDef and related component definitions. These are
 used by components to define the apps which can then be launched via a TorchX

torchx/specs/api.py

@@ -5,6 +5,8 @@
 # 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
+
 import copy
 import json
 import re
@@ -637,11 +639,11 @@ class AppStatusError(Exception):
         self.status = status
 
 
-# valid run cfg values; only support primitives (str, int, float, bool, List[str])
+# valid run cfg values; only support primitives (str, int, float, bool, List[str], Dict[str, str])
 # TODO(wilsonhong): python 3.9+ supports list[T] in typing, which can be used directly
 # in isinstance(). Should replace with that.
 # see: https://docs.python.org/3/library/stdtypes.html#generic-alias-type
-CfgVal = Union[str, int, float, bool, List[str], None]
+CfgVal = Union[str, int, float, bool, List[str], Dict[str, str], None]
 
 
 T = TypeVar("T")
@@ -755,6 +757,10 @@ class runopts:
         except TypeError:
             if isinstance(obj, list):
                 return all(isinstance(e, str) for e in obj)
+            elif isinstance(obj, dict):
+                return all(
+                    isinstance(k, str) and isinstance(v, str) for k, v in obj.items()
+                )
             else:
                 return False
 
@@ -863,8 +869,13 @@ class runopts:
                 # lists may be ; or , delimited
                 # also deal with trailing "," by removing empty strings
                 return [v for v in value.replace(";", ",").split(",") if v]
+            elif opt_type == Dict[str, str]:
+                return {
+                    s.split(":", 1)[0]: s.split(":", 1)[1]
+                    for s in value.replace(";", ",").split(",")
+                }
             else:
-                # pyre-ignore[19]
+                # pyre-ignore[19, 6] type won't be dict here as we handled it above
                 return opt_type(value)
 
         cfg: Dict[str, CfgVal] = {}

torchx/specs/builders.py

@@ -4,25 +4,27 @@
 # 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
+
 import argparse
 import inspect
+import os
+from argparse import Namespace
 from typing import Any, Callable, Dict, List, Mapping, Optional, Union
 
 from torchx.specs.api import BindMount, MountType, VolumeMount
 from torchx.specs.file_linter import get_fn_docstring, TorchXArgumentHelpFormatter
-from torchx.util.types import (
-    decode_from_string,
-    decode_optional,
-    get_argparse_param_type,
-    is_bool,
-    is_primitive,
-)
+from torchx.util.types import decode, decode_optional, get_argparse_param_type, is_bool
 
 from .api import AppDef, DeviceMount
 
+ENV_TORCHX_COMPONENT_ARGS = "TORCHX_COMPONENT_ARGS"
+
 
 def _create_args_parser(
-    cmpnt_fn: Callable[..., AppDef], cmpnt_defaults: Optional[Dict[str, str]] = None
+    cmpnt_fn: Callable[..., AppDef],
+    cmpnt_defaults: Optional[Dict[str, str]] = None,
+    config: Optional[Dict[str, Any]] = None,
 ) -> argparse.ArgumentParser:
     parameters = inspect.signature(cmpnt_fn).parameters
     function_desc, args_desc = get_fn_docstring(cmpnt_fn)
@@ -85,15 +87,56 @@ def _create_args_parser(
             if len(param_name) == 1:
                 arg_names = [f"-{param_name}"] + arg_names
             if "default" not in args:
-                args["required"] = True
+                if (config and param_name not in config) or not config:
+                    args["required"] = True
+
             script_parser.add_argument(*arg_names, **args)
     return script_parser
 
 
+def _merge_config_values_with_args(
+    parsed_args: argparse.Namespace, config: Dict[str, Any]
+) -> None:
+    for key, val in config.items():
+        if key in parsed_args:
+            setattr(parsed_args, key, val)
+
+
+def parse_args(
+    cmpnt_fn: Callable[..., AppDef],
+    cmpnt_args: List[str],
+    cmpnt_defaults: Optional[Dict[str, Any]] = None,
+    config: Optional[Dict[str, Any]] = None,
+) -> Namespace:
+    """
+    Parse passed arguments, defaults, and config values into a namespace for
+    a component function.
+
+    Args:
+    cmpnt_fn: Component function
+    cmpnt_args: Function args
+    cmpnt_defaults: Additional default values for parameters of ``app_fn``
+                        (overrides the defaults set on the fn declaration)
+    config: Optional dict containing additional configuration for the component from a passed config file
+
+    Returns:
+    A Namespace object with the args, defaults, and config values incorporated.
+    """
+
+    script_parser = _create_args_parser(cmpnt_fn, cmpnt_defaults, config)
+    parsed_args = script_parser.parse_args(cmpnt_args)
+    if config:
+        _merge_config_values_with_args(parsed_args, config)
+
+    return parsed_args
+
+
 def materialize_appdef(
     cmpnt_fn: Callable[..., AppDef],
     cmpnt_args: List[str],
-    cmpnt_defaults: Optional[Dict[str, str]] = None,
+    cmpnt_defaults: Optional[Dict[str, Any]] = None,
+    config: Optional[Dict[str, Any]] = None,
+    component_args_string: Optional[str] = None,
 ) -> AppDef:
     """
     Creates an application by running user defined ``app_fn``.
@@ -118,26 +161,23 @@ def materialize_appdef(
         cmpnt_args: Function args
         cmpnt_defaults: Additional default values for parameters of ``app_fn``
                           (overrides the defaults set on the fn declaration)
+        config: Optional dict containing additional configuration for the component from a passed config file
     Returns:
         An application spec
     """
 
-    script_parser = _create_args_parser(cmpnt_fn, cmpnt_defaults)
-    parsed_args = script_parser.parse_args(cmpnt_args)
-
     function_args = []
     var_arg = []
     kwargs = {}
 
+    parsed_args = parse_args(cmpnt_fn, cmpnt_args, cmpnt_defaults, config)
+
     parameters = inspect.signature(cmpnt_fn).parameters
     for param_name, parameter in parameters.items():
         arg_value = getattr(parsed_args, param_name)
         parameter_type = parameter.annotation
         parameter_type = decode_optional(parameter_type)
-        if is_bool(parameter_type):
-            arg_value = arg_value and arg_value.lower() == "true"
-        elif not is_primitive(parameter_type):
-            arg_value = decode_from_string(arg_value, parameter_type)
+        arg_value = decode(arg_value, parameter_type)
         if parameter.kind == inspect.Parameter.VAR_POSITIONAL:
             var_arg = arg_value
         elif parameter.kind == inspect.Parameter.KEYWORD_ONLY:
@@ -149,7 +189,13 @@ def materialize_appdef(
     if len(var_arg) > 0 and var_arg[0] == "--":
         var_arg = var_arg[1:]
 
-    return cmpnt_fn(*function_args, *var_arg, **kwargs)
+    appdef = cmpnt_fn(*function_args, *var_arg, **kwargs)
+
+    if component_args_string:
+        for role in appdef.roles:
+            role.env[ENV_TORCHX_COMPONENT_ARGS] = component_args_string
+
+    return appdef
 
 
 def make_app_handle(scheduler_backend: str, session_name: str, app_id: str) -> str:
@@ -205,9 +251,12 @@ def parse_mounts(opts: List[str]) -> List[Union[BindMount, VolumeMount, DeviceMo
     for opts in mount_opts:
         typ = opts.get("type")
         if typ == MountType.BIND:
+            src_path = opts["src"]
+            if src_path.startswith("~"):
+                src_path = os.path.expanduser(src_path)
             mounts.append(
                 BindMount(
-                    src_path=opts["src"],
+                    src_path=src_path,
                     dst_path=opts["dst"],
                     read_only="readonly" in opts,
                 )

torchx/specs/file_linter.py

@@ -5,6 +5,8 @@
 # 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
+
 import abc
 import argparse
 import ast
@@ -29,7 +31,11 @@ def _get_default_arguments_descriptions(fn: Callable[..., object]) -> Dict[str,
     return args_decs
 
 
-class TorchXArgumentHelpFormatter(argparse.HelpFormatter):
+class TorchXArgumentHelpFormatter(
+    argparse.RawDescriptionHelpFormatter,
+    argparse.ArgumentDefaultsHelpFormatter,
+    argparse.MetavarTypeHelpFormatter,
+):
     """Help message formatter which adds default values and required to argument help.
 
     If the argument is required, the class appends `(required)` at the end of the help message.
@@ -79,7 +85,7 @@ to your component (see: https://pytorch.org/torchx/latest/component_best_practic
             args_description[param.arg_name] = param.description
     short_func_description = docstring.short_description or default_fn_desc
     if docstring.long_description:
-        short_func_description += " ..."
+        short_func_description += "\n" + docstring.long_description
     return (short_func_description or default_fn_desc, args_description)
 
 

torchx/specs/finder.py

@@ -4,6 +4,8 @@
 # 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
+
 import abc
 import importlib
 import inspect

torchx/specs/named_resources_aws.py

@@ -4,6 +4,8 @@
 # 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
+
 r"""
 `torchx.specs.named_resources_aws` contains resource definitions that represent corresponding AWS instance types
 taken from https://aws.amazon.com/ec2/instance-types/. The resources are exposed

torchx/specs/named_resources_generic.py

@@ -4,6 +4,8 @@
 # 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
+
 """
 Defines generic named resources that are not specific to any cloud provider's
 instance types. These generic named resources are meant to be used as

torchx/specs/test/components/__init__.py

@@ -3,3 +3,5 @@
 #
 # 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

torchx/specs/test/components/a/__init__.py

@@ -3,6 +3,8 @@
 #
 # 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
 import torchx
 from torchx import specs
 

torchx/specs/test/components/a/b/__init__.py

@@ -3,3 +3,5 @@
 #
 # 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

torchx/specs/test/components/a/b/c.py

@@ -3,6 +3,8 @@
 #
 # 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
 import torchx
 from torchx import specs
 

torchx/specs/test/components/c/__init__.py

@@ -4,3 +4,5 @@
 #
 # 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

torchx/specs/test/components/c/d.py

@@ -3,6 +3,8 @@
 #
 # 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
 import torchx
 from torchx import specs
 

torchx/tracker/__init__.py

@@ -4,6 +4,8 @@
 # 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
+
 """
 .. note:: PROTOTYPE, USE AT YOUR OWN RISK, APIs SUBJECT TO CHANGE
 

torchx/tracker/api.py

@@ -4,6 +4,8 @@
 # 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
+
 from __future__ import annotations
 
 import logging
@@ -67,8 +69,7 @@ class AppRunTrackableSource:
     artifact_name: Optional[str]
 
 
-class Lineage:
-    ...
+class Lineage: ...
 
 
 class TrackerBase(ABC):
@@ -332,5 +333,4 @@ class AppRun:
 
         return model_run_sources
 
-    def children(self) -> Iterable[AppRun]:
-        ...
+    def children(self) -> Iterable[AppRun]: ...

torchx/tracker/backend/fsspec.py

@@ -4,6 +4,8 @@
 # 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
+
 from __future__ import annotations
 
 import json

torchx/util/cuda.py

@@ -4,6 +4,8 @@
 # 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
+
 import torch
 
 

torchx/util/datetime.py

@@ -4,6 +4,8 @@
 # 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
+
 from datetime import datetime, timedelta
 
 

torchx/util/entrypoints.py

@@ -4,6 +4,8 @@
 # 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
+
 from typing import Any, Dict, Optional
 
 import importlib_metadata as metadata

torchx/util/io.py

@@ -4,6 +4,8 @@
 # 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
+
 from os import path
 from pathlib import Path
 from typing import Optional

torchx/util/modules.py

@@ -4,6 +4,8 @@
 # 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
+
 import importlib
 from types import ModuleType
 from typing import Callable, Optional, Union

torchx/util/shlex.py

@@ -4,6 +4,8 @@
 # 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
+
 import shlex
 from typing import Iterable
 

torchx/util/strings.py

@@ -4,6 +4,8 @@
 # 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
+
 import re
 
 

torchx/util/types.py

@@ -4,6 +4,8 @@
 # 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
+
 import inspect
 from typing import Any, Callable, Dict, List, Optional, Tuple, Type, TypeVar, Union
 
@@ -120,6 +122,16 @@ def _decode_string_to_list(
     return arg_values
 
 
+def decode(encoded_value: Any, annotation: Any):
+    if encoded_value is None:
+        return None
+    if is_bool(annotation):
+        return encoded_value and encoded_value.lower() == "true"
+    if not is_primitive(annotation) and type(encoded_value) == str:
+        return decode_from_string(encoded_value, annotation)
+    return encoded_value
+
+
 def decode_from_string(
     encoded_value: str, annotation: Any
 ) -> Union[Dict[Any, Any], List[Any], None]: