from __future__ import annotations import os import gzip 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 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_checkpoint_recent_window_and_gzip_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, archive_pending_dir=archive, ) 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"]) shard = next(archive.rglob("*.jsonl.gz")) import gzip with gzip.open(shard, "rt", encoding="utf-8") as handle: archived_text = handle.read() self.assertNotIn("logCosUrl", archived_text) self.assertNotIn("token=hidden", archived_text) restored = OutcomeTracker( outcomes, checkpoint_path=checkpoint, recent_path=recent, archive_pending_dir=archive, ) 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, archive_pending_dir=archive, ) 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}], ) pending_archive = ( project / ".modelhub_state" / "archive_pending" / "outcomes" / "2026-08" / "shard.jsonl.gz" ) pending_archive.parent.mkdir(parents=True, exist_ok=True) with gzip.open(pending_archive, "wt", encoding="utf-8") as handle: handle.write('{"taskId":"archived"}\n') 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()) self.assertFalse(pending_archive.exists()) archive_refs = porcelain.ls_remote(str(remote)).refs self.assertIn(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_archived(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(250) ] + [{"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(201, len(retained)) self.assertEqual(1, sum(row.get("status") == "pending" for row in retained)) shard = next((root / ".modelhub_state" / "archive_pending" / "attempts").rglob("*.jsonl.gz")) with gzip.open(shard, "rt", encoding="utf-8") as handle: self.assertEqual(50, len(handle.readlines())) 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()