feat: secure metrics ingestion

.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+.pytest_cache/
+.ruff_cache/

README.md

@@ -12,5 +12,35 @@ hf_oauth_scopes:
 ---
 
 Every [OpenEnv Arena](https://huggingface.co/spaces/openenvarena/arena) training run streams its numbers here while it
-trains: reward, loss and how long each step took. Runs are named after the participant and their submission, and the
+trains: reward, loss and how long each step took. Each participant has a project with one run per training run, and the
 arena's own submissions page links each run to its curves. Built with [Trackio](https://github.com/gradio-app/trackio).
+
+## Worker ingestion
+
+Set the Space secret `TRACKIO_INGEST_TOKEN` to the same dedicated random credential
+configured on Arena's workers. It only authorizes `POST /ingest`; it is not a Hugging
+Face token or Trackio's dashboard administration token. When unset, ingestion is
+disabled and the dashboard continues serving existing runs.
+
+Send `Authorization: Bearer <credential>` with a JSON body:
+
+```json
+{"logs":[{"project":"burtenshaw","run":"arena-run-id","run_id":"arena-run-id","step":1,"log_id":"unique-event-id","metrics":{"reward":0.75,"arena/step_seconds":1.2},"config":{"submission_id":"submission-id"}}]}
+```
+
+A batch contains 1–128 events for one project and run, up to 256 KiB. Each event
+contains 1–128 finite numeric metrics. Project, run and event IDs contain only
+ASCII letters, digits, `_` and `-`, up to 128 characters. `run_id` may be omitted;
+when present it must equal `run`. Preserve each `log_id` when retrying delivery;
+duplicates cannot replace existing metrics. Distinct events at the same step need
+distinct IDs.
+
+Optional config accepts strings `participant`, `submission_id`, `run_id`, `model`,
+`model_revision`, `recipe`, and nonnegative numbers `num_updates`, `train_tasks`.
+Supply it on the first batch; it cannot be changed after a run exists. Repeating
+the same config on subsequent events is supported.
+
+The existing `openenvarena/training-dashboard` bucket stays mounted at `/data`, with
+`TRACKIO_DIR=/data/trackio`. The app uses this mounted storage directly and does not
+provision Spaces, buckets or credentials. Run one server process for Trackio's
+bucket-backed SQLite storage. Trackio remains pinned to `0.40.0`.

app.py

@@ -1,2 +1,22 @@
-import trackio
-trackio.show()
+import os
+
+import uvicorn
+from starlette.routing import Route
+from trackio.server import build_starlette_app_only
+
+from ingest import ingest
+
+
+def create_app():
+    # Keep Trackio's dashboard, authentication and mounted bucket configuration.
+    app, _ = build_starlette_app_only(mcp_server=False)
+    app.router.routes.insert(0, Route("/ingest", ingest, methods=["POST"]))
+    return app
+
+
+if __name__ == "__main__":
+    uvicorn.run(
+        create_app(),
+        host="0.0.0.0",
+        port=int(os.environ.get("GRADIO_SERVER_PORT", "7860")),
+    )

ingest.py

@@ -0,0 +1,150 @@
+"""Append-only numeric metric ingestion for trusted Arena workers."""
+
+import json
+import math
+import os
+import re
+import secrets
+from threading import Lock
+
+from starlette.concurrency import run_in_threadpool
+from starlette.exceptions import HTTPException
+from starlette.responses import JSONResponse
+from trackio.sqlite_storage import SQLiteStorage
+
+MAX_BODY_BYTES = 256 * 1024
+MAX_LOGS = 128
+MAX_METRICS = 128
+IDENTIFIER = re.compile(r"[A-Za-z0-9_-]{1,128}\Z")
+STRING_CONFIG = {
+    "participant",
+    "submission_id",
+    "run_id",
+    "model",
+    "model_revision",
+    "recipe",
+}
+NUMBER_CONFIG = {"num_updates", "train_tasks"}
+LOG_FIELDS = {"project", "run", "run_id", "step", "log_id", "metrics", "config"}
+# Trackio's bucket-backed SQLite storage supports one server process. Serialize
+# the configuration check with ingestion so two first batches cannot replace it.
+_write_lock = Lock()
+
+
+def _number(value):
+    if type(value) is int:
+        return -(2**63) <= value <= 2**63 - 1
+    return type(value) is float and math.isfinite(value)
+
+
+def _identifier(value):
+    return isinstance(value, str) and IDENTIFIER.fullmatch(value) is not None
+
+
+def _config(value):
+    if value is None:
+        return None
+    if not isinstance(value, dict) or value.keys() - STRING_CONFIG - NUMBER_CONFIG:
+        raise ValueError("config contains unsupported fields")
+    for key, item in value.items():
+        if key in STRING_CONFIG:
+            if not isinstance(item, str) or len(item) > 512:
+                raise ValueError(
+                    "config string values must have at most 512 characters"
+                )
+            item.encode("utf-8")
+        elif not _number(item) or item < 0:
+            raise ValueError("config counts must be finite nonnegative numbers")
+    return value or None
+
+
+def validate(payload):
+    if not isinstance(payload, dict) or set(payload) != {"logs"}:
+        raise ValueError("expected an object containing only logs")
+    logs = payload["logs"]
+    if not isinstance(logs, list) or not 1 <= len(logs) <= MAX_LOGS:
+        raise ValueError(f"logs must contain 1 to {MAX_LOGS} entries")
+    identity = None
+    config = None
+    for entry in logs:
+        if not isinstance(entry, dict) or entry.keys() - LOG_FIELDS:
+            raise ValueError("log contains unsupported fields")
+        for key in ("project", "run", "log_id"):
+            if not _identifier(entry.get(key)):
+                raise ValueError(f"{key} must contain 1 to 128 letters, digits, _ or -")
+        if entry.get("run_id", entry["run"]) != entry["run"]:
+            raise ValueError("run_id must equal run")
+        current = (entry["project"], entry["run"])
+        if identity is not None and current != identity:
+            raise ValueError("a batch must contain one project and run")
+        identity = current
+        step = entry.get("step")
+        if type(step) is not int or not 0 <= step <= 2**63 - 1:
+            raise ValueError("step must be a nonnegative 64-bit integer")
+        metrics = entry.get("metrics")
+        if not isinstance(metrics, dict) or not 1 <= len(metrics) <= MAX_METRICS:
+            raise ValueError(f"metrics must contain 1 to {MAX_METRICS} numeric values")
+        for key, value in metrics.items():
+            if not isinstance(key, str) or not 1 <= len(key) <= 128:
+                raise ValueError("metric names must contain 1 to 128 characters")
+            key.encode("utf-8")
+            if not _number(value):
+                raise ValueError("metrics must contain only finite numbers")
+        entry_config = _config(entry.get("config"))
+        if entry_config is not None:
+            if config is not None and entry_config != config:
+                raise ValueError("config must be consistent within a batch")
+            config = entry_config
+    return logs, identity, config
+
+
+def append(logs, identity, config):
+    project, run = identity
+    with _write_lock:
+        existing = next(
+            (
+                record
+                for record in SQLiteStorage.get_run_records(project)
+                if record["id"] == run
+            ),
+            None,
+        )
+        stored = SQLiteStorage.get_run_config(project, run, run_id=run)
+        if existing is not None or stored is not None:
+            if config is not None and config != stored:
+                raise HTTPException(
+                    409, "configuration of an existing run is immutable"
+                )
+            config = None
+        SQLiteStorage.bulk_log(
+            project=project,
+            run=run,
+            run_id=run,
+            metrics_list=[entry["metrics"] for entry in logs],
+            steps=[entry["step"] for entry in logs],
+            # Scope event IDs to the run; retries cannot replace existing rows.
+            log_ids=[f"{run}:{entry['log_id']}" for entry in logs],
+            config=config,
+        )
+
+
+async def ingest(request):
+    expected = os.environ.get("TRACKIO_INGEST_TOKEN", "")
+    if not expected:
+        return JSONResponse(
+            {"error": "metric ingestion is not configured"}, status_code=503
+        )
+    supplied = request.headers.get("authorization", "")
+    if not secrets.compare_digest(supplied.encode(), f"Bearer {expected}".encode()):
+        return JSONResponse({"error": "invalid ingest credential"}, status_code=401)
+    body = bytearray()
+    async for chunk in request.stream():
+        if len(body) + len(chunk) > MAX_BODY_BYTES:
+            return JSONResponse({"error": "metric batch is too large"}, status_code=413)
+        body.extend(chunk)
+    try:
+        logs, identity, config = validate(json.loads(body))
+    except (ValueError, OverflowError, RecursionError):
+        return JSONResponse({"error": "invalid metric batch"}, status_code=400)
+    await run_in_threadpool(append, logs, identity, config)
+    return JSONResponse({"accepted": len(logs)})

test_ingest.py

@@ -0,0 +1,256 @@
+import importlib
+import json
+import os
+import re
+import subprocess
+import sys
+
+import pytest
+from starlette.applications import Starlette
+from starlette.routing import Route
+from starlette.testclient import TestClient
+
+
+@pytest.fixture
+def client(tmp_path, monkeypatch):
+    import trackio.sqlite_storage
+    import trackio.utils
+
+    monkeypatch.setattr(trackio.sqlite_storage, "TRACKIO_DIR", tmp_path)
+    monkeypatch.setattr(trackio.utils, "TRACKIO_DIR", tmp_path)
+    monkeypatch.setenv("TRACKIO_INGEST_TOKEN", "test-ingest-token")
+    module = importlib.import_module("ingest")
+    app = Starlette(routes=[Route("/ingest", module.ingest, methods=["POST"])])
+    with TestClient(app) as test_client:
+        test_client.headers["Authorization"] = "Bearer test-ingest-token"
+        yield test_client
+
+
+def batch(**updates):
+    entry = {
+        "project": "burtenshaw",
+        "run": "arena-run-123",
+        "run_id": "arena-run-123",
+        "step": 1,
+        "log_id": "event-1",
+        "metrics": {"reward/mean": 0.75, "arena/step_seconds": 1.2},
+        "config": {
+            "participant": "space:burtenshaw",
+            "submission_id": "submission-123",
+        },
+    }
+    entry.update(updates)
+    return {"logs": [entry]}
+
+
+def test_missing_or_invalid_auth_cannot_write(client, monkeypatch):
+    for authorization in ("", "Bearer wrong", "test-ingest-token"):
+        assert (
+            client.post(
+                "/ingest", json=batch(), headers={"Authorization": authorization}
+            ).status_code
+            == 401
+        )
+    monkeypatch.delenv("TRACKIO_INGEST_TOKEN")
+    assert client.post("/ingest", json=batch()).status_code == 503
+
+
+@pytest.mark.parametrize(
+    "updates",
+    [
+        {"project": "../another-project"},
+        {"run_id": "another-run"},
+        {"step": -1},
+        {"step": True},
+        {"log_id": ""},
+        {"metrics": {"loss": "private text"}},
+        {"metrics": {"loss": True}},
+        {"metrics": {"loss": 2**80}},
+        {"metrics": {"loss": {"_type": "html", "html": "<script>"}}},
+        {"metrics": {str(i): i for i in range(129)}},
+        {"config": {"HF_TOKEN": "not-allowed"}},
+        {"config": {"model": ["not-a-string"]}},
+        {"config": {"num_updates": 2**80}},
+        {"delete": True},
+    ],
+)
+def test_invalid_payload_rejected(client, updates):
+    assert client.post("/ingest", json=batch(**updates)).status_code == 400
+
+
+def test_nonfinite_and_oversized_payload_rejected(client):
+    assert (
+        client.post(
+            "/ingest", content=json.dumps(batch(metrics={"loss": float("nan")}))
+        ).status_code
+        == 400
+    )
+    assert (
+        client.post(
+            "/ingest", content=json.dumps(batch(metrics={"loss": float("inf")}))
+        ).status_code
+        == 400
+    )
+    assert client.post("/ingest", content=b" " * (256 * 1024 + 1)).status_code == 413
+    assert (
+        client.post("/ingest", json={"logs": batch()["logs"] * 129}).status_code == 400
+    )
+
+
+@pytest.mark.parametrize(
+    "updates",
+    [
+        {"metrics": {"\ud800": 0.2}},
+        {"config": {"model": "\ud800"}},
+    ],
+)
+def test_unpaired_surrogates_rejected(client, updates):
+    response = client.post("/ingest", content=json.dumps(batch(**updates)))
+    assert response.status_code == 400
+
+
+def test_retry_is_idempotent_and_cannot_overwrite_metric(client):
+    from trackio.sqlite_storage import SQLiteStorage
+
+    assert client.post("/ingest", json=batch()).json() == {"accepted": 1}
+    assert client.post("/ingest", json=batch()).status_code == 200
+    assert (
+        client.post("/ingest", json=batch(metrics={"reward/mean": 999})).status_code
+        == 200
+    )
+    with SQLiteStorage._get_connection(
+        SQLiteStorage.get_project_db_path("burtenshaw")
+    ) as conn:
+        rows = conn.execute("SELECT metrics FROM metrics").fetchall()
+    assert len(rows) == 1
+    assert json.loads(rows[0][0])["reward/mean"] == 0.75
+    assert (
+        client.post("/ingest", json=batch(log_id="event-2", config=None)).status_code
+        == 200
+    )
+
+
+def test_existing_config_cannot_be_changed(client):
+    from trackio.sqlite_storage import SQLiteStorage
+
+    assert client.post("/ingest", json=batch()).status_code == 200
+    assert (
+        client.post(
+            "/ingest", json=batch(config={"model": "changed"}, log_id="event-2")
+        ).status_code
+        == 409
+    )
+    assert (
+        SQLiteStorage.get_run_config(
+            "burtenshaw", "arena-run-123", run_id="arena-run-123"
+        )
+        == batch()["logs"][0]["config"]
+    )
+
+
+def test_config_cannot_be_added_to_an_existing_run(client):
+    assert client.post("/ingest", json=batch(config=None)).status_code == 200
+    assert client.post("/ingest", json=batch(log_id="event-2")).status_code == 409
+
+
+def test_config_without_metric_rows_cannot_be_overwritten(client):
+    from trackio.sqlite_storage import SQLiteStorage
+
+    assert client.post("/ingest", json=batch()).status_code == 200
+    with SQLiteStorage._get_connection(
+        SQLiteStorage.get_project_db_path("burtenshaw")
+    ) as conn:
+        conn.execute("DELETE FROM metrics")
+    assert SQLiteStorage.get_run_records("burtenshaw") == []
+    assert (
+        client.post("/ingest", json=batch(config={"model": "changed"})).status_code
+        == 409
+    )
+    assert (
+        SQLiteStorage.get_all_run_configs("burtenshaw")["arena-run-123"]
+        == batch()["logs"][0]["config"]
+    )
+
+
+def test_batches_cannot_mix_runs_or_mutate_partial_data(client):
+    from trackio.sqlite_storage import SQLiteStorage
+
+    payload = {"logs": batch()["logs"] + batch(run="another", run_id="another")["logs"]}
+    assert client.post("/ingest", json=payload).status_code == 400
+    assert SQLiteStorage.get_run_records("burtenshaw") == []
+
+
+def test_dashboard_still_serves_without_ingest_credentials(client, monkeypatch):
+    import huggingface_hub
+    from app import create_app
+
+    def unexpected_mutation(*args, **kwargs):
+        pytest.fail("starting the dashboard must not provision or upload Hub resources")
+
+    for name in (
+        "create_repo",
+        "create_bucket",
+        "set_space_volumes",
+        "upload_file",
+        "add_space_secret",
+    ):
+        monkeypatch.setattr(huggingface_hub, name, unexpected_mutation)
+        monkeypatch.setattr(huggingface_hub.HfApi, name, unexpected_mutation)
+    monkeypatch.delenv("TRACKIO_INGEST_TOKEN")
+    monkeypatch.setenv("TRACKIO_DIR", "/data/trackio")
+    with TestClient(create_app()) as dashboard:
+        page = dashboard.get("/")
+        assert page.status_code == 200
+        asset = re.search(r'<script[^>]+src="([^"]+)"', page.text)
+        assert asset is not None
+        assert dashboard.get(asset.group(1)).status_code == 200
+        assert dashboard.post("/ingest", json=batch()).status_code == 503
+    assert os.environ["TRACKIO_DIR"] == "/data/trackio"
+
+
+def test_metrics_survive_fresh_process_with_same_storage(client, tmp_path):
+    assert client.post("/ingest", json=batch()).status_code == 200
+    result = subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            """
+import json
+from trackio.sqlite_storage import SQLiteStorage
+rows = SQLiteStorage.get_run_records('burtenshaw')
+assert rows[0]['id'] == 'arena-run-123'
+with SQLiteStorage._get_connection(SQLiteStorage.get_project_db_path('burtenshaw')) as conn:
+    row = conn.execute('SELECT metrics FROM metrics').fetchone()
+assert json.loads(row[0])['reward/mean'] == 0.75
+print('persisted metrics recovered')
+""",
+        ],
+        env={**os.environ, "TRACKIO_DIR": str(tmp_path)},
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    assert "persisted metrics recovered" in result.stdout
+
+
+def test_dashboard_does_not_accept_ingest_token_for_delete_or_upload(
+    client, monkeypatch
+):
+    monkeypatch.setenv("SYSTEM", "spaces")
+    from app import create_app
+
+    assert client.post("/ingest", json=batch()).status_code == 200
+    with TestClient(create_app()) as dashboard:
+        dashboard.headers["Authorization"] = "Bearer test-ingest-token"
+        response = dashboard.post(
+            "/api/delete_run",
+            json={"kwargs": {"project": "burtenshaw", "run": "arena-run-123"}},
+        )
+        assert response.status_code >= 400 or response.json().get("error")
+        response = dashboard.post(
+            "/api/upload_db_to_space",
+            json={
+                "kwargs": {"project": "burtenshaw", "uploaded_db": {}, "hf_token": None}
+            },
+        )
+        assert response.status_code >= 400 or response.json().get("error")