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

torchx/__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
+
 from .version import (  # noqa F401; noqa F401
     __version__ as __version__,
     TORCHX_IMAGE as IMAGE,

torchx/apps/serve/serve.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 argparse
 import binascii
 import os

torchx/apps/utils/booth_main.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 argparse
 import sys
 from typing import List

torchx/apps/utils/copy_main.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 argparse
 import os
 import shutil

torchx/apps/utils/process_monitor.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 argparse
 import subprocess
 import sys

torchx/cli/__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
+
 """
 The ``torchx`` CLI is a commandline tool around :py:class:`torchx.runner.Runner`.
 It allows users to launch :py:class:`torchx.specs.AppDef` directly onto

torchx/cli/argparse_util.py

@@ -4,19 +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 logging
+import sys
 from argparse import Action, ArgumentParser, Namespace
-from typing import Any, Dict, Optional, Sequence, Text
+from typing import Any, Dict, List, Optional, Sequence, Set, Text
 
 from torchx.runner import config
 
+logger: logging.Logger = logging.getLogger(__name__)
+
 
-class _torchxconfig(Action):
+class torchxconfig(Action):
     """
     Custom argparse action that loads default torchx CLI options
     from .torchxconfig file.
 
     """
 
+    called_args: Set[str] = set()
+
     # since this action is used for each argparse argument
     # load the config section for the subcmd once
     _subcmd_configs: Dict[str, Dict[str, str]] = {}
@@ -64,13 +72,18 @@ class _torchxconfig(Action):
         values: Any,  # pyre-ignore[2] declared as Any in superclass Action
         option_string: Optional[str] = None,
     ) -> None:
+        if option_string is not None:
+            if option_string in self.called_args:
+                logger.error(f"{option_string} is specified more than once")
+                sys.exit(1)
+            self.called_args.add(option_string)
         setattr(namespace, self.dest, values)
 
 
 # argparse takes the action as a Type[Action] so we can't have custom constructors
 # hence for each subcommand we need to subclass the base _torchxconfig Action
 # this is also how store_true and store_false builtin actions are implemented in argparse
-class torchxconfig_run(_torchxconfig):
+class torchxconfig_run(torchxconfig):
     """
     Custom action that gets the default argument from .torchxconfig.
     """
@@ -92,3 +105,25 @@ class torchxconfig_run(_torchxconfig):
             option_strings=option_strings,
             **kwargs,
         )
+
+
+class ArgOnceAction(Action):
+    """
+    Custom argparse action only allows argument to be specified once
+    """
+
+    called_args: Set[str] = set()
+
+    def __call__(
+        self,
+        parser: ArgumentParser,
+        namespace: Namespace,
+        values: List[str],
+        option_string: Optional[str] = None,
+    ) -> None:
+        if option_string is not None:
+            if option_string in self.called_args:
+                logger.error(f"{option_string} is specified more than once")
+                sys.exit(1)
+            self.called_args.add(option_string)
+        setattr(namespace, self.dest, values)

torchx/cli/cmd_base.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 argparse
 

torchx/cli/cmd_cancel.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 argparse
 import logging
 

torchx/cli/cmd_configure.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 argparse
 import logging
 import sys

torchx/cli/cmd_describe.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 argparse
 import dataclasses
 import logging

torchx/cli/cmd_list.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 argparse
 import logging
 

torchx/cli/cmd_log.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 argparse
 import logging
 import re

torchx/cli/cmd_run.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 argparse
 import logging
 import os
@@ -15,7 +17,7 @@ from pprint import pformat
 from typing import Dict, List, Optional, Tuple
 
 import torchx.specs as specs
-from torchx.cli.argparse_util import torchxconfig_run
+from torchx.cli.argparse_util import ArgOnceAction, torchxconfig_run
 from torchx.cli.cmd_base import SubCommand
 from torchx.cli.cmd_log import get_logs
 from torchx.runner import config, get_runner, Runner
@@ -131,6 +133,7 @@ class CmdRun(SubCommand):
             "-cfg",
             "--scheduler_args",
             type=str,
+            action=ArgOnceAction,
             help="Arguments to pass to the scheduler (Ex:`cluster=foo,user=bar`)."
             " For a list of scheduler run options run: `torchx runopts`",
         )
@@ -163,6 +166,7 @@ class CmdRun(SubCommand):
         subparser.add_argument(
             "--parent_run_id",
             type=str,
+            action=ArgOnceAction,
             help="optional parent run ID that this run belongs to."
             " It can be used to group runs for experiment tracking purposes",
         )

torchx/cli/cmd_runopts.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 argparse
 import logging
 

torchx/cli/cmd_status.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 argparse
 import logging
 import sys

torchx/cli/cmd_tracker.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 argparse
 import logging
 

torchx/cli/colors.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 sys
 
 # only print colors if outputting directly to a terminal

torchx/cli/main.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 logging
 import os
 import sys

torchx/components/__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
+
 """
 This module contains a collection of builtin TorchX components. The directory
 structure is organized by component category. Components are simply

torchx/components/component_test_base.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
+
 """
 You can unit test the component definitions as you would normal Python code
 since they are valid Python definitions.

torchx/components/dist.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
+
 """
 For distributed training, TorchX relies on the scheduler's gang scheduling
 capabilities to schedule ``n`` copies of nodes. Once launched, the application

torchx/components/integration_tests/component_provider.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 os
 import tempfile
 from abc import ABC, abstractmethod

torchx/components/integration_tests/integ_tests.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
 import logging
 import sys

torchx/components/serve.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
+
 """
 These components aim to make it easier to interact with inference and serving
 tools such as `torchserve <https://pytorch.org/serve/>`_.

torchx/components/structured_arg.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 methods for structured (higher order) component argument parsing.
 Use the functionalities defined in this module to author components

torchx/components/utils.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
+
 """
 This contains TorchX utility components that are `ready-to-use` out of the box. These are
 components that simply execute well known binaries (e.g. ``cp``)

torchx/examples/apps/datapreproc/datapreproc.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
+
 """
 Data Preprocessing App Example
 ====================================

torchx/examples/apps/lightning/data.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
+
 """
 Trainer Datasets Example
 ========================

torchx/examples/apps/lightning/model.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
+
 """
 Tiny ImageNet Model
 ====================

torchx/examples/apps/lightning/profiler.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
+
 """
 Simple Logging Profiler
 ===========================

torchx/examples/apps/lightning/train.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
+
 """
 Trainer Example
 =============================================

torchx/examples/pipelines/kfp/advanced_pipeline.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
+
 """
 Advanced KubeFlow Pipelines Example
 ===================================

torchx/examples/pipelines/kfp/dist_pipeline.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
+
 """
 Distributed KubeFlow Pipelines Example
 ======================================

torchx/examples/pipelines/kfp/intro_pipeline.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
+
 """
 Intro KubeFlow Pipelines Example
 ================================

torchx/notebook.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 TorchX utilities for creating and running components and apps from
 an Jupyter/IPython Notebook.

torchx/pipelines/kfp/__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 module contains adapters for converting TorchX components into KubeFlow
 Pipeline components.

torchx/pipelines/kfp/adapter.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 json
 import os
 import os.path
@@ -74,8 +76,7 @@ class ContainerFactory(Protocol):
     kfp.dsl.ContainerOp.
     """
 
-    def __call__(self, *args: object, **kwargs: object) -> dsl.ContainerOp:
-        ...
+    def __call__(self, *args: object, **kwargs: object) -> dsl.ContainerOp: ...
 
 
 class KFPContainerFactory(ContainerFactory, Protocol):

torchx/pipelines/kfp/version.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
+
 # Follows PEP-0440 version scheme guidelines
 # https://www.python.org/dev/peps/pep-0440/#version-scheme
 #

torchx/runner/__init__.py

@@ -5,4 +5,6 @@
 # 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 torchx.runner.api import get_runner, Runner  # noqa: F401 F403

torchx/runner/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 json
 import logging
 import os
@@ -131,7 +133,7 @@ class Runner:
         It is ok to call this method multiple times on the same runner object.
         """
 
-        for name, scheduler in self._scheduler_instances.items():
+        for scheduler in self._scheduler_instances.values():
             scheduler.close()
 
     def run_component(
@@ -654,7 +656,7 @@ class Runner:
     def _scheduler_app_id(
         self,
         app_handle: AppHandle,
-        check_session: bool = True
+        check_session: bool = True,
         # pyre-fixme[24]: SchedulerOpts is a generic, and we don't have access to the corresponding type
     ) -> Tuple[Scheduler, str, str]:
         """

torchx/runner/config.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
+
 """
 Status: Beta
 
@@ -266,6 +268,13 @@ def dump(
                         val = ";".join(opt.default)
                     else:
                         val = _NONE
+                elif opt.opt_type == Dict[str, str]:
+                    # deal with empty or None default lists
+                    if opt.default:
+                        # pyre-ignore[16] opt.default type checked already as Dict[str, str]
+                        val = ";".join([f"{k}:{v}" for k, v in opt.default.items()])
+                    else:
+                        val = _NONE
                 else:
                     val = f"{opt.default}"
 
@@ -525,6 +534,11 @@ def load(scheduler: str, f: TextIO, cfg: Dict[str, CfgVal]) -> None:
                         cfg[name] = config.getboolean(section, name)
                     elif runopt.opt_type is List[str]:
                         cfg[name] = value.split(";")
+                    elif runopt.opt_type is Dict[str, str]:
+                        cfg[name] = {
+                            s.split(":", 1)[0]: s.split(":", 1)[1]
+                            for s in value.replace(",", ";").split(";")
+                        }
                     else:
                         # pyre-ignore[29]
                         cfg[name] = runopt.opt_type(value)

torchx/runner/events/__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
+
 """
 Module contains events processing mechanisms that are integrated with the standard python logging.
 

torchx/runner/events/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 json
 from dataclasses import asdict, dataclass
 from enum import Enum

torchx/runner/events/handlers.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 logging
 from typing import Dict
 

torchx/runtime/tracking/__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
+
 """
 .. note:: EXPERIMENTAL, USE AT YOUR OWN RISK, APIs SUBJECT TO CHANGE
 

torchx/runtime/tracking/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 abc
 import json
 from typing import Dict, Union

torchx/schedulers/__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
+
 import importlib
 from typing import Dict, Mapping
 
@@ -19,6 +21,7 @@ DEFAULT_SCHEDULER_MODULES: Mapping[str, str] = {
     "kubernetes": "torchx.schedulers.kubernetes_scheduler",
     "kubernetes_mcad": "torchx.schedulers.kubernetes_mcad_scheduler",
     "aws_batch": "torchx.schedulers.aws_batch_scheduler",
+    "aws_sagemaker": "torchx.schedulers.aws_sagemaker_scheduler",
     "gcp_batch": "torchx.schedulers.gcp_batch_scheduler",
     "ray": "torchx.schedulers.ray_scheduler",
     "lsf": "torchx.schedulers.lsf_scheduler",
@@ -27,8 +30,7 @@ DEFAULT_SCHEDULER_MODULES: Mapping[str, str] = {
 
 class SchedulerFactory(Protocol):
     # pyre-fixme: Scheduler opts
-    def __call__(self, session_name: str, **kwargs: object) -> Scheduler:
-        ...
+    def __call__(self, session_name: str, **kwargs: object) -> Scheduler: ...
 
 
 def _defer_load_scheduler(path: str) -> SchedulerFactory:

torchx/schedulers/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 abc
 import re
 from dataclasses import dataclass, field

torchx/schedulers/aws_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 AWS Batch scheduler which can be used to run TorchX