fix: prioritize preferred GPU routes with fallback
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user