submmit/tests/test_poll_policy.py

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