torchx-nightly 2024.1.6__py3-none-any.whl → 2025.12.24__py3-none-any.whl

torchx/cli/cmd_run.py

@@ -4,23 +4,29 @@
 # 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 json
 import logging
 import os
 import sys
 import threading
-from dataclasses import asdict
+from collections import Counter
+from dataclasses import asdict, dataclass, field, fields, MISSING as DATACLASS_MISSING
+from itertools import groupby
 from pathlib import Path
 from pprint import pformat
-from typing import Dict, List, Optional, Tuple
+from typing import Any, 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
+from torchx.specs import CfgVal, Workspace
 from torchx.specs.finder import (
     _Component,
     ComponentNotFoundException,
@@ -28,6 +34,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
 
 
@@ -35,10 +42,81 @@ MISSING_COMPONENT_ERROR_MSG = (
     "missing component name, either provide it from the CLI or in .torchxconfig"
 )
 
+LOCAL_SCHEDULER_WARNING_MSG = (
+    "`local` scheduler is deprecated and will be"
+    " removed in the near future,"
+    " please use other variants of the local scheduler"
+    " (e.g. `local_cwd`)"
+)
 
 logger: logging.Logger = logging.getLogger(__name__)
 
 
+@dataclass
+class TorchXRunArgs:
+    component_name: str
+    scheduler: str
+    scheduler_args: Dict[str, Any]
+    scheduler_cfg: Dict[str, CfgVal] = field(default_factory=dict)
+    dryrun: bool = False
+    wait: bool = False
+    log: bool = False
+    workspace: str = ""
+    parent_run_id: Optional[str] = None
+    tee_logs: bool = False
+    component_args: Dict[str, Any] = field(default_factory=dict)
+    component_args_str: List[str] = field(default_factory=list)
+
+
+def torchx_run_args_from_json(json_data: Dict[str, Any]) -> TorchXRunArgs:
+    all_fields = [f.name for f in fields(TorchXRunArgs)]
+    required_fields = {
+        f.name
+        for f in fields(TorchXRunArgs)
+        if f.default is DATACLASS_MISSING and f.default_factory is DATACLASS_MISSING
+    }
+    missing_fields = required_fields - json_data.keys()
+    if missing_fields:
+        raise ValueError(
+            f"The following required fields are missing: {', '.join(missing_fields)}"
+        )
+
+    # Fail if there are fields that aren't part of the run command
+    filtered_json_data = {k: v for k, v in json_data.items() if k in all_fields}
+    extra_fields = set(json_data.keys()) - set(all_fields)
+    if extra_fields:
+        raise ValueError(
+            f"The following fields are not part of the run command: {', '.join(extra_fields)}.",
+            "Please check your JSON and try launching again.",
+        )
+
+    torchx_args = TorchXRunArgs(**filtered_json_data)
+    if torchx_args.workspace == "":
+        torchx_args.workspace = f"{Path.cwd()}"
+    return torchx_args
+
+
+def torchx_run_args_from_argparse(
+    args: argparse.Namespace,
+    component_name: str,
+    component_args: List[str],
+    scheduler_cfg: Dict[str, CfgVal],
+) -> TorchXRunArgs:
+    return TorchXRunArgs(
+        component_name=component_name,
+        scheduler=args.scheduler,
+        scheduler_args={},
+        scheduler_cfg=scheduler_cfg,
+        dryrun=args.dryrun,
+        wait=args.wait,
+        log=args.log,
+        workspace=args.workspace,
+        parent_run_id=args.parent_run_id,
+        tee_logs=args.tee_logs,
+        component_args_str=component_args,
+    )
+
+
 def _parse_component_name_and_args(
     component_name_and_args: List[str],
     subparser: argparse.ArgumentParser,
@@ -82,6 +160,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)
 
@@ -114,6 +206,7 @@ class CmdBuiltins(SubCommand):
 class CmdRun(SubCommand):
     def __init__(self) -> None:
         self._subparser: Optional[argparse.ArgumentParser] = None
+        self._stdin_data_json: Optional[Dict[str, Any]] = None
 
     def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
         scheduler_names = get_scheduler_factories().keys()
@@ -131,6 +224,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`",
         )
@@ -156,46 +250,54 @@ class CmdRun(SubCommand):
         subparser.add_argument(
             "--workspace",
             "--buck-target",
-            default=f"file://{Path.cwd()}",
+            default=f"{Path.cwd()}",
             action=torchxconfig_run,
             help="local workspace to build/patch (buck-target of main binary if using buck)",
         )
         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",
         )
+        subparser.add_argument(
+            "--tee_logs",
+            action="store_true",
+            default=False,
+            help="Add additional prefix to log lines to indicate which replica is printing the log",
+        )
+        subparser.add_argument(
+            "--stdin",
+            action="store_true",
+            default=False,
+            help="Read JSON input from stdin to parse into torchx run args and run the component.",
+        )
         subparser.add_argument(
             "component_name_and_args",
             nargs=argparse.REMAINDER,
         )
 
-    def _run(self, runner: Runner, args: argparse.Namespace) -> None:
+    def _run_inner(self, runner: Runner, args: TorchXRunArgs) -> None:
         if args.scheduler == "local":
-            logger.warning(
-                "`local` scheduler is deprecated and will be"
-                " removed in the near future,"
-                " please use other variants of the local scheduler"
-                " (e.g. `local_cwd`)"
-            )
-
-        scheduler_opts = runner.scheduler_run_opts(args.scheduler)
-        cfg = scheduler_opts.cfg_from_str(args.scheduler_args)
-        config.apply(scheduler=args.scheduler, cfg=cfg)
+            logger.warning(LOCAL_SCHEDULER_WARNING_MSG)
 
-        component, component_args = _parse_component_name_and_args(
-            args.component_name_and_args,
-            none_throws(self._subparser),
+        config.apply(scheduler=args.scheduler, cfg=args.scheduler_cfg)
+        component_args = (
+            args.component_args_str
+            if args.component_args_str != []
+            else args.component_args
         )
         try:
+            workspace = Workspace.from_str(args.workspace) if args.workspace else None
+
             if args.dryrun:
                 dryrun_info = runner.dryrun_component(
-                    component,
+                    args.component_name,
                     component_args,
                     args.scheduler,
-                    workspace=args.workspace,
-                    cfg=cfg,
+                    workspace=workspace,
+                    cfg=args.scheduler_cfg,
                     parent_run_id=args.parent_run_id,
                 )
                 print(
@@ -206,40 +308,143 @@ class CmdRun(SubCommand):
                 print("\n=== SCHEDULER REQUEST ===\n" f"{dryrun_info}")
             else:
                 app_handle = runner.run_component(
-                    component,
+                    args.component_name,
                     component_args,
                     args.scheduler,
                     workspace=args.workspace,
-                    cfg=cfg,
+                    cfg=args.scheduler_cfg,
                     parent_run_id=args.parent_run_id,
                 )
                 # DO NOT delete this line. It is used by slurm tests to retrieve the app id
                 print(app_handle)
 
                 if args.scheduler.startswith("local"):
-                    self._wait_and_exit(runner, app_handle, log=True)
+                    self._wait_and_exit(
+                        runner, app_handle, log=True, tee_logs=args.tee_logs
+                    )
                 else:
                     logger.info(f"Launched app: {app_handle}")
                     app_status = runner.status(app_handle)
                     if app_status:
                         logger.info(app_status.format())
                     if args.wait or args.log:
-                        self._wait_and_exit(runner, app_handle, log=args.log)
+                        self._wait_and_exit(
+                            runner, app_handle, log=args.log, tee_logs=args.tee_logs
+                        )
 
         except (ComponentValidationException, ComponentNotFoundException) as e:
-            error_msg = f"\nFailed to run component `{component}` got errors: \n {e}"
+            error_msg = (
+                f"\nFailed to run component `{args.component_name}` got errors: \n {e}"
+            )
             logger.error(error_msg)
             sys.exit(1)
         except specs.InvalidRunConfigException as e:
             error_msg = (
-                f"Scheduler arg is incorrect or missing required option: `{e.cfg_key}`\n"
-                f"Run `torchx runopts` to check configuration for `{args.scheduler}` scheduler\n"
-                f"Use `-cfg` to specify run cfg as `key1=value1,key2=value2` pair\n"
-                "of setup `.torchxconfig` file, see: https://pytorch.org/torchx/main/experimental/runner.config.html"
+                "Invalid scheduler configuration: %s\n"
+                "To configure scheduler options, either:\n"
+                "  1. Use the `-cfg` command-line argument, e.g., `-cfg key1=value1,key2=value2`\n"
+                "  2. Set up a `.torchxconfig` file. For more details, visit: https://meta-pytorch.org/torchx/main/runner.config.html\n"
+                "Run `torchx runopts %s` to check all available configuration options for the "
+                "`%s` scheduler."
             )
-            logger.error(error_msg)
+            print(error_msg % (e, args.scheduler, args.scheduler), file=sys.stderr)
             sys.exit(1)
 
+    def _run_from_cli_args(self, runner: Runner, args: argparse.Namespace) -> None:
+        scheduler_opts = runner.scheduler_run_opts(args.scheduler)
+        cfg = scheduler_opts.cfg_from_str(args.scheduler_args)
+
+        component, component_args = _parse_component_name_and_args(
+            args.component_name_and_args,
+            none_throws(self._subparser),
+        )
+        torchx_run_args = torchx_run_args_from_argparse(
+            args, component, component_args, cfg
+        )
+        self._run_inner(runner, torchx_run_args)
+
+    def _run_from_stdin_args(self, runner: Runner, stdin_data: Dict[str, Any]) -> None:
+        torchx_run_args = torchx_run_args_from_json(stdin_data)
+        scheduler_opts = runner.scheduler_run_opts(torchx_run_args.scheduler)
+        cfg = scheduler_opts.cfg_from_json_repr(
+            json.dumps(torchx_run_args.scheduler_args)
+        )
+        torchx_run_args.scheduler_cfg = cfg
+        self._run_inner(runner, torchx_run_args)
+
+    def _get_torchx_stdin_args(
+        self, args: argparse.Namespace
+    ) -> Optional[Dict[str, Any]]:
+        if not args.stdin:
+            return None
+        if self._stdin_data_json is None:
+            self._stdin_data_json = self.torchx_json_from_stdin(args)
+        return self._stdin_data_json
+
+    def torchx_json_from_stdin(
+        self, args: Optional[argparse.Namespace] = None
+    ) -> Dict[str, Any]:
+        try:
+            stdin_data_json = json.load(sys.stdin)
+            if args and args.dryrun:
+                stdin_data_json["dryrun"] = True
+            if not isinstance(stdin_data_json, dict):
+                logger.error(
+                    "Invalid JSON input for `torchx run` command. Expected a dictionary."
+                )
+                sys.exit(1)
+            return stdin_data_json
+        except (json.JSONDecodeError, EOFError):
+            logger.error(
+                "Unable to parse JSON input for `torchx run` command, please make sure it's a valid JSON input."
+            )
+            sys.exit(1)
+
+    def verify_no_extra_args(self, args: argparse.Namespace) -> None:
+        """
+        Verifies that only --stdin was provided when using stdin mode.
+        """
+        if not args.stdin:
+            return
+
+        subparser = none_throws(self._subparser)
+        conflicting_args = []
+
+        # Check each argument against its default value
+        for action in subparser._actions:
+            if action.dest == "stdin":  # Skip stdin itself
+                continue
+            if action.dest == "help":  # Skip help
+                continue
+            if action.dest == "dryrun":  # Skip dryrun
+                continue
+
+            current_value = getattr(args, action.dest, None)
+            default_value = action.default
+
+            # For arguments that differ from default
+            if current_value != default_value:
+                # Handle special cases where non-default doesn't mean explicitly set
+                if action.dest == "component_name_and_args" and current_value == []:
+                    continue  # Empty list is still default
+                print(f"*********\n {default_value} = {current_value}")
+                conflicting_args.append(f"--{action.dest.replace('_', '-')}")
+
+        if conflicting_args:
+            subparser.error(
+                f"Cannot specify {', '.join(conflicting_args)} when using --stdin. "
+                "All configuration should be provided in JSON input."
+            )
+
+    def _run(self, runner: Runner, args: argparse.Namespace) -> None:
+        self.verify_no_extra_args(args)
+        if args.stdin:
+            stdin_data_json = self._get_torchx_stdin_args(args)
+            if stdin_data_json is not None:
+                self._run_from_stdin_args(runner, stdin_data_json)
+        else:
+            self._run_from_cli_args(runner, args)
+
     def run(self, args: argparse.Namespace) -> None:
         os.environ["TORCHX_CONTEXT_NAME"] = os.getenv("TORCHX_CONTEXT_NAME", "cli_run")
         component_defaults = load_sections(prefix="component")
@@ -247,10 +452,16 @@ class CmdRun(SubCommand):
         with get_runner(component_defaults=component_defaults) as runner:
             self._run(runner, args)
 
-    def _wait_and_exit(self, runner: Runner, app_handle: str, log: bool) -> None:
+    def _wait_and_exit(
+        self, runner: Runner, app_handle: str, log: bool, tee_logs: bool = False
+    ) -> None:
         logger.info("Waiting for the app to finish...")
 
-        log_thread = self._start_log_thread(runner, app_handle) if log else None
+        log_thread = (
+            self._start_log_thread(runner, app_handle, tee_logs_enabled=tee_logs)
+            if log
+            else None
+        )
 
         status = runner.wait(app_handle, wait_interval=1)
         if not status:
@@ -267,17 +478,30 @@ class CmdRun(SubCommand):
         else:
             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.daemon = True
+    def _start_log_thread(
+        self, runner: Runner, app_handle: str, tee_logs_enabled: bool = False
+    ) -> threading.Thread:
+        if tee_logs_enabled:
+            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(),
+            )
+        else:
+            thread = threading.Thread(
+                target=get_logs,
+                kwargs={
+                    "file": sys.stderr,
+                    "runner": runner,
+                    "identifier": app_handle,
+                    "regex": None,
+                    "should_tail": True,
+                },
+            )
+            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,7 +5,10 @@
 # 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 json
 import logging
 import sys
 from typing import List, Optional
@@ -44,6 +47,11 @@ class CmdStatus(SubCommand):
         subparser.add_argument(
             "--roles", type=str, default="", help="comma separated roles to filter"
         )
+        subparser.add_argument(
+            "--json",
+            action="store_true",
+            help="output the status in JSON format",
+        )
 
     def run(self, args: argparse.Namespace) -> None:
         app_handle = args.app_handle
@@ -52,7 +60,10 @@ class CmdStatus(SubCommand):
         app_status = runner.status(app_handle)
         filter_roles = parse_list_arg(args.roles)
         if app_status:
-            print(app_status.format(filter_roles))
+            if args.json:
+                print(json.dumps(app_status.to_json(filter_roles)))
+            else:
+                print(app_status.format(filter_roles))
         else:
             logger.error(
                 f"AppDef: {app_id},"

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
 
@@ -43,7 +45,7 @@ class CmdTracker(SubCommand):
         else:
             raise RuntimeError(
                 "No trackers configured."
-                " See: https://pytorch.org/torchx/latest/runtime/tracking.html"
+                " See: https://meta-pytorch.org/torchx/latest/runtime/tracking.html"
             )
 
     def add_list_job_arguments(self, subparser: argparse.ArgumentParser) -> None:

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
@@ -14,6 +16,7 @@ import torchx
 from torchx.cli.cmd_base import SubCommand
 from torchx.cli.cmd_cancel import CmdCancel
 from torchx.cli.cmd_configure import CmdConfigure
+from torchx.cli.cmd_delete import CmdDelete
 from torchx.cli.cmd_describe import CmdDescribe
 from torchx.cli.cmd_list import CmdList
 from torchx.cli.cmd_log import CmdLog
@@ -35,6 +38,7 @@ def get_default_sub_cmds() -> Dict[str, SubCommand]:
         "builtins": CmdBuiltins(),
         "cancel": CmdCancel(),
         "configure": CmdConfigure(),
+        "delete": CmdDelete(),
         "describe": CmdDescribe(),
         "list": CmdList(),
         "log": CmdLog(),

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
@@ -179,7 +181,7 @@ To validate that you've defined your component correctly you can either:
 
 1. (easiest) Dryrun your component's ``--help`` with the cli: ``torchx run --dryrun ~/component.py:train --help``
 2. Use the component :ref:`linter<specs:Component Linter>`
-   (see `dist_test.py <https://github.com/pytorch/torchx/blob/main/torchx/components/test/dist_test.py>`_ as an example)
+   (see `dist_test.py <https://github.com/meta-pytorch/torchx/blob/main/torchx/components/test/dist_test.py>`_ as an example)
 
 
 Running as a Job
@@ -296,13 +298,6 @@ imagine the component is defined as:
    * ``*args=["--help"]``: ``torchx run comp.py:f -- --help``
    * ``*args=["--i", "2"]``: ``torchx run comp.py:f --i 1 -- --i 2``
 
-Run in a Pipeline
---------------------------------
-
-The :ref:`torchx.pipelines<pipelines:torchx.pipelines>` define adapters that
-convert a torchx component into the object that represents a pipeline "stage" in the
-target pipeline platform (see :ref:`Pipelines` for a list of supported pipeline orchestrators).
-
 Additional Resources
 -----------------------
 

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.