457 lines
20 KiB
Python
457 lines
20 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import sys
|
|
|
|
from dulwich import porcelain
|
|
|
|
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_json, 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 outcome_tracker import OutcomeTracker # 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_outcome_history_compacts_to_decision_checkpoint_without_raw_archive(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
root = Path(temporary_dir)
|
|
outcomes = root / "outcomes.jsonl"
|
|
checkpoint = root / "checkpoint.json"
|
|
recent = root / "recent.jsonl"
|
|
archive = root / "archive"
|
|
rows = [
|
|
{
|
|
"taskId": str(index),
|
|
"modelId": f"owner/model-{index}",
|
|
"targetGpu": "gpu-a",
|
|
"framework": "vllm",
|
|
"taskType": "text-generation",
|
|
"submitTime": f"2026-08-{1 + index // 100:02d}T00:00:{index % 60:02d}+00:00",
|
|
"lastSyncTime": "2026-08-21T00:00:00+00:00",
|
|
"outcome": "success" if index % 2 == 0 else "failed",
|
|
"failureCategory": "model_runtime" if index % 2 else None,
|
|
"failureScope": "model" if index % 2 else None,
|
|
"logCosUrl": "https://secret.invalid/signed?token=hidden",
|
|
}
|
|
for index in range(600)
|
|
]
|
|
write_jsonl(outcomes, rows)
|
|
with patch.dict(
|
|
os.environ,
|
|
{
|
|
"MODELHUB_AGENT_OUTCOME_COMPACT_THRESHOLD": "500",
|
|
"MODELHUB_AGENT_RECENT_OUTCOME_LIMIT": "100",
|
|
},
|
|
):
|
|
tracker = OutcomeTracker(
|
|
outcomes,
|
|
checkpoint_path=checkpoint,
|
|
recent_path=recent,
|
|
)
|
|
self.assertTrue(tracker.has_durable_checkpoint)
|
|
self.assertEqual([], read_jsonl(outcomes))
|
|
self.assertEqual(100, len(read_jsonl(recent)))
|
|
report = tracker.get_stats_report()
|
|
self.assertEqual(600, report["terminalRecords"])
|
|
self.assertEqual(300, report["totals"]["successCount"])
|
|
self.assertFalse(archive.exists())
|
|
self.assertNotIn("logCosUrl", read_jsonl(recent)[0])
|
|
self.assertEqual("decision_state_only", read_json(checkpoint)["storageMode"])
|
|
|
|
restored = OutcomeTracker(
|
|
outcomes,
|
|
checkpoint_path=checkpoint,
|
|
recent_path=recent,
|
|
)
|
|
self.assertEqual(600, restored.get_stats_report()["terminalRecords"])
|
|
|
|
restored.record_submission(
|
|
model_id="owner/new-model",
|
|
target_gpu="gpu-a",
|
|
framework="vllm",
|
|
task_type="text-generation",
|
|
task_id="new-task",
|
|
submit_time="2026-08-22T00:00:00+00:00",
|
|
)
|
|
restored._update_record_from_task(
|
|
restored._by_task_id["new-task"],
|
|
{"status": "success", "verifyResult": 1},
|
|
)
|
|
restored.save()
|
|
updated = restored.get_stats_report()
|
|
self.assertEqual(601, updated["terminalRecords"])
|
|
self.assertEqual(301, updated["totals"]["successCount"])
|
|
|
|
restarted = OutcomeTracker(
|
|
outcomes,
|
|
checkpoint_path=checkpoint,
|
|
recent_path=recent,
|
|
)
|
|
self.assertEqual(601, restarted.get_stats_report()["terminalRecords"])
|
|
self.assertIn("599", restarted._by_task_id)
|
|
|
|
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_transformers_is_fallback_when_an_executable_framework_exists(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
engine = SuccessFirstRoutingEngine(Path(temporary_dir) / "routing.json", log_fn=lambda _: None)
|
|
candidates = [
|
|
{
|
|
"repoId": "owner/model",
|
|
"targetGpu": "gpu-a",
|
|
"framework": "transformers",
|
|
"taskType": "text-generation",
|
|
"frameworkMarketSamples": 1000,
|
|
"frameworkMarketSuccessRate": 0.95,
|
|
},
|
|
{
|
|
"repoId": "owner/model",
|
|
"targetGpu": "gpu-b",
|
|
"framework": "vllm",
|
|
"taskType": "text-generation",
|
|
"frameworkMarketSamples": 1000,
|
|
"frameworkMarketSuccessRate": 0.40,
|
|
},
|
|
]
|
|
ordered = engine.order_candidates(candidates)
|
|
self.assertEqual(1, len(ordered))
|
|
self.assertEqual("vllm", ordered[0]["framework"])
|
|
|
|
candidates[0]["transformersPrerequisiteSatisfied"] = True
|
|
unlocked = engine.order_candidates(candidates)
|
|
self.assertEqual("transformers", unlocked[0]["framework"])
|
|
|
|
def test_outcome_context_repairs_framework_missing_from_history_api(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
root = Path(temporary_dir)
|
|
tracker = OutcomeTracker(root / "outcomes.jsonl")
|
|
tracker._records = [
|
|
{
|
|
"taskId": "42",
|
|
"modelId": "owner/model",
|
|
"targetGpu": "gpu-a",
|
|
"framework": "",
|
|
"taskType": "",
|
|
"outcome": "failed",
|
|
}
|
|
]
|
|
tracker._rebuild_indexes()
|
|
repaired = tracker.merge_task_contexts(
|
|
{
|
|
"42": {
|
|
"framework": "vllm",
|
|
"taskType": "text-generation",
|
|
"modelProfile": {"modelType": "Qwen2ForCausalLM"},
|
|
}
|
|
}
|
|
)
|
|
self.assertEqual(1, repaired)
|
|
self.assertEqual("vllm", tracker._by_task_id["42"]["framework"])
|
|
self.assertEqual("Qwen2ForCausalLM", tracker._by_task_id["42"]["modelProfile"]["modelType"])
|
|
|
|
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()
|
|
porcelain.init(remote, bare=True)
|
|
write_json(project / ".modelhub_state" / "account_capacity.json", {"version": 1})
|
|
write_jsonl(
|
|
project / ".modelhub_state" / "worker_crashes.jsonl",
|
|
[{"at": "2026-08-21T00:00:00+00:00", "exitCode": 137}],
|
|
)
|
|
write_jsonl(
|
|
restored_project / ".modelhub_state" / "worker_crashes.jsonl",
|
|
[{"at": "2026-08-21T01:00:00+00:00", "exitCode": 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)
|
|
generation = manager.generation
|
|
remote_head = manager._remote_oid()
|
|
self.assertTrue(manager.sync("unchanged_cycle"))
|
|
self.assertEqual(generation, manager.generation)
|
|
self.assertEqual(remote_head, manager._remote_oid())
|
|
archive_refs = porcelain.ls_remote(str(remote)).refs
|
|
self.assertNotIn(b"refs/heads/agent-archive-2026-08", archive_refs)
|
|
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"])
|
|
crashes = read_jsonl(restored_project / ".modelhub_state" / "worker_crashes.jsonl")
|
|
self.assertEqual([137, 1], [row["exitCode"] for row in crashes])
|
|
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_state_push_retries_the_same_commit_and_generation(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
root = Path(temporary_dir)
|
|
remote = root / "remote.git"
|
|
project = root / "project"
|
|
project.mkdir()
|
|
porcelain.init(remote, bare=True)
|
|
manager = StateGitSync(
|
|
project_root=project,
|
|
credentials={"username": "u", "email": "e@example.com", "password": "p"},
|
|
remote=str(remote),
|
|
log_fn=lambda _: None,
|
|
)
|
|
manager.acquire_process_lock()
|
|
self.assertTrue(manager.restore())
|
|
write_json(project / ".modelhub_state" / "account_capacity.json", {"version": 1})
|
|
with patch("state_sync.porcelain.push", side_effect=RuntimeError("temporary failure")):
|
|
self.assertFalse(manager.sync("cycle"))
|
|
pending_oid = manager._pending_push_oid
|
|
self.assertIsNotNone(pending_oid)
|
|
self.assertEqual(0, manager.generation)
|
|
self.assertTrue(manager.sync("retry"))
|
|
self.assertEqual(1, manager.generation)
|
|
self.assertEqual(pending_oid, manager._remote_oid())
|
|
manager.close()
|
|
|
|
def test_terminal_intents_are_bounded_and_reduced_to_decision_fields(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
root = Path(temporary_dir)
|
|
intent_path = root / ".modelhub_state" / "recovery_intents.jsonl"
|
|
write_jsonl(
|
|
intent_path,
|
|
[
|
|
{
|
|
"intentId": str(index),
|
|
"status": "failed",
|
|
"createdAt": f"2026-08-01T00:{index % 60:02d}:00+00:00",
|
|
"completedAt": f"2026-08-02T00:{index % 60:02d}:00+00:00",
|
|
}
|
|
for index in range(350)
|
|
]
|
|
+ [{"intentId": "pending", "status": "pending"}],
|
|
)
|
|
manager = StateGitSync(
|
|
project_root=root,
|
|
credentials={"username": "u", "email": "e@example.com", "password": "p"},
|
|
remote="unused",
|
|
log_fn=lambda _: None,
|
|
)
|
|
self.assertEqual(50, manager._compact_intents())
|
|
retained = read_jsonl(intent_path)
|
|
self.assertEqual(301, len(retained))
|
|
self.assertEqual(1, sum(row.get("status") == "pending" for row in retained))
|
|
self.assertFalse((root / ".modelhub_state" / "archive_pending").exists())
|
|
|
|
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()
|