Spaces:
Paused
Paused
| 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 | |
| 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 alert_batch(**updates): | |
| entry = { | |
| "project": "burtenshaw", | |
| "run": "arena-run-123", | |
| "run_id": "arena-run-123", | |
| "step": 39, | |
| "alert_id": "failure-39", | |
| "title": "Sandbox cleanup timed out", | |
| "text": "Training ended after 39 updates. Evaluation will use the last checkpoint.", | |
| "level": "error", | |
| "config": batch()["logs"][0]["config"], | |
| } | |
| entry.update(updates) | |
| return {"alerts": [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 | |
| 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 | |
| ) | |
| 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") | |
| def test_alert_retry_is_idempotent_and_does_not_replace_the_original(client): | |
| from trackio.sqlite_storage import SQLiteStorage | |
| assert client.post("/ingest", json=alert_batch()).json() == {"accepted": 1} | |
| assert client.post("/ingest", json=alert_batch()).status_code == 200 | |
| assert ( | |
| client.post("/ingest", json=alert_batch(text="replacement")).status_code == 200 | |
| ) | |
| assert ( | |
| client.post("/ingest", json=alert_batch(alert_id="failure-40")).status_code | |
| == 200 | |
| ) | |
| alerts = SQLiteStorage.get_alerts("burtenshaw", run_id="arena-run-123") | |
| assert len(alerts) == 2 | |
| assert all(alert["text"] == alert_batch()["alerts"][0]["text"] for alert in alerts) | |
| assert all(alert["level"] == "error" and alert["step"] == 39 for alert in alerts) | |
| def test_early_alert_registers_a_visible_run_without_fabricating_metrics(client): | |
| from trackio.sqlite_storage import SQLiteStorage | |
| from app import create_app | |
| assert client.post("/ingest", json=alert_batch(step=None)).status_code == 200 | |
| records = SQLiteStorage.get_run_records("burtenshaw") | |
| assert [record["id"] for record in records] == ["arena-run-123"] | |
| with SQLiteStorage._get_connection( | |
| SQLiteStorage.get_project_db_path("burtenshaw") | |
| ) as conn: | |
| metrics = conn.execute("SELECT metrics FROM metrics").fetchall() | |
| assert [json.loads(row[0]) for row in metrics] == [{}] | |
| with TestClient(create_app()) as dashboard: | |
| assert ( | |
| dashboard.get("/reports?project=burtenshaw&runs=arena-run-123").status_code | |
| == 200 | |
| ) | |
| response = dashboard.post( | |
| "/api/get_alerts", | |
| json={ | |
| "kwargs": { | |
| "project": "burtenshaw", | |
| "run": "arena-run-123", | |
| "level": "error", | |
| } | |
| }, | |
| ) | |
| assert response.status_code == 200 | |
| alerts = response.json()["data"] | |
| assert len(alerts) == 1 | |
| assert alerts[0]["title"] == "Sandbox cleanup timed out" | |
| assert alerts[0]["step"] is None | |
| # The next real metric can use the same immutable configuration. | |
| assert client.post("/ingest", json=batch()).status_code == 200 | |
| def test_metrics_and_alerts_share_run_identity_and_immutable_config(client): | |
| from trackio.sqlite_storage import SQLiteStorage | |
| payload = {**batch(), **alert_batch()} | |
| assert client.post("/ingest", json=payload).json() == {"accepted": 2} | |
| assert ( | |
| client.post( | |
| "/ingest", json=alert_batch(config={"model": "changed"}) | |
| ).status_code | |
| == 409 | |
| ) | |
| assert len(SQLiteStorage.get_alerts("burtenshaw", run_id="arena-run-123")) == 1 | |
| another = alert_batch(run="another", run_id="another") | |
| assert client.post("/ingest", json=another).status_code == 200 | |
| assert len(SQLiteStorage.get_alerts("burtenshaw", run_id="another")) == 1 | |
| def test_invalid_alert_rejected_before_any_metric_write(client, updates): | |
| from trackio.sqlite_storage import SQLiteStorage | |
| payload = {**batch(), **alert_batch(**updates)} | |
| assert client.post("/ingest", json=payload).status_code == 400 | |
| assert SQLiteStorage.get_run_records("burtenshaw") == [] | |
| def test_alerts_have_the_same_auth_batch_and_unicode_limits(client): | |
| for authorization in ("", "Bearer wrong"): | |
| assert ( | |
| client.post( | |
| "/ingest", json=alert_batch(), headers={"Authorization": authorization} | |
| ).status_code | |
| == 401 | |
| ) | |
| payload = {**batch(), "alerts": alert_batch()["alerts"] * 128} | |
| assert client.post("/ingest", json=payload).status_code == 400 | |
| for payload in ( | |
| {}, | |
| {"alerts": []}, | |
| {"logs": [], "alerts": []}, | |
| {"alerts": {}}, | |
| {**batch(), **alert_batch(run="another", run_id="another")}, | |
| ): | |
| assert client.post("/ingest", json=payload).status_code == 400 | |
| for field in ("title", "text"): | |
| assert ( | |
| client.post( | |
| "/ingest", content=json.dumps(alert_batch(**{field: "\ud800"})) | |
| ).status_code | |
| == 400 | |
| ) | |
| def test_alerts_survive_fresh_process_with_same_storage(client, tmp_path): | |
| assert client.post("/ingest", json=alert_batch()).status_code == 200 | |
| result = subprocess.run( | |
| [ | |
| sys.executable, | |
| "-c", | |
| ( | |
| "from trackio.sqlite_storage import SQLiteStorage; " | |
| "alerts = SQLiteStorage.get_alerts('burtenshaw', run_id='arena-run-123'); " | |
| "assert len(alerts) == 1 and alerts[0]['step'] == 39" | |
| ), | |
| ], | |
| env={**os.environ, "TRACKIO_DIR": str(tmp_path)}, | |
| capture_output=True, | |
| text=True, | |
| check=False, | |
| ) | |
| assert result.returncode == 0, result.stderr | |