fix: prioritize preferred GPU routes with fallback

This commit is contained in:
CoolBoy
2026-09-18 00:25:49 +08:00
parent 2c748dcec2
commit 3077b169c5
4 changed files with 217 additions and 85 deletions

View File

@@ -68,6 +68,51 @@ class ExactGpuProcessedClient:
return {"code": 0, "data": {"id": "should-not-submit"}}
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}
@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 UniquenessRejectingClient:
@staticmethod
def add_task(_payload: dict) -> dict:
@@ -107,6 +152,50 @@ def candidate(target_gpu: str = "Vastai_va16") -> dict:
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,