torchx-nightly 2025.2.14__py3-none-any.whl → 2025.2.16__py3-none-any.whl

torchx/cli/cmd_run.py

@@ -21,6 +21,7 @@ from typing import Dict, List, Optional, Tuple
 import torchx.specs as specs
 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
@@ -186,6 +187,12 @@ class CmdRun(SubCommand):
             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(
             "component_name_and_args",
             nargs=argparse.REMAINDER,
@@ -237,14 +244,18 @@ class CmdRun(SubCommand):
                 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}"
@@ -267,10 +278,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:
@@ -287,15 +304,30 @@ class CmdRun(SubCommand):
         else:
             logger.debug(status)
 
-    def _start_log_thread(self, runner: Runner, app_handle: str) -> threading.Thread:
-        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(),
-        )
+    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_nightly-2025.2.14.dist-info → torchx_nightly-2025.2.16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: torchx-nightly
-Version: 2025.2.14
+Version: 2025.2.16
 Summary: TorchX SDK and Components
 Home-page: https://github.com/pytorch/torchx
 Author: TorchX Devs

{torchx_nightly-2025.2.14.dist-info → torchx_nightly-2025.2.16.dist-info}/RECORD

@@ -16,7 +16,7 @@ torchx/cli/cmd_configure.py,sha256=1kTv0qbsbV44So74plAySwWu56pQrqjhfW_kbfdC3Rw,1
 torchx/cli/cmd_describe.py,sha256=E5disbHoKTsqYKp2s3DaFW9GDLCCOgdOc3pQoHKoyCs,1283
 torchx/cli/cmd_list.py,sha256=BVqHEW2oTEJ3GqcFK7c1K-i2R-DUjaXQ-WBr0meeIGM,1429
 torchx/cli/cmd_log.py,sha256=v-EZYUDOcG95rEgTnrsmPJMUyxM9Mk8YFAJtUxtgViE,5475
-torchx/cli/cmd_run.py,sha256=OYAp0Rp_YxdS8FLrRxoWnYX24VIsNXaGiUHAyisEXYI,11185
+torchx/cli/cmd_run.py,sha256=4M1JJc7YmEa5T_2OFakCwCwiP0Ibpy-3zcLp1arrj9w,12203
 torchx/cli/cmd_runopts.py,sha256=NWZiP8XpQjfTDJgays2c6MgL_8wxFoeDge6NstaZdKk,1302
 torchx/cli/cmd_status.py,sha256=ubtmCp4PylrIh_kC3ZJ5QJm7lzXRt_aRPmY7j-sZu_0,1836
 torchx/cli/cmd_tracker.py,sha256=RfLxE4Cq1wfk7k051RtZ8RPJp0pEKSCa3KmTeRs3LF8,5218
@@ -115,9 +115,9 @@ torchx/workspace/__init__.py,sha256=FqN8AN4VhR1C_SBY10MggQvNZmyanbbuPuE-JCjkyUY,
 torchx/workspace/api.py,sha256=PtDkGTC5lX03pRoYpuMz2KCmM1ZOycRP1UknqvNb97Y,6341
 torchx/workspace/dir_workspace.py,sha256=npNW_IjUZm_yS5r-8hrRkH46ndDd9a_eApT64m1S1T4,2268
 torchx/workspace/docker_workspace.py,sha256=PFu2KQNVC-0p2aKJ-W_BKA9ZOmXdCY2ABEkCExp3udQ,10269
-torchx_nightly-2025.2.14.dist-info/LICENSE,sha256=WVHfXhFC0Ia8LTKt_nJVYobdqTJVg_4J3Crrfm2A8KQ,1721
-torchx_nightly-2025.2.14.dist-info/METADATA,sha256=V_4xZJYmKyLdAAzX_4cv4m7G3qSUcJiqOUUeTnxaH2E,6169
-torchx_nightly-2025.2.14.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-torchx_nightly-2025.2.14.dist-info/entry_points.txt,sha256=T328AMXeKI3JZnnxfkEew2ZcMN1oQDtkXjMz7lkV-P4,169
-torchx_nightly-2025.2.14.dist-info/top_level.txt,sha256=pxew3bc2gsiViS0zADs0jb6kC5v8o_Yy_85fhHj_J1A,7
-torchx_nightly-2025.2.14.dist-info/RECORD,,
+torchx_nightly-2025.2.16.dist-info/LICENSE,sha256=WVHfXhFC0Ia8LTKt_nJVYobdqTJVg_4J3Crrfm2A8KQ,1721
+torchx_nightly-2025.2.16.dist-info/METADATA,sha256=5EK5Tjh6Ju12-2gDZvHVNu3CP-7xBMPLDF8i6HfmqOI,6169
+torchx_nightly-2025.2.16.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+torchx_nightly-2025.2.16.dist-info/entry_points.txt,sha256=T328AMXeKI3JZnnxfkEew2ZcMN1oQDtkXjMz7lkV-P4,169
+torchx_nightly-2025.2.16.dist-info/top_level.txt,sha256=pxew3bc2gsiViS0zADs0jb6kC5v8o_Yy_85fhHj_J1A,7
+torchx_nightly-2025.2.16.dist-info/RECORD,,