from __future__ import annotations import argparse import json import sys import tempfile import unittest from datetime import datetime, timezone from pathlib import Path 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 _bootstrap_architecture_history, _load_task_compatibility_contexts, resolve_age_cleanup_policy, ) class PollPolicyTests(unittest.TestCase): def test_architecture_bootstrap_falls_back_when_public_logs_are_hidden(self) -> None: class HistoryClient: @staticmethod def list_tasks_page(**_kwargs): # noqa: ANN003 return { "data": { "records": [ { "taskId": "public-failure", "modelId": "public/model", "gpuType": "gpu", "status": "success", "verifyResult": -1, "logCosUrl": None, } ] } } @staticmethod def list_tasks(**_kwargs): # noqa: ANN003 return [ { "taskId": "owned-success", "modelId": "owner/model", "gpuType": "gpu", "modelTaskLevelId": 23, "status": "success", "verifyResult": 1, "updateTime": datetime.now(timezone.utc).isoformat(), } ] with tempfile.TemporaryDirectory() as temporary_dir: root = Path(temporary_dir) tracker = OutcomeTracker(root / "outcomes.jsonl") summary = _bootstrap_architecture_history( modelhub_client=HistoryClient(), # type: ignore[arg-type] outcome_tracker=tracker, ledger_path=root / "ledger.jsonl", now=datetime.now(timezone.utc), ) self.assertEqual("owned_full_history", summary["source"]) self.assertEqual(1, summary["terminalRecords"]) self.assertEqual(0, summary["communityUsableFailureDetails"]) 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_age_cleanup_uses_minus_ten_once_then_minus_five(self) -> None: args = argparse.Namespace( recent_model_reserve_slots=10, dynamic_old_model_cleanup_reserve_slots=5, ) self.assertEqual( ("initial", 10), resolve_age_cleanup_policy(args, initial_cleanup_pending=True), ) self.assertEqual( ("dynamic", 5), resolve_age_cleanup_policy(args, initial_cleanup_pending=False), ) def test_dynamic_cleanup_cannot_be_stricter_than_initial_cleanup(self) -> None: args = argparse.Namespace( recent_model_reserve_slots=10, dynamic_old_model_cleanup_reserve_slots=20, ) self.assertEqual( ("dynamic", 10), resolve_age_cleanup_policy(args, initial_cleanup_pending=False), ) if __name__ == "__main__": unittest.main()