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()