torchx-nightly 2024.2.11__py3-none-any.whl → 2025.1.14__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
@@ -21,6 +23,10 @@ from torchx.runner import get_runner, Runner
 from torchx.schedulers.api import Stream
 from torchx.specs.api import is_started
 from torchx.specs.builders import make_app_handle
+from torchx.util.log_tee_helpers import (
+    _find_role_replicas as find_role_replicas,
+    _prefix_line,
+)
 
 from torchx.util.types import none_throws
 
@@ -37,19 +43,6 @@ def validate(job_identifier: str) -> None:
         sys.exit(1)
 
 
-def _prefix_line(prefix: str, line: str) -> str:
-    """
-    _prefix_line ensure the prefix is still present even when dealing with return characters
-    """
-    if "\r" in line:
-        line = line.replace("\r", f"\r{prefix}")
-    if "\n" in line[:-1]:
-        line = line[:-1].replace("\n", f"\n{prefix}") + line[-1:]
-    if not line.startswith("\r"):
-        line = f"{prefix}{line}"
-    return line
-
-
 def print_log_lines(
     file: TextIO,
     runner: Runner,
@@ -165,17 +158,6 @@ def get_logs(
         raise threads_exceptions[0]
 
 
-def find_role_replicas(
-    app: specs.AppDef, role_name: Optional[str]
-) -> List[Tuple[str, int]]:
-    role_replicas = []
-    for role in app.roles:
-        if role_name is None or role_name == role.name:
-            for i in range(role.num_replicas):
-                role_replicas.append((role.name, i))
-    return role_replicas
-
-
 class CmdLog(SubCommand):
     def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
         subparser.add_argument(

torchx/cli/cmd_run.py

@@ -4,20 +4,23 @@
 # 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
 import sys
 import threading
+from collections import Counter
 from dataclasses import asdict
+from itertools import groupby
 from pathlib import Path
 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
 from torchx.runner.config import load_sections
 from torchx.schedulers import get_default_scheduler_name, get_scheduler_factories
@@ -28,6 +31,7 @@ from torchx.specs.finder import (
     get_builtin_source,
     get_components,
 )
+from torchx.util.log_tee_helpers import tee_logs
 from torchx.util.types import none_throws
 
 
@@ -82,6 +86,20 @@ def _parse_component_name_and_args(
             component = args[0]
             component_args = args[1:]
 
+    # Error if there are repeated command line arguments each group of arguments,
+    # where the groups are separated by "--"
+    arg_groups = [list(g) for _, g in groupby(component_args, key=lambda x: x == "--")]
+    for arg_group in arg_groups:
+        all_options = [
+            x
+            for x in arg_group
+            if x.startswith("-") and x.strip() != "-" and x.strip() != "--"
+        ]
+        arg_count = Counter(all_options)
+        duplicates = [arg for arg, count in arg_count.items() if count > 1]
+        if len(duplicates) > 0:
+            subparser.error(f"Repeated Command Line Arguments: {duplicates}")
+
     if not component:
         subparser.error(MISSING_COMPONENT_ERROR_MSG)
 
@@ -131,6 +149,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 +182,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",
         )
@@ -268,16 +288,14 @@ class CmdRun(SubCommand):
             logger.debug(status)
 
     def _start_log_thread(self, runner: Runner, app_handle: str) -> threading.Thread:
-        thread = threading.Thread(
-            target=get_logs,
-            kwargs={
-                "file": sys.stderr,
-                "runner": runner,
-                "identifier": app_handle,
-                "regex": None,
-                "should_tail": True,
-            },
+        thread = tee_logs(
+            dst=sys.stderr,
+            app_handle=app_handle,
+            regex=None,
+            runner=runner,
+            should_tail=True,
+            streams=None,
+            colorize=not sys.stderr.closed and sys.stderr.isatty(),
         )
-        thread.daemon = True
         thread.start()
         return thread

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
 ========================
@@ -62,17 +64,17 @@ class ImageFolderSamplesDataset(datasets.ImageFolder):
 # our trainer and other components that need to load data.
 
 
-# pyre-fixme[13]: Attribute `test_ds` is never initialized.
-# pyre-fixme[13]: Attribute `train_ds` is never initialized.
-# pyre-fixme[13]: Attribute `val_ds` is never initialized.
 class TinyImageNetDataModule(pl.LightningDataModule):
     """
     TinyImageNetDataModule is a pytorch LightningDataModule for the tiny
     imagenet dataset.
     """
 
+    # pyre-fixme[13]: Attribute `test_ds` is never initialized.
     train_ds: ImageFolderSamplesDataset
+    # pyre-fixme[13]: Attribute `train_ds` is never initialized.
     val_ds: ImageFolderSamplesDataset
+    # pyre-fixme[13]: Attribute `val_ds` is never initialized.
     test_ds: ImageFolderSamplesDataset
 
     def __init__(

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
 ===========================
@@ -17,18 +19,19 @@ output is used for HPO optimization with Ax.
 import time
 from typing import Dict
 
-from pytorch_lightning.loggers.base import LightningLoggerBase
-from pytorch_lightning.profiler.base import BaseProfiler
+from pytorch_lightning.loggers.logger import Logger
+
+from pytorch_lightning.profilers.profiler import Profiler
 
 
-class SimpleLoggingProfiler(BaseProfiler):
+class SimpleLoggingProfiler(Profiler):
     """
     This profiler records the duration of actions (in seconds) and reports the
     mean duration of each action to the specified logger. Reported metrics are
     in the format `duration_<event>`.
     """
 
-    def __init__(self, logger: LightningLoggerBase) -> None:
+    def __init__(self, logger: Logger) -> None:
         super().__init__()
 
         self.current_actions: Dict[str, float] = {}

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
 ======================================
@@ -58,7 +60,7 @@ with open("pipeline.yaml", "rt") as f:
 # a kfp.Client.
 #
 # See the
-# `KFP SDK Examples <https://www.kubeflow.org/docs/components/pipelines/tutorials/sdk-examples/#examples>`_
+# `KFP SDK Examples <https://www.kubeflow.org/docs/components/pipelines/legacy-v1/tutorials/sdk-examples/#examples>`_
 # for more info on launching KFP pipelines.
 
 # %%

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
 ================================
@@ -71,7 +73,7 @@ with open("pipeline.yaml", "rt") as f:
 # a kfp.Client.
 #
 # See the
-# `KFP SDK Examples <https://www.kubeflow.org/docs/components/pipelines/tutorials/sdk-examples/#examples>`_
+# `KFP SDK Examples <https://www.kubeflow.org/docs/components/pipelines/legacy-v1/tutorials/sdk-examples/#examples>`_
 # for more info on launching KFP pipelines.
 
 # %%

torchx/examples/torchx_out_of_sync_training.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# 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 file is for lex team to ramp up on torchX OSS. It's for training purpose.
+# Please DO NOT modify this file unless you know what you are doing.

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.