from __future__ import annotations import json import sys import tempfile import unittest from datetime import datetime, timezone from pathlib import Path from unittest.mock import patch PACKAGE_DIR = Path(__file__).resolve().parents[1] / "modelhub_submmit_api" if str(PACKAGE_DIR) in sys.path: sys.path.remove(str(PACKAGE_DIR)) sys.path.insert(0, str(PACKAGE_DIR)) from outcome_tracker import OutcomeTracker # noqa: E402 from poll_runner import ( # noqa: E402 _advance_architecture_history_backfill, _complete_architecture_history_backfill, _load_task_compatibility_contexts, build_parser, ) class PollPolicyTests(unittest.TestCase): def test_architecture_history_backfill_resumes_and_keeps_only_decision_state(self) -> None: class PagedHistoryClient: calls: list[int] = [] def list_tasks_page(self, **kwargs): # noqa: ANN003, ANN201 current = int(kwargs["current"]) self.calls.append(current) records = { 1: [ { "taskId": "history-1", "modelId": "owner/model-1", "gpuType": "gpu-a", "status": "success", "verifyResult": 1, }, { "taskId": "history-2", "modelId": "owner/model-2", "gpuType": "gpu-a", "status": "failed", "verifyResult": -1, }, ], 2: [ { "taskId": "history-3", "modelId": "owner/model-3", "gpuType": "gpu-b", "status": "success", "verifyResult": 1, } ], }[current] return {"data": {"records": records, "pages": 2}} with tempfile.TemporaryDirectory() as temporary_dir, patch.dict( "os.environ", {"MODELHUB_ARCHITECTURE_BACKFILL_PAGES_PER_BATCH": "1"}, ): root = Path(temporary_dir) outcomes = root / "outcomes.jsonl" checkpoint = root / "checkpoint.json" recent = root / "recent.jsonl" progress_path = root / "backfill.json" tracker = OutcomeTracker( outcomes, checkpoint_path=checkpoint, recent_path=recent, ) client = PagedHistoryClient() first = _advance_architecture_history_backfill( modelhub_client=client, # type: ignore[arg-type] outcome_tracker=tracker, ledger_path=root / "ledger.jsonl", progress_path=progress_path, now=datetime.now(timezone.utc), ) self.assertFalse(first["complete"]) self.assertTrue(tracker.has_durable_checkpoint) restored = OutcomeTracker( outcomes, checkpoint_path=checkpoint, recent_path=recent, ) second = _advance_architecture_history_backfill( modelhub_client=client, # type: ignore[arg-type] outcome_tracker=restored, ledger_path=root / "ledger.jsonl", progress_path=progress_path, now=datetime.now(timezone.utc), ) self.assertTrue(second["complete"]) self.assertNotIn("seenTaskIds", second) self.assertEqual([1, 2], client.calls) self.assertEqual(3, restored.get_stats_report()["terminalRecords"]) self.assertEqual([], json.loads(outcomes.read_text(encoding="utf-8") or "[]")) persisted_progress = progress_path.read_text(encoding="utf-8") self.assertNotIn("owner/model", persisted_progress) self.assertNotIn("logs.invalid", persisted_progress) def test_cold_start_backfill_finishes_all_pages_in_one_phase(self) -> None: class ThreePageClient: calls: list[int] = [] def list_tasks_page(self, **kwargs): # noqa: ANN003, ANN201 current = int(kwargs["current"]) self.calls.append(current) return { "data": { "records": [ { "taskId": f"task-{current}", "modelId": f"owner/model-{current}", "gpuType": "gpu-a", "status": "success", "verifyResult": 1, } ], "pages": 3, } } with tempfile.TemporaryDirectory() as temporary_dir, patch.dict( "os.environ", {"MODELHUB_ARCHITECTURE_BACKFILL_PAGES_PER_BATCH": "1"}, ): root = Path(temporary_dir) tracker = OutcomeTracker( root / "outcomes.jsonl", checkpoint_path=root / "checkpoint.json", recent_path=root / "recent.jsonl", ) client = ThreePageClient() durable_phases: list[str] = [] progress = _complete_architecture_history_backfill( modelhub_client=client, # type: ignore[arg-type] outcome_tracker=tracker, ledger_path=root / "ledger.jsonl", progress_path=root / "backfill.json", sync_callback=lambda phase: durable_phases.append(phase) or True, ) self.assertTrue(progress["complete"]) self.assertEqual([1, 2, 3], client.calls) self.assertEqual(["history_backfill"], durable_phases) self.assertEqual(3, tracker.get_stats_report()["terminalRecords"]) def test_cleanup_contexts_merge_outcomes_with_older_ledger_entries(self) -> None: with tempfile.TemporaryDirectory() as temporary_dir: root = Path(temporary_dir) tracker = OutcomeTracker(root / "outcomes.jsonl") tracker.record_submission( "owner/new", "gpu-a", "vllm", "text-generation", "task-new", datetime.now(timezone.utc).isoformat(), model_profile={"architectures": ["Qwen2ForCausalLM"]}, ) ledger_path = root / "ledger.jsonl" ledger_path.write_text( json.dumps( { "taskId": "task-old", "modelId": "owner/old", "targetGpu": "gpu-b", "framework": "mindie", "taskType": "text-generation", } ) + "\n", encoding="utf-8", ) contexts = _load_task_compatibility_contexts( tracker, ledger_path=ledger_path, ) self.assertEqual(["Qwen2ForCausalLM"], contexts["task-new"]["modelProfile"]["architectures"]) self.assertEqual("mindie", contexts["task-old"]["framework"]) self.assertEqual({}, contexts["task-old"]["modelProfile"]) def test_history_batch_is_classified_before_checkpoint_compaction(self) -> None: tasks = [ { "taskId": f"failed-{index}", "modelId": f"owner/model-{index}", "gpuType": "gpu-a", "status": "failed", "verifyResult": -1, "logCosUrl": f"https://logs.invalid/{index}", } for index in range(600) ] def classify(records, **kwargs): # noqa: ANN001, ANN003, ANN202 for record in records: record["failureCategory"] = "model_runtime" record["failureScope"] = "model" record["failureDeterministic"] = True progress = kwargs.get("progress") if progress is not None: progress(len(records), len(records), len(records), 0) return len(records) with tempfile.TemporaryDirectory() as temporary_dir: root = Path(temporary_dir) tracker = OutcomeTracker( root / "outcomes.jsonl", checkpoint_path=root / "checkpoint.json", recent_path=root / "recent.jsonl", ) with patch.object(tracker, "_enrich_failure_records", side_effect=classify): summary = tracker.bootstrap_from_history_tasks( tasks, enrichment_limit=600, ) report = tracker.get_stats_report() self.assertEqual(600, summary["enrichmentAttempts"]) self.assertEqual(600, report["totals"]["attributableFailureCount"]) self.assertEqual(600, report["totals"]["failureBreakdown"]["model_runtime"]) def test_age_policy_defaults_to_admission_only_reserve_five(self) -> None: args = build_parser().parse_args([]) self.assertEqual(5, args.recent_model_reserve_slots) self.assertFalse(hasattr(args, "dynamic_old_model_cleanup_reserve_slots")) if __name__ == "__main__": unittest.main()