feat: add central slot scheduler for model requests

codai/admin/routes.py

@@ -307,9 +307,25 @@ async def api_status(username: str = Depends(require_auth)):
     # Request stats from queue manager
     req_total = 0
     req_active = 0
+    req_waiting = 0
+    req_metrics = {
+        "max_parallel_requests": 0,
+        "queue_max_size": 0,
+        "active_by_model": {},
+        "waiting_by_model": {},
+    }
     try:
         from codai.queue.manager import queue_manager
-        req_active = 1 if queue_manager._processing else 0
+        metrics = queue_manager.get_metrics()
+        req_active = int(metrics.get("active", 0))
+        req_waiting = int(metrics.get("waiting", 0))
+        req_total = req_active + req_waiting
+        req_metrics = {
+            "max_parallel_requests": metrics.get("max_parallel_requests", 0),
+            "queue_max_size": metrics.get("queue_max_size", 0),
+            "active_by_model": metrics.get("active_by_model", {}),
+            "waiting_by_model": metrics.get("waiting_by_model", {}),
+        }
     except Exception:
         pass
 
@@ -364,7 +380,15 @@ async def api_status(username: str = Depends(require_auth)):
         "enabled_models": enabled_models,
         "vram": vram,
         "cuda": is_cuda,
-        "requests": {"total": req_total, "active": req_active},
+        "requests": {
+            "total": req_total,
+            "active": req_active,
+            "waiting": req_waiting,
+            "max_parallel_requests": req_metrics["max_parallel_requests"],
+            "queue_max_size": req_metrics["queue_max_size"],
+            "active_by_model": req_metrics["active_by_model"],
+            "waiting_by_model": req_metrics["waiting_by_model"],
+        },
         "recent_activity": recent_activity,
         "whisper_server": whisper_status,
     }
@@ -1423,6 +1447,7 @@ async def api_get_settings(username: str = Depends(require_admin)):
             "https_key_path": c.server.https_key_path,
             "https_cert_path": c.server.https_cert_path,
             "queue_max_size": c.server.queue_max_size,
+            "max_parallel_requests": c.server.max_parallel_requests,
         },
         "backend": {
             "type": c.backend.type,
@@ -1478,6 +1503,10 @@ async def api_save_settings(request: Request, username: str = Depends(require_ad
             c.server.queue_max_size = max(1, int(srv["queue_max_size"]))
             from codai.queue.manager import queue_manager
             queue_manager.max_size = c.server.queue_max_size
+        if "max_parallel_requests" in srv:
+            c.server.max_parallel_requests = int(srv["max_parallel_requests"])
+            from codai.queue.manager import queue_manager
+            queue_manager.max_parallel_requests = c.server.max_parallel_requests
 
     if "backend" in data:
         bk = data["backend"]

codai/api/app.py

@@ -92,6 +92,8 @@ from codai.api.tts import router as tts_router
 from codai.api.text import router as text_router
 from codai.api.video import router as video_router
 from codai.api.audio_gen import router as audio_gen_router
+from codai.api.audio_stems import router as audio_stems_router
+from codai.api.audio_clean import router as audio_clean_router
 from codai.api.embeddings import router as embeddings_router
 from codai.api.pipelines import router as pipelines_router
 from codai.api.custom_pipelines import router as custom_pipelines_router

codai/api/custom_pipelines.py

@@ -274,6 +274,42 @@ async def _run_step(step: Dict, context: Dict, http_request) -> Dict:
     return _extract_output(step_type, result)
 
 
+def _infer_step_model_key(step: Dict) -> Optional[str]:
+    step_type = step.get('type')
+    params = step.get('params', {})
+
+    if step_type == 'stt':
+        model = params.get('model') or params.get('audio_model')
+        return f"audio:{model}" if model else None
+    if step_type == 'text_gen':
+        return params.get('model')
+    if step_type in {'image_gen', 'image_edit', 'image_upscale', 'image_depth', 'image_segment'}:
+        model = params.get('model')
+        return f"image:{model}" if model else None
+    if step_type in {'embed', 'embedding'}:
+        model = params.get('model')
+        return f"embedding:{model}" if model else None
+    if step_type in {'video_gen', 'video'}:
+        model = params.get('model')
+        return f"video:{model}" if model else None
+    return None
+
+
+async def _run_scheduled_step(step: Dict, context: Dict, http_request) -> Dict:
+    from codai.queue.manager import queue_manager
+
+    model_key = _infer_step_model_key(step)
+    if not model_key:
+        return await _run_step(step, context, http_request)
+
+    request_id = f"pipeline-step-{uuid.uuid4().hex[:8]}"
+    lease = await queue_manager.acquire(request_id, model_key)
+    try:
+        return await _run_step(step, context, http_request)
+    finally:
+        await queue_manager.release(lease)
+
+
 async def _execute_pipeline(pipeline_def: Dict, pipeline_input: str, http_request) -> Dict:
     """Execute all steps of a pipeline definition."""
     context = {'input': pipeline_input}
@@ -281,7 +317,7 @@ async def _execute_pipeline(pipeline_def: Dict, pipeline_input: str, http_reques
 
     for i, step in enumerate(pipeline_def.get('steps', [])):
         try:
-            out = await _run_step(step, context, http_request)
+            out = await _run_scheduled_step(step, context, http_request)
             context[f'step{i}'] = out
             steps_output.append({'step': i, 'type': step['type'],
                                   'label': step.get('label', step['type']), **out})
@@ -446,7 +482,7 @@ async def run_audio_understanding(request: AudioUnderstandRequest, http_request:
             'response_format': 'json',
         },
     }
-    stt_out = await _run_step(stt_step, {'input': request.input or ''}, http_request)
+    stt_out = await _run_scheduled_step(stt_step, {'input': request.input or ''}, http_request)
     transcript = stt_out.get('text') or stt_out.get('output') or ''
     steps.append({'step': 0, 'type': 'stt', 'label': 'Transcribe audio', **stt_out})
 
@@ -460,7 +496,7 @@ async def run_audio_understanding(request: AudioUnderstandRequest, http_request:
                 'prompt': f"{request.input or 'Summarize this audio transcript clearly.'}\n\nTranscript:\n{{{{step0.output}}}}",
             },
         }
-        text_out = await _run_step(text_step, {'input': request.input or '', 'step0': {'output': transcript, 'text': transcript}}, http_request)
+        text_out = await _run_scheduled_step(text_step, {'input': request.input or '', 'step0': {'output': transcript, 'text': transcript}}, http_request)
         summary = text_out.get('output')
         steps.append({'step': 1, 'type': 'text_gen', 'label': 'Reason over transcript', **text_out})
 
@@ -486,7 +522,7 @@ async def run_full_music_dub(request: AudioMusicDubRequest, http_request: Reques
             'response_format': 'json',
         },
     }
-    stt_out = await _run_step(stt_step, {'input': request.notes or ''}, http_request)
+    stt_out = await _run_scheduled_step(stt_step, {'input': request.notes or ''}, http_request)
     transcript = stt_out.get('text') or stt_out.get('output') or ''
     translated = transcript if not request.target_lang else f"[{request.target_lang}] {transcript}"
     steps = [

codai/config.py

@@ -31,6 +31,7 @@ class ServerConfig:
     https_key_path: Optional[str] = None
     https_cert_path: Optional[str] = None
     queue_max_size: int = 6
+    max_parallel_requests: int = 2
 
 
 @dataclass
@@ -302,6 +303,7 @@ class ConfigManager:
                 "https_key_path": self.config.server.https_key_path,
                 "https_cert_path": self.config.server.https_cert_path,
                 "queue_max_size": self.config.server.queue_max_size,
+                "max_parallel_requests": self.config.server.max_parallel_requests,
             },
             "backend": {
                 "type": self.config.backend.type,

codai/main.py

@@ -646,9 +646,10 @@ def main():
 
 
     
-    # Apply queue max size from config
+    # Apply queue scheduler settings from config
     from codai.queue.manager import queue_manager
     queue_manager.max_size = config.server.queue_max_size
+    queue_manager.max_parallel_requests = config.server.max_parallel_requests
 
     # Start the server
     import uvicorn

codai/models/manager.py

@@ -428,6 +428,14 @@ class MultiModelManager:
         """Return the first image model or None."""
         return self.image_models[0] if self.image_models else None
 
+    def get_loaded_model_keys(self) -> set:
+        """Return the set of currently loaded model keys."""
+        return set(self.models.keys())
+
+    def has_loaded_model(self, model_key: str) -> bool:
+        """Return True when the given model key is currently loaded."""
+        return model_key in self.models
+    
     def cleanup(self):
         """Cleanup all models and resources."""
         # Cleanup all model managers

tests/test_studio_composed_surfaces.py

@@ -104,6 +104,35 @@ def test_audio_understanding_returns_transcript_only_without_text_model(monkeypa
     assert len(body["steps"]) == 1
 
 
+def test_audio_understanding_pipeline_steps_release_scheduler_slots(monkeypatch, studio_client):
+    from codai.api import custom_pipelines
+    from codai.queue.manager import queue_manager
+
+    observed = []
+
+    async def fake_run_step(step, context, http_request):
+        observed.append(queue_manager.get_metrics()["active"])
+        return {"output": step["type"], "text": step["type"]}
+
+    monkeypatch.setattr(custom_pipelines, "_run_step", fake_run_step)
+    queue_manager.reset_for_tests()
+    queue_manager.set_loaded_models({"audio:whisper-small", "qwen-text"})
+
+    response = studio_client.post(
+        "/v1/pipelines/audio-understand",
+        json={
+            "input": "Summarize",
+            "audio_model": "whisper-small",
+            "text_model": "qwen-text",
+            "audio": "ZmFrZQ==",
+        },
+    )
+
+    assert response.status_code == 200
+    assert observed == [1, 1]
+    assert queue_manager.get_metrics()["active"] == 0
+
+
 def test_audio_understanding_requires_audio_source(studio_client):
     response = studio_client.post(
         "/v1/pipelines/audio-understand",