torchx-nightly 2024.9.8__py3-none-any.whl → 2024.9.13__py3-none-any.whl

torchx/cli/cmd_log.py

@@ -23,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
 
@@ -39,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,
@@ -167,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

@@ -21,7 +21,6 @@ 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
@@ -32,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
 
 
@@ -288,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/runner/events/__init__.py

@@ -27,6 +27,7 @@ from types import TracebackType
 from typing import Dict, Optional, Type
 
 from torchx.runner.events.handlers import get_logging_handler
+from torchx.util.session import get_session_id_or_create_new
 
 from .api import SourceType, TorchxEvent  # noqa F401
 
@@ -136,7 +137,7 @@ class log_event:
         workspace: Optional[str] = None,
     ) -> TorchxEvent:
         return TorchxEvent(
-            session=app_id or "",
+            session=get_session_id_or_create_new(),
             scheduler=scheduler,
             api=api,
             app_id=app_id,

torchx/runner/events/api.py

@@ -25,7 +25,7 @@ class TorchxEvent:
     The class represents the event produced by ``torchx.runner`` api calls.
 
     Arguments:
-        session: Session id that was used to execute request.
+        session: Session id of the current run
         scheduler: Scheduler that is used to execute request
         api: Api name
         app_id: Unique id that is set by the underlying scheduler

torchx/util/log_tee_helpers.py

@@ -0,0 +1,210 @@
+# 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
+
+"""
+If you're wrapping the TorchX API with your own CLI, these functions can
+help show the logs of the job within your CLI, just like
+`torchx log`
+"""
+
+import logging
+import threading
+from queue import Queue
+from typing import List, Optional, TextIO, Tuple, TYPE_CHECKING
+
+from torchx.util.types import none_throws
+
+if TYPE_CHECKING:
+    from torchx.runner.api import Runner
+    from torchx.schedulers.api import Stream
+    from torchx.specs.api import AppDef
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+# A torchX job can have stderr/stdout for many replicas, of many roles
+# The scheduler API has functions that allow us to get,
+# with unspecified detail, the log lines of a given replica of
+# a given role.
+#
+# So, to neatly tee the results, we:
+# 1) Determine every role ID / replica ID pair we want to monitor
+# 2) Request the given stderr / stdout / combined streams from them (1 thread each)
+# 3) Concatenate each of those streams to a given destination file
+
+
+def _find_role_replicas(
+    app: "AppDef",
+    role_name: Optional[str],
+) -> List[Tuple[str, int]]:
+    """
+    Enumerate all (role, replica id) pairs in the given AppDef.
+    Replica IDs are 0-indexed, and range up to num_replicas,
+    for each role.
+    If role_name is provided, filters to only that name.
+    """
+    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
+
+
+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_for_role_replica(
+    dst: TextIO,
+    app_handle: str,
+    regex: Optional[str],
+    runner: "Runner",
+    which_role: str,
+    which_replica: int,
+    exceptions: "Queue[Exception]",
+    should_tail: bool,
+    streams: Optional["Stream"],
+    colorize: bool = False,
+) -> None:
+    """
+    Helper function that'll run in parallel - one
+    per monitored replica of a given role.
+
+    Based on print_log_lines .. but not designed for TTY
+    """
+    try:
+        for line in runner.log_lines(
+            app_handle,
+            which_role,
+            which_replica,
+            regex,
+            should_tail=should_tail,
+            streams=streams,
+        ):
+            if colorize:
+                color_begin = "\033[32m"
+                color_end = "\033[0m"
+            else:
+                color_begin = ""
+                color_end = ""
+            prefix = f"{color_begin}{which_role}/{which_replica}{color_end} "
+            print(_prefix_line(prefix, line), file=dst, end="", flush=True)
+    except Exception as e:
+        exceptions.put(e)
+        raise
+
+
+def _start_threads_to_monitor_role_replicas(
+    dst: TextIO,
+    app_handle: str,
+    regex: Optional[str],
+    runner: "Runner",
+    which_role: Optional[str] = None,
+    should_tail: bool = False,
+    streams: Optional["Stream"] = None,
+    colorize: bool = False,
+) -> None:
+    threads = []
+
+    app = none_throws(runner.describe(app_handle))
+    replica_ids = _find_role_replicas(app, role_name=which_role)
+
+    # Holds exceptions raised by all threads, in a thread-safe
+    # object
+    exceptions = Queue()
+
+    if not replica_ids:
+        valid_roles = [role.name for role in app.roles]
+        raise ValueError(
+            f"{which_role} is not a valid role name. Available: {valid_roles}"
+        )
+
+    for role_name, replica_id in replica_ids:
+        threads.append(
+            threading.Thread(
+                target=_print_log_lines_for_role_replica,
+                kwargs={
+                    "dst": dst,
+                    "runner": runner,
+                    "app_handle": app_handle,
+                    "which_role": role_name,
+                    "which_replica": replica_id,
+                    "regex": regex,
+                    "should_tail": should_tail,
+                    "exceptions": exceptions,
+                    "streams": streams,
+                    "colorize": colorize,
+                },
+                daemon=True,
+            )
+        )
+
+    for t in threads:
+        t.start()
+
+    for t in threads:
+        t.join()
+
+    # Retrieve all exceptions, print all except one and raise the first recorded exception
+    threads_exceptions = []
+    while not exceptions.empty():
+        threads_exceptions.append(exceptions.get())
+
+    if len(threads_exceptions) > 0:
+        for i in range(1, len(threads_exceptions)):
+            logger.error(threads_exceptions[i])
+
+        raise threads_exceptions[0]
+
+
+def tee_logs(
+    dst: TextIO,
+    app_handle: str,
+    regex: Optional[str],
+    runner: "Runner",
+    should_tail: bool = False,
+    streams: Optional["Stream"] = None,
+    colorize: bool = False,
+) -> threading.Thread:
+    """
+    Makes a thread, which in turn will start 1 thread per replica
+    per role, that tees that role-replica's logs to the given
+    destination buffer.
+
+    You'll need to start and join with this parent thread.
+
+    dst:  TextIO to tee the logs into
+    app_handle: The return value of runner.run() or runner.schedule()
+    regex: Regex to filter the logs that are tee-d
+    runner: The Runner you used to schedule the job
+    should_tail: If true, continue until we run out of logs. Otherwise, just fetch
+                 what's available
+    streams: Whether to fetch STDERR, STDOUT, or the temporally COMBINED (default) logs
+    """
+    thread = threading.Thread(
+        target=_start_threads_to_monitor_role_replicas,
+        kwargs={
+            "dst": dst,
+            "runner": runner,
+            "app_handle": app_handle,
+            "regex": None,
+            "should_tail": True,
+            "colorize": colorize,
+        },
+        daemon=True,
+    )
+    return thread

torchx/util/session.py

@@ -0,0 +1,26 @@
+#!/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
+
+import uuid
+from typing import Optional
+
+CURRENT_SESSION_ID: Optional[str] = None
+
+
+def get_session_id_or_create_new() -> str:
+    """
+    Returns the current session ID, or creates a new one if none exists.
+    The session ID remains the same as long as it is in the same process.
+    """
+    global CURRENT_SESSION_ID
+    if CURRENT_SESSION_ID:
+        return CURRENT_SESSION_ID
+    session_id = str(uuid.uuid4())
+    CURRENT_SESSION_ID = session_id
+    return session_id

{torchx_nightly-2024.9.8.dist-info → torchx_nightly-2024.9.13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: torchx-nightly
-Version: 2024.9.8
+Version: 2024.9.13
 Summary: TorchX SDK and Components
 Home-page: https://github.com/pytorch/torchx
 Author: TorchX Devs
@@ -62,6 +62,7 @@ Requires-Dist: torchtext==0.17.1; extra == "dev"
 Requires-Dist: torchvision==0.17.1; extra == "dev"
 Requires-Dist: ts==0.5.1; extra == "dev"
 Requires-Dist: ray[default]; extra == "dev"
+Requires-Dist: wheel; extra == "dev"
 Requires-Dist: lintrunner; extra == "dev"
 Requires-Dist: lintrunner-adapters; extra == "dev"
 Requires-Dist: grpcio==1.62.1; extra == "dev"

{torchx_nightly-2024.9.8.dist-info → torchx_nightly-2024.9.13.dist-info}/RECORD

@@ -15,8 +15,8 @@ torchx/cli/cmd_cancel.py,sha256=NKfOCu_44Lch9vliGSQ0Uv6BVqpUqj7Tob652TI-ua4,835
 torchx/cli/cmd_configure.py,sha256=1kTv0qbsbV44So74plAySwWu56pQrqjhfW_kbfdC3Rw,1722
 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=Xh5vrsbwyV_ppwurrENGBNKxc1XLVbFC6YH1b8jlHAM,6104
-torchx/cli/cmd_run.py,sha256=T2TmkVZNbIljCEmVQHAQjzVIkHkOH_BV5dlY_4ErrUs,11220
+torchx/cli/cmd_log.py,sha256=v-EZYUDOcG95rEgTnrsmPJMUyxM9Mk8YFAJtUxtgViE,5475
+torchx/cli/cmd_run.py,sha256=OYAp0Rp_YxdS8FLrRxoWnYX24VIsNXaGiUHAyisEXYI,11185
 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
@@ -58,8 +58,8 @@ torchx/pipelines/kfp/version.py,sha256=mYBxd6bm4MeR34D--xo-JLQ9wHeAl_ZQLwbItCf9t
 torchx/runner/__init__.py,sha256=x8Sz7s_tLxPgJgvWIhK4ju9BNZU61uBFywGwDY6CqJs,315
 torchx/runner/api.py,sha256=17ZFEYeuIK0kUrTUysdgsFaEVc6_-ggvBuSoabQnrcA,28625
 torchx/runner/config.py,sha256=fTdCcf-MKlBg6MzXopF4W0hYyDDoPAuvZs2v2bKzwG0,17849
-torchx/runner/events/__init__.py,sha256=KxvodAlq9ZDUIKzdlzHuyQWa0aoIGKTtBfE7vjM6GRk,4570
-torchx/runner/events/api.py,sha256=UTdExV7cMR5kC0nQIViwTlyblhBF3D5bbuqd4-kWrkE,2542
+torchx/runner/events/__init__.py,sha256=TtzBLZ9oaHKfr689R4NnjCii3G8kxiRafe7Q0jRE5_k,4649
+torchx/runner/events/api.py,sha256=_6mjS5B5FcSMEea1U0M-fD4jErmbpZ1gQOyUHV-4dqY,2527
 torchx/runner/events/handlers.py,sha256=ThHCIJW21BfBgB7b6ftyjASJmD1KdizpjuTtsyqnvJs,522
 torchx/runtime/__init__.py,sha256=Wxje2BryzeQneFu5r6P9JJiEKG-_C9W1CcZ_JNrKT6g,593
 torchx/runtime/tracking/__init__.py,sha256=dYnAPnrXYREfPXkpHhdOFkcYIODWEbA13PdD-wLQYBo,3055
@@ -105,7 +105,9 @@ torchx/util/cuda.py,sha256=-ZTa1WCLnY2WtSWAdWufLQqZSDCZfZsloBuiS84LIkU,1099
 torchx/util/datetime.py,sha256=hV6Sg0u5KTBe68yrmy_RGCC5su0i4Tb_mAYphWamiXI,405
 torchx/util/entrypoints.py,sha256=4rqmA81XYLj4Kk7GboJi0z78h4NIQxSrcOzDuuTwCkw,2725
 torchx/util/io.py,sha256=HNpWLcFUX0WTAP3CsdamHz--FR5A4kSdLCPfNqa2UkA,1807
+torchx/util/log_tee_helpers.py,sha256=yJ4ODU3QE0t7ZgYHrE-HiwlOhfg8BK1e4oVaw_zMa2E,6320
 torchx/util/modules.py,sha256=LRTuZRH5bbRr0ZaCtCtvKbgwhMoPsTx-GokWbCLGPdk,1131
+torchx/util/session.py,sha256=__4DFoJqpTzqc99FPlKyLPDQFn2g3rtu0dFb5b-Z-7E,721
 torchx/util/shlex.py,sha256=eXEKu8KC3zIcd8tEy9_s8Ds5oma8BORr-0VGWNpG2dk,463
 torchx/util/strings.py,sha256=GkLWCmYS89Uv6bWc5hH0XwvHy7oQmprv2U7axC4A2e8,678
 torchx/util/types.py,sha256=een55pV-N8aBc3qUBjHRc1llJcX10JVa19pB8dBE8No,7564
@@ -113,9 +115,9 @@ torchx/workspace/__init__.py,sha256=FqN8AN4VhR1C_SBY10MggQvNZmyanbbuPuE-JCjkyUY,
 torchx/workspace/api.py,sha256=1heBmPgB-W5Zf9gwViM7NrqvHpZlVYeMN7jpY8Qkytc,5479
 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-2024.9.8.dist-info/LICENSE,sha256=WVHfXhFC0Ia8LTKt_nJVYobdqTJVg_4J3Crrfm2A8KQ,1721
-torchx_nightly-2024.9.8.dist-info/METADATA,sha256=GDLO7VCHMTXs4ZwHgCRkSf6pWFBw3p08VMPhfqL7dEA,6131
-torchx_nightly-2024.9.8.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-torchx_nightly-2024.9.8.dist-info/entry_points.txt,sha256=T328AMXeKI3JZnnxfkEew2ZcMN1oQDtkXjMz7lkV-P4,169
-torchx_nightly-2024.9.8.dist-info/top_level.txt,sha256=pxew3bc2gsiViS0zADs0jb6kC5v8o_Yy_85fhHj_J1A,7
-torchx_nightly-2024.9.8.dist-info/RECORD,,
+torchx_nightly-2024.9.13.dist-info/LICENSE,sha256=WVHfXhFC0Ia8LTKt_nJVYobdqTJVg_4J3Crrfm2A8KQ,1721
+torchx_nightly-2024.9.13.dist-info/METADATA,sha256=8rzgUn_Ac9ga2suLgjBnQwrS86iCpGmcao1CC9Oixjk,6169
+torchx_nightly-2024.9.13.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+torchx_nightly-2024.9.13.dist-info/entry_points.txt,sha256=T328AMXeKI3JZnnxfkEew2ZcMN1oQDtkXjMz7lkV-P4,169
+torchx_nightly-2024.9.13.dist-info/top_level.txt,sha256=pxew3bc2gsiViS0zADs0jb6kC5v8o_Yy_85fhHj_J1A,7
+torchx_nightly-2024.9.13.dist-info/RECORD,,