Metadata-Version: 2.4
Name: sagemaker-hyperpod-observability-sdk
Version: 1.0.0
Summary: Python library for emitting training events from machine learning workloads on Amazon SageMaker HyperPod clusters
Project-URL: Homepage, https://github.com/aws/sagemaker-hyperpod-observability-sdk
Project-URL: Repository, https://github.com/aws/sagemaker-hyperpod-observability-sdk
Project-URL: Issues, https://github.com/aws/sagemaker-hyperpod-observability-sdk/issues
Project-URL: Documentation, https://github.com/aws/sagemaker-hyperpod-observability-sdk/tree/main/packages/sagemaker-hyperpod-observability-sdk#readme
Author: Amazon Web Services
License-Expression: Apache-2.0
License-File: LICENSE
Keywords: hyperpod,machine-learning,observability,opentelemetry,sagemaker,training-metrics
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: System :: Monitoring
Classifier: Typing :: Typed
Requires-Python: >=3.11
Provides-Extra: dev
Requires-Dist: keras>=3.0.0; extra == 'dev'
Requires-Dist: lightning>=2.0.0; extra == 'dev'
Requires-Dist: mypy>=1.0.0; extra == 'dev'
Requires-Dist: opentelemetry-api>=1.20.0; extra == 'dev'
Requires-Dist: opentelemetry-exporter-otlp>=1.20.0; extra == 'dev'
Requires-Dist: opentelemetry-sdk>=1.20.0; extra == 'dev'
Requires-Dist: pre-commit>=3.0.0; extra == 'dev'
Requires-Dist: pytest-cov>=4.0.0; extra == 'dev'
Requires-Dist: pytest-mock>=3.10.0; extra == 'dev'
Requires-Dist: pytest-xdist>=3.0.0; extra == 'dev'
Requires-Dist: pytest>=7.0.0; extra == 'dev'
Requires-Dist: ruff>=0.1.0; extra == 'dev'
Requires-Dist: tensorflow>=2.16.0; extra == 'dev'
Provides-Extra: frameworks
Requires-Dist: keras>=3.0.0; extra == 'frameworks'
Requires-Dist: lightning>=2.0.0; extra == 'frameworks'
Requires-Dist: tensorflow>=2.16.0; extra == 'frameworks'
Provides-Extra: otlp
Requires-Dist: opentelemetry-api>=1.20.0; extra == 'otlp'
Requires-Dist: opentelemetry-exporter-otlp>=1.20.0; extra == 'otlp'
Requires-Dist: opentelemetry-sdk>=1.20.0; extra == 'otlp'
Provides-Extra: quality
Requires-Dist: mypy>=1.0.0; extra == 'quality'
Requires-Dist: pre-commit>=3.0.0; extra == 'quality'
Requires-Dist: ruff>=0.1.0; extra == 'quality'
Provides-Extra: test
Requires-Dist: pytest-cov>=4.0.0; extra == 'test'
Requires-Dist: pytest-mock>=3.10.0; extra == 'test'
Requires-Dist: pytest-xdist>=3.0.0; extra == 'test'
Requires-Dist: pytest>=7.0.0; extra == 'test'
Description-Content-Type: text/markdown

# Observability SDK for Training Performance Metrics

The SageMaker HyperPod Observability SDK is a Python library for emitting training events from machine learning workloads running on Amazon SageMaker HyperPod clusters. It provides OpenTelemetry-compliant event structures and flexible export capabilities to give you deep visibility into your training jobs.

---

## Overview

HyperPod's observability add-on delivers infrastructure-level training metrics out of the box. The Observability SDK extends this capability by adding application-level visibility from inside your training process, capturing batch timing, checkpoint duration, validation phases, and training lifecycle boundaries.

When the observability add-on is enabled with advanced training metrics, events emitted by the SDK are processed into metrics and surfaced on the Grafana Training dashboard under the **Core Training Performance** and **Overall Training Health** metrics sections. Together, the infrastructure-level and SDK-level metrics provide a complete, end-to-end picture of your training performance.

---

## Requirements

Before using the SDK, ensure the following prerequisites are met:

- Python 3.11 or later
- A HyperPod EKS cluster with the [observability add-on](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-hyperpod-observability-addon.html) installed and advanced training metrics enabled
- The latest version of the observability add-on (required for SDK metrics dashboard panels)

---

## Key Features

- **Zero external dependencies** for the core SDK. OTLP export requires the `[otlp]` extra.
- **Built-in framework adapters** for PyTorch Lightning and Keras. Any other framework, including PyTorch, XGBoost, JAX, or custom training loops, is supported through manual instrumentation.
- **OpenTelemetry-compliant events** emitted as LogRecords with trace and span correlation for paired START/END events.

---

## Frameworks Supported

The following table summarizes instrumentation support across frameworks:

| Framework | Auto | Programmatic | Manual |
|---|---|---|---|
| PyTorch Lightning | ✅ | ✅ | ✅ |
| Keras (TF / PyTorch / JAX backend) | - | ✅ | ✅ |
| PyTorch | - | - | ✅ |
| XGBoost, JAX, custom frameworks | - | - | ✅ |

## Installation

Install the core SDK (stdout export only) or include the OTLP exporter for full pipeline integration:

```bash
pip install sagemaker-hyperpod-observability-sdk           # core (stdout only)
pip install sagemaker-hyperpod-observability-sdk[otlp]     # with OTLP exporter
```

---

## Quick Start

### Auto-instrumentation (PyTorch Lightning)

No code changes are required. Set the following environment variable and run your training script:

```bash
export HYPERPOD_OBSERVABILITY_AUTO_INSTRUMENT=true
python train.py
```

The SDK automatically detects PyTorch Lightning and registers callbacks for batch, validation, checkpoint, and training lifecycle events.

### Programmatic Instrumentation

If you prefer to register adapters explicitly, use the following approach for your framework.

**PyTorch Lightning:**

```python
from sagemaker_hyperpod_observability_sdk import AutoInstrumentor, Config, FrameworkType
from sagemaker_hyperpod_observability_sdk.common.types import ExporterType

config = Config(exporters=[ExporterType.OTLP], auto_instrument=True)
instrumentor = AutoInstrumentor.create(config)
instrumentor.start()

callback = instrumentor.get_adapter(FrameworkType.PYTORCH_LIGHTNING).get_callback()
trainer = Trainer(callbacks=[callback])
trainer.fit(model, train_loader, val_loader)
```

**Keras:**

```python
from sagemaker_hyperpod_observability_sdk import Config, KerasAdapter
from sagemaker_hyperpod_observability_sdk.common.types import ExporterType
from sagemaker_hyperpod_observability_sdk.core.event_processor import EventProcessor

config = Config(exporters=[ExporterType.OTLP])
event_processor = EventProcessor.create(config)
callback = KerasAdapter.create(event_processor).get_callback()

model.fit(X, y, callbacks=[callback])
```

### Manual Instrumentation (PyTorch, XGBoost, JAX, and custom frameworks)

For frameworks without a built-in adapter, use `ManualInstrumentor` to emit events directly from your training loop.

**Using a context manager:**

```python
from sagemaker_hyperpod_observability_sdk import Config
from sagemaker_hyperpod_observability_sdk.common.types import ExporterType
from sagemaker_hyperpod_observability_sdk.instrumentors.manual_instrumentor import ManualInstrumentor

config = Config(exporters=[ExporterType.OTLP])
instrumentor = ManualInstrumentor.create(config)

for epoch in range(num_epochs):
    for batch in train_loader:
        with instrumentor.training_context_manager("BATCH"):
            loss = model(batch)
            loss.backward()
            optimizer.step()

    with instrumentor.training_context_manager("VALIDATION"):
        validate(model, val_loader)

instrumentor.close()
```

**Using a decorator:**

```python
@instrumentor.training_decorator("BATCH")
def train_step(batch):
    loss = model(batch)
    loss.backward()
    optimizer.step()
    return loss
```

**Emitting events directly:**

```python
instrumentor.emit_event("CHECKPOINT_START", global_step=1000)
save_checkpoint(model, path)
instrumentor.emit_event("CHECKPOINT_SYNC_END", global_step=1000)
```

---

## Configuration

Configuration follows this precedence order: **programmatic config > environment variables > JSON file > defaults**.

**JSON configuration file** (`sagemaker-hyperpod-observability.json`):

```json
{
  "exporters": ["otlp"],
  "frameworks": ["pytorch_lightning"],
  "auto_instrument": true,
  "otlp_config": {
    "endpoint": "hyperpod-otel-collector.hyperpod-observability:4317",
    "timeout": 30,
    "insecure": true
  }
}
```

**Programmatic configuration:**

```python
from sagemaker_hyperpod_observability_sdk import Config
from sagemaker_hyperpod_observability_sdk.common.types import ExporterType
from sagemaker_hyperpod_observability_sdk.config.config_models import OTLPConfig

config = Config(
    exporters=[ExporterType.OTLP],
    auto_instrument=True,
    otlp_config=OTLPConfig(
        endpoint="hyperpod-otel-collector.hyperpod-observability:4317",
        insecure=True,
    ),
)
```

---

## Supported Events

The SDK emits the following event types. Paired START/END events carry trace and span correlation for accurate duration measurement.

| Event | Description | Lightning | Keras | Manual |
|---|---|---|---|---|
| BATCH_START / BATCH_END | Training batch boundaries | ✅ | ✅ | ✅ |
| TRAINING_START / TRAINING_END | Training session lifecycle | ✅ | ✅ | ✅ |
| VALIDATION_START / VALIDATION_END | Validation phase boundaries | ✅ | ✅ | ✅ |
| CHECKPOINT_START / CHECKPOINT_SYNC_END | Synchronous checkpoint | - | - | ✅ |
| CHECKPOINT_SYNC_END / CHECKPOINT_ASYNC_END | Asynchronous checkpoint | - | - | ✅ |
| CHECKPOINT_SAVED | Checkpoint saved (unpaired event) | ✅ | - | ✅ |
| EXCEPTION_DETECTED | Application-level fault | ✅ | - | ✅ |
| Custom events | User-defined events | - | - | ✅ |

---

## Dashboard Metrics

When the observability add-on is enabled with advanced training metrics, events emitted by the SDK are processed into metrics and surfaced on the Grafana Training dashboard under the **Core Training Performance** and **Overall Training Health** metrics sections. If you are on the older version of the observability add-on, upgrade the add-on to the latest version to make these panels visible.

| Metric | Description |
|---|---|
| training_batch_count | The number of training batches completed |
| training_batch_time_ms | Cumulative time spent executing training batches, in milliseconds |
| training_validation_count | Total number of validation phases successfully completed |
| training_validation_time_ms | Cumulative time spent running validation phases, in milliseconds |
| training_checkpoint_count | Total number of checkpoint operations performed during training |
| training_checkpoint_time_ms | Total time taken to complete all checkpoint operations, including both synchronous and asynchronous phases, in milliseconds |
| training_checkpoint_sync_time_ms | Time spent in the synchronous portion of checkpoint operations, in milliseconds |
| training_checkpoint_async_time_ms | Time spent in the asynchronous portion of checkpoint operations, in milliseconds |
| training_checkpoint_saved_count | Total number of checkpoints recorded via CHECKPOINT_SAVED events |
| training_lost_step_count | Number of training steps that were lost due to faults or unexpected restarts |
| training_reprocessing_time_ms | Time spent re-running training steps after recovering from a fault, in milliseconds |
| training_fault_detection_time_ms | Elapsed time between the last successfully completed batch and the point at which a fault was detected, in milliseconds |
| training_system_restart_time_ms | Time elapsed between fault recovery and the resumption of training, in milliseconds |

### Accessing dashboards

To view your SageMaker HyperPod cluster's metrics in Amazon Managed Grafana, perform the following steps:

- Open the Amazon SageMaker AI console at https://console.aws.amazon.com/sagemaker/.
- Go to your cluster's details page.
- On the **Dashboard** tab, locate the **HyperPod Observability** section, and choose **Open dashboard in Grafana**.

#### Adding new users to an Amazon Managed Grafana workspace

For information about how to add users to an Amazon Managed Grafana workspace, see ["Use AWS IAM Identity Center with your Amazon Managed Grafana workspace"](https://docs.aws.amazon.com/grafana/latest/userguide/authentication-in-AMG-SSO.html) in the *Amazon Managed Grafana User Guide*.

---

## Additional Resources

- [Examples](./examples) of runnable scripts covering all instrumentation approaches
- [Development guide](./DEVELOPING.md) for build system, testing, and project structure
- [Contributing guidelines](./CONTRIBUTING.md) for contributing and reporting issues

---

## Security

For information on reporting security issues, see [CONTRIBUTING](./CONTRIBUTING.md#security-issue-notifications).

---

## License

This SDK is licensed under the Apache-2.0 License. See [LICENSE](./LICENSE) for full details.
