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

torchx/tracker/__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
+
 """
 .. note:: PROTOTYPE, USE AT YOUR OWN RISK, APIs SUBJECT TO CHANGE
 
@@ -30,14 +32,14 @@ implementation.
 
 Example usage
 -------------
-Sample `code <https://github.com/pytorch/torchx/blob/main/torchx/examples/apps/tracker/main.py>`__ using tracker API.
+Sample `code <https://github.com/meta-pytorch/torchx/blob/main/torchx/examples/apps/tracker/main.py>`__ using tracker API.
 
 
 Tracker Setup
 -------------
 To enable tracking it requires:
 
-1. Defining tracker backends (entrypoints and configuration) on launcher side using :doc:`runner.config`
+1. Defining tracker backends (entrypoints/modules and configuration) on launcher side using :doc:`runner.config`
 2. Adding entrypoints within a user job using entry_points (`specification`_)
 
 .. _specification: https://packaging.python.org/en/latest/specifications/entry-points/
@@ -49,13 +51,13 @@ To enable tracking it requires:
 User can define any number of tracker backends under **torchx:tracker** section in :doc:`runner.config`, where:
    * Key: is an arbitrary name for the tracker, where the name will be used to configure its properties
         under [tracker:<TRACKER_NAME>]
-   * Value: is *entrypoint/factory method* that must be available within user job. The value will be injected into a
+   * Value: is *entrypoint* or *module* factory method that must be available within user job. The value will be injected into a
         user job and used to construct tracker implementation.
 
 .. code-block:: ini
 
     [torchx:tracker]
-    tracker_name=<entry_point>
+    tracker_name=<entry_point_or_module_factory_method>
 
 
 Each tracker can be additionally configured (currently limited to `config` parameter) under `[tracker:<TRACKER NAME>]` section:
@@ -71,11 +73,15 @@ For example, ~/.torchxconfig may be setup as:
 
     [torchx:tracker]
     tracker1=tracker1
-    tracker12=backend_2_entry_point
+    tracker2=backend_2_entry_point
+    tracker3=torchx.tracker.mlflow:create_tracker
 
     [tracker:tracker1]
     config=s3://my_bucket/config.json
 
+    [tracker:tracker3]
+    config=my_config.json
+
 
 2. User job configuration (Advanced)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -105,7 +111,7 @@ Use :py:meth:`~torchx.tracker.app_run_from_env`:
 Reference :py:class:`~torchx.tracker.api.TrackerBase` implementation
 --------------------------------------------------------------------
 :py:class:`~torchx.tracker.backend.fsspec.FsspecTracker` provides reference implementation of a tracker backend.
-GitHub example `directory <https://github.com/pytorch/torchx/blob/main/torchx/examples/apps/tracker/>`__ provides example on how to
+GitHub example `directory <https://github.com/meta-pytorch/torchx/blob/main/torchx/examples/apps/tracker/>`__ provides example on how to
 configure and use it in user application.
 
 

torchx/tracker/api.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
+
 from __future__ import annotations
 
 import logging
@@ -14,6 +16,7 @@ from functools import lru_cache
 from typing import Iterable, Mapping, Optional
 
 from torchx.util.entrypoints import load_group
+from torchx.util.modules import load_module
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -66,8 +69,7 @@ class AppRunTrackableSource:
     artifact_name: Optional[str]
 
 
-class Lineage:
-    ...
+class Lineage: ...
 
 
 class TrackerBase(ABC):
@@ -177,30 +179,26 @@ def _extract_tracker_name_and_config_from_environ() -> Mapping[str, Optional[str
 
 
 def build_trackers(
-    entrypoint_and_config: Mapping[str, Optional[str]]
+    factory_and_config: Mapping[str, Optional[str]]
 ) -> Iterable[TrackerBase]:
     trackers = []
 
-    entrypoint_factories = load_group("torchx.tracker")
+    entrypoint_factories = load_group("torchx.tracker") or {}
     if not entrypoint_factories:
-        logger.warning(
-            "No 'torchx.tracker' entry_points are defined. Tracking will not capture any data."
-        )
-        return trackers
+        logger.warning("No 'torchx.tracker' entry_points are defined.")
 
-    for entrypoint_key, config in entrypoint_and_config.items():
-        if entrypoint_key not in entrypoint_factories:
+    for factory_name, config in factory_and_config.items():
+        factory = entrypoint_factories.get(factory_name) or load_module(factory_name)
+        if not factory:
             logger.warning(
-                f"Could not find `{entrypoint_key}` tracker entrypoint. Skipping..."
+                f"No tracker factory `{factory_name}` found in entry_points or modules. See https://meta-pytorch.org/torchx/main/tracker.html#module-torchx.tracker"
             )
             continue
-        factory = entrypoint_factories[entrypoint_key]
         if config:
-            logger.info(f"Tracker config found for `{entrypoint_key}` as `{config}`")
-            tracker = factory(config)
+            logger.info(f"Tracker config found for `{factory_name}` as `{config}`")
         else:
-            logger.info(f"No tracker config specified for `{entrypoint_key}`")
-            tracker = factory(None)
+            logger.info(f"No tracker config specified for `{factory_name}`")
+        tracker = factory(config)
         trackers.append(tracker)
     return trackers
 
@@ -335,5 +333,4 @@ class AppRun:
 
         return model_run_sources
 
-    def children(self) -> Iterable[AppRun]:
-        ...
+    def children(self) -> Iterable[AppRun]: ...

torchx/tracker/backend/fsspec.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
+
 from __future__ import annotations
 
 import json

torchx/util/cuda.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 torch
 
 

torchx/util/datetime.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
+
 from datetime import datetime, timedelta
 
 

torchx/util/entrypoints.py

@@ -4,13 +4,14 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Any, Dict, Optional
+# pyre-strict
+# pyre-ignore-all-errors[3, 2, 16]
 
-import importlib_metadata as metadata
-from importlib_metadata import EntryPoint
+from importlib import metadata
+from importlib.metadata import EntryPoint
+from typing import Any, Dict, Optional
 
 
-# pyre-ignore-all-errors[3, 2]
 def load(group: str, name: str, default=None):
     """
     Loads the entry point specified by
@@ -28,13 +29,34 @@ def load(group: str, name: str, default=None):
     raises an error.
     """
 
-    entrypoints = metadata.entry_points().select(group=group)
+    # [note_on_entrypoints]
+    # return type of importlib.metadata.entry_points() is different between python-3.9 and python-3.10
+    # https://docs.python.org/3.9/library/importlib.metadata.html#importlib.metadata.entry_points
+    # https://docs.python.org/3.10/library/importlib.metadata.html#importlib.metadata.entry_points
+    if hasattr(metadata.entry_points(), "select"):
+        # python>=3.10
+        entrypoints = metadata.entry_points().select(group=group)
 
-    if name not in entrypoints.names and default is not None:
-        return default
+        if name not in entrypoints.names and default is not None:
+            return default
+
+        ep = entrypoints[name]
+        return ep.load()
 
-    ep = entrypoints[name]
-    return ep.load()
+    else:
+        # python<3.10 (e.g. 3.9)
+        # metadata.entry_points() returns dict[str, tuple[EntryPoint]] (not EntryPoints) in python-3.9
+        entrypoints = metadata.entry_points().get(group, ())
+
+        for ep in entrypoints:
+            if ep.name == name:
+                return ep.load()
+
+        # [group].name not found
+        if default is not None:
+            return default
+        else:
+            raise KeyError(f"entrypoint {group}.{name} not found")
 
 
 def _defer_load_ep(ep: EntryPoint) -> object:
@@ -47,11 +69,7 @@ def _defer_load_ep(ep: EntryPoint) -> object:
     return run
 
 
-# pyre-ignore-all-errors[3, 2]
-def load_group(
-    group: str,
-    default: Optional[Dict[str, Any]] = None,
-):
+def load_group(group: str, default: Optional[Dict[str, Any]] = None):
     """
     Loads all the entry points specified by ``group`` and returns
     the entry points as a map of ``name (str) -> deferred_load_fn``.
@@ -85,7 +103,13 @@ def load_group(
 
     """
 
-    entrypoints = metadata.entry_points().select(group=group)
+    # see [note_on_entrypoints] above
+    if hasattr(metadata.entry_points(), "select"):
+        # python>=3.10
+        entrypoints = metadata.entry_points().select(group=group)
+    else:
+        # python<3.10 (e.g. 3.9)
+        entrypoints = metadata.entry_points().get(group, ())
 
     if len(entrypoints) == 0:
         return default

torchx/util/io.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
+
 from os import path
 from pathlib import Path
 from typing import Optional

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.strip()), file=dst, end="\n", 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/modules.py

@@ -0,0 +1,65 @@
+# 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 importlib
+from types import ModuleType
+from typing import Callable, Optional, TypeVar, Union
+
+
+def load_module(path: str) -> Union[ModuleType, Optional[Callable[..., object]]]:
+    """
+    Loads and returns the module/module attr represented by the ``path``: ``full.module.path:optional_attr``
+
+    1. ``load_module("this.is.a_module:fn")`` -> equivalent to ``this.is.a_module.fn``
+    1. ``load_module("this.is.a_module")`` -> equivalent to ``this.is.a_module``
+    """
+    parts = path.split(":", 2)
+    module_path, method = parts[0], parts[1] if len(parts) > 1 else None
+    module = None
+    i, n = -1, len(module_path)
+    try:
+        while i < n:
+            i = module_path.find(".", i + 1)
+            i = i if i >= 0 else n
+            module = importlib.import_module(module_path[:i])
+        return getattr(module, method) if method else module
+    except Exception:
+        return None
+
+
+T = TypeVar("T")
+
+
+def import_attr(name: str, attr: str, default: T) -> T:
+    """
+    Imports ``name.attr`` and returns it if the module is found.
+    Otherwise, returns the specified ``default``.
+    Useful when getting an attribute from an optional dependency.
+
+    Note that the ``default`` parameter is intentionally not an optional
+    since this function is intended to be used with modules that may not be
+    installed as a dependency. Therefore the caller must ALWAYS provide a
+    sensible default.
+
+    Usage:
+
+    .. code-block:: python
+
+        aws_resources = import_attr("torchx.specs.named_resources_aws", "NAMED_RESOURCES", default={})
+        all_resources.update(aws_resources)
+
+    Raises:
+        AttributeError: If the module exists (e.g. can be imported)
+            but does not have an attribute with name ``attr``.
+    """
+    try:
+        mod = importlib.import_module(name)
+    except ModuleNotFoundError:
+        return default
+    else:
+        return getattr(mod, attr)

torchx/util/session.py

@@ -0,0 +1,42 @@
+#!/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 os
+import uuid
+from typing import Optional
+
+TORCHX_INTERNAL_SESSION_ID = "TORCHX_INTERNAL_SESSION_ID"
+
+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.
+    Please DO NOT use this function out of torchx codebase.
+    """
+    global CURRENT_SESSION_ID
+    if CURRENT_SESSION_ID:
+        return CURRENT_SESSION_ID
+    env_session_id = os.getenv(TORCHX_INTERNAL_SESSION_ID)
+    if env_session_id:
+        CURRENT_SESSION_ID = env_session_id
+        return CURRENT_SESSION_ID
+    session_id = str(uuid.uuid4())
+    CURRENT_SESSION_ID = session_id
+    return session_id
+
+
+def get_torchx_session_id() -> Optional[str]:
+    """
+    Returns the torchx session ID.
+    Please use this function to get the session ID out of torchx codebase.
+    """
+    return CURRENT_SESSION_ID

torchx/util/shlex.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 shlex
 from typing import Iterable
 

torchx/util/strings.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 re
 
 
@@ -11,7 +13,7 @@ def normalize_str(data: str) -> str:
     """
     Invokes ``lower`` on thes string and removes all
     characters that do not satisfy ``[a-z0-9\\-]`` pattern.
-    This method is mostly used to make sure kubernetes and gcp_batch scheduler gets
+    This method is mostly used to make sure kubernetes scheduler gets
     the job name that does not violate its restrictions.
     """
     if data.startswith("-"):