Databricks + carbon-llm : pull les audit logs Mosaic AI Gateway

Console Databricks, Account Admin

# 1. Account Console → Previews → "system.ai schema" : Enable
# 2. Workspace → Catalog → system → ai → usage_log : Verify Read access

# 3. Pour chaque Serving Endpoint qui sert des LLM :
#    Serving → <endpoint_name> → Edit → AI Gateway tab :
#    [x] Enable usage tracking
#    [x] Enable payload logging (optionnel, utile pour debug, pas requis pour carbon)

# 4. Vérification SQL :
SELECT * FROM system.ai.usage_log
WHERE request_time > current_timestamp() - INTERVAL 1 HOUR
LIMIT 10;

CLI databricks-cli

# Créer un scope si pas déjà fait
databricks secrets create-scope carbon-llm

# Stocker la clé (lecture interactive, pas dans l'historique shell)
databricks secrets put-secret carbon-llm carbon-llm-api-key
# (paste live_xxx en stdin, Ctrl+D)

# Optionnel, base URL si vous routez par un proxy
databricks secrets put-secret carbon-llm carbon-llm-base-url
# (paste https://carbon-llm.com)

# Vérifier
databricks secrets list-secrets carbon-llm

Notebook : carbon_llm_sync_hourly.py

# Databricks notebook source
# COMMAND ----------

import requests
import json
from datetime import datetime, timedelta, timezone

# COMMAND ----------

API_KEY = dbutils.secrets.get(scope="carbon-llm", key="carbon-llm-api-key")
BASE_URL = "https://carbon-llm.com"

# Workspace + catalog → tenant_id naturel pour le multi-tenant
WORKSPACE_ID = spark.conf.get("spark.databricks.workspaceUrl", "unknown")

# COMMAND ----------

# Lire la dernière heure depuis system.ai.usage_log
# Le filtre is null sur error garantit qu'on ne compte que les calls réussis
df = spark.sql("""
    SELECT
        request_time,
        endpoint_name,
        served_entity_name AS model,
        request_token_count AS prompt_tokens,
        response_token_count AS completion_tokens,
        request_id,
        identity_metadata.run_as AS run_as_user,
        request_metadata
    FROM system.ai.usage_log
    WHERE request_time > current_timestamp() - INTERVAL 1 HOUR + INTERVAL 5 MINUTES
      AND request_time <= current_timestamp() - INTERVAL 5 MINUTES
      AND served_entity_name IS NOT NULL
      AND request_token_count IS NOT NULL
      AND error_code IS NULL
""")

print(f"Rows à forwarder : {df.count()}")

# COMMAND ----------

def forward_event(row):
    """POST vers /api/v1/track avec event_id idempotent dérivé de request_id."""
    payload = {
        "event_id": f"databricks:{WORKSPACE_ID}:{row.request_id}",
        "model": row.model,
        "prompt_tokens": int(row.prompt_tokens or 0),
        "completion_tokens": int(row.completion_tokens or 0),
        "tenant_id": f"databricks:{row.endpoint_name}",
        "ingestion_source": "databricks_ai_gateway",
        "client_instance_id": row.run_as_user or "unknown",
    }

    try:
        r = requests.post(
            f"{BASE_URL}/api/v1/track",
            headers={
                "Authorization": f"Bearer {API_KEY}",
                "Content-Type": "application/json",
            },
            json=payload,
            timeout=10.0,
        )
        # 202 Accepted ou 409 Conflict (replay event_id) sont OK
        return r.status_code in (200, 202, 409)
    except Exception as e:
        print(f"forward_event failed: {e}")
        return False

# Foreach séquentiel, pour les volumes Databricks (1k-100k events/h),
# pas besoin de paralléliser. Le bottleneck est carbon-llm (rate-limit Pro
# = 600 req/min) pas Databricks.
results = [forward_event(row) for row in df.collect()]
print(f"Successfully forwarded: {sum(results)} / {len(results)}")

Job JSON spec (workflows.create)

{
  "name": "carbon-llm-sync-hourly",
  "tasks": [{
    "task_key": "sync_to_carbon_llm",
    "notebook_task": {
      "notebook_path": "/Users/data-eng/carbon_llm_sync_hourly",
      "source": "WORKSPACE"
    },
    "existing_cluster_id": "0123-456789-abcdef"
  }],
  "schedule": {
    "quartz_cron_expression": "0 0 * * * ?",
    "timezone_id": "Europe/Paris",
    "pause_status": "UNPAUSED"
  },
  "email_notifications": {
    "on_failure": ["data-eng@company.com"]
  },
  "max_concurrent_runs": 1
}