perf(rollout): six independent rollout-path fixes and knobs

miles/backends/fsdp_utils/checkpoint.py

@@ -85,13 +85,25 @@ def state_dict(self):
 
     def load_state_dict(self, state_dict):
         options = StateDictOptions(strict=False) if self.lora_only else None
-        set_state_dict(
+        incompatible = set_state_dict(
             self.model,
             optimizers=[],
             model_state_dict=state_dict["model"],
             optim_state_dict=None,
             options=options,
         )
+        # strict=False swallows mismatches, so a checkpoint whose keys match
+        # nothing "loads" as a silent no-op. In lora_only mode missing base keys
+        # are expected; unclaimed checkpoint keys or unfilled LoRA params are not.
+        unclaimed = list(incompatible.unexpected_keys)
+        unfilled = [k for k in incompatible.missing_keys if not self.lora_only or _is_lora_param_name(k)]
+        if unclaimed or unfilled:
+            logger.error(
+                f"[FSDP] Checkpoint mismatch: {len(unclaimed)} unclaimed checkpoint keys, "
+                f"{len(unfilled)} unfilled model params (e.g. {(unclaimed + unfilled)[:3]})"
+            )
+        else:
+            logger.info(f"[FSDP] Applied {len(state_dict['model'])} params from the checkpoint")
 
 
 class OptimizerState(Stateful):

miles/rollout/rm_hub/core.py

@@ -82,11 +82,16 @@ def _next_actor(self):
         return actor
 
     async def score(self, images: list, prompts: list[str]) -> list[float]:
-        refs = []
-        for start in range(0, len(images), self._batch_size):
-            end = start + self._batch_size
-            refs.append(self._next_actor().score_batch.remote(images[start:end], prompts[start:end]))
+        chunks = [
+            (self._next_actor(), images[start:start + self._batch_size], prompts[start:start + self._batch_size])
+            for start in range(0, len(images), self._batch_size)
+        ]
+
+        def submit_and_collect():
+            # .remote() copies each chunk into plasma in the calling thread, so
+            # submitting on the loop stalls every concurrent request for that copy.
+            return ray.get([actor.score_batch.remote(img, prm) for actor, img, prm in chunks])
 
         loop = asyncio.get_running_loop()
-        chunked_scores = await loop.run_in_executor(None, ray.get, refs)
+        chunked_scores = await loop.run_in_executor(None, submit_and_collect)
         return [float(score) for chunk in chunked_scores for score in chunk]

miles/rollout/rm_hub/pickscore.py

@@ -8,7 +8,7 @@
 from PIL import Image
 
 from miles.utils.misc import SingletonMeta
-from miles.utils.processing_utils import cfhw_to_fhwc, image_or_video_to_uint8
+from miles.utils.processing_utils import cfhw_to_fhwc
 from miles.utils.types import Sample
 
 from .core import AsyncRewardActorPool
@@ -39,9 +39,17 @@ def _sample_to_rgb_hwc_uint8_frames(sample: Sample, num_frames: int | None) -> l
     if cfhw is None:
         raise ValueError("generated_output is None")
 
-    fhwc = image_or_video_to_uint8(cfhw_to_fhwc(cfhw.detach().cpu()))
-    indices = sample_frame_indices(fhwc.shape[0], num_frames)
-    return [np.ascontiguousarray(fhwc[i].numpy()) for i in indices]
+    # Convert only the frames that survive: a 107-frame clip yields 8 here, and
+    # converting the whole clip first costs several full-size float32 copies of
+    # it. The unit-scale decision still comes from the whole clip, so selecting
+    # first cannot change how a frame is quantised.
+    indices = sample_frame_indices(cfhw.shape[1], num_frames)
+    unit_scale = float(cfhw.max()) <= 1.0 + 1e-3
+    selected = cfhw[:, indices].detach().cpu().float()
+    if unit_scale:
+        selected = selected * 255.0
+    fhwc = cfhw_to_fhwc(selected.clamp(0, 255).to(torch.uint8))
+    return [np.ascontiguousarray(fhwc[i].numpy()) for i in range(len(indices))]
 
 
 class PickScoreScorer(torch.nn.Module):

miles/rollout/sglang_diffusion_rollout.py

@@ -2,6 +2,7 @@
 import copy
 import inspect
 import logging
+import random
 from argparse import Namespace
 from collections.abc import Callable
 from contextlib import contextmanager
@@ -60,6 +61,7 @@ def build_rollout_sampling_params(
                 "rollout_noise_level": args.diffusion_noise_level,
                 "rollout_log_prob_no_const": args.diffusion_log_prob_no_const,
                 "rollout_debug_mode": args.diffusion_debug_mode,
+                "rollout_video_dtype": args.rollout_video_dtype,
                 "rollout_return_denoising_env": True,
                 "rollout_return_dit_trajectory": True,
             }
@@ -225,6 +227,8 @@ async def generate_and_rm_microgroup(
     state = GenerateState(args)
 
     # generate
+    if args.rollout_request_stagger > 0:
+        await asyncio.sleep(random.uniform(0, args.rollout_request_stagger))
     async with state.semaphore:
         with state.dp_rank_context() as _:
             if args.custom_generate_function_path is not None:

miles/router/router.py

@@ -64,6 +64,7 @@ def _setup_routes(self):
         """Setup all the HTTP routes"""
         # sglang-router api
         self.app.post("/add_worker")(self.add_worker)
+        self.app.post("/remove_worker")(self.remove_worker)
         self.app.get("/list_workers")(self.list_workers)
         # Catch-all route for proxying to SGLang - must be registered LAST
         self.app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])(self.proxy)
@@ -177,6 +178,30 @@ async def add_worker(self, request: Request):
 
         return {"status": "success", "worker_urls": self.worker_request_counts}
 
+    async def remove_worker(self, request: Request):
+        """Drop a worker from the routing pool, same URL forms as add_worker.
+
+        Engines call this on shutdown. Without the route it fell through to the
+        catch-all and was proxied to a worker as if it were a generate request.
+        """
+        worker_url = request.query_params.get("url") or request.query_params.get("worker_url")
+        if not worker_url:
+            body = await request.body()
+            payload = json.loads(body) if body else {}
+            worker_url = payload.get("url") or payload.get("worker_url")
+
+        if not worker_url:
+            return JSONResponse(
+                status_code=400, content={"error": "worker_url is required (use query ?url=... or JSON body)"}
+            )
+
+        self.worker_request_counts.pop(worker_url, None)
+        self.worker_failure_counts.pop(worker_url, None)
+        self.dead_workers.discard(worker_url)
+        if self.verbose:
+            print(f"[miles-router] Removed worker: {worker_url}")
+        return {"status": "success", "worker_urls": self.worker_request_counts}
+
     async def list_workers(self, request: Request):
         """List all registered workers"""
         return {"urls": list(self.worker_request_counts.keys())}

miles/utils/arguments.py

@@ -252,6 +252,30 @@ def add_rollout_arguments(parser):
                 default=1,
                 help="Samples per prompt sent in one rollout request (sub-batch of the group).",
             )
+            parser.add_argument(
+                "--rollout-video-dtype",
+                type=str,
+                choices=["keep", "uint8"],
+                default="keep",
+                help=(
+                    "Dtype of the decoded video in rollout responses. 'keep' returns the "
+                    "engine's raw float tensor; 'uint8' quantises engine-side with the same "
+                    "formula the reward path applies, cutting the response body ~4x. Use "
+                    "'keep' for any consumer that needs the unquantised tensor."
+                ),
+            )
+            parser.add_argument(
+                "--rollout-request-stagger",
+                type=float,
+                default=0.0,
+                help=(
+                    "Max uniform random delay (s) before a request first tries to acquire a "
+                    "concurrency slot. Requests admitted together stay in phase for the whole "
+                    "rollout because generation time is near-constant, so their response bodies "
+                    "land on the transfer path simultaneously; a one-shot jitter breaks that. "
+                    "Applied before the acquire, so it never idles a held slot. 0 disables."
+                ),
+            )
             parser.add_argument(
                 "--diffusion-fps",
                 type=float,
@@ -1482,6 +1506,9 @@ def miles_validate_args(args):
     if args.wandb_log_image_interval < 1:
         raise ValueError(f"wandb_log_image_interval must be >= 1, got {args.wandb_log_image_interval}")
 
+    if args.rollout_request_stagger < 0:
+        raise ValueError(f"--rollout-request-stagger must be non-negative, got {args.rollout_request_stagger}")
+
     args.rollout_patch_groups = [name for name in (args.rollout_patch_group or "").split(",") if name]
 
     if not args.hf_checkpoint:

tests/fast/rollout/test_rollout_negative_prompt.py

@@ -30,6 +30,7 @@ def _args(**overrides):
         diffusion_noise_level=0.7,
         diffusion_log_prob_no_const=False,
         diffusion_debug_mode=False,
+        rollout_video_dtype="keep",
         train_pipeline_config_path=None,
     )
     values.update(overrides)