From d0e85aab265139fb51e0d03c493028c6d08d2e95 Mon Sep 17 00:00:00 2001 From: Rachel Li Date: Thu, 15 Jan 2026 16:14:16 -0800 Subject: [PATCH] Make Prometheus init robust to duplicates --- unsloth/metrics/prometheus.py | 74 +++++++++++++++++++++++++---------- 1 file changed, 54 insertions(+), 20 deletions(-) diff --git a/unsloth/metrics/prometheus.py b/unsloth/metrics/prometheus.py index b076fccdb2..581a7f16ff 100644 --- a/unsloth/metrics/prometheus.py +++ b/unsloth/metrics/prometheus.py @@ -83,6 +83,20 @@ def _get_existing_collector(metric_name: str): return REGISTRY._names_to_collectors.get(metric_name) # type: ignore[attr-defined] +def _get_or_create_metric(metric_cls, name: str, *args, **kwargs): + """Return existing collector if present, otherwise create one safely.""" + existing = _get_existing_collector(name) + if existing is not None: + return existing + try: + return metric_cls(name, *args, **kwargs) + except ValueError: + existing = _get_existing_collector(name) + if existing is not None: + return existing + raise + + def _init_metrics(): """Initialize Prometheus metrics if available.""" global _metrics_registry @@ -96,55 +110,66 @@ def _init_metrics(): # Inference metrics inference_metrics = { # Counters - "request_total": _get_existing_collector("unsloth_request_total") or Counter( + "request_total": _get_or_create_metric( + Counter, "unsloth_request_total", "Total number of inference requests", ["finish_reason"], ), - "prompt_tokens_total": _get_existing_collector("unsloth_prompt_tokens_total") or Counter( + "prompt_tokens_total": _get_or_create_metric( + Counter, "unsloth_prompt_tokens_total", "Total number of prompt tokens processed", ), - "generation_tokens_total": _get_existing_collector("unsloth_generation_tokens_total") or Counter( + "generation_tokens_total": _get_or_create_metric( + Counter, "unsloth_generation_tokens_total", "Total number of generation tokens produced", ), # Gauges - "requests_active": _get_existing_collector("unsloth_requests_active") or Gauge( + "requests_active": _get_or_create_metric( + Gauge, "unsloth_requests_active", "Number of currently active inference requests", ), - "tokens_per_second": _get_existing_collector("unsloth_tokens_per_second") or Gauge( + "tokens_per_second": _get_or_create_metric( + Gauge, "unsloth_tokens_per_second", "Current tokens per second throughput", ), # Histograms - "request_latency_seconds": _get_existing_collector("unsloth_request_latency_seconds") or Histogram( + "request_latency_seconds": _get_or_create_metric( + 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( + "prefill_latency_seconds": _get_or_create_metric( + 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( + "decode_latency_seconds": _get_or_create_metric( + 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( + "time_per_output_token_seconds": _get_or_create_metric( + 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( + "prompt_tokens": _get_or_create_metric( + 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( + "generation_tokens": _get_or_create_metric( + Histogram, "unsloth_generation_tokens", "Number of generation tokens per request", buckets = [10, 50, 100, 500, 1000, 2000, 4000, 8000, 16000, 32000], @@ -154,43 +179,52 @@ def _init_metrics(): # Training metrics training_metrics = { # Counters - "training_steps_total": _get_existing_collector("unsloth_training_steps_total") or Counter( + "training_steps_total": _get_or_create_metric( + Counter, "unsloth_training_steps_total", "Total number of training steps", ), - "training_samples_total": _get_existing_collector("unsloth_training_samples_total") or Counter( + "training_samples_total": _get_or_create_metric( + Counter, "unsloth_training_samples_total", "Total number of training samples processed", ), # Gauges - "training_loss": _get_existing_collector("unsloth_training_loss") or Gauge( + "training_loss": _get_or_create_metric( + Gauge, "unsloth_training_loss", "Current training loss", ), - "learning_rate": _get_existing_collector("unsloth_learning_rate") or Gauge( + "learning_rate": _get_or_create_metric( + Gauge, "unsloth_learning_rate", "Current learning rate", ), - "samples_per_second": _get_existing_collector("unsloth_training_samples_per_second") or Gauge( + "samples_per_second": _get_or_create_metric( + Gauge, "unsloth_training_samples_per_second", "Training throughput in samples per second", ), - "gradient_norm": _get_existing_collector("unsloth_gradient_norm") or Gauge( + "gradient_norm": _get_or_create_metric( + Gauge, "unsloth_gradient_norm", "Current gradient norm", ), # Histograms - "forward_time_seconds": _get_existing_collector("unsloth_training_forward_time_seconds") or Histogram( + "forward_time_seconds": _get_or_create_metric( + 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( + "backward_time_seconds": _get_or_create_metric( + 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( + "batch_size": _get_or_create_metric( + Histogram, "unsloth_training_batch_size", "Training batch size", buckets = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512],