From 59369b31522cd9cfe62f474ded994d27a648f29b Mon Sep 17 00:00:00 2001 From: Rachel Li Date: Thu, 15 Jan 2026 15:24:44 -0800 Subject: [PATCH] Add opt-in telemetry forwarding for metrics --- unsloth/__init__.py | 3 + unsloth/metrics/README.md | 10 ++ unsloth/metrics/__init__.py | 10 ++ unsloth/metrics/stats.py | 16 +++ unsloth/metrics/telemetry.py | 198 +++++++++++++++++++++++++++++++++++ 5 files changed, 237 insertions(+) create mode 100644 unsloth/metrics/telemetry.py diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 95f3ae4fbc..650c37cd8a 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -330,6 +330,9 @@ from .metrics import ( 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 diff --git a/unsloth/metrics/README.md b/unsloth/metrics/README.md index 28e610d006..a2f4c1ed19 100644 --- a/unsloth/metrics/README.md +++ b/unsloth/metrics/README.md @@ -169,6 +169,10 @@ The metrics system works without `prometheus_client` - it gracefully degrades an ## 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_INTERVAL` - Telemetry interval seconds (default: 300) ## API Reference @@ -180,6 +184,12 @@ The metrics system works without `prometheus_client` - it gracefully degrades an - `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 + ### HTTP Server Functions - `start_metrics_server(host="0.0.0.0", port=9090)` - Start metrics HTTP server diff --git a/unsloth/metrics/__init__.py b/unsloth/metrics/__init__.py index a6a4274056..a4c990e61c 100644 --- a/unsloth/metrics/__init__.py +++ b/unsloth/metrics/__init__.py @@ -36,6 +36,12 @@ from unsloth.metrics.server import ( is_metrics_server_running, test_metrics_server, ) +from unsloth.metrics.telemetry import ( + enable_telemetry, + disable_telemetry, + is_telemetry_enabled, + schedule_telemetry, +) __all__ = [ "InferenceStats", @@ -51,4 +57,8 @@ __all__ = [ "stop_metrics_server", "is_metrics_server_running", "test_metrics_server", + "enable_telemetry", + "disable_telemetry", + "is_telemetry_enabled", + "schedule_telemetry", ] diff --git a/unsloth/metrics/stats.py b/unsloth/metrics/stats.py index 75cf3bb62d..446d30d562 100644 --- a/unsloth/metrics/stats.py +++ b/unsloth/metrics/stats.py @@ -186,6 +186,14 @@ class InferenceStats: # 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: @@ -304,6 +312,14 @@ class TrainingStats: 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: diff --git a/unsloth/metrics/telemetry.py b/unsloth/metrics/telemetry.py new file mode 100644 index 0000000000..34ed551a03 --- /dev/null +++ b/unsloth/metrics/telemetry.py @@ -0,0 +1,198 @@ +# 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_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() + + +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 + _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 + + 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 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()