submmit/tests/test_submission_safety.py

from __future__ import annotations

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 main import build_parser, one_candidate_per_model, process_model_for_candidates, run_submission, submit_candidate  # noqa: E402
from candidate_preflight import CandidatePreflightAdvisor  # noqa: E402
from modelhub_client import (  # noqa: E402
    ModelHubAPIError,
    ModelHubClientPool,
    is_model_uniqueness_error,
    parse_model_submission_precheck,
)
from models import HFModelSummary, ModelInspection  # noqa: E402
from submission_exclusions import SubmissionExclusionStore  # noqa: E402
from template_selector import TemplateSelector  # noqa: E402


class SafeDiscovery:
    def __init__(self) -> None:
        self.model = HFModelSummary(
            repo_id="owner/model",
            downloads=100,
            last_modified=datetime(2026, 1, 1, tzinfo=timezone.utc),
            pipeline_tag="text-generation",
        )

    def list_recent_models(self, **_kwargs) -> list[HFModelSummary]:  # noqa: ANN003
        return [self.model]

    @staticmethod
    def inspect_model(model: HFModelSummary) -> ModelInspection:
        return ModelInspection(repo_id=model.repo_id, weight_files=["model.safetensors"])


class OtherGpuProcessedClient:
    @staticmethod
    def model_submission_precheck(_model_id: str, *, force_refresh: bool = False) -> dict:
        del force_refresh
        return {"processedGpus": {"Biren_166m"}, "isInDB": True}


class FailingLookupClient:
    @staticmethod
    def search_by_model_id(_model_id: str) -> dict:
        raise ModelHubAPIError("temporary lookup outage")


class ExactGpuProcessedClient:
    def __init__(self) -> None:
        self.add_calls = 0

    @staticmethod
    def model_submission_precheck(_model_id: str, *, force_refresh: bool = False) -> dict:
        del force_refresh
        return {"processedGpus": {"Vastai_va16"}, "isInDB": True}

    def add_task(self, _payload: dict) -> dict:
        self.add_calls += 1
        return {"code": 0, "data": {"id": "should-not-submit"}}


class FailedFrameworkClient:
    def __init__(self) -> None:
        self.add_calls = 0

    @staticmethod
    def model_submission_precheck(_model_id: str, *, force_refresh: bool = False) -> dict:
        del force_refresh
        return {
            "processedGpus": set(),
            "failedFrameworksByGpu": {"Cambricon_mlu-370-x4": {"vllm"}},
        }

    def add_task(self, _payload: dict) -> dict:
        self.add_calls += 1
        return {"code": 0, "data": {"id": str(self.add_calls)}}


class AllGpuOfficialRegistry:
    @staticmethod
    def task_types_for(_client, *, model_address: str, model_last_modified: str | None, gpu: str) -> list[str]:
        del model_address, model_last_modified, gpu
        return ["text-generation"]


class AllGpuVettedMarket:
    @staticmethod
    def has_framework_route(_task_type: str, _target_gpu: str) -> bool:
        return True

    @staticmethod
    def ensure_framework_route(_client, _task_type: str, _target_gpu: str) -> None:
        return None

    @staticmethod
    def selectable_frameworks(
        *, task_type: str, target_gpu: str, incumbent_frameworks: list[str], inspection
    ) -> list[str]:
        del task_type, target_gpu, incumbent_frameworks, inspection
        return ["vllm"]

    @staticmethod
    def official_config(
        *, task_type: str, target_gpu: str, framework: str, gguf_filename: str | None = None
    ) -> str:
        del task_type, target_gpu, framework, gguf_filename
        return "framework: vllm"

    @staticmethod
    def framework_metadata(_task_type: str, _target_gpu: str, _framework: str) -> dict:
        return {
            "frameworkCombinedScore": 1.0,
            "frameworkOfficialConfigValid": True,
            "frameworkEvidenceQualified": True,
        }

    @staticmethod
    def gpu_metadata(_target_gpu: str) -> dict:
        return {}


class SeventhGpuFallbackMarket(AllGpuVettedMarket):
    @staticmethod
    def has_framework_route(_task_type: str, target_gpu: str) -> bool:
        return target_gpu == "gpu-7"


class TwoFrameworkMarket(AllGpuVettedMarket):
    @staticmethod
    def selectable_frameworks(
        *, task_type: str, target_gpu: str, incumbent_frameworks: list[str], inspection
    ) -> list[str]:
        del task_type, target_gpu, incumbent_frameworks, inspection
        return ["vllm", "vllm-mlu"]


class UniquenessRejectingClient:
    @staticmethod
    def add_task(_payload: dict) -> dict:
        raise ModelHubAPIError("模型唯一性检查没有通过，无法进行同步")


class UniquenessRunClient(UniquenessRejectingClient):
    def __init__(self) -> None:
        self.add_calls = 0

    @staticmethod
    def available_submit_slots() -> int:
        return 1

    @staticmethod
    def list_tasks_page(**_kwargs) -> dict:  # noqa: ANN003
        return {"code": 0, "data": {"records": []}}

    @staticmethod
    def processed_gpus_for_model(_model_id: str) -> set[str]:
        return set()

    def add_task(self, _payload: dict) -> dict:
        self.add_calls += 1
        return super().add_task(_payload)


def candidate(target_gpu: str = "Vastai_va16") -> dict:
    return {
        "repoId": "owner/model",
        "modelAddress": "https://modelscope.cn/models/owner/model",
        "taskType": "text-generation",
        "targetGpu": target_gpu,
        "framework": "vllm",
        "configParams": "framework: vllm",
    }


class SubmissionSafetyTests(unittest.TestCase):
    def test_official_routes_prefer_first_six_gpus(self) -> None:
        target_gpus = [f"gpu-{index}" for index in range(1, 8)]
        model = HFModelSummary(
            repo_id="owner/model",
            downloads=100,
            last_modified=datetime(2026, 1, 1, tzinfo=timezone.utc),
            pipeline_tag="text-generation",
        )
        candidates, skipped, failed = process_model_for_candidates(
            model=model,
            hf_discovery=SafeDiscovery(),  # type: ignore[arg-type]
            modelhub_client=OtherGpuProcessedClient(),  # type: ignore[arg-type]
            template_selector=TemplateSelector(),
            target_gpus=target_gpus,
            allowed_task_types=["text-generation"],
            market_intelligence=AllGpuVettedMarket(),  # type: ignore[arg-type]
            official_registry=AllGpuOfficialRegistry(),  # type: ignore[arg-type]
        )
        self.assertEqual(target_gpus[:6], [item["targetGpu"] for item in candidates])
        self.assertEqual([], skipped)
        self.assertEqual([], failed)

    def test_official_routes_fall_back_after_preferred_six_are_unusable(self) -> None:
        target_gpus = [f"gpu-{index}" for index in range(1, 8)]
        model = HFModelSummary(
            repo_id="owner/model",
            downloads=100,
            last_modified=datetime(2026, 1, 1, tzinfo=timezone.utc),
            pipeline_tag="text-generation",
        )
        candidates, skipped, failed = process_model_for_candidates(
            model=model,
            hf_discovery=SafeDiscovery(),  # type: ignore[arg-type]
            modelhub_client=OtherGpuProcessedClient(),  # type: ignore[arg-type]
            template_selector=TemplateSelector(),
            target_gpus=target_gpus,
            allowed_task_types=["text-generation"],
            market_intelligence=SeventhGpuFallbackMarket(),  # type: ignore[arg-type]
            official_registry=AllGpuOfficialRegistry(),  # type: ignore[arg-type]
        )
        self.assertEqual(["gpu-7"], [item["targetGpu"] for item in candidates])
        self.assertEqual(6, len(skipped))
        self.assertEqual([], failed)

    def test_precheck_tracks_processed_gpus_without_blocking_other_gpus(self) -> None:
        payload = {
            "code": 0,
            "data": {
                "isInDB": True,
                "verifyResult": {
                    "Biren_166m": {
                        "result": "已验证",
                        "records": [{"verifyResult": 1}],
                    }
                },
            },
        }
        precheck = parse_model_submission_precheck(payload)
        self.assertEqual({"Biren_166m"}, precheck["processedGpus"])

        model = HFModelSummary(
            repo_id="owner/model",
            downloads=100,
            last_modified=datetime(2026, 1, 1, tzinfo=timezone.utc),
            pipeline_tag="text-generation",
        )
        candidates, skipped, failed = process_model_for_candidates(
            model=model,
            hf_discovery=SafeDiscovery(),  # type: ignore[arg-type]
            modelhub_client=OtherGpuProcessedClient(),  # type: ignore[arg-type]
            template_selector=TemplateSelector(),
            target_gpus=["Biren_166m", "Vastai_va16"],
            allowed_task_types=["text-generation"],
        )
        self.assertEqual(["Vastai_va16"], [item["targetGpu"] for item in candidates])
        self.assertEqual("already_processed_for_gpu", skipped[0]["reason"])
        self.assertEqual([], failed)

    def test_precheck_distinguishes_unverified_failed_and_active_gpus(self) -> None:
        payload = {
            "data": {
                "isInDB": True,
                "verifyResult": {
                    "free": {"result": "未验证", "records": []},
                    "failed": {
                        "result": "验证失败",
                        "records": [{"framework": "vllm", "verifyResult": -1}],
                    },
                    "running": {"result": "验证中", "records": []},
                    "success": {
                        "result": "已验证",
                        "records": [{"framework": "vllm", "verifyResult": 1}],
                    },
                    "unknown": {"result": "未知状态", "records": []},
                    "failed_without_framework": {
                        "result": "验证失败",
                        "records": [{"verifyResult": -1}],
                    },
                },
            },
        }
        precheck = parse_model_submission_precheck(payload)
        self.assertEqual(
            {"running", "success", "unknown", "failed_without_framework"},
            precheck["processedGpus"],
        )
        self.assertEqual({"failed": {"vllm"}}, precheck["failedFrameworksByGpu"])

    def test_unverified_gpu_becomes_candidate_but_validating_gpu_does_not(self) -> None:
        class StatusClient:
            def __init__(self, status: str) -> None:
                self.status = status

            def model_submission_precheck(self, _model_id: str) -> dict:
                return parse_model_submission_precheck({
                    "data": {
                        "verifyResult": {
                            "Vastai_va16": {"result": self.status, "records": []}
                        }
                    }
                })

        def collect(status: str):
            return process_model_for_candidates(
                model=SafeDiscovery().model,
                hf_discovery=SafeDiscovery(),  # type: ignore[arg-type]
                modelhub_client=StatusClient(status),  # type: ignore[arg-type]
                template_selector=TemplateSelector(),
                target_gpus=["Vastai_va16"],
                allowed_task_types=["text-generation"],
            )

        open_candidates, _, _ = collect("未验证")
        active_candidates, active_skips, _ = collect("验证中")
        self.assertEqual(["Vastai_va16"], [item["targetGpu"] for item in open_candidates])
        self.assertEqual([], active_candidates)
        self.assertEqual("already_processed_for_gpu", active_skips[0]["reason"])

    def test_failed_gpu_requires_an_alternative_framework(self) -> None:
        model = SafeDiscovery().model
        client = FailedFrameworkClient()
        candidates, skipped, failed = process_model_for_candidates(
            model=model,
            hf_discovery=SafeDiscovery(),  # type: ignore[arg-type]
            modelhub_client=client,  # type: ignore[arg-type]
            template_selector=TemplateSelector(),
            target_gpus=["Cambricon_mlu-370-x4"],
            allowed_task_types=["text-generation"],
            market_intelligence=TwoFrameworkMarket(),  # type: ignore[arg-type]
            official_registry=AllGpuOfficialRegistry(),  # type: ignore[arg-type]
        )
        self.assertEqual([], skipped)
        self.assertEqual([], failed)
        self.assertEqual(["vllm-mlu"], [item["framework"] for item in candidates])
        self.assertEqual(["vllm"], candidates[0]["previousFailedFrameworks"])

        blocked = submit_candidate({**candidates[0], "framework": "vllm"}, client)  # type: ignore[arg-type]
        accepted = submit_candidate(candidates[0], client)  # type: ignore[arg-type]
        self.assertEqual("duplicate", blocked["outcome"])
        self.assertEqual("prior_failed_framework_for_gpu", blocked["reason"])
        self.assertEqual("submitted", accepted["outcome"])
        self.assertEqual(1, client.add_calls)

        unvetted = submit_candidate(
            {**candidates[0], "frameworkEvidenceQualified": False}, client  # type: ignore[arg-type]
        )
        self.assertEqual("framework_evidence_deferred", unvetted["outcome"])
        self.assertEqual(1, client.add_calls)

        no_routes, no_route_skips, _ = process_model_for_candidates(
            model=model,
            hf_discovery=SafeDiscovery(),  # type: ignore[arg-type]
            modelhub_client=client,  # type: ignore[arg-type]
            template_selector=TemplateSelector(),
            target_gpus=["Cambricon_mlu-370-x4"],
            allowed_task_types=["text-generation"],
            market_intelligence=AllGpuVettedMarket(),  # type: ignore[arg-type]
            official_registry=AllGpuOfficialRegistry(),  # type: ignore[arg-type]
        )
        self.assertEqual([], no_routes)
        self.assertEqual("prior_failed_framework_no_alternative", no_route_skips[-1]["reason"])

        no_market_routes, no_market_skips, _ = process_model_for_candidates(
            model=model,
            hf_discovery=SafeDiscovery(),  # type: ignore[arg-type]
            modelhub_client=client,  # type: ignore[arg-type]
            template_selector=TemplateSelector(),
            target_gpus=["Cambricon_mlu-370-x4"],
            allowed_task_types=["text-generation"],
        )
        self.assertEqual([], no_market_routes)
        self.assertEqual(
            "prior_failed_gpu_requires_public_framework_evidence",
            no_market_skips[-1]["reason"],
        )

    def test_published_size_avoids_official_calls_for_certain_oom(self) -> None:
        class CountingRegistry(AllGpuOfficialRegistry):
            calls = 0

            def task_types_for(self, *_args, **_kwargs) -> list[str]:
                self.calls += 1
                return ["text-generation"]

        registry = CountingRegistry()
        model = HFModelSummary(
            repo_id="owner/oversized",
            downloads=100,
            last_modified=datetime(2026, 1, 1, tzinfo=timezone.utc),
            pipeline_tag="text-generation",
            file_size=100 * 1024**3,
        )
        candidates, skipped, failed = process_model_for_candidates(
            model=model,
            hf_discovery=SafeDiscovery(),  # type: ignore[arg-type]
            modelhub_client=OtherGpuProcessedClient(),  # type: ignore[arg-type]
            template_selector=TemplateSelector(),
            target_gpus=["MetaX_c-500"],
            allowed_task_types=["text-generation"],
            preflight_advisor=CandidatePreflightAdvisor(),
            official_registry=registry,  # type: ignore[arg-type]
        )
        self.assertEqual([], candidates)
        self.assertEqual("preflight_predicted_oom", skipped[0]["reason"])
        self.assertEqual([], failed)
        self.assertEqual(0, registry.calls)

    def test_lookup_failure_is_fail_closed(self) -> None:
        pool = ModelHubClientPool([FailingLookupClient()])  # type: ignore[arg-type]
        with self.assertRaises(ModelHubAPIError):
            pool.model_submission_precheck("owner/model")

    def test_submit_precheck_stops_an_exact_processed_gpu(self) -> None:
        client = ExactGpuProcessedClient()
        result = submit_candidate(candidate(), client)  # type: ignore[arg-type]
        self.assertEqual("duplicate", result["outcome"])
        self.assertEqual(0, client.add_calls)

    def test_uniqueness_rejection_is_non_retryable_for_that_combination(self) -> None:
        error = ModelHubAPIError("模型唯一性检查没有通过，无法进行同步")
        self.assertTrue(is_model_uniqueness_error(error))
        result = submit_candidate(candidate(), UniquenessRejectingClient())  # type: ignore[arg-type]
        self.assertEqual("uniqueness_rejected", result["outcome"])

        with tempfile.TemporaryDirectory() as temporary_dir:
            path = Path(temporary_dir) / "exclusions.jsonl"
            store = SubmissionExclusionStore(path)
            store.block("owner/model", "Vastai_va16", reason=str(error))
            reloaded = SubmissionExclusionStore(path)
            self.assertTrue(reloaded.is_blocked("owner/model", "Vastai_va16"))
            self.assertFalse(reloaded.is_blocked("owner/model", "Biren_166m"))

    def test_run_persists_uniqueness_rejection_and_does_not_retry_it(self) -> None:
        with tempfile.TemporaryDirectory() as temporary_dir:
            root = Path(temporary_dir)
            args = build_parser().parse_args(
                [
                    "--gpus",
                    "Vastai_va16",
                    "--task-types",
                    "text-generation",
                    "--limit",
                    "1",
                    "--max-scan-models",
                    "1",
                    "--skip-outcome-sync",
                    "--skip-history-archive",
                ]
            )
            args.runs_dir = str(root / "runs")
            args.ledger_path = str(root / "ledger.jsonl")
            args.outcomes_path = str(root / "outcomes.jsonl")
            args.claims_path = str(root / "claims.jsonl")
            args.submission_exclusions_path = str(root / "exclusions.jsonl")
            args.history_archive_path = str(root / "history.jsonl")
            client = UniquenessRunClient()
            discovery = SafeDiscovery()

            first = run_submission(
                args,
                now=datetime(2026, 1, 1, 12, tzinfo=timezone.utc),
                hf_discovery=discovery,  # type: ignore[arg-type]
                modelhub_client=client,  # type: ignore[arg-type]
            )
            second = run_submission(
                args,
                now=datetime(2026, 1, 1, 12, 1, tzinfo=timezone.utc),
                hf_discovery=discovery,  # type: ignore[arg-type]
                modelhub_client=client,  # type: ignore[arg-type]
            )

            self.assertEqual(1, first["modelGpuUniquenessRejectedCount"])
            self.assertEqual(1, client.add_calls)
            self.assertEqual(0, second["candidateCount"])
            self.assertEqual(1, second["skipReasonCounts"]["model_gpu_uniqueness_blocklist"])

    def test_concurrent_batch_uses_at_most_one_gpu_per_model(self) -> None:
        selected = one_candidate_per_model(
            [
                candidate("Vastai_va16"),
                candidate("Biren_166m"),
                {**candidate("Biren_166m"), "repoId": "owner/other"},
            ]
        )
        self.assertEqual(["owner/model", "owner/other"], [item["repoId"] for item in selected])


if __name__ == "__main__":
    unittest.main()