ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Make Prometheus init robust to duplicates

Browse source
This commit is contained in:
Rachel Li 2026-01-15 16:14:16 -08:00 • committed by Daniel Han
commit d0e85aab26
1 changed files with 54 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

74
unsloth/metrics/prometheus.py
View file

@ -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],