Merge pull request zhaochenyang20#14 from zhaochenyang20/validate_video_generation

[router] Validate video generation before slection; Support selctive worker choice

Author: zhaochenyang20

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "sglang"]
+	path = sglang
+	url = https://github.com/sgl-project/sglang.git

README.md

@@ -20,15 +20,19 @@ From repository root:
 # python3 -m venv .venv
 # source .venv/bin/activate
 # pip install uv
-git clone https://github.com/sglang/sglang-diffusion-routing.git
+git clone --recursive https://github.com/sglang/sglang-diffusion-routing.git
 cd sglang-diffusion-routing
 uv pip install .
 ```
 
 Workers require SGLang diffusion support:
 
 ```bash
+# If cloned sglang-diffusion-routing without --recursive, run:
+# git submodule update --init --recursive
+cd sglang
 uv pip install "sglang[diffusion]" --prerelease=allow
+cd ..
 ```
 
 ## Quick Start
@@ -114,16 +118,10 @@ with open('output.png', 'wb') as f:
 print('Saved to output.png')
 "
 
-# Video generation request
-curl -X POST http://localhost:30081/generate_video \
-    -H "Content-Type: application/json" \
-    -d '{
-        "model": "Qwen/Qwen-Image",
-        "prompt": "a flowing river"
-    }'
 
-# Check per-worker health and load
-curl http://localhost:30081/health_workers
+curl -X POST http://localhost:30081/update_weights_from_disk \
+    -H "Content-Type: application/json" \
+    -d '{"model_path": "Qwen/Qwen-Image-2512"}'
 ```
 
 ### Python requests examples
@@ -177,66 +175,10 @@ print(resp.json())
 - `GET /health`: aggregated router health.
 - `GET /health_workers`: per-worker health and active request counts.
 - `POST /generate`: forwards to worker `/v1/images/generations`.
-- `POST /generate_video`: forwards to worker `/v1/videos`.
+- `POST /generate_video`: forwards to worker `/v1/videos`; rejects image-only workers (`T2I`/`I2I`/`TI2I`) with `400`.
 - `POST /update_weights_from_disk`: broadcast to healthy workers.
 - `GET|POST|PUT|DELETE /{path}`: catch-all proxy forwarding.
-
-## `update_weights_from_disk` behavior
-
-Full details: [docs/update_weights_from_disk.md](docs/update_weights_from_disk.md)
-
-- The router forwards request payloads as-is to each healthy worker.
-- The router does not validate payload schema; payload semantics are worker-defined.
-- Worker servers must implement `POST /update_weights_from_disk`.
-
-Example:
-
-```bash
-curl -X POST http://localhost:30081/update_weights_from_disk \
-    -H "Content-Type: application/json" \
-    -d '{"model_path": "/path/to/new/weights"}'
-```
-
-Response shape:
-
-```json
-{
-  "results": [
-    {
-      "worker_url": "http://localhost:30000",
-      "status_code": 200,
-      "body": {
-        "ok": true
-      }
-    }
-  ]
-}
-```
-
-## Benchmark Scripts
-
-Benchmark scripts are available under `tests/benchmarks/diffusion_router/` and are intended for manual runs.
-They are not part of default unit test collection (`pytest tests/unit -v`).
-
-Single benchmark:
-
-```bash
-SGLANG_USE_MODELSCOPE=TRUE python tests/benchmarks/diffusion_router/bench_router.py \
-    --model Qwen/Qwen-Image \
-    --num-workers 2 \
-    --num-prompts 20 \
-    --max-concurrency 4
-```
-
-Algorithm comparison:
-
-```bash
-SGLANG_USE_MODELSCOPE=TRUE python tests/benchmarks/diffusion_router/bench_routing_algorithms.py \
-    --model Qwen/Qwen-Image \
-    --num-workers 2 \
-    --num-prompts 20 \
-    --max-concurrency 4
-```
+- `POST /update_weights_from_disk`: broadcast to all healthy workers.
 
 ## Project Layout
 

development.md

@@ -1,4 +1,6 @@
-Development install:
+# Development
+
+## Development Install
 
 ```bash
 pip install -e .
@@ -10,3 +12,28 @@ Run tests:
 pip install pytest
 pytest tests/unit -v
 ```
+
+## Benchmark Scripts
+
+Benchmark scripts are available under `tests/benchmarks/diffusion_router/` and are intended for manual runs.
+They are not part of default unit test collection (`pytest tests/unit -v`).
+
+Single benchmark:
+
+```bash
+python tests/benchmarks/diffusion_router/bench_router.py \
+    --model Qwen/Qwen-Image \
+    --num-workers 2 \
+    --num-prompts 20 \
+    --max-concurrency 4
+```
+
+Algorithm comparison:
+
+```bash
+python tests/benchmarks/diffusion_router/bench_routing_algorithms.py \
+    --model Qwen/Qwen-Image \
+    --num-workers 2 \
+    --num-prompts 20 \
+    --max-concurrency 4
+```

docs/update_weights_from_disk.md

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import argparse
+import asyncio
 import sys
 
 from sglang_diffusion_routing import DiffusionRouter
@@ -25,8 +26,18 @@ def _run_router_server(
         worker_urls if worker_urls is not None else args.worker_urls or []
     )
     router = DiffusionRouter(args, verbose=args.verbose)
+    refresh_tasks = []
     for url in worker_urls:
-        router.register_worker(url)
+        normalized_url = router.normalize_worker_url(url)
+        router.register_worker(normalized_url)
+        refresh_tasks.append(router.refresh_worker_video_support(normalized_url))
+
+    if refresh_tasks:
+
+        async def _refresh_all_worker_video_support() -> None:
+            await asyncio.gather(*refresh_tasks)
+
+        asyncio.run(_refresh_all_worker_video_support())
 
     print(f"{log_prefix} starting router on {args.host}:{args.port}", flush=True)
     print(

sglang

@@ -16,9 +16,11 @@
 logger = logging.getLogger(__name__)
 
 _METADATA_HOSTS = {"169.254.169.254", "metadata.google.internal"}
+_IMAGE_TASK_TYPES = {"T2I", "I2I", "TI2I"}
 
 
 class DiffusionRouter:
+
     def __init__(self, args, verbose: bool = False):
         """Initialize the router for load-balancing sglang-diffusion workers."""
         self.args = args
@@ -32,6 +34,9 @@ def __init__(self, args, verbose: bool = False):
         self.worker_request_counts: dict[str, int] = {}
         # URL -> consecutive health check failures
         self.worker_failure_counts: dict[str, int] = {}
+        # URL -> whether worker supports video generation
+        # True: supports, False: image-only, None: unknown/unprobed
+        self.worker_video_support: dict[str, bool | None] = {}
         # quarantined workers excluded from routing
         self.dead_workers: set[str] = set()
         self._health_task: asyncio.Task | None = None
@@ -139,14 +144,23 @@ async def _health_check_loop(self) -> None:
                 )
                 await asyncio.sleep(5)
 
-    def _use_url(self) -> str:
-        """Select a worker URL based on the configured routing algorithm."""
+    def _select_worker_by_routing(self, worker_urls: list[str] | None = None) -> str:
+        """Select a worker URL based on routing algorithm and optional candidates.
+
+        Args:
+            worker_urls: Optional list of worker URLs to consider. If provided,
+            only these workers will be considered for selection. If not provided,
+            all registered workers will be considered.
+        """
         if not self.worker_request_counts:
             raise RuntimeError("No workers registered in the pool")
 
         valid_workers = [
             w for w in self.worker_request_counts if w not in self.dead_workers
         ]
+        if worker_urls is not None:
+            allowed = {w for w in worker_urls if w in self.worker_request_counts}
+            valid_workers = [w for w in valid_workers if w in allowed]
         if not valid_workers:
             raise RuntimeError("No healthy workers available in the pool")
 
@@ -202,13 +216,14 @@ def _build_proxy_response(
             media_type=content_type,
         )
 
-    async def _forward_to_worker(self, request: Request, path: str) -> Response:
-        """Forward a request to a selected worker and return the response."""
+    async def _forward_to_worker(
+        self, request: Request, path: str, worker_urls: list[str] | None = None
+    ) -> Response:
+        """Forward request to a selected worker (optionally from candidate URLs)."""
         try:
-            worker_url = self._use_url()
+            worker_url = self._select_worker_by_routing(worker_urls=worker_urls)
         except RuntimeError as exc:
             return JSONResponse(status_code=503, content={"error": str(exc)})
-
         try:
             query = request.url.query
             url = (
@@ -243,6 +258,29 @@ async def _forward_to_worker(self, request: Request, path: str) -> Response:
         finally:
             self._finish_url(worker_url)
 
+    async def _probe_worker_video_support(self, worker_url: str) -> bool | None:
+        """Probe /v1/models and infer if this worker supports video generation."""
+        try:
+            response = await self.client.get(f"{worker_url}/v1/models", timeout=5.0)
+            if response.status_code == 200:
+                payload = response.json()
+                data = payload.get("data")
+                task_type = (
+                    data[0].get("task_type")
+                    if isinstance(data, list) and data
+                    else None
+                )
+                if isinstance(task_type, str):
+                    return task_type.upper() not in _IMAGE_TASK_TYPES
+        except (httpx.RequestError, json.JSONDecodeError):
+            return None
+
+    async def refresh_worker_video_support(self, worker_url: str) -> None:
+        """Refresh cached video capability for a single worker."""
+        self.worker_video_support[worker_url] = await self._probe_worker_video_support(
+            worker_url
+        )
+
     async def _broadcast_to_workers(
         self, path: str, body: bytes, headers: dict
     ) -> list[dict]:
@@ -297,7 +335,7 @@ def _sanitize_response_headers(headers) -> dict:
         }
 
     @staticmethod
-    def _normalize_worker_url(url: str) -> str:
+    def normalize_worker_url(url: str) -> str:
         if not isinstance(url, str):
             raise ValueError("worker_url must be a string")
 
@@ -345,7 +383,22 @@ async def generate(self, request: Request):
 
     async def generate_video(self, request: Request):
         """Route video generation to /v1/videos."""
-        return await self._forward_to_worker(request, "v1/videos")
+        candidate_workers = [
+            worker_url
+            for worker_url, support in self.worker_video_support.items()
+            if support
+        ]
+
+        if not candidate_workers:
+            return JSONResponse(
+                status_code=400,
+                content={
+                    "error": "No video-capable workers available in current worker pool.",
+                },
+            )
+        return await self._forward_to_worker(
+            request, "v1/videos", worker_urls=candidate_workers
+        )
 
     async def health(self, request: Request):
         """Aggregated health status: healthy if at least one worker is alive."""
@@ -388,10 +441,11 @@ async def update_weights_from_disk(self, request: Request):
 
     def register_worker(self, url: str) -> None:
         """Register a worker URL if not already known."""
-        normalized_url = self._normalize_worker_url(url)
+        normalized_url = self.normalize_worker_url(url)
         if normalized_url not in self.worker_request_counts:
             self.worker_request_counts[normalized_url] = 0
             self.worker_failure_counts[normalized_url] = 0
+            self.worker_video_support[normalized_url] = None
             if self.verbose:
                 print(f"[diffusion-router] Added new worker: {normalized_url}")
 
@@ -422,6 +476,7 @@ async def add_worker(self, request: Request):
             self.register_worker(worker_url)
         except ValueError as exc:
             return JSONResponse(status_code=400, content={"error": str(exc)})
+        await self.refresh_worker_video_support(worker_url)
         return {
             "status": "success",
             "worker_urls": list(self.worker_request_counts.keys()),