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