diff --git a/pyproject.toml b/pyproject.toml index 5ee66153f0..74c816e4fc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,6 +51,9 @@ triton = [ "triton>=3.0.0 ; ('linux' in sys_platform)", "triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] +metrics = [ + "prometheus_client>=0.20.0", +] huggingfacenotorch = [ "wheel>=0.42.0", diff --git a/tests/metrics/README.md b/tests/metrics/README.md new file mode 100644 index 0000000000..a9b604ff90 --- /dev/null +++ b/tests/metrics/README.md @@ -0,0 +1,50 @@ +# Metrics Tests + +## Running Tests + +```bash +python3 tests/metrics/test_metrics_standalone.py +``` + +## Test Coverage + +The test suite verifies: + +1. **InferenceStats** - Request tracking, token counts, latencies +2. **TrainingStats** - Batch tracking, loss, throughput +3. **StatsCollector** - Singleton pattern, enable/disable +4. **Prometheus Export** - Client availability check + +## Test Results + +See `TEST_RESULTS.txt` for the latest test run output. + +## Expected Output + +``` +============================================================ +Unsloth Metrics Collection - Standalone Test +============================================================ + +1. Testing InferenceStats... + ✅ Total requests: 1 + ✅ Prompt tokens: 10 + ✅ Generation tokens: 5 + ✅ Avg latency: 0.087s + +2. Testing TrainingStats... + ✅ Total steps: 3 + ✅ Total samples: 12 + ✅ Avg loss: 0.4500 + +3. Testing StatsCollector... + ✅ StatsCollector working + +4. Testing Prometheus export... + ✅ Prometheus client available + ℹ️ Full Prometheus export test requires GPU environment + +============================================================ +✅ All metrics tests passed! +============================================================ +``` diff --git a/tests/metrics/__init__.py b/tests/metrics/__init__.py new file mode 100644 index 0000000000..5ade23931b --- /dev/null +++ b/tests/metrics/__init__.py @@ -0,0 +1 @@ +# Metrics tests diff --git a/tests/metrics/test_metrics_standalone.py b/tests/metrics/test_metrics_standalone.py new file mode 100644 index 0000000000..5527b6285f --- /dev/null +++ b/tests/metrics/test_metrics_standalone.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python3 +""" +Standalone test for metrics collection - tests modules directly without importing unsloth package. +""" + +import sys +import os +import time + +# Add project root to path +project_root = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +) +sys.path.insert(0, os.path.join(project_root, "unsloth", "metrics")) + +# Import metrics modules directly +import importlib.util + + +def load_module(filepath, module_name): + """Load a Python module directly from file.""" + spec = importlib.util.spec_from_file_location(module_name, filepath) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +# Load modules +metrics_dir = os.path.join(project_root, "unsloth", "metrics") +stats = load_module(os.path.join(metrics_dir, "stats.py"), "stats") + +print("=" * 60) +print("Unsloth Metrics Collection - Standalone Test") +print("=" * 60) + +# Test 1: InferenceStats +print("\n1. Testing InferenceStats...") +inference_stats = stats.InferenceStats() +request_id = "test_1" + +inference_stats.start_request(request_id, num_prompt_tokens = 10, max_tokens = 20) +time.sleep(0.01) +inference_stats.record_scheduled(request_id) +time.sleep(0.01) +inference_stats.record_first_token(request_id) +for i in range(5): + inference_stats.record_token(request_id) + time.sleep(0.01) +inference_stats.finish_request( + request_id, finish_reason = "stop", num_generation_tokens = 5 +) + +stats_dict = inference_stats.get_stats() +print(f" ✅ Total requests: {stats_dict['total_requests']}") +print(f" ✅ Prompt tokens: {stats_dict['total_prompt_tokens']}") +print(f" ✅ Generation tokens: {stats_dict['total_generation_tokens']}") +print(f" ✅ Avg latency: {stats_dict['avg_e2e_latency']:.3f}s") +assert stats_dict["total_requests"] == 1 +assert stats_dict["total_prompt_tokens"] == 10 +assert stats_dict["total_generation_tokens"] == 5 + +# Test 2: TrainingStats +print("\n2. Testing TrainingStats...") +training_stats = stats.TrainingStats() + +for step in range(3): + training_stats.record_batch( + step = step, + batch_size = 4, + forward_time = 0.1, + backward_time = 0.15, + loss = 0.5 - (step * 0.05), + learning_rate = 2e-4, + grad_norm = 1.0, + ) + +stats_dict = training_stats.get_stats() +print(f" ✅ Total steps: {stats_dict['total_steps']}") +print(f" ✅ Total samples: {stats_dict['total_samples']}") +print(f" ✅ Avg loss: {stats_dict['avg_loss']:.4f}") +assert stats_dict["total_steps"] == 3 +assert stats_dict["total_samples"] == 12 + +# Test 3: StatsCollector +print("\n3. Testing StatsCollector...") +collector = stats.StatsCollector() +collector.enable() +assert collector.is_enabled() + +all_stats = collector.get_all_stats() +assert "inference" in all_stats +assert "training" in all_stats +print(" ✅ StatsCollector working") + +# Test 4: Prometheus (if available) +print("\n4. Testing Prometheus export...") +try: + # Check if prometheus_client is available + try: + import prometheus_client + + prometheus_available = True + except ImportError: + prometheus_available = False + print(f" ℹ️ prometheus_client not installed (optional dependency)") + print(f" ✅ Metrics collection works without it") + + if prometheus_available: + # Create a minimal test without importing the full prometheus module + # (since it imports stats which triggers package init) + print(f" ✅ Prometheus client available") + print(f" ℹ️ Full Prometheus export test requires GPU environment") +except Exception as e: + print(f" ⚠️ Prometheus test skipped: {e}") + import traceback + + traceback.print_exc() + +print("\n" + "=" * 60) +print("✅ All metrics tests passed!") +print("=" * 60) diff --git a/unsloth/__init__.py b/unsloth/__init__.py index dbacd551c4..650c37cd8a 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -315,6 +315,26 @@ from .chat_templates import * from .tokenizer_utils import * from .trainer import * +# Export metrics functionality (optional, requires prometheus_client for full features) +from .metrics import ( + InferenceStats, + TrainingStats, + StatsCollector, + get_stats_collector, + get_metrics_registry, + generate_prometheus_metrics, + enable_prometheus_metrics, + disable_prometheus_metrics, + is_prometheus_available, + start_metrics_server, + stop_metrics_server, + is_metrics_server_running, + test_metrics_server, + enable_telemetry, + disable_telemetry, + is_telemetry_enabled, +) + # Export dataprep utilities for CLI and downstream users from .dataprep.raw_text import RawTextDataLoader, TextPreprocessor from unsloth_zoo.rl_environments import ( diff --git a/unsloth/metrics/README.md b/unsloth/metrics/README.md new file mode 100644 index 0000000000..a87bf2afd7 --- /dev/null +++ b/unsloth/metrics/README.md @@ -0,0 +1,287 @@ +# Unsloth Metrics Collection + +Comprehensive runtime performance metrics collection for Unsloth, inspired by vLLM's metrics system. + +## Overview + +This module provides detailed tracking of both inference and training metrics, with optional Prometheus-compatible export and HTTP server support. Metrics are collected automatically when enabled, requiring no changes to your existing code. + +## Quick Start + +```python +from unsloth import enable_prometheus_metrics, get_stats_collector + +# Enable metrics (once at start) +enable_prometheus_metrics() + +# Use Unsloth normally - metrics collected automatically +model, tokenizer = FastLanguageModel.from_pretrained(...) +output = model.generate(...) + +# Access metrics programmatically +stats = get_stats_collector().get_all_stats() +print("Inference:", stats['inference']) +print("Training:", stats['training']) +``` + +## Features + +### Inference Metrics +- **Request tracking**: Total requests, active requests, finish reasons +- **Latency metrics**: End-to-end latency (measured), prefill latency (estimated), decode latency (estimated), time per token +- **Token metrics**: Prompt tokens, generation tokens, tokens per second +- **Throughput**: Real-time throughput calculations + +**Note**: Prefill and decode latencies are estimated after generation completes. For more accurate metrics, consider hooking into the generation process itself (e.g., via LogitsProcessor or StoppingCriteria). + +### Training Metrics +- **Step tracking**: Total steps, samples processed +- **Performance**: Forward/backward pass times (estimated), samples per second +- **Training state**: Loss, learning rate, gradient norm +- **Batch metrics**: Batch size tracking + +**Note**: Forward and backward pass times are estimated as a 40/60 split of the total step duration. For precise timings, more intrusive instrumentation would be required. + +### Prometheus Integration +- Prometheus-compatible metrics export +- Standard metric types (Counter, Gauge, Histogram) +- Optional dependency (`prometheus_client`) +- Works without Prometheus (graceful degradation) + +### HTTP Server (Optional) +- Background HTTP server for metrics scraping +- Standard `/metrics` endpoint +- Health check endpoint +- Configurable host/port + +## Usage + +### Programmatic Access (Recommended) + +```python +from unsloth import enable_prometheus_metrics, get_stats_collector, generate_prometheus_metrics + +enable_prometheus_metrics() + +# ... your code ... + +# Get stats +stats = get_stats_collector().get_all_stats() +inference = stats['inference'] +print(f"Requests: {inference['total_requests']}, Tokens/sec: {inference['tokens_per_second']:.2f}") + +# Get Prometheus format +prometheus_text = generate_prometheus_metrics() +print(prometheus_text.decode()) +``` + +### Optional: HTTP Server + +```python +from unsloth import start_metrics_server + +start_metrics_server(port=9090) +# Access at http://localhost:9090/metrics +``` + +## Collected Metrics + +### Inference Metrics + +**Counters:** +- `unsloth_request_total` - Total number of requests (labeled by finish_reason) +- `unsloth_prompt_tokens_total` - Total prompt tokens processed +- `unsloth_generation_tokens_total` - Total generation tokens produced + +**Gauges:** +- `unsloth_requests_active` - Number of currently active requests +- `unsloth_tokens_per_second` - Current tokens per second throughput + +**Histograms:** +- `unsloth_request_latency_seconds` - End-to-end request latency +- `unsloth_prefill_latency_seconds` - Prefill (prompt processing) latency +- `unsloth_decode_latency_seconds` - Decode (generation) latency +- `unsloth_time_per_output_token_seconds` - Time per output token +- `unsloth_prompt_tokens` - Prompt tokens per request +- `unsloth_generation_tokens` - Generation tokens per request + +### Training Metrics + +**Counters:** +- `unsloth_training_steps_total` - Total training steps +- `unsloth_training_samples_total` - Total samples processed + +**Gauges:** +- `unsloth_training_loss` - Current training loss +- `unsloth_learning_rate` - Current learning rate +- `unsloth_training_samples_per_second` - Training throughput +- `unsloth_gradient_norm` - Current gradient norm + +**Histograms:** +- `unsloth_training_forward_time_seconds` - Forward pass time +- `unsloth_training_backward_time_seconds` - Backward pass time +- `unsloth_training_batch_size` - Batch size + +## How It Works + +### Automatic Instrumentation + +Metrics are collected automatically when enabled: + +1. **Inference**: `unsloth_base_fast_generate()` is instrumented to track: + - Request lifecycle (start, scheduled, first token, tokens, finish) + - Latencies (E2E, prefill, decode) + - Token counts and throughput + +2. **Training**: `Trainer.training_step()` is patched to track: + - Forward/backward pass times + - Loss, learning rate, gradient norm + - Batch size and samples per second + +### Integration Points + +- **Inference hook**: `unsloth/models/vision.py` - `unsloth_base_fast_generate()` +- **Training hook**: `unsloth/models/_utils.py` - `_patch_training_metrics()` +- **Public API**: `unsloth/__init__.py` - Exports metrics functions + +### Design Decisions + +1. **Non-intrusive**: Metrics are opt-in via `enable_prometheus_metrics()` +2. **Graceful degradation**: Works without Prometheus client installed +3. **Thread-safe**: Singleton pattern with proper locking +4. **Low overhead**: Minimal performance impact when enabled +5. **Modular**: Separate modules for stats, Prometheus, server +6. **vLLM-inspired**: Similar architecture to vLLM's metrics system + +## Dependencies + +**Optional**: `prometheus_client>=0.20.0` for Prometheus export + +Install via: +```bash +pip install prometheus_client +# or +pip install unsloth[metrics] +``` + +The metrics system works without `prometheus_client` - it gracefully degrades and only provides programmatic access. + +## Environment Variables + +- `UNSLOTH_ENABLE_METRICS=1` - Enable metrics collection (default: disabled) +- `UNSLOTH_ENABLE_METRICS_TELEMETRY=1` - Enable metrics telemetry (opt-in) +- `UNSLOTH_DISABLE_METRICS_TELEMETRY=1` - Disable metrics telemetry (opt-out) +- `UNSLOTH_METRICS_TELEMETRY_ENDPOINT` - Telemetry endpoint (default: https://api.unsloth.ai/metrics) +- `UNSLOTH_METRICS_TELEMETRY_SETTINGS_ENDPOINT` - Optional settings endpoint for server-side opt-out polling (default: unset) +- `UNSLOTH_METRICS_TELEMETRY_INTERVAL` - Telemetry interval seconds (default: 300) + +## API Reference + +### Core Functions + +- `enable_prometheus_metrics()` - Enable metrics collection +- `disable_prometheus_metrics()` - Disable metrics collection +- `get_stats_collector()` - Get the global stats collector singleton +- `generate_prometheus_metrics()` - Generate Prometheus-format metrics +- `is_prometheus_available()` - Check if Prometheus client is available + +### Telemetry Functions (Server-Side Forwarding) + +- `enable_telemetry()` - Enable telemetry (opt-in) +- `disable_telemetry()` - Disable telemetry +- `is_telemetry_enabled()` - Check if telemetry is enabled + +If `UNSLOTH_METRICS_TELEMETRY_SETTINGS_ENDPOINT` is set, Unsloth will poll it once per session and disable telemetry if it returns `{"enabled": false}`. + +### HTTP Server Functions + +- `start_metrics_server(host="0.0.0.0", port=9090)` - Start metrics HTTP server +- `stop_metrics_server()` - Stop metrics HTTP server +- `is_metrics_server_running()` - Check if server is running +- `test_metrics_server(port=9090)` - Test server connectivity + +### Stats Classes + +- `StatsCollector` - Global singleton that manages inference and training stats +- `InferenceStats` - Inference metrics collection +- `TrainingStats` - Training metrics collection + +## Examples + +### Basic Usage + +```python +from unsloth import enable_prometheus_metrics, get_stats_collector + +enable_prometheus_metrics() + +# Run inference/training +# ... + +# Get metrics +stats = get_stats_collector().get_all_stats() +print(f"Inference requests: {stats['inference']['total_requests']}") +print(f"Training steps: {stats['training']['total_steps']}") +``` + +### Prometheus Export + +```python +from unsloth import generate_prometheus_metrics + +# Get Prometheus format +metrics_text = generate_prometheus_metrics() + +# Print or save +print(metrics_text.decode()) +# or +with open("metrics.prom", "wb") as f: + f.write(metrics_text) +``` + +### HTTP Server + +```python +from unsloth import start_metrics_server, test_metrics_server + +# Start server +start_metrics_server(port=9090) + +# Test connection +test_metrics_server(port=9090) + +# Prometheus can now scrape from http://localhost:9090/metrics +``` + +## Comparison with vLLM + +This implementation is inspired by vLLM's metrics collection: + +- Similar metrics structure (request stats, latency breakdowns, token counts) +- Prometheus-compatible export format +- HTTP endpoint for metrics scraping +- Sliding window aggregation for recent metrics + +The main difference is that Unsloth's metrics are integrated into the Transformers-based training/inference pipeline rather than a custom engine. + +## Testing + +Comprehensive test suite available: +```bash +python3 tests/metrics/test_metrics_standalone.py +``` + +Tests cover: +- Inference metrics tracking +- Training metrics tracking +- StatsCollector singleton pattern +- Prometheus export (when available) + +## Future Enhancements + +Potential future improvements: +- More detailed latency breakdowns +- Per-model metrics +- Metrics aggregation windows +- Export to other formats (JSON, CSV) +- Metrics dashboard integration diff --git a/unsloth/metrics/__init__.py b/unsloth/metrics/__init__.py new file mode 100644 index 0000000000..a4c990e61c --- /dev/null +++ b/unsloth/metrics/__init__.py @@ -0,0 +1,64 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Metrics collection module for Unsloth. +Provides comprehensive runtime performance metrics similar to vLLM's metrics system. +""" + +from unsloth.metrics.stats import ( + InferenceStats, + TrainingStats, + StatsCollector, + get_stats_collector, +) +from unsloth.metrics.prometheus import ( + get_metrics_registry, + generate_prometheus_metrics, + enable_prometheus_metrics, + disable_prometheus_metrics, + is_prometheus_available, +) +from unsloth.metrics.server import ( + start_metrics_server, + stop_metrics_server, + is_metrics_server_running, + test_metrics_server, +) +from unsloth.metrics.telemetry import ( + enable_telemetry, + disable_telemetry, + is_telemetry_enabled, + schedule_telemetry, +) + +__all__ = [ + "InferenceStats", + "TrainingStats", + "StatsCollector", + "get_stats_collector", + "get_metrics_registry", + "generate_prometheus_metrics", + "enable_prometheus_metrics", + "disable_prometheus_metrics", + "is_prometheus_available", + "start_metrics_server", + "stop_metrics_server", + "is_metrics_server_running", + "test_metrics_server", + "enable_telemetry", + "disable_telemetry", + "is_telemetry_enabled", + "schedule_telemetry", +] diff --git a/unsloth/metrics/prometheus.py b/unsloth/metrics/prometheus.py new file mode 100644 index 0000000000..372c37010b --- /dev/null +++ b/unsloth/metrics/prometheus.py @@ -0,0 +1,320 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Prometheus metrics export module for Unsloth. +Provides Prometheus-compatible metrics for monitoring. +""" + +import os +from typing import Optional, Dict, Any + +from unsloth.metrics.stats import get_stats_collector + +try: + from prometheus_client import ( + Counter, + Gauge, + Histogram, + REGISTRY, + generate_latest, + CONTENT_TYPE_LATEST, + ) + + PROMETHEUS_AVAILABLE = True +except ImportError: + PROMETHEUS_AVAILABLE = False + + # Mock classes for when prometheus_client is not available + class Counter: + def __init__(self, *args, **kwargs): + pass + + def inc(self, *args, **kwargs): + pass + + def observe(self, *args, **kwargs): + pass + + class Gauge: + def __init__(self, *args, **kwargs): + pass + + def set(self, *args, **kwargs): + pass + + def inc(self, *args, **kwargs): + pass + + def dec(self, *args, **kwargs): + pass + + class Histogram: + def __init__(self, *args, **kwargs): + pass + + def observe(self, *args, **kwargs): + pass + + REGISTRY = None + generate_latest = None + CONTENT_TYPE_LATEST = None + + +# Prometheus metrics (initialized if available) +_metrics_registry: Optional[Dict[str, Any]] = None +_metrics_enabled = False + + +def _get_existing_collector(metric_name: str): + if REGISTRY is None: + return None + return REGISTRY._names_to_collectors.get(metric_name) # type: ignore[attr-defined] + + +def _init_metrics(): + """Initialize Prometheus metrics if available.""" + global _metrics_registry + + if not PROMETHEUS_AVAILABLE: + return None + + if _metrics_registry is not None: + return _metrics_registry + + # Inference metrics + inference_metrics = { + # Counters + "request_total": _get_existing_collector("unsloth_request_total") + or Counter( + "unsloth_request_total", + "Total number of inference requests", + ["finish_reason"], + ), + "prompt_tokens_total": _get_existing_collector("unsloth_prompt_tokens_total") + or Counter( + "unsloth_prompt_tokens_total", + "Total number of prompt tokens processed", + ), + "generation_tokens_total": _get_existing_collector( + "unsloth_generation_tokens_total" + ) + or Counter( + "unsloth_generation_tokens_total", + "Total number of generation tokens produced", + ), + # Gauges + "requests_active": _get_existing_collector("unsloth_requests_active") + or Gauge( + "unsloth_requests_active", + "Number of currently active inference requests", + ), + "tokens_per_second": _get_existing_collector("unsloth_tokens_per_second") + or Gauge( + "unsloth_tokens_per_second", + "Current tokens per second throughput", + ), + # Histograms + "request_latency_seconds": _get_existing_collector( + "unsloth_request_latency_seconds" + ) + or Histogram( + "unsloth_request_latency_seconds", + "End-to-end request latency in seconds", + buckets = [0.1, 0.5, 1.0, 2.0, 5.0, 10.0, 30.0, 60.0, 120.0], + ), + "prefill_latency_seconds": _get_existing_collector( + "unsloth_prefill_latency_seconds" + ) + or Histogram( + "unsloth_prefill_latency_seconds", + "Prefill (prompt processing) latency in seconds", + buckets = [0.01, 0.05, 0.1, 0.5, 1.0, 2.0, 5.0, 10.0], + ), + "decode_latency_seconds": _get_existing_collector( + "unsloth_decode_latency_seconds" + ) + or Histogram( + "unsloth_decode_latency_seconds", + "Decode (generation) latency in seconds", + buckets = [0.01, 0.05, 0.1, 0.5, 1.0, 2.0, 5.0, 10.0], + ), + "time_per_output_token_seconds": _get_existing_collector( + "unsloth_time_per_output_token_seconds" + ) + or Histogram( + "unsloth_time_per_output_token_seconds", + "Time per output token in seconds", + buckets = [0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1.0], + ), + "prompt_tokens": _get_existing_collector("unsloth_prompt_tokens") + or Histogram( + "unsloth_prompt_tokens", + "Number of prompt tokens per request", + buckets = [10, 50, 100, 500, 1000, 2000, 4000, 8000, 16000, 32000], + ), + "generation_tokens": _get_existing_collector("unsloth_generation_tokens") + or Histogram( + "unsloth_generation_tokens", + "Number of generation tokens per request", + buckets = [10, 50, 100, 500, 1000, 2000, 4000, 8000, 16000, 32000], + ), + } + + # Training metrics + training_metrics = { + # Counters + "training_steps_total": _get_existing_collector("unsloth_training_steps_total") + or Counter( + "unsloth_training_steps_total", + "Total number of training steps", + ), + "training_samples_total": _get_existing_collector( + "unsloth_training_samples_total" + ) + or Counter( + "unsloth_training_samples_total", + "Total number of training samples processed", + ), + # Gauges + "training_loss": _get_existing_collector("unsloth_training_loss") + or Gauge( + "unsloth_training_loss", + "Current training loss", + ), + "learning_rate": _get_existing_collector("unsloth_learning_rate") + or Gauge( + "unsloth_learning_rate", + "Current learning rate", + ), + "samples_per_second": _get_existing_collector( + "unsloth_training_samples_per_second" + ) + or Gauge( + "unsloth_training_samples_per_second", + "Training throughput in samples per second", + ), + "gradient_norm": _get_existing_collector("unsloth_gradient_norm") + or Gauge( + "unsloth_gradient_norm", + "Current gradient norm", + ), + # Histograms + "forward_time_seconds": _get_existing_collector( + "unsloth_training_forward_time_seconds" + ) + or Histogram( + "unsloth_training_forward_time_seconds", + "Forward pass time in seconds", + buckets = [0.01, 0.05, 0.1, 0.5, 1.0, 2.0, 5.0], + ), + "backward_time_seconds": _get_existing_collector( + "unsloth_training_backward_time_seconds" + ) + or Histogram( + "unsloth_training_backward_time_seconds", + "Backward pass time in seconds", + buckets = [0.01, 0.05, 0.1, 0.5, 1.0, 2.0, 5.0], + ), + "batch_size": _get_existing_collector("unsloth_training_batch_size") + or Histogram( + "unsloth_training_batch_size", + "Training batch size", + buckets = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512], + ), + } + + _metrics_registry = { + "inference": inference_metrics, + "training": training_metrics, + } + + return _metrics_registry + + +def get_metrics_registry() -> Optional[Dict[str, Any]]: + """Get the Prometheus metrics registry.""" + return _init_metrics() + + +def update_prometheus_metrics(): + """Update Prometheus metrics from stats collector.""" + if not _metrics_enabled or not PROMETHEUS_AVAILABLE: + return + + registry = get_metrics_registry() + if registry is None: + return + + collector = get_stats_collector() + if not collector.is_enabled(): + return + + # Update inference metrics + inference_stats = collector.inference_stats.get_stats() + inference_metrics = registry["inference"] + + inference_metrics["requests_active"].set(inference_stats.get("active_requests", 0)) + inference_metrics["tokens_per_second"].set( + inference_stats.get("tokens_per_second", 0.0) + ) + + # Note: Counters and histograms are updated when events occur, + # not from aggregated stats. They're updated in the integration hooks. + + # Update training metrics + training_stats = collector.training_stats.get_stats() + training_metrics = registry["training"] + + training_metrics["training_loss"].set(training_stats.get("avg_loss", 0.0)) + training_metrics["learning_rate"].set(training_stats.get("current_lr", 0.0)) + training_metrics["samples_per_second"].set( + training_stats.get("samples_per_second", 0.0) + ) + + +def generate_prometheus_metrics() -> bytes: + """Generate Prometheus metrics output in text format.""" + if not PROMETHEUS_AVAILABLE: + return b"# Prometheus metrics not available (prometheus_client not installed)\n" + + update_prometheus_metrics() + return generate_latest(REGISTRY) + + +def enable_prometheus_metrics(): + """Enable Prometheus metrics collection and export.""" + global _metrics_enabled + _metrics_enabled = True + _init_metrics() + get_stats_collector().enable() + + +def disable_prometheus_metrics(): + """Disable Prometheus metrics collection.""" + global _metrics_enabled + _metrics_enabled = False + get_stats_collector().disable() + + +def is_prometheus_available() -> bool: + """Check if prometheus_client is available.""" + return PROMETHEUS_AVAILABLE + + +def get_metrics_content_type() -> str: + """Get the Content-Type header for Prometheus metrics.""" + if PROMETHEUS_AVAILABLE: + return CONTENT_TYPE_LATEST + return "text/plain; charset=utf-8" diff --git a/unsloth/metrics/server.py b/unsloth/metrics/server.py new file mode 100644 index 0000000000..5569a849b3 --- /dev/null +++ b/unsloth/metrics/server.py @@ -0,0 +1,253 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Optional HTTP server for exposing Prometheus metrics. +""" + +import threading +import socket +from typing import Optional +from http.server import HTTPServer, BaseHTTPRequestHandler + +from unsloth.metrics.prometheus import ( + generate_prometheus_metrics, + get_metrics_content_type, + enable_prometheus_metrics, +) + + +class MetricsHandler(BaseHTTPRequestHandler): + """HTTP request handler for metrics endpoint.""" + + def do_GET(self): + """Handle GET requests to /metrics endpoint.""" + try: + if self.path == "/metrics": + try: + metrics_output = generate_prometheus_metrics() + self.send_response(200) + self.send_header("Content-Type", get_metrics_content_type()) + self.send_header("Content-Length", str(len(metrics_output))) + self.end_headers() + self.wfile.write(metrics_output) + self.wfile.flush() + except Exception as e: + error_msg = f"Error generating metrics: {str(e)}\n" + self.send_response(500) + self.send_header("Content-Type", "text/plain") + self.send_header("Content-Length", str(len(error_msg))) + self.end_headers() + self.wfile.write(error_msg.encode()) + self.wfile.flush() + elif self.path == "/" or self.path == "": + # Simple health check endpoint + response = ( + b"Unsloth Metrics Server\n/metrics - Prometheus metrics endpoint\n" + ) + self.send_response(200) + self.send_header("Content-Type", "text/plain") + self.send_header("Content-Length", str(len(response))) + self.end_headers() + self.wfile.write(response) + self.wfile.flush() + else: + response = b"Not Found\n" + self.send_response(404) + self.send_header("Content-Type", "text/plain") + self.send_header("Content-Length", str(len(response))) + self.end_headers() + self.wfile.write(response) + self.wfile.flush() + except Exception as e: + # Handle any errors in request processing + try: + error_msg = f"Internal server error: {str(e)}\n" + self.send_response(500) + self.send_header("Content-Type", "text/plain") + self.send_header("Content-Length", str(len(error_msg))) + self.end_headers() + self.wfile.write(error_msg.encode()) + self.wfile.flush() + except Exception: + pass # Connection might be closed + + def log_message(self, format, *args): + """Suppress default logging.""" + + +_metrics_server: Optional[HTTPServer] = None +_server_thread: Optional[threading.Thread] = None + + +def start_metrics_server(host: str = "0.0.0.0", port: int = 9090): + """ + Start a background HTTP server to expose Prometheus metrics. + + Args: + host: Host to bind to (default: "0.0.0.0") + port: Port to bind to (default: 9090) + + Returns: + Thread object running the server + """ + global _metrics_server, _server_thread + + if _metrics_server is not None: + print(f"📊 Metrics server already running at http://{host}:{port}/metrics") + return _server_thread + + # Enable Prometheus metrics + enable_prometheus_metrics() + + def run_server(): + global _metrics_server + try: + _metrics_server = HTTPServer((host, port), MetricsHandler) + # Set socket options to allow reuse + _metrics_server.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + # Give the server a moment to bind + import time + + time.sleep(0.1) + _metrics_server.serve_forever() + except OSError as e: + if "Address already in use" in str(e) or "already in use" in str(e).lower(): + print( + f"⚠️ Port {port} is already in use. Please use a different port or stop the other service." + ) + else: + print(f"⚠️ Failed to start metrics server: {e}") + _metrics_server = None + except Exception as e: + print(f"⚠️ Error starting metrics server: {e}") + import traceback + + traceback.print_exc() + _metrics_server = None + + # Use daemon=False in some cases to keep server alive + # But daemon=True is better for cleanup when main program exits + _server_thread = threading.Thread( + target = run_server, daemon = True, name = "UnslothMetricsServer" + ) + _server_thread.start() + + # Give the thread a moment to start and bind + import time + + max_wait = 2.0 + waited = 0.0 + while _metrics_server is None and waited < max_wait: + time.sleep(0.1) + waited += 0.1 + + # Check if server started successfully + if _metrics_server is not None: + # Verify the server is actually listening + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + result = sock.connect_ex(("localhost", port)) + sock.close() + if result == 0: + print( + f"📊 Unsloth metrics server started at http://localhost:{port}/metrics" + ) + print(f" (Also accessible at http://{host}:{port}/metrics)") + else: + print(f"⚠️ Server thread started but port {port} is not accessible") + print(f" Waiting a bit longer for server to bind...") + time.sleep(0.5) + # Try one more time + sock2 = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + result2 = sock2.connect_ex(("localhost", port)) + sock2.close() + if result2 == 0: + print( + f"📊 Server is now accessible at http://localhost:{port}/metrics" + ) + else: + print( + f"⚠️ Server still not accessible. Try restarting or use a different port." + ) + except Exception as e: + print(f"⚠️ Could not verify server: {e}") + else: + print(f"⚠️ Failed to start metrics server. Check if port {port} is available.") + print(f" Try: start_metrics_server(port=9091) to use a different port") + + return _server_thread + + +def stop_metrics_server(): + """Stop the metrics server.""" + global _metrics_server, _server_thread + + if _metrics_server is not None: + _metrics_server.shutdown() + _metrics_server = None + _server_thread = None + print("📊 Metrics server stopped") + + +def is_metrics_server_running() -> bool: + """Check if the metrics server is currently running.""" + global _metrics_server + return _metrics_server is not None + + +def test_metrics_server(port: int = 9090): + """Test if the metrics server is accessible.""" + import urllib.request + import urllib.error + + # First check if server object exists + if not is_metrics_server_running(): + print(f"❌ Metrics server is not running") + print(f" Call start_metrics_server() first") + return False + + # Check if port is listening + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + result = sock.connect_ex(("localhost", port)) + sock.close() + if result != 0: + print(f"❌ Port {port} is not listening") + print(f" Server object exists but port is not accessible") + return False + except Exception as e: + print(f"❌ Could not check port {port}: {e}") + sock.close() + return False + + # Try HTTP request + try: + url = f"http://localhost:{port}/metrics" + response = urllib.request.urlopen(url, timeout = 2) + print(f"✅ Metrics server is running and accessible at {url}") + print(f" Status: {response.getcode()}") + print(f" Content-Length: {response.headers.get('Content-Length', 'unknown')}") + return True + except urllib.error.URLError as e: + print( + f"❌ Could not connect to metrics server at http://localhost:{port}/metrics" + ) + print(f" Error: {e}") + print(f" Error code: {getattr(e, 'code', 'unknown')}") + print(f" Error reason: {getattr(e, 'reason', 'unknown')}") + return False + except Exception as e: + print(f"❌ Error testing metrics server: {e}") + return False diff --git a/unsloth/metrics/stats.py b/unsloth/metrics/stats.py new file mode 100644 index 0000000000..13650dd7bb --- /dev/null +++ b/unsloth/metrics/stats.py @@ -0,0 +1,436 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Statistics tracking module for Unsloth. +Tracks runtime performance metrics during inference and training. +""" + +import time +from collections import defaultdict, deque +from dataclasses import dataclass, field +from typing import Optional, Dict, List, Any +from threading import Lock +import os + +import torch + + +@dataclass +class RequestStats: + """Statistics for a single inference request.""" + + request_id: str + arrival_time: float + scheduled_time: Optional[float] = None + first_token_time: Optional[float] = None + last_token_time: Optional[float] = None + finish_time: Optional[float] = None + num_prompt_tokens: int = 0 + num_generation_tokens: int = 0 + max_tokens_param: Optional[int] = None + finish_reason: Optional[str] = None # "stop", "length", "error" + + +@dataclass +class TrainingBatchStats: + """Statistics for a single training batch.""" + + step: int + batch_size: int + forward_time: float + backward_time: float + loss: float + learning_rate: float + grad_norm: Optional[float] = None + + +class InferenceStats: + """Collects and aggregates inference statistics.""" + + def __init__(self, max_recent_requests: int = 1000): + self.max_recent_requests = max_recent_requests + self._lock = Lock() + + # Per-request tracking + self._active_requests: Dict[str, RequestStats] = {} + + # Aggregated statistics + self.total_requests: int = 0 + self.total_prompt_tokens: int = 0 + self.total_generation_tokens: int = 0 + self.total_e2e_latency: float = 0.0 + self.total_prefill_latency: float = 0.0 + self.total_decode_latency: float = 0.0 + + # Finished requests (for sliding window) + self._finished_requests: deque = deque(maxlen = max_recent_requests) + + # Finish reason counts + self.finish_reasons: Dict[str, int] = defaultdict(int) + + # Timing breakdowns + self._queued_times: deque = deque(maxlen = max_recent_requests) + self._prefill_times: deque = deque(maxlen = max_recent_requests) + self._decode_times: deque = deque(maxlen = max_recent_requests) + self._e2e_times: deque = deque(maxlen = max_recent_requests) + + def start_request( + self, + request_id: str, + num_prompt_tokens: int, + max_tokens: Optional[int] = None, + ): + """Record the start of an inference request.""" + with self._lock: + self._active_requests[request_id] = RequestStats( + request_id = request_id, + arrival_time = time.time(), + num_prompt_tokens = num_prompt_tokens, + max_tokens_param = max_tokens, + ) + + def record_scheduled(self, request_id: str): + """Record when a request was scheduled for processing.""" + with self._lock: + if request_id in self._active_requests: + if self._active_requests[request_id].scheduled_time is None: + self._active_requests[request_id].scheduled_time = time.time() + + def record_first_token(self, request_id: str, timestamp: Optional[float] = None): + """Record when the first token was generated. + + Args: + request_id: Unique identifier for the request + timestamp: Optional timestamp. If None, uses current time. + """ + with self._lock: + if request_id in self._active_requests: + req = self._active_requests[request_id] + if req.first_token_time is None: + req.first_token_time = ( + timestamp if timestamp is not None else time.time() + ) + if req.scheduled_time is None: + req.scheduled_time = req.first_token_time + + def record_token(self, request_id: str): + """Record each token generation (updates last_token_time).""" + with self._lock: + if request_id in self._active_requests: + self._active_requests[request_id].last_token_time = time.time() + req = self._active_requests[request_id] + if req.num_generation_tokens == 0 and req.first_token_time is None: + req.first_token_time = req.last_token_time + req.num_generation_tokens += 1 + + def finish_request( + self, + request_id: str, + finish_reason: str = "stop", + num_generation_tokens: Optional[int] = None, + ): + """Record the completion of an inference request.""" + with self._lock: + if request_id not in self._active_requests: + return + + req = self._active_requests.pop(request_id) + req.finish_time = time.time() + req.finish_reason = finish_reason + + if num_generation_tokens is not None: + req.num_generation_tokens = num_generation_tokens + + # Calculate latencies + if req.scheduled_time is None: + req.scheduled_time = req.arrival_time + if req.first_token_time is None: + req.first_token_time = req.finish_time + + e2e_latency = req.finish_time - req.arrival_time + queued_time = req.scheduled_time - req.arrival_time + prefill_time = req.first_token_time - req.scheduled_time + decode_time = ( + req.last_token_time - req.first_token_time + if req.last_token_time + else 0.0 + ) + + # Update aggregates + self.total_requests += 1 + self.total_prompt_tokens += req.num_prompt_tokens + self.total_generation_tokens += req.num_generation_tokens + self.total_e2e_latency += e2e_latency + self.total_prefill_latency += prefill_time + self.total_decode_latency += decode_time + + # Store in sliding window + self._finished_requests.append(req) + self._queued_times.append(queued_time) + self._prefill_times.append(prefill_time) + self._decode_times.append(decode_time) + self._e2e_times.append(e2e_latency) + + # Count finish reasons + self.finish_reasons[finish_reason] += 1 + + # Schedule telemetry if enabled (non-blocking) + try: + from unsloth.metrics.telemetry import schedule_telemetry + + schedule_telemetry() + except Exception: + pass + + def get_stats(self) -> Dict[str, Any]: + """Get current aggregated statistics.""" + with self._lock: + num_finished = len(self._finished_requests) + if num_finished == 0: + return { + "total_requests": 0, + "active_requests": len(self._active_requests), + "avg_e2e_latency": 0.0, + "avg_prefill_latency": 0.0, + "avg_decode_latency": 0.0, + "avg_time_per_output_token": 0.0, + "total_prompt_tokens": 0, + "total_generation_tokens": 0, + "tokens_per_second": 0.0, + "finish_reasons": {}, + } + + # Calculate averages from recent requests + avg_e2e = sum(self._e2e_times) / num_finished + avg_prefill = sum(self._prefill_times) / num_finished + avg_decode = sum(self._decode_times) / num_finished + + # Calculate tokens per second from recent requests + total_recent_time = sum(self._e2e_times) + total_recent_tokens = sum( + req.num_generation_tokens for req in self._finished_requests + ) + tokens_per_second = ( + total_recent_tokens / total_recent_time + if total_recent_time > 0 + else 0.0 + ) + + # Average time per output token + avg_time_per_token = ( + avg_decode / max(1, total_recent_tokens / num_finished) + if total_recent_tokens > 0 + else 0.0 + ) + + return { + "total_requests": self.total_requests, + "active_requests": len(self._active_requests), + "avg_e2e_latency": avg_e2e, + "avg_prefill_latency": avg_prefill, + "avg_decode_latency": avg_decode, + "avg_time_per_output_token": avg_time_per_token, + "total_prompt_tokens": self.total_prompt_tokens, + "total_generation_tokens": self.total_generation_tokens, + "tokens_per_second": tokens_per_second, + "finish_reasons": dict(self.finish_reasons), + } + + def reset(self): + """Reset all statistics.""" + with self._lock: + self._active_requests.clear() + self._finished_requests.clear() + self.total_requests = 0 + self.total_prompt_tokens = 0 + self.total_generation_tokens = 0 + self.total_e2e_latency = 0.0 + self.total_prefill_latency = 0.0 + self.total_decode_latency = 0.0 + self.finish_reasons.clear() + self._queued_times.clear() + self._prefill_times.clear() + self._decode_times.clear() + self._e2e_times.clear() + + +class TrainingStats: + """Collects and aggregates training statistics.""" + + def __init__(self, max_recent_batches: int = 1000): + self.max_recent_batches = max_recent_batches + self._lock = Lock() + + # Aggregated statistics + self.total_steps: int = 0 + self.total_samples: int = 0 + self.total_forward_time: float = 0.0 + self.total_backward_time: float = 0.0 + self.total_loss: float = 0.0 + + # Recent batches for sliding window + self._recent_batches: deque = deque(maxlen = max_recent_batches) + + def record_batch( + self, + step: int, + batch_size: int, + forward_time: float, + backward_time: float, + loss: float, + learning_rate: float, + grad_norm: Optional[float] = None, + ): + """Record statistics for a training batch.""" + with self._lock: + batch_stats = TrainingBatchStats( + step = step, + batch_size = batch_size, + forward_time = forward_time, + backward_time = backward_time, + loss = loss, + learning_rate = learning_rate, + grad_norm = grad_norm, + ) + self._recent_batches.append(batch_stats) + + self.total_steps += 1 + self.total_samples += batch_size + self.total_forward_time += forward_time + self.total_backward_time += backward_time + self.total_loss += loss + + # Schedule telemetry if enabled (non-blocking) + try: + from unsloth.metrics.telemetry import schedule_telemetry + + schedule_telemetry() + except Exception: + pass + + def get_stats(self) -> Dict[str, Any]: + """Get current aggregated statistics.""" + with self._lock: + num_batches = len(self._recent_batches) + if num_batches == 0: + return { + "total_steps": 0, + "total_samples": 0, + "avg_loss": 0.0, + "avg_forward_time": 0.0, + "avg_backward_time": 0.0, + "samples_per_second": 0.0, + "current_lr": 0.0, + } + + recent_loss = sum(b.loss for b in self._recent_batches) / num_batches + recent_forward = ( + sum(b.forward_time for b in self._recent_batches) / num_batches + ) + recent_backward = ( + sum(b.backward_time for b in self._recent_batches) / num_batches + ) + recent_samples = sum(b.batch_size for b in self._recent_batches) + recent_time = sum( + b.forward_time + b.backward_time for b in self._recent_batches + ) + samples_per_second = ( + recent_samples / recent_time if recent_time > 0 else 0.0 + ) + + current_lr = ( + self._recent_batches[-1].learning_rate if self._recent_batches else 0.0 + ) + + return { + "total_steps": self.total_steps, + "total_samples": self.total_samples, + "avg_loss": recent_loss, + "avg_forward_time": recent_forward, + "avg_backward_time": recent_backward, + "samples_per_second": samples_per_second, + "current_lr": current_lr, + } + + def reset(self): + """Reset all statistics.""" + with self._lock: + self._recent_batches.clear() + self.total_steps = 0 + self.total_samples = 0 + self.total_forward_time = 0.0 + self.total_backward_time = 0.0 + self.total_loss = 0.0 + + +class StatsCollector: + """Global statistics collector that manages both inference and training stats.""" + + _instance: Optional["StatsCollector"] = None + _lock = Lock() + + def __new__(cls): + if cls._instance is None: + with cls._lock: + if cls._instance is None: + cls._instance = super().__new__(cls) + cls._instance._initialized = False + return cls._instance + + def __init__(self): + if self._initialized: + return + + self.inference_stats = InferenceStats() + self.training_stats = TrainingStats() + self._enabled = os.environ.get("UNSLOTH_ENABLE_METRICS", "0") == "1" + self._initialized = True + + def enable(self): + """Enable metrics collection.""" + self._enabled = True + + def disable(self): + """Disable metrics collection.""" + self._enabled = False + + def is_enabled(self) -> bool: + """Check if metrics collection is enabled.""" + return self._enabled + + def get_all_stats(self) -> Dict[str, Any]: + """Get all statistics (inference + training).""" + return { + "inference": self.inference_stats.get_stats(), + "training": self.training_stats.get_stats(), + "enabled": self._enabled, + } + + def reset_all(self): + """Reset all statistics.""" + self.inference_stats.reset() + self.training_stats.reset() + + +# Global singleton instance +_stats_collector = None + + +def get_stats_collector() -> StatsCollector: + """Get the global stats collector instance.""" + global _stats_collector + if _stats_collector is None: + _stats_collector = StatsCollector() + return _stats_collector diff --git a/unsloth/metrics/telemetry.py b/unsloth/metrics/telemetry.py new file mode 100644 index 0000000000..cb7f6a5271 --- /dev/null +++ b/unsloth/metrics/telemetry.py @@ -0,0 +1,235 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Optional telemetry module for forwarding metrics to Unsloth's server. +Disabled by default. Opt-in via environment variable or explicit enable. +""" + +from __future__ import annotations + +import json +import os +import threading +import time +from queue import Empty, Queue +from typing import Any, Dict, Optional +import urllib.error +import urllib.request + +from unsloth.metrics.stats import get_stats_collector + + +_TELEMETRY_ENABLED = os.environ.get("UNSLOTH_ENABLE_METRICS_TELEMETRY", "0") == "1" +_TELEMETRY_DISABLED = os.environ.get("UNSLOTH_DISABLE_METRICS_TELEMETRY", "0") == "1" +_TELEMETRY_ENDPOINT = os.environ.get( + "UNSLOTH_METRICS_TELEMETRY_ENDPOINT", + "https://api.unsloth.ai/metrics", +) +_TELEMETRY_SETTINGS_ENDPOINT = os.environ.get( + "UNSLOTH_METRICS_TELEMETRY_SETTINGS_ENDPOINT", + "", +) +_TELEMETRY_INTERVAL = int(os.environ.get("UNSLOTH_METRICS_TELEMETRY_INTERVAL", "300")) + +_telemetry_queue: Optional[Queue] = None +_telemetry_thread: Optional[threading.Thread] = None +_telemetry_lock = threading.Lock() +_settings_checked = False + + +def is_telemetry_enabled() -> bool: + """Check if telemetry is enabled.""" + return _TELEMETRY_ENABLED and not _TELEMETRY_DISABLED + + +def enable_telemetry() -> None: + """Enable metrics telemetry (if not disabled via env var).""" + global _TELEMETRY_ENABLED + if _TELEMETRY_DISABLED: + return + _TELEMETRY_ENABLED = True + _check_server_opt_out_once() + _start_telemetry_thread() + + +def disable_telemetry() -> None: + """Disable metrics telemetry.""" + global _TELEMETRY_ENABLED + _TELEMETRY_ENABLED = False + _stop_telemetry_thread() + + +def _start_telemetry_thread() -> None: + """Start background thread for sending telemetry.""" + global _telemetry_queue, _telemetry_thread + + with _telemetry_lock: + if _telemetry_thread is not None and _telemetry_thread.is_alive(): + return + + _telemetry_queue = Queue() + _telemetry_thread = threading.Thread( + target = _telemetry_worker, + daemon = True, + name = "UnslothMetricsTelemetry", + ) + _telemetry_thread.start() + + +def _stop_telemetry_thread() -> None: + """Stop background telemetry thread.""" + global _telemetry_queue, _telemetry_thread + + with _telemetry_lock: + if _telemetry_queue is not None: + _telemetry_queue.put(None) + _telemetry_thread = None + _telemetry_queue = None + + +def _telemetry_worker() -> None: + """Background worker that sends telemetry data periodically.""" + global _telemetry_queue + if _telemetry_queue is None: + return + + while True: + try: + item = _telemetry_queue.get(timeout = _TELEMETRY_INTERVAL) + except Empty: + item = "timeout" + + if item is None: + break + + _check_server_opt_out_once() + if is_telemetry_enabled(): + _send_telemetry_batch() + + +def _get_package_version() -> str: + try: + from importlib.metadata import version + + return version("unsloth") + except Exception: + return "unknown" + + +def _send_telemetry_batch() -> None: + """Send a batch of metrics to Unsloth's server.""" + if not is_telemetry_enabled(): + return + + try: + collector = get_stats_collector() + if not collector.is_enabled(): + return + + stats = collector.get_all_stats() + payload = { + "timestamp": time.time(), + "version": _get_package_version(), + "metrics": { + "inference": { + "total_requests": stats["inference"].get("total_requests", 0), + "avg_e2e_latency": stats["inference"].get("avg_e2e_latency", 0.0), + "tokens_per_second": stats["inference"].get( + "tokens_per_second", 0.0 + ), + "total_prompt_tokens": stats["inference"].get( + "total_prompt_tokens", 0 + ), + "total_generation_tokens": stats["inference"].get( + "total_generation_tokens", 0 + ), + }, + "training": { + "total_steps": stats["training"].get("total_steps", 0), + "avg_loss": stats["training"].get("avg_loss", 0.0), + "samples_per_second": stats["training"].get( + "samples_per_second", 0.0 + ), + "total_samples": stats["training"].get("total_samples", 0), + }, + }, + } + + _send_to_server(payload) + except Exception: + # Telemetry should never break user code + pass + + +def _send_to_server(payload: Dict[str, Any]) -> None: + """Send payload to Unsloth's telemetry endpoint.""" + try: + data = json.dumps(payload).encode("utf-8") + request = urllib.request.Request( + _TELEMETRY_ENDPOINT, + data = data, + headers = { + "Content-Type": "application/json", + "User-Agent": "Unsloth-Metrics/1.0", + }, + ) + urllib.request.urlopen(request, timeout = 5) + except (urllib.error.URLError, urllib.error.HTTPError, Exception): + pass + + +def _check_server_opt_out_once() -> None: + """Optional one-time poll for server-side opt-out setting.""" + global _settings_checked, _TELEMETRY_ENABLED + if _settings_checked: + return + _settings_checked = True + + if not _TELEMETRY_SETTINGS_ENDPOINT: + return + + try: + request = urllib.request.Request( + _TELEMETRY_SETTINGS_ENDPOINT, + headers = {"User-Agent": "Unsloth-Metrics/1.0"}, + ) + with urllib.request.urlopen(request, timeout = 5) as response: + data = response.read().decode("utf-8") + # Expect JSON payload like {"enabled": true} + enabled = json.loads(data).get("enabled", True) + if not enabled: + _TELEMETRY_ENABLED = False + except Exception: + # If polling fails, keep telemetry enabled (opt-in still required) + pass + + +def schedule_telemetry() -> None: + """Schedule telemetry to be sent (non-blocking).""" + if not is_telemetry_enabled(): + return + + global _telemetry_queue + if _telemetry_queue is not None: + try: + _telemetry_queue.put_nowait("send") + except Exception: + pass + + +# Opt-in via UNSLOTH_ENABLE_METRICS_TELEMETRY=1 +# Opt-out via UNSLOTH_DISABLE_METRICS_TELEMETRY=1 +if is_telemetry_enabled(): + _start_telemetry_thread() diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 19b5fe0574..e8801c9121 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -1877,6 +1877,143 @@ def _unsloth_pre_compute_loss(self, model, inputs, *args, **kwargs): return outputs +def _patch_training_metrics(Trainer): + """Patch Trainer.training_step to collect training metrics.""" + if hasattr(Trainer, "_unsloth_metrics_patched"): + return + + # Get the current training_step (which might already be _unsloth_training_step) + original_training_step = Trainer.training_step + + @functools.wraps(original_training_step) + def training_step_with_metrics(self, model, inputs, *args, **kwargs): + # Try to collect metrics if enabled + try: + from unsloth.metrics.stats import get_stats_collector + from unsloth.metrics.prometheus import ( + get_metrics_registry, + _metrics_enabled, + ) + import time + + collector = get_stats_collector() + track_metrics = collector.is_enabled() + except Exception: + track_metrics = False + + if not track_metrics: + return original_training_step(self, model, inputs, *args, **kwargs) + + # Get current step + step = getattr(self.state, "global_step", 0) + + # Get batch size + batch_size = 0 + if isinstance(inputs, dict): + # Try to get batch size from input_ids or input tensor + for key in ["input_ids", "inputs", "input_features"]: + if key in inputs and inputs[key] is not None: + tensor = inputs[key] + if hasattr(tensor, "shape") and len(tensor.shape) > 0: + batch_size = tensor.shape[0] + break + + # Track step duration (includes both forward and backward passes) + step_start = time.time() + + # Call original training_step + try: + result = original_training_step(self, model, inputs, *args, **kwargs) + except Exception as e: + # Re-raise exception but don't track metrics on error + raise + + step_end = time.time() + step_duration = step_end - step_start + + # Extract loss and other info from result + loss_value = None + if isinstance(result, (int, float)): + loss_value = float(result) + elif hasattr(result, "numel") and result.numel() == 1: # torch.Tensor + loss_value = float(result.item()) + elif isinstance(result, dict): + loss_value = result.get("loss") + if loss_value is not None and hasattr(loss_value, "item"): + loss_value = float(loss_value.item()) + + # Get learning rate + learning_rate = 0.0 + if hasattr(self, "lr_scheduler") and self.lr_scheduler is not None: + try: + if hasattr(self.lr_scheduler, "get_last_lr"): + lrs = self.lr_scheduler.get_last_lr() + learning_rate = float(lrs[0]) if lrs else 0.0 + elif hasattr(self.lr_scheduler, "get_lr"): + lrs = self.lr_scheduler.get_lr() + learning_rate = float(lrs[0]) if lrs else 0.0 + except Exception: + pass + + # Estimate backward time (simplified - backward happens inside training_step) + # We'll approximate it as a fraction of step duration + backward_time = step_duration * 0.6 # Rough estimate + forward_time = step_duration * 0.4 # Estimate forward time + # Get gradient norm if available + grad_norm = None + if hasattr(self, "accelerator"): + try: + # Gradient norm might be stored after clipping + if hasattr(self.state, "grad_norm"): + grad_norm = float(self.state.grad_norm) + except Exception: + pass + + # Record to stats collector + if loss_value is not None: + try: + collector.training_stats.record_batch( + step = step, + batch_size = batch_size, + forward_time = forward_time, + backward_time = backward_time, + loss = loss_value, + learning_rate = learning_rate, + grad_norm = grad_norm, + ) + + # Update Prometheus metrics + if _metrics_enabled: + registry = get_metrics_registry() + if registry: + training_metrics = registry["training"] + training_metrics["training_steps_total"].inc() + training_metrics["training_samples_total"].inc(batch_size) + training_metrics["training_loss"].set(loss_value) + training_metrics["learning_rate"].set(learning_rate) + training_metrics["forward_time_seconds"].observe(forward_time) + training_metrics["backward_time_seconds"].observe(backward_time) + training_metrics["batch_size"].observe(batch_size) + if grad_norm is not None: + training_metrics["gradient_norm"].set(grad_norm) + + # Update samples per second + total_time = forward_time + backward_time + if total_time > 0: + samples_per_second = batch_size / total_time + training_metrics["samples_per_second"].set( + samples_per_second + ) + except Exception: + # Metrics collection failed, continue without metrics + pass + + return result + + Trainer.training_step = training_step_with_metrics + Trainer._unsloth_metrics_patched = True + + def patch_gradient_accumulation_fix(Trainer): # Fixes gradient accumulation # Fixes Output 0 of UnslothFusedLossBackward is a view and is being modified inplace. @@ -1999,6 +2136,9 @@ def patch_gradient_accumulation_fix(Trainer): exec(function, globals()) Trainer.training_step = _unsloth_training_step + # Add metrics tracking to training_step + _patch_training_metrics(Trainer) + # Prevent double scaling gradient accumulation # https://github.com/huggingface/transformers/pull/37208 # Patch model_accepts_loss_kwargs detection in Trainer.__init__ diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index a8adba99e7..f648199e3f 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -69,6 +69,7 @@ import functools import os import gc import math +import time from typing import Optional, Tuple, List, Union import re, inspect, sys import contextlib @@ -292,9 +293,156 @@ def unsloth_base_fast_generate( pass # DO INFERENCE + # Track metrics if enabled + collector = None + request_id = None + prompt_tokens_per_sequence = input_ids.shape[-1] if input_ids is not None else 0 + prompt_batch_size = input_ids.shape[0] if input_ids.dim() > 1 else 1 + num_prompt_tokens = prompt_tokens_per_sequence * prompt_batch_size + max_tokens = kwargs.get("max_new_tokens") or kwargs.get("max_length") + + try: + from unsloth.metrics.stats import get_stats_collector + from unsloth.metrics.prometheus import get_metrics_registry, _metrics_enabled + import uuid + + collector = get_stats_collector() + if collector.is_enabled(): + request_id = str(uuid.uuid4()) + collector.inference_stats.start_request( + request_id = request_id, + num_prompt_tokens = num_prompt_tokens, + max_tokens = max_tokens, + ) + collector.inference_stats.record_scheduled(request_id) + + # Update Prometheus counters + if _metrics_enabled: + registry = get_metrics_registry() + if registry: + registry["inference"]["prompt_tokens_total"].inc(num_prompt_tokens) + registry["inference"]["prompt_tokens"].observe(num_prompt_tokens) + except Exception: + # Metrics collection is optional, continue even if it fails + collector = None + request_id = None + + start_time = time.time() + first_token_seen = False + with torch.inference_mode(), autocaster: output = self._old_generate(*args, **kwargs) + # Record metrics after generation + if collector and request_id: + try: + end_time = time.time() + e2e_latency = end_time - start_time + + # Calculate generated tokens + num_generation_tokens = 0 + effective_prompt_tokens = num_prompt_tokens + if isinstance(output, torch.Tensor): + output_batch_size = output.shape[0] if output.dim() > 1 else 1 + if output_batch_size != prompt_batch_size: + effective_prompt_tokens = ( + prompt_tokens_per_sequence * output_batch_size + ) + if output.dim() > 1: + total_tokens = output.shape[-1] * output.shape[0] + else: + total_tokens = output.shape[-1] + num_generation_tokens = max(0, total_tokens - effective_prompt_tokens) + elif isinstance(output, dict) and "sequences" in output: + # Handle ModelOutput when return_dict_in_generate=True + sequences = output["sequences"] + if isinstance(sequences, torch.Tensor): + output_batch_size = sequences.shape[0] if sequences.dim() > 1 else 1 + if output_batch_size != prompt_batch_size: + effective_prompt_tokens = ( + prompt_tokens_per_sequence * output_batch_size + ) + if sequences.dim() > 1: + total_tokens = sequences.shape[-1] * sequences.shape[0] + else: + total_tokens = sequences.shape[-1] + num_generation_tokens = max( + 0, total_tokens - effective_prompt_tokens + ) + elif hasattr(output, "sequences"): + # Handle ModelOutput object directly + sequences = output.sequences + if isinstance(sequences, torch.Tensor): + output_batch_size = sequences.shape[0] if sequences.dim() > 1 else 1 + if output_batch_size != prompt_batch_size: + effective_prompt_tokens = ( + prompt_tokens_per_sequence * output_batch_size + ) + if sequences.dim() > 1: + total_tokens = sequences.shape[-1] * sequences.shape[0] + else: + total_tokens = sequences.shape[-1] + num_generation_tokens = max( + 0, total_tokens - effective_prompt_tokens + ) + + # Estimate timing (simplified) + # Note: These are estimations. For more accurate metrics, consider hooking into + # the generation process itself (e.g., via LogitsProcessor or StoppingCriteria) + if num_generation_tokens > 0: + # Estimate first token time + estimated_first_token_time = start_time + ( + e2e_latency / (num_generation_tokens + 1) + ) + collector.inference_stats.record_first_token( + request_id, timestamp = estimated_first_token_time + ) + + # Record tokens (simplified - records all at once after generation) + for _ in range(num_generation_tokens): + collector.inference_stats.record_token(request_id) + + # Determine finish reason (simplified - could be improved) + finish_reason = "stop" # Default + if hasattr(output, "finish_reason"): + output_reason = getattr(output, "finish_reason") + if output_reason is not None: + finish_reason = output_reason + elif isinstance(output, dict): + output_reason = output.get("finish_reason") + if output_reason is not None: + finish_reason = output_reason + collector.inference_stats.finish_request( + request_id = request_id, + finish_reason = finish_reason, + num_generation_tokens = num_generation_tokens, + ) + + # Update Prometheus metrics + if _metrics_enabled: + registry = get_metrics_registry() + if registry: + registry["inference"]["request_total"].labels( + finish_reason = finish_reason + ).inc() + registry["inference"]["generation_tokens_total"].inc( + num_generation_tokens + ) + if num_generation_tokens > 0: + registry["inference"]["generation_tokens"].observe( + num_generation_tokens + ) + registry["inference"]["request_latency_seconds"].observe( + e2e_latency + ) + time_per_token = e2e_latency / num_generation_tokens + registry["inference"]["time_per_output_token_seconds"].observe( + time_per_token + ) + except Exception: + # Metrics collection failed, continue + pass + # Delete cached Flex Attention masks to reset inference for name, module in self.named_modules(): if hasattr(module, "_flex_attention_cache"):