torchx.tracker

Tip

Trackers record artifacts, metadata, and lineage for training runs. Use AppRun inside your job and register custom backends via entry points.

Prerequisites: Basic Concepts (core concepts). For registering custom tracker backends, see Registering Custom Trackers in the Advanced Usage guide.

Overview & Usage

Note

PROTOTYPE, USE AT YOUR OWN RISK, APIs SUBJECT TO CHANGE

Practitioners running ML jobs often need to track information such as:

• Job inputs:

  • configuration

    • model configuration

    • HPO parameters

  • data

    • version

    • sources

• Job results:

  • metrics

  • model location

• Conceptual job groupings

AppRun provides a uniform interface as an experiment and artifact tracking solution that supports wrapping pluggable tracking implementations by providing TrackerBase adapter implementation.

Example usage

Sample code using tracker API.

Tracker Setup

To enable tracking it requires:

1. Defining tracker backends (entrypoints/modules and configuration) on launcher side using .torchxconfig

2. Adding entrypoints within a user job using entry_points (specification)

1. Launcher side configuration

User can define any number of tracker backends under torchx:tracker section in .torchxconfig, 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 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.

[torchx:tracker]
tracker_name=<entry_point_or_module_factory_method>

Each tracker can be additionally configured (currently limited to config parameter) under [tracker:<TRACKER NAME>] section:

[tracker:<TRACKER NAME>]
config=configvalue

For example, ~/.torchxconfig may be setup as:

[torchx:tracker]
tracker1=tracker1
tracker2=backend_2_entry_point
tracker3=my_module.trackers:create_tracker

[tracker:tracker1]
config=s3://my_bucket/config.json

[tracker:tracker3]
config=my_config.json

2. User job configuration (Advanced)

Entrypoint value defined in the previous step must be discoverable under [torchx.tracker] group and callable within user job (depending on packaging/distribution mechanism) to create an instance of the TrackerBase.

To accomplish that define entrypoint in the distribution in entry_points.txt as:

[torchx.tracker]
entry_point_name=my_module:create_tracker_fn

Acquiring AppRun instance

Use app_run_from_env():

>>> import os; os.environ["TORCHX_JOB_ID"] = "scheduler://session/job_id" # Simulate running job first
>>> from torchx.tracker import app_run_from_env
>>> app_run = app_run_from_env()

Reference TrackerBase implementation

FsspecTracker provides reference implementation of a tracker backend. GitHub example directory provides example on how to configure and use it in user application.

Querying data

• CmdTracker exposes operations available to users at the CLI level:

  • torchx tracker list jobs [–parent-run-id RUN_ID]

  • torchx tracker list metadata RUN_ID

  • torchx tracker list artifacts [–artifact ARTIFACT_NAME] RUN_ID

• Alternatively, backend implementations may expose UI for user consumption.

Trackers operate at two levels:

• Backend level (TrackerBase) – the storage implementation. TorchX ships with FsspecTracker (filesystem-based). You can implement your own backend.

• Job level (AppRun) – the user-facing API that delegates to one or more TrackerBase backends. AppRun is constructed automatically from environment variables set by the TorchX runner (TORCHX_JOB_ID, TORCHX_TRACKERS).

Typical usage inside a training job:

from torchx.tracker.api import AppRun

# Singleton created from TORCHX_JOB_ID and TORCHX_TRACKERS env vars
run = AppRun.run_from_env()

# Store metadata (key-value pairs)
run.add_metadata(lr=0.001, epochs=10, model="resnet50")

# Store an artifact (named path + optional metadata)
run.add_artifact("checkpoint", "s3://bucket/checkpoints/epoch_10.pt")

# Link a parent run for lineage tracking
run.add_source("local_cwd://torchx/parent_job_123")

Implementing a Custom Tracker

Subclass TrackerBase and implement its eight abstract methods. Then provide a factory function and register it as a torchx.tracker entry point. See Registering Custom Trackers for the full walkthrough.

Testing Your Tracker

Study torchx/tracker/test/api_test.py for a complete in-memory example. A minimal test writes metadata and artifacts, then reads them back:

import unittest

class MyTrackerTest(unittest.TestCase):
    def test_round_trip_metadata(self) -> None:
        tracker = MyTracker(connection_str="test://localhost")
        tracker.add_metadata("run-1", lr=0.01, epochs=10)
        md = tracker.metadata("run-1")
        self.assertEqual(md["lr"], 0.01)
        self.assertEqual(md["epochs"], 10)

    def test_round_trip_artifact(self) -> None:
        tracker = MyTracker(connection_str="test://localhost")
        tracker.add_artifact("run-1", "checkpoint", "/path/to/ckpt.pt")
        arts = tracker.artifacts("run-1")
        self.assertIn("checkpoint", arts)
        self.assertEqual(arts["checkpoint"].path, "/path/to/ckpt.pt")

Test factory wiring with patch.dict:

import os
from unittest.mock import patch

@patch.dict(os.environ, {
    "TORCHX_TRACKERS": "my_tracker",
    "TORCHX_TRACKER_MY_TRACKER_CONFIG": "test://localhost",
    "TORCHX_JOB_ID": "test-run-1",
})
def test_tracker_from_env(self) -> None:
    from torchx.tracker.api import trackers_from_environ
    trackers = list(trackers_from_environ())
    self.assertEqual(len(trackers), 1)

Common Pitfalls

• Entry point targets the class, not the factory: The entry point must reference a factory function (config: str | None) -> TrackerBase, not the class itself.

• Factory signature mismatch: The factory receives config: str | None, not keyword arguments. Parse connection strings or JSON inside the factory.

• Forgetting to handle ``None`` config: When no TORCHX_TRACKER_<NAME>_CONFIG env var is set, config is None. Provide a sensible default or raise a clear error.

API Reference

class torchx.tracker.AppRun(id: str, backends: Iterable[TrackerBase])

Job-level tracker API that delegates to one or more TrackerBase backends.

Warning

This API is experimental and may change significantly.

>>> from torchx.tracker.api import AppRun
>>> run = AppRun(id="my_job_123", backends=[])
>>> run.add_metadata(lr=0.01, epochs=10)  # no-op with empty backends
>>> run.job_id()
'my_job_123'

add_artifact(name: str, path: str, metadata: Mapping[str, object] | None = None) → None

Store an artifact (name, path, optional metadata) for this run.

add_metadata(**kwargs: object) → None

Store key-value metadata for this run.

add_source(source_id: str, artifact_name: str | None = None) → None

Link a source (TorchX run or external entity) to this run for lineage tracking.

job_id() → str

Return the run ID.

static run_from_env() → AppRun

Create a singleton AppRun from environment variables.

Reads TORCHX_JOB_ID and TORCHX_TRACKERS (set by the torchx runner). Returns a cached singleton so all callers share the same tracker backends.

Note

When not launched via torchx, returns an empty AppRun with job_id="<UNDEFINED>" and no backends (write methods become no-ops).

>>> from torchx.tracker.api import AppRun
>>> apprun = AppRun.run_from_env()
>>> apprun.add_metadata(md_1="foo", md_2="bar")

sources() → Iterable[AppRunTrackableSource]

Return source links for this run (queries the first backend).

class torchx.tracker.api.TrackerBase

Abstract base for tracker backend implementations.

Warning

This API is experimental and may change significantly.

abstractmethod add_artifact(run_id: str, name: str, path: str, metadata: Mapping[str, object] | None = None) → None

Add an artifact with the given name, path, and optional metadata.

abstractmethod add_metadata(run_id: str, **kwargs: object) → None

Store arbitrary key-value metadata for the given run.

abstractmethod add_source(run_id: str, source_id: str, artifact_name: str | None = None) → None

Link a source run (lineage) to the given run.

abstractmethod artifacts(run_id: str) → Mapping[str, TrackerArtifact]

Return all artifacts for the given run.

abstractmethod lineage(run_id: str) → Lineage

Return full lineage (parents and consumers) for the given run.

abstractmethod metadata(run_id: str) → Mapping[str, object]

Return metadata for the given run.

abstractmethod run_ids(**kwargs: str) → Iterable[str]

Return run IDs, optionally filtered by keyword arguments.

abstractmethod sources(run_id: str, artifact_name: str | None = None) → Iterable[TrackerSource]

Return sources for the given run, optionally filtered by artifact_name.

Data Types

class torchx.tracker.api.TrackerArtifact(name: str, path: str, metadata: Mapping[str, object] | None)

An artifact stored by a backend tracker (name, path, and optional metadata).

class torchx.tracker.api.TrackerSource(source_run_id: str, artifact_name: str | None)

A source link at the backend tracker level.

source_run_id is a TorchX handle or external entity ID. artifact_name classifies the relationship (used for filtering).

Built-in Backends

class torchx.tracker.backend.fsspec.FsspecTracker(fs: AbstractFileSystem, root_dir: str)

Bases: TrackerBase

Implements TrackerBase using Fsspec abstraction and has an advantage of using various storage options for persisting the data.

Important: torchx.tracker.api API is still experimental, hence there are no backwards compatibility gurantees with future releases yet.

Each run will have a directory with subdirs for metadata, artifact, source and descendants data.

add_artifact(run_id: str, name: str, path: str, metadata: Mapping[str, object] | None = None) → None

Add an artifact with the given name, path, and optional metadata.

add_metadata(run_id: str, **metadata: object) → None

Store arbitrary key-value metadata for the given run.

add_source(run_id: str, source_id: str, artifact_name: str | None = None) → None

Link a source run (lineage) to the given run.

artifacts(run_id: str) → Mapping[str, TrackerArtifact]

Return all artifacts for the given run.

lineage(run_id: str) → Lineage

Return full lineage (parents and consumers) for the given run.

metadata(run_id: str) → Mapping[str, object]

Return metadata for the given run.

run_ids(parent_run_id: str | None = None) → Iterable[str]

Return run IDs, optionally filtered by keyword arguments.

sources(run_id: str, artifact_name: str | None = None) → Iterable[TrackerSource]

Return sources for the given run, optionally filtered by artifact_name.

CLI Command

class torchx.cli.cmd_tracker.CmdTracker

Prototype TorchX tracker subcommand that allows querying data by interacting with tracker implementation.

Important: commands and the arguments may be modified in the future.

Supported commands:

• tracker list jobs [–parent-run-id RUN_ID]

• tracker list metadata RUN_ID

• tracker list artifacts [–artifact ARTIFACT_NAME] RUN_ID

add_arguments(subparser: ArgumentParser) → None

Adds the arguments to this sub command

run(args: Namespace) → None

Runs the sub command. Parsed arguments are available as args.

Environment Variables

Set automatically by the runner when trackers are configured.

Variable	Purpose
TORCHX_JOB_ID	The AppHandle for the current job. Set by the runner and used by AppRun.run_from_env() to identify the run.
TORCHX_TRACKERS	Comma-separated list of tracker entry-point keys to activate (e.g. fsspec,my_tracker).
TORCHX_TRACKER_<NAME>_CONFIG	Per-tracker configuration string passed to the factory function. <NAME> is the upper-cased entry-point key.
TORCHX_PARENT_RUN_ID	Optional parent run ID for lineage tracking. Set by the runner; read by AppRun.run_from_env() which automatically calls tracker.add_source() on each backend to record the lineage link.

See also

Advanced Usage

Entry-point registration for custom trackers, schedulers, and components.

Tracking

Runtime tracking utilities for use within applications.