from __future__ import annotations 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_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() 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) 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_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()