Spaces:
Paused
Paused
File size: 7,004 Bytes
935d4b2 a456b58 935d4b2 a456b58 935d4b2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """A generation-only Studio cannot advertise or admit the other generation."""
import os
import unittest
from unittest.mock import patch
from fastapi.testclient import TestClient
from app import create_app
from model_registry import PROFILES, model_registry
from tetris_arena import HTTPDecisionAdapter, LOCAL_MODELS
TOKEN = "t" * 32
QUESTIONS = {"accept": {"type": "noul", "instructions": "Is this accepted?"}}
def generation_registry(version):
return [
{
"id": wire,
"label": profile["repo_id"].rsplit("/", 1)[-1].split("-")[2],
"version": version,
"manifest_sha256": f"{index + 1:064x}",
"revision": f"{index + 1:040x}",
}
for index, (wire, profile) in enumerate(PROFILES.items())
if f"/Decision-{version}-" in profile["repo_id"]
]
class GenerationSplitTests(unittest.TestCase):
def test_each_generation_is_an_independent_valid_registry(self):
for version in ("1.0", "2.0"):
with self.subTest(version=version):
registry = model_registry(generation_registry(version))
self.assertEqual(len(registry), 6)
self.assertTrue(all(f"/Decision-{version}-" in row["repo_id"]
for row in registry.values()))
def test_native_rejects_single_models_other_than_decision1_kai(self):
for wire in ("decision2-kai", "decision-lux"):
version = "2.0" if wire.startswith("decision2-") else "1.0"
entry = next(row for row in generation_registry(version) if row["id"] == wire)
with self.subTest(wire=wire), self.assertRaisesRegex(ValueError, "Native mode requires"):
create_app(mode="native", registry=[entry])
def test_pull_queue_discovery_readiness_and_admission_follow_generation(self):
for version, other in (("1.0", "2.0"), ("2.0", "1.0")):
with self.subTest(version=version), patch.dict(os.environ, {
"DECISION_WORKER_TOKEN": TOKEN,
"TETRIS_LOCAL_API_URL": "http://127.0.0.1:7860/v1/systemone",
"TETRIS_JEV_API_URL": "https://cloud.example.test/v1/systemone",
"TETRIS_JEV_API_KEY": "test-cloud-key",
"TETRIS_SYSTEMONE_API_URL": "https://mirror.example.test/v1/systemone",
"TETRIS_SYSTEMONE_API_KEY": "test-mirror-key",
}, clear=True):
entries = generation_registry(version)
api = create_app(mode="pull_queue", registry=entries)
with TestClient(api) as client:
expected = {PROFILES[row["id"]]["repo_id"] for row in entries}
self.assertEqual(set(api.state.relays), {row["id"] for row in entries})
self.assertEqual({row["id"] for row in client.get("/v1/models").json()["models"]}, expected)
self.assertEqual(client.get("/api/ready").status_code, 503)
for index, row in enumerate(entries):
heartbeat = client.post("/internal/worker/heartbeat", headers={
"Authorization": "Bearer " + TOKEN,
}, json={
"model": row["id"],
"manifest_sha256": row["manifest_sha256"],
"worker_id": f"{index + 1:032x}",
"phase": "ready",
"capabilities": ["context_batch_v1"],
})
self.assertEqual(heartbeat.status_code, 200, heartbeat.text)
for data in ({"state": "A request"}, {"states": [{"id": "a", "state": "A request"}]}):
job = client.post("/api/jobs", json={
"model": row["id"], "questions": QUESTIONS, **data,
})
self.assertEqual(job.status_code, 202, job.text)
self.assertEqual(job.json()["model"], row["id"])
client.delete(f"/api/jobs/{job.json()['id']}", params={"model": row["id"]})
ready = client.get("/api/ready")
self.assertEqual(ready.status_code, 200)
self.assertEqual(ready.json()["models"], dict.fromkeys(expected, True))
catalog = client.get("/api/tetris/config").json()["competitors"]
allowed_competitors = {row.id for row in LOCAL_MODELS if row.request_model in expected}
self.assertEqual({row["id"] for row in catalog if row["family"] == "decision"}, allowed_competitors)
self.assertTrue(all(row["ready"] for row in catalog if row["family"] == "decision"))
self.assertFalse(any(row["family"] == "cloud" for row in catalog))
self.assertNotIn("jev-cloud", api.state.tetris.adapter._endpoints)
self.assertNotIn("jev-cloud-mirror", api.state.tetris.adapter._endpoints)
for cloud in ("jev-cloud", "jev-cloud-mirror"):
race = client.post("/api/tetris/races", json={
"left": cloud, "right": cloud, "mode": "steps", "max_steps": 1,
})
self.assertEqual(race.status_code, 422, race.text)
self.assertEqual(api.state.tetris.sessions, {})
excluded = generation_registry(other)[0]
canonical = PROFILES[excluded["id"]]["repo_id"]
for route, data in (("/v1/systemone", {"state": "A request"}),
("/v1/systemone/batches", {"states": [{"id": "a", "state": "A request"}]})):
self.assertEqual(client.post(route, json={
"model": canonical, "questions": QUESTIONS, **data,
}).status_code, 422)
self.assertEqual(client.post("/api/jobs", json={
"model": excluded["id"], "state": "A request", "questions": QUESTIONS,
}).status_code, 422)
excluded_competitor = next(row.id for row in LOCAL_MODELS if row.request_model == canonical)
race = client.post("/api/tetris/races", json={
"left": excluded_competitor,
"right": next(iter(allowed_competitors)),
"mode": "steps", "max_steps": 1,
})
self.assertEqual(race.status_code, 422, race.text)
self.assertEqual(api.state.tetris.sessions, {})
def test_tetris_allowlist_rejects_unknown_model(self):
with self.assertRaisesRegex(ValueError, "Unknown local Decision model allowlist"):
HTTPDecisionAdapter.from_environment({}, allowed_local_models={"unknown/model"})
if __name__ == "__main__":
unittest.main()
|