543 lines
22 KiB
Python
543 lines
22 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 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 outcome_tracker import OutcomeTracker # 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(
|
|
{"failed", "running", "success", "unknown", "failed_without_framework"},
|
|
precheck["processedGpus"],
|
|
)
|
|
self.assertEqual({"success"}, precheck["successfulGpus"])
|
|
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_remains_closed_even_when_an_alternative_framework_exists(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("prior_failed_gpu_closed_accuracy_first", skipped[-1]["reason"])
|
|
self.assertEqual([], failed)
|
|
self.assertEqual([], candidates)
|
|
|
|
blocked = submit_candidate(candidate("Cambricon_mlu-370-x4"), client) # type: ignore[arg-type]
|
|
self.assertEqual("duplicate", blocked["outcome"])
|
|
self.assertEqual("prior_failed_gpu_closed_accuracy_first", blocked["reason"])
|
|
self.assertEqual(0, 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_gpu_closed_accuracy_first", 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_closed_accuracy_first",
|
|
no_market_skips[-1]["reason"],
|
|
)
|
|
|
|
def test_model_route_limit_is_global_across_poll_cycles(self) -> None:
|
|
class SixRouteClient:
|
|
@staticmethod
|
|
def model_submission_precheck(_model_id: str) -> dict:
|
|
return {
|
|
"processedGpus": {f"gpu-{index}" for index in range(1, 7)},
|
|
"successfulGpus": {"gpu-1"},
|
|
"failedFrameworksByGpu": {},
|
|
}
|
|
|
|
candidates, skipped, failed = process_model_for_candidates(
|
|
model=SafeDiscovery().model,
|
|
hf_discovery=SafeDiscovery(), # type: ignore[arg-type]
|
|
modelhub_client=SixRouteClient(), # type: ignore[arg-type]
|
|
template_selector=TemplateSelector(),
|
|
target_gpus=["gpu-7"],
|
|
allowed_task_types=["text-generation"],
|
|
)
|
|
self.assertEqual([], candidates)
|
|
self.assertEqual([], failed)
|
|
self.assertEqual("model_route_limit_reached", skipped[0]["reason"])
|
|
self.assertEqual(6, skipped[0]["processedGpuCount"])
|
|
|
|
def test_model_waits_for_feedback_after_two_pending_routes(self) -> None:
|
|
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",
|
|
)
|
|
tracker.record_submission(
|
|
"owner/model", "gpu-1", "vllm", "text-generation", "task-1",
|
|
"2026-09-22T00:00:00+00:00",
|
|
)
|
|
tracker.record_submission(
|
|
"owner/model", "gpu-2", "vllm", "text-generation", "task-2",
|
|
"2026-09-22T00:01:00+00:00",
|
|
)
|
|
|
|
candidates, skipped, failed = process_model_for_candidates(
|
|
model=SafeDiscovery().model,
|
|
hf_discovery=SafeDiscovery(), # type: ignore[arg-type]
|
|
modelhub_client=OtherGpuProcessedClient(), # type: ignore[arg-type]
|
|
template_selector=TemplateSelector(),
|
|
target_gpus=["gpu-3"],
|
|
allowed_task_types=["text-generation"],
|
|
outcome_tracker=tracker,
|
|
)
|
|
|
|
self.assertEqual([], candidates)
|
|
self.assertEqual([], failed)
|
|
self.assertEqual("model_pending_feedback_limit", skipped[0]["reason"])
|
|
self.assertEqual(2, skipped[0]["pendingRouteCount"])
|
|
|
|
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()
|