Metadata-Version: 2.3
Name: aidot-training-telemetry
Version: 1.1.2
Summary: Utilities for monitoring training of large foundation models
License: NVIDIA Mission Control, License no 744-SW7001+P1CMI36
Author: Stefania Alborghetti
Author-email: salborghetti@nvidia.com
Requires-Python: >=3.10,<3.15
Classifier: Development Status :: 3 - Alpha
Classifier: Environment :: Console
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Information Technology
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Natural Language :: English
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Classifier: Topic :: Software Development :: Libraries
Classifier: Topic :: Utilities
Provides-Extra: jax-deps
Provides-Extra: torch-deps
Requires-Dist: jax[cpu] (>=0.4.0,<1.0.0) ; extra == "jax-deps"
Requires-Dist: lightning (>=2.5.6,<3.0.0) ; extra == "torch-deps"
Requires-Dist: numpy (>=1.24.0,<3.0.0) ; extra == "torch-deps" or extra == "jax-deps"
Requires-Dist: nvtx (>=0.2.5,<1.0.0)
Requires-Dist: opentelemetry-api (>=1.27.0,<2.0.0)
Requires-Dist: opentelemetry-exporter-otlp-proto-grpc (>=1.27.0,<2.0.0)
Requires-Dist: opentelemetry-exporter-otlp-proto-http (>=1.27.0,<2.0.0)
Requires-Dist: opentelemetry-sdk (>=1.27.0,<2.0.0)
Requires-Dist: pyyaml (>=6.0.2,<7.0.0)
Requires-Dist: torch (>=2.9.1,<3.0.0) ; extra == "torch-deps"
Requires-Dist: tzlocal (>=5.3,<6.0)
Description-Content-Type: text/markdown

# Training Telemetry

A Python library that records events, metrics, and errors during model training in standardized formats:

* structured key=value logs
* JSON for text files
* Open Telemetry (OTEL) traces and logs
* NVTX code markers for Nvidia Nsight Systems

## Overview

The objective of this library is to provide a standard format for logging events, metrics and errors that can be adopted by existing frameworks and applications for training large AI models. The end result is that the runtime performance and errors of these training models can be monitored in a consistent manner, without impacting the training performance. Time spans provide detailed information on how each training process spends its time during startup, training and checkpoint saving. Errors can be analyzed and correlated with infrastructure events once the application fails, in order to provide users with more actionable information.

This library is lightweight and intentionally has very few dependencies, so as to facilitate integration with training frameworks that normally have a long list of dependencies. The API is provided on two levels:

* A context-based API, where monitoring can be done via context managers or function decorators
* A low-level recorder API with start/stop/event/error functions for callback implementations and other low-level requirements

The following events are currently supported:

- Application runtime and application-specific metrics
- Training loop progress and timing
- Individual iteration metrics, including loss, accuracy, TFLOPS, consumed samples, forward and backward times
- Checkpoint saves, including global and local checkpoints, async and sync checkpoint strategies
- Errors and exceptions  
- Model validation and testing
- Custom metrics and events

Events are logged by one or more of the following backends:

* A Python logger backend, logging events as messages using a logger at INFO level with structured log format
* A file logger backend, where each event is logged as a one-line JSON object
* An OpenTelemetry backend, where each event is converted to a span and sent to the OTEL collector

Events have metrics attached to them. A special class of events, error events, captures error messages and stack traces.

## Key Features

- Context managers for timing code blocks
- Event recording with customizable metrics
- Exception handling and error reporting
- Flexible backend system for storing/analyzing telemetry data as log messages, JSON objects or OTEL traces
- Low overhead monitoring

## Installation

The library package is available on the following pypi public repositories:

* https://pypi.org/project/aidot-training-telemetry/
* https://pypi.nvidia.com/aidot-training-telemetry/

Install with:

```
pip install aidot-training-telemetry
```

If using Poetry, run the following command:

```bash
poetry add aidot-training-telemetry
```


## Usage

Using the context API, initialize the main function with:

```
def get_application_metrics():
    return ApplicationMetrics.create(
        rank=get_rank_index(),
        world_size=get_rank_count(),
        node_name="localhost",
        timezone=str(get_localzone()),
        total_iterations=num_epochs * len(dataloader),
        checkpoint_enabled=True,
        checkpoint_strategy="sync",
    )


@application_running(metrics=get_application_metrics())
def main():
    [...]
```

This will capture any exceptions not handled by the application, and log them as an error event before re-raising them.

For the training loop and iterations:

```
with training_iteration() as training_iteration_span:
    [...]
    training_iteration_span.add_metrics(
        IterationMetrics.create(
            current_iteration=current_iteration,
            num_iterations=len(dataloader),
            loss=loss.item(),
            accuracy=accuracy.item(),
        )
    )
```

For checkpoint monitoring:

```
with checkpoint_save() as checkpoint_save_span:
    [...]
    checkpoint_save_span.add_metrics(
        CheckpointSaveMetrics.create(
            checkpoint_type=CheckPointType.LOCAL,
            current_iteration=current_iteration,
            checkpoint_directory=temp_dir,
            checkpoint_filename=os.path.basename(checkpoint_file_name),
        )
    )

```

For a concrete example refer to the [torch example](training_telemetry/torch/example.py) or usage examples.

It's also possible to manually create spans and events, refer to the [recorder](training_telemetry/recorder.py) API for how to do this.









