submmit/tests/test_poll_policy.py

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