# higgsmeter

Prometheus-based activity and state tracking for GPU worker services.

## Overview

`higgsmeter` provides a lightweight, thread-safe state tracker for GPU workers, enabling detailed observability of worker operations through Prometheus metrics.

## Features

- **Mutually Exclusive States**: Ensures workers are in exactly one state at a time
- **Automatic Time Tracking**: Records time spent in each state and operation
- **Context Manager Support**: Easy state transitions with automatic reversion
- **Decorator Support**: Track state for sync and async functions
- **Operation Labels**: Track specific operations within each state
- **Prometheus Integration**: Exposes standard Prometheus metrics
- **Multi-rank Aware**: Enable tracking only on rank 0 in distributed setups

## Installation

```bash
# From PyPI
pip install higgsmeter

# Or with uv
uv pip install higgsmeter
```

## Quick Start

```python
from higgsmeter import enable_tracker, state_tracker, AppState, track_state

# Initialize tracker (call once at startup)
enable_tracker(rank=0)

# Manual state tracking
state_tracker.set_state(AppState.IDLE, operation="waiting_for_work")

# Context manager (auto-reverts)
with state_tracker.scope(AppState.GPU, operation="inference"):
    # GPU work here
    run_inference()

# Decorator (most convenient)
@track_state(AppState.IO)
def download_data(url: str):
    # I/O operation
    fetch(url)

# Cleanup at shutdown
state_tracker.set_state(AppState.SHUTDOWN)
state_tracker.finalize()
```

## Available States

- `AppState.IDLE` - Worker idle, waiting for tasks
- `AppState.IO` - I/O operations (network, disk)
- `AppState.CPU` - CPU-intensive processing
- `AppState.GPU` - GPU computation
- `AppState.WARMUP` - Model warmup/initialization
- `AppState.STARTUP` - Service startup
- `AppState.SHUTDOWN` - Graceful shutdown

## Prometheus Metrics

### `worker_activity_time_seconds_total` (Counter)

Total cumulative seconds spent in each state/operation combination.

**Labels:**
- `state`: The application state (idle, io, cpu, gpu, warmup, startup, shutdown)
- `operation`: Specific operation name (function name or custom label)

**Example:**
```promql
worker_activity_time_seconds_total{state="gpu",operation="inference"} 1234.5
worker_activity_time_seconds_total{state="io",operation="SqsClient.receive_message"} 45.2
```

### `worker_current_state` (Gauge)

Current state indicator (numeric value for easy alerting).

**Values:**
- `0` = Idle
- `1` = I/O
- `2` = CPU
- `3` = GPU
- `4` = Warmup
- `5` = Startup
- `6` = Shutdown

## Usage Patterns

### Pattern 1: Service Lifecycle

```python
from higgsmeter import enable_tracker, state_tracker, AppState

# Initialize
rank = get_rank()  # From your distributed framework
enable_tracker(rank)

# Startup
state_tracker.set_state(AppState.STARTUP)

# Load models
with state_tracker.scope(AppState.WARMUP, "model_loading"):
    load_models()

# Main loop
state_tracker.set_state(AppState.IDLE, "consuming_messages")
while not shutdown:
    message = queue.receive()
    if message:
        with state_tracker.scope(AppState.CPU, "processing"):
            process(message)

# Shutdown
state_tracker.set_state(AppState.SHUTDOWN)
state_tracker.finalize()
```

### Pattern 2: Nested Operations

```python
@track_state(AppState.CPU)
def process_job(data):
    # Parse input (CPU)
    parsed = parse(data)

    # Fetch additional data (I/O)
    with state_tracker.scope(AppState.IO, "fetch_metadata"):
        metadata = fetch_metadata(parsed.id)

    # Run inference (GPU)
    with state_tracker.scope(AppState.GPU, "inference"):
        result = model.predict(parsed)

    # Upload result (I/O)
    with state_tracker.scope(AppState.IO, "upload_result"):
        upload(result)
```

### Pattern 3: Client Decorators

```python
from higgsmeter import track_state, AppState

class S3Client:
    @track_state(AppState.IO)  # Auto-uses function name as operation
    def upload_file(self, path, key):
        self._client.upload_file(path, key)

    @track_state(AppState.IO, operation="s3_download")  # Custom operation name
    def download_file(self, key, path):
        self._client.download_file(key, path)
```

## Prometheus Query Examples

```promql
# GPU utilization percentage
rate(worker_activity_time_seconds_total{state="gpu"}[5m])
/
sum(rate(worker_activity_time_seconds_total{state!~"startup|warmup|shutdown"}[5m]))
* 100

# Idle time percentage (wasted cost)
rate(worker_activity_time_seconds_total{state="idle"}[5m])
/
sum(rate(worker_activity_time_seconds_total{state!~"startup|warmup|shutdown"}[5m]))
* 100

# I/O operations breakdown
sum by (operation) (rate(worker_activity_time_seconds_total{state="io"}[5m]))

# Current state
worker_current_state

# Top 5 slowest operations
topk(5, sum by (operation, state) (rate(worker_activity_time_seconds_total[5m])))
```

## Multi-GPU / Distributed Training

```python
import torch.distributed as dist
from higgsmeter import enable_tracker

# Get rank from your distributed framework
if dist.is_initialized():
    rank = dist.get_rank()
else:
    rank = 0

# Only rank 0 will track and record metrics
enable_tracker(rank)

# All ranks can call tracking methods (no-op on rank != 0)
state_tracker.set_state(AppState.GPU)
```

## Development

```bash
# Clone repository
git clone git@github.com:higgsfield/higgsmeter.git
cd higgsmeter

# Install with dev dependencies
uv pip install -e ".[dev]"

# Run tests
uv run pytest

# Run linter
uv run ruff check src tests

# Format code
uv run ruff format src tests
```

## License

MIT License - see [LICENSE](LICENSE) file for details
