feat: add durable success-first modelhub agent
This commit is contained in:
CoolBoy
@@ -3,6 +3,7 @@ from __future__ import annotations
|
||||
import importlib.util
|
||||
import os
|
||||
import unittest
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
@@ -28,6 +29,16 @@ class HostedAgentEntrypointTests(unittest.TestCase):
|
||||
command = ENTRYPOINT._worker_command()
|
||||
|
||||
self.assertEqual(["--max-submits-per-run", "0"], command[-2:])
|
||||
self.assertIn("--state-sync", command)
|
||||
|
||||
def test_readiness_file_is_separate_from_liveness(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||||
path = Path(temporary_dir) / "readiness.json"
|
||||
path.write_text('{"ready": false, "reason": "state_sync_unhealthy"}', encoding="utf-8")
|
||||
with patch.object(ENTRYPOINT, "READINESS_PATH", path):
|
||||
readiness = ENTRYPOINT._readiness()
|
||||
self.assertFalse(readiness["ready"])
|
||||
self.assertEqual("state_sync_unhealthy", readiness["reason"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
246
tests/test_super_agent.py
Normal file
246
tests/test_super_agent.py
Normal file
@@ -0,0 +1,246 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
import tempfile
|
||||
import unittest
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import sys
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
MODULE_ROOT = ROOT / "modelhub_submmit_api"
|
||||
if str(MODULE_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(MODULE_ROOT))
|
||||
|
||||
from common import read_jsonl, write_json, write_jsonl # noqa: E402
|
||||
from config_optimizer import SafeConfigOptimizer # noqa: E402
|
||||
from hf_discovery import HuggingFaceDiscovery, parse_model_card_front_matter # noqa: E402
|
||||
from official_capabilities import OfficialCapabilityRegistry # noqa: E402
|
||||
from routing_engine import SuccessFirstRoutingEngine # noqa: E402
|
||||
from state_sync import StateGitSync # noqa: E402
|
||||
|
||||
|
||||
class OfficialClient:
|
||||
def __init__(self, *, fail_catalog: bool = False) -> None:
|
||||
self.fail_catalog = fail_catalog
|
||||
|
||||
def list_machine_info(self): # noqa: ANN201
|
||||
if self.fail_catalog:
|
||||
raise RuntimeError("offline")
|
||||
return [
|
||||
{"gpuType": "gpu-fast", "canVerify": True, "maxConcurrentTasks": 2},
|
||||
{"gpuType": "gpu-disabled", "canVerify": False, "maxConcurrentTasks": 8},
|
||||
]
|
||||
|
||||
def list_task_levels(self): # noqa: ANN201
|
||||
return {"data": [{"taskType": "text-generation"}, {"taskType": "new-task"}]}
|
||||
|
||||
def list_model_task_types(self, target_gpu, model_address): # noqa: ANN001, ANN201
|
||||
del target_gpu, model_address
|
||||
return {"data": [{"taskType": "text-generation"}]}
|
||||
|
||||
|
||||
class SuperAgentTests(unittest.TestCase):
|
||||
def test_official_registry_discovers_catalog_and_exact_model_routes(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||||
registry = OfficialCapabilityRegistry(Path(temporary_dir) / "official.json", log_fn=lambda _: None)
|
||||
client = OfficialClient()
|
||||
registry.prepare(
|
||||
client,
|
||||
fallback_gpus=["legacy"],
|
||||
task_types=["text-generation"],
|
||||
now=datetime(2026, 8, 15, tzinfo=timezone.utc),
|
||||
)
|
||||
self.assertTrue(registry.ready)
|
||||
self.assertEqual(["gpu-fast"], registry.eligible_gpus())
|
||||
self.assertEqual(
|
||||
["text-generation"],
|
||||
registry.task_types_for(
|
||||
client,
|
||||
model_address="https://modelscope.cn/models/owner/model",
|
||||
model_last_modified="2026-08-15T00:00:00+00:00",
|
||||
gpu="gpu-fast",
|
||||
),
|
||||
)
|
||||
|
||||
def test_official_registry_fails_closed_without_catalog_cache(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||||
registry = OfficialCapabilityRegistry(Path(temporary_dir) / "official.json", log_fn=lambda _: None)
|
||||
registry.prepare(
|
||||
OfficialClient(fail_catalog=True),
|
||||
fallback_gpus=["legacy"],
|
||||
task_types=["text-generation"],
|
||||
now=datetime(2026, 8, 15, tzinfo=timezone.utc),
|
||||
)
|
||||
self.assertFalse(registry.ready)
|
||||
self.assertEqual("critical_official_signal_unavailable", registry.pause_reason)
|
||||
|
||||
def test_success_band_beats_shorter_queue(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||||
engine = SuccessFirstRoutingEngine(Path(temporary_dir) / "routing.json", log_fn=lambda _: None)
|
||||
candidates = [
|
||||
{
|
||||
"repoId": "owner/model-a",
|
||||
"targetGpu": "reliable",
|
||||
"framework": "vllm",
|
||||
"taskType": "text-generation",
|
||||
"frameworkMarketSamples": 1000,
|
||||
"frameworkMarketSuccessRate": 0.9,
|
||||
"queueBacklogHours": 12,
|
||||
},
|
||||
{
|
||||
"repoId": "owner/model-b",
|
||||
"targetGpu": "fast",
|
||||
"framework": "vllm",
|
||||
"taskType": "text-generation",
|
||||
"frameworkMarketSamples": 1000,
|
||||
"frameworkMarketSuccessRate": 0.4,
|
||||
"queueBacklogHours": 0.25,
|
||||
},
|
||||
]
|
||||
ordered = engine.order_candidates(candidates)
|
||||
self.assertEqual("reliable", ordered[0]["targetGpu"])
|
||||
self.assertGreater(ordered[0]["routingSuccessBand"], ordered[1]["routingSuccessBand"])
|
||||
|
||||
def test_modelscope_metadata_and_model_card_lineage_are_structured(self) -> None:
|
||||
item = {
|
||||
"id": "owner/model",
|
||||
"downloads": 123,
|
||||
"params": 7_000_000_000,
|
||||
"file_size": 14_000_000_000,
|
||||
"tags": ["qwen", "chat"],
|
||||
"tasks": ["text-generation"],
|
||||
"license": "apache-2.0",
|
||||
"likes": 9,
|
||||
}
|
||||
model = HuggingFaceDiscovery._parse_model(
|
||||
item,
|
||||
fallback_pipeline_tag="text-generation",
|
||||
min_downloads=0,
|
||||
)
|
||||
self.assertIsNotNone(model)
|
||||
assert model is not None
|
||||
self.assertEqual(7_000_000_000, model.params)
|
||||
self.assertEqual(("qwen", "chat"), model.tags)
|
||||
metadata = parse_model_card_front_matter(
|
||||
"---\nbase_model: Qwen/base\nframeworks:\n - transformers\ntasks:\n - text-generation\n---\nbody"
|
||||
)
|
||||
self.assertEqual("Qwen/base", metadata["base_model"])
|
||||
self.assertEqual(["transformers"], metadata["frameworks"])
|
||||
|
||||
def test_state_branch_round_trip_persists_intent_without_config_or_secret(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||||
root = Path(temporary_dir)
|
||||
remote = root / "remote.git"
|
||||
project = root / "project"
|
||||
restored_project = root / "restored"
|
||||
project.mkdir()
|
||||
restored_project.mkdir()
|
||||
subprocess.run(["git", "init", "--bare", str(remote)], check=True, stdout=subprocess.DEVNULL)
|
||||
write_json(project / ".modelhub_state" / "account_capacity.json", {"version": 1})
|
||||
credentials = {"username": "tester", "email": "tester@example.com", "password": "secret-value"}
|
||||
manager = StateGitSync(
|
||||
project_root=project,
|
||||
credentials=credentials,
|
||||
remote=str(remote),
|
||||
log_fn=lambda _: None,
|
||||
)
|
||||
manager.acquire_process_lock()
|
||||
self.assertTrue(manager.restore())
|
||||
batch_id = manager.begin_batch(
|
||||
[
|
||||
{
|
||||
"repoId": "owner/model",
|
||||
"modelAddress": "https://modelscope.cn/models/owner/model",
|
||||
"targetGpu": "gpu-a",
|
||||
"taskType": "text-generation",
|
||||
"framework": "vllm",
|
||||
"configParams": "password: must-not-be-copied",
|
||||
}
|
||||
]
|
||||
)
|
||||
self.assertIsNotNone(batch_id)
|
||||
manager.close()
|
||||
|
||||
restored = StateGitSync(
|
||||
project_root=restored_project,
|
||||
credentials=credentials,
|
||||
remote=str(remote),
|
||||
log_fn=lambda _: None,
|
||||
)
|
||||
restored.acquire_process_lock()
|
||||
self.assertTrue(restored.restore())
|
||||
intents = read_jsonl(restored_project / ".modelhub_state" / "recovery_intents.jsonl")
|
||||
self.assertEqual("owner/model", intents[0]["repoId"])
|
||||
state_text = "\n".join(
|
||||
path.read_text(encoding="utf-8")
|
||||
for path in restored._workspace.rglob("*")
|
||||
if path.is_file() and ".git" not in path.parts
|
||||
)
|
||||
self.assertNotIn("secret-value", state_text)
|
||||
self.assertNotIn("must-not-be-copied", state_text)
|
||||
restored.close()
|
||||
|
||||
def test_failed_intent_push_returns_no_batch_id(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||||
manager = StateGitSync(
|
||||
project_root=Path(temporary_dir),
|
||||
credentials={"username": "u", "email": "e@example.com", "password": "p"},
|
||||
remote="unused",
|
||||
log_fn=lambda _: None,
|
||||
)
|
||||
manager.healthy = True
|
||||
with patch.object(manager, "sync", return_value=False):
|
||||
self.assertIsNone(
|
||||
manager.begin_batch(
|
||||
[
|
||||
{
|
||||
"repoId": "owner/model",
|
||||
"targetGpu": "gpu",
|
||||
"taskType": "text-generation",
|
||||
"framework": "vllm",
|
||||
"configParams": "safe",
|
||||
}
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
def test_config_patch_requires_repeated_cross_model_success(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||||
root = Path(temporary_dir)
|
||||
intents = []
|
||||
outcomes = []
|
||||
for index in range(5):
|
||||
intents.append(
|
||||
{
|
||||
"taskId": str(index),
|
||||
"taskType": "text-generation",
|
||||
"targetGpu": "gpu-a",
|
||||
"framework": "vllm",
|
||||
"repoId": f"owner/model-{index % 2}",
|
||||
"configFingerprint": "proven",
|
||||
"safeConfigVector": {"gpuNum": 1, "tensorParallel": 1},
|
||||
}
|
||||
)
|
||||
outcomes.append({"taskId": str(index), "outcome": "success"})
|
||||
write_jsonl(root / "intents.jsonl", intents)
|
||||
write_jsonl(root / "outcomes.jsonl", outcomes)
|
||||
optimizer = SafeConfigOptimizer(
|
||||
intents_path=root / "intents.jsonl",
|
||||
outcomes_path=root / "outcomes.jsonl",
|
||||
)
|
||||
config, metadata = optimizer.optimize(
|
||||
task_type="text-generation",
|
||||
target_gpu="gpu-a",
|
||||
framework="vllm",
|
||||
official_config="framework: vllm\nsut_config:\n gpu_num: 2\nref_config:\n gpu_num: 2\n",
|
||||
official_lower_bound=0.40,
|
||||
)
|
||||
self.assertTrue(metadata["applied"])
|
||||
self.assertNotIn("gpu_num: 2", config)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user