feat: add queue-aware adaptive scheduling

modelhub_submmit_api/task_registry.py

@@ -63,58 +63,69 @@ def task_specs_for_model(model: HFModelSummary) -> list[TaskSpec]:
     return [task for task in TASK_SPECS if pipeline_tag in task.pipeline_tags]
 
 
-def choose_text_generation_framework(target_gpu: str, supported_frameworks: set[str], inspection: ModelInspection) -> str:
+def compatible_text_generation_frameworks(
+    target_gpu: str,
+    supported_frameworks: set[str],
+    inspection: ModelInspection,
+) -> list[str]:
     can_llamacpp = "llamacpp" in supported_frameworks
     vllm_like = [framework for framework in VLLM_LIKE_FRAMEWORKS if framework in supported_frameworks]
     can_transformers = "transformers" in supported_frameworks
+    compatible: list[str] = []
 
-    if can_llamacpp and vllm_like:
-        if inspection.has_gguf:
-            return "llamacpp"
-        if inspection.has_vllm_weights:
-            return vllm_like[0]
-        raise ValueError(f"{target_gpu} supports both llamacpp and vLLM-like frameworks, but the model lacks usable weights")
-    if can_llamacpp:
-        if inspection.has_gguf:
-            return "llamacpp"
-        raise ValueError(f"{target_gpu} only supports llamacpp for this workflow, but no usable GGUF was found")
-    if vllm_like:
-        if inspection.has_vllm_weights:
-            return vllm_like[0]
-        raise ValueError(f"{target_gpu} only supports vLLM-like frameworks for this workflow, but no standard weights were found")
-    if can_transformers:
-        if inspection.has_vllm_weights:
-            return "transformers"
-        raise ValueError(f"{target_gpu} only supports transformers for this workflow, but no standard weights were found")
-    raise ValueError(f"No supported LLM framework template is available for {target_gpu}")
+    if can_llamacpp and inspection.has_gguf:
+        compatible.append("llamacpp")
+    if inspection.has_vllm_weights:
+        compatible.extend(vllm_like)
+        if can_transformers:
+            compatible.append("transformers")
+    if not compatible:
+        raise ValueError(f"No compatible LLM weights/framework combination is available for {target_gpu}")
+    return compatible
 
 
-def choose_framework_for_task(task_type: str, target_gpu: str, supported_frameworks: set[str], inspection: ModelInspection) -> str:
+def choose_text_generation_framework(target_gpu: str, supported_frameworks: set[str], inspection: ModelInspection) -> str:
+    return compatible_text_generation_frameworks(target_gpu, supported_frameworks, inspection)[0]
+
+
+def compatible_frameworks_for_task(
+    task_type: str,
+    target_gpu: str,
+    supported_frameworks: set[str],
+    inspection: ModelInspection,
+) -> list[str]:
     if task_type in {"text-generation", "visual-multi-modal", "reinforcement_learning"}:
-        return choose_text_generation_framework(target_gpu, supported_frameworks, inspection)
+        return compatible_text_generation_frameworks(target_gpu, supported_frameworks, inspection)
 
     if task_type == "asr":
+        compatible: list[str] = []
         if "sherpa-onnx" in supported_frameworks and inspection.has_onnx_weights:
-            return "sherpa-onnx"
-        for framework in ("transformers", "funasr"):
-            if framework in supported_frameworks and inspection.has_standard_weights:
-                return framework
+            compatible.append("sherpa-onnx")
+        if inspection.has_standard_weights:
+            compatible.extend(framework for framework in ("transformers", "funasr") if framework in supported_frameworks)
+        if compatible:
+            return compatible
         raise ValueError(f"No compatible ASR framework found for {target_gpu}")
 
     if task_type == "feature_emb":
-        for framework in ("sentence-transformers", "transformers"):
-            if framework in supported_frameworks and inspection.has_standard_weights:
-                return framework
+        if inspection.has_standard_weights:
+            compatible = [framework for framework in ("sentence-transformers", "transformers") if framework in supported_frameworks]
+            if compatible:
+                return compatible
         raise ValueError(f"No compatible embedding framework found for {target_gpu}")
 
     if task_type in {"question_answering", "vision_classification", "text_classification"}:
         if "transformers" in supported_frameworks and inspection.has_standard_weights:
-            return "transformers"
+            return ["transformers"]
         raise ValueError(f"No compatible transformers template found for {task_type} on {target_gpu}")
 
     if task_type == "text-to-image-generation":
         if "diffusers" in supported_frameworks and inspection.has_standard_weights:
-            return "diffusers"
+            return ["diffusers"]
         raise ValueError(f"No compatible diffusers template found for {target_gpu}")
 
     raise ValueError(f"Unsupported task type for auto framework selection: {task_type}")
+
+
+def choose_framework_for_task(task_type: str, target_gpu: str, supported_frameworks: set[str], inspection: ModelInspection) -> str:
+    return compatible_frameworks_for_task(task_type, target_gpu, supported_frameworks, inspection)[0]