226 lines
8.3 KiB
Python
226 lines
8.3 KiB
Python
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()
|