fix: prevent repeated model GPU submissions
This commit is contained in:
CoolBoy
225
tests/test_submission_safety.py
Normal file
225
tests/test_submission_safety.py
Normal file
@@ -0,0 +1,225 @@
|
||||
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 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 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_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_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()
|
||||
Reference in New Issue
Block a user