support identical method

Author: zhaochenyang20

README.md

@@ -55,7 +55,7 @@ sglang-d-router --port 30081 --launcher-config examples/local_launcher.yaml
 ```yaml
 launcher:
   backend: local
-  model: stabilityai/stable-diffusion-3-medium-diffusers
+  model: Qwen/Qwen-Image
   num_workers: 8
   num_gpus_per_worker: 1
   worker_base_port: 10090
@@ -70,14 +70,14 @@ launcher:
 
 # worker 1
 CUDA_VISIBLE_DEVICES=0 sglang serve \
-    --model-path stabilityai/stable-diffusion-3-medium-diffusers \
+    --model-path Qwen/Qwen-Image \
     --num-gpus 1 \
     --host 127.0.0.1 \
     --port 30000
 
 # worker 2
 CUDA_VISIBLE_DEVICES=1 sglang serve \
-    --model-path stabilityai/stable-diffusion-3-medium-diffusers \
+    --model-path Qwen/Qwen-Image \
     --num-gpus 1 \
     --host 127.0.0.1 \
     --port 30002
@@ -139,7 +139,7 @@ print("Saved to output.png")
 # so this request will fail. Use a video-capable model instead.
 
 resp = requests.post(f"{ROUTER}/v1/videos", json={
-    "model": "stabilityai/stable-diffusion-3-medium-diffusers",
+    "model": "Qwen/Qwen-Image",
     "prompt": "a flowing river",
 })
 print(resp.json())
@@ -149,9 +149,17 @@ if video_id:
 
 # Update weights from disk
 resp = requests.post(f"{ROUTER}/update_weights_from_disk", json={
-    "model_path": "/path/to/new/checkpoint",
+    "model_path": "Qwen/Qwen-Image-2512",
 })
 print(resp.json())
+
+# sleep and wake up
+resp = requests.post(f"{ROUTER}/release_memory_occupation", json={})
+print(resp.json())
+
+
+resp = requests.post(f"{ROUTER}/resume_memory_occupation", json={})
+print(resp.json())
 ```
 
 ### Native Diffusion Generate Endpoint (with Trajectory & Log-Prob)

src/sglang_diffusion_routing/router/diffusion_router.py

@@ -771,30 +771,53 @@ async def release_memory_occupation(self, request: Request):
         status, payload = await self._broadcast_to_pool(
             request, "release_memory_occupation"
         )
-        if status != 200:
-            return JSONResponse(status_code=status, content=payload)
-
-        for item in payload["results"]:
-            if item.get("status_code") == 200:
-                self.sleeping_workers.add(item["worker_url"])
+        if status == 200:
+            for item in payload["results"]:
+                if item.get("status_code") == 200:
+                    self.sleeping_workers.add(item["worker_url"])
+            return JSONResponse(content=payload)
+
+        all_already_sleeping = bool(self.worker_request_counts) and all(
+            url in self.sleeping_workers
+            for url in self.worker_request_counts
+            if url not in self.dead_workers
+        )
+        if all_already_sleeping:
+            return JSONResponse(
+                content={
+                    "message": "All workers are already sleeping",
+                    "sleeping_workers": len(self.sleeping_workers),
+                }
+            )
 
-        return JSONResponse(content=payload)
+        return JSONResponse(status_code=status, content=payload)
 
     async def resume_memory_occupation(self, request: Request):
         status, payload = await self._broadcast_to_pool(
             request, "resume_memory_occupation"
         )
-        if status != 200:
-            return JSONResponse(status_code=status, content=payload)
-
-        for item in payload["results"]:
-            if item.get("status_code") == 200:
-                url = item["worker_url"]
-                self.sleeping_workers.discard(url)
-                self.worker_failure_counts[url] = 0
-                self.dead_workers.discard(url)  # wake success => recover
+        if status == 200:
+            for item in payload["results"]:
+                if item.get("status_code") == 200:
+                    url = item["worker_url"]
+                    self.sleeping_workers.discard(url)
+                    self.worker_failure_counts[url] = 0
+                    self.dead_workers.discard(url)
+            return JSONResponse(content=payload)
+
+        has_no_sleeping_workers = not self.sleeping_workers and bool(
+            self.worker_request_counts
+        )
+        if has_no_sleeping_workers:
+            return JSONResponse(
+                content={
+                    "message": "All workers are already active",
+                    "active_workers": len(self.worker_request_counts)
+                    - len(self.dead_workers),
+                }
+            )
 
-        return JSONResponse(content=payload)
+        return JSONResponse(status_code=status, content=payload)
 
     def register_worker(self, url: str) -> None:
         """Register a worker URL if not already known."""

tests/unit/test_router_endpoints.py

@@ -415,6 +415,31 @@ async def fake_broadcast(path: str, body: bytes, headers: dict):
     assert router.worker_failure_counts["http://localhost:10091"] == 2
 
 
+def test_release_memory_occupation_idempotent_when_all_already_sleeping():
+    router = DiffusionRouter(make_router_args())
+    router.register_worker("http://localhost:10090")
+    router.register_worker("http://localhost:10091")
+    router.sleeping_workers.update(["http://localhost:10090", "http://localhost:10091"])
+
+    with TestClient(router.app) as client:
+        resp = client.post("/release_memory_occupation", json={})
+        assert resp.status_code == 200
+        assert "already sleeping" in resp.json()["message"].lower()
+        assert resp.json()["sleeping_workers"] == 2
+
+
+def test_resume_memory_occupation_idempotent_when_all_already_active():
+    router = DiffusionRouter(make_router_args())
+    router.register_worker("http://localhost:10090")
+    router.register_worker("http://localhost:10091")
+
+    with TestClient(router.app) as client:
+        resp = client.post("/resume_memory_occupation", json={})
+        assert resp.status_code == 200
+        assert "already active" in resp.json()["message"].lower()
+        assert resp.json()["active_workers"] == 2
+
+
 def test_select_worker_excludes_sleeping_workers():
     router = DiffusionRouter(make_router_args(routing_algorithm="least-request"))
     router.register_worker("http://localhost:10090")