Add one-time server opt-out polling for telemetry

This commit is contained in:
Rachel Li 2026-01-15 17:01:25 -08:00 committed by Daniel Han
commit 755bd0f678
2 changed files with 36 additions and 0 deletions

View file

@ -172,6 +172,7 @@ The metrics system works without `prometheus_client` - it gracefully degrades an
- `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
@ -190,6 +191,8 @@ The metrics system works without `prometheus_client` - it gracefully degrades an
- `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

View file

@ -37,11 +37,16 @@ _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:
@ -55,6 +60,7 @@ def enable_telemetry() -> None:
if _TELEMETRY_DISABLED:
return
_TELEMETRY_ENABLED = True
_check_server_opt_out_once()
_start_telemetry_thread()
@ -108,6 +114,7 @@ def _telemetry_worker() -> None:
if item is None:
break
_check_server_opt_out_once()
if is_telemetry_enabled():
_send_telemetry_batch()
@ -179,6 +186,32 @@ def _send_to_server(payload: Dict[str, Any]) -> None:
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():