test(e2e): run every RL e2e with the output store on, generated output included (CI only, do not merge)

.github/workflows/_run-ci.yml

@@ -62,6 +62,10 @@ jobs:
     # (e.g. the label name '["h200", "2gpu"]' with brackets), which no
     # runner carries, and the job queues forever.
     runs-on: ${{ fromJSON(inputs.runs_on) }}
+    # The CI node's RDMA devices are RoCE, and a RoCE connection needs the device's network interface, which the
+    # runner's bridge network does not have. `ci-roce` is an ipvlan network on bond0 that adds it. Created once on
+    # the node; without it the container is not created:
+    #   docker network create -d ipvlan --internal --subnet 192.168.250.0/24 -o parent=bond0 ci-roce
     container:
       image: ${{ inputs.container_image }}
       options: >
@@ -80,6 +84,7 @@ jobs:
         -v /data/miles_ci/paddleocr:/github/home/.paddleocr
         --privileged
         --ulimit nofile=65535:65535
+        --network ci-roce
     defaults:
       run:
         working-directory: ${{ github.workspace }}
@@ -90,6 +95,9 @@ jobs:
       # Gated models (e.g. stabilityai/stable-diffusion-3.5-medium for the SD3
       # e2e) need an HF token with repo access; unset for public-model tests.
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
+      # Mooncake's whitelist of RDMA devices: only bond0's is reachable through ci-roce, and a transfer that
+      # tries any of the other seven fails.
+      MC_MS_FILTERS: mlx5_bond_0
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4

.github/workflows/pr-test.yml

@@ -319,7 +319,7 @@ jobs:
         --labels ${{ join(github.event.pull_request.labels.*.name, ' ') }}
         ${{ (github.event_name == 'workflow_dispatch' || github.event_name == 'schedule' || needs.ci-gate.outputs.docs_only == 'true' || contains(github.event.pull_request.labels.*.name, 'run-ci-image') || contains(github.event.pull_request.labels.*.name, 'run-ci-all')) && '--match-all-labels' || '' }}
         ${{ (github.event_name == 'schedule' || inputs.nightly) && '--nightly' || '' }}
-        ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }}
+        ${{ (github.event_name == 'schedule' || github.event_name == 'pull_request' || inputs.continue_on_error) && '--continue-on-error' || '' }}
     secrets: inherit
 
   # Stage B: CPU-only slower bucket (currently empty; reserved for future
@@ -337,7 +337,7 @@ jobs:
         --labels ${{ join(github.event.pull_request.labels.*.name, ' ') }}
         ${{ (github.event_name == 'workflow_dispatch' || github.event_name == 'schedule' || needs.ci-gate.outputs.docs_only == 'true' || contains(github.event.pull_request.labels.*.name, 'run-ci-image') || contains(github.event.pull_request.labels.*.name, 'run-ci-all')) && '--match-all-labels' || '' }}
         ${{ (github.event_name == 'schedule' || inputs.nightly) && '--nightly' || '' }}
-        ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }}
+        ${{ (github.event_name == 'schedule' || github.event_name == 'pull_request' || inputs.continue_on_error) && '--continue-on-error' || '' }}
     secrets: inherit
 
   # Stage B GPU: fast-gpu bucket, runs once stage-a-cpu is green (fast-fail:
@@ -364,7 +364,7 @@ jobs:
         --labels ${{ join(github.event.pull_request.labels.*.name, ' ') }}
         ${{ (github.event_name == 'workflow_dispatch' || github.event_name == 'schedule' || contains(github.event.pull_request.labels.*.name, 'run-ci-image') || contains(github.event.pull_request.labels.*.name, 'run-ci-all')) && '--match-all-labels' || '' }}
         ${{ (github.event_name == 'schedule' || inputs.nightly) && '--nightly' || '' }}
-        ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }}
+        ${{ (github.event_name == 'schedule' || github.event_name == 'pull_request' || inputs.continue_on_error) && '--continue-on-error' || '' }}
     secrets: inherit
 
   # Stage B GPU: 4-GPU distributed tests on the 5gpu runner.
@@ -387,7 +387,7 @@ jobs:
         --labels ${{ join(github.event.pull_request.labels.*.name, ' ') }}
         ${{ (github.event_name == 'workflow_dispatch' || github.event_name == 'schedule' || contains(github.event.pull_request.labels.*.name, 'run-ci-image') || contains(github.event.pull_request.labels.*.name, 'run-ci-all')) && '--match-all-labels' || '' }}
         ${{ (github.event_name == 'schedule' || inputs.nightly) && '--nightly' || '' }}
-        ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }}
+        ${{ (github.event_name == 'schedule' || github.event_name == 'pull_request' || inputs.continue_on_error) && '--continue-on-error' || '' }}
     secrets: inherit
 
   # Fan-in so each stage-c needs one job instead of every stage-b.
@@ -422,7 +422,7 @@ jobs:
         --labels ${{ join(github.event.pull_request.labels.*.name, ' ') }}
         ${{ (github.event_name == 'workflow_dispatch' || github.event_name == 'schedule' || contains(github.event.pull_request.labels.*.name, 'run-ci-image') || contains(github.event.pull_request.labels.*.name, 'run-ci-all')) && '--match-all-labels' || '' }}
         ${{ (github.event_name == 'schedule' || inputs.nightly) && '--nightly' || '' }}
-        ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }}
+        ${{ (github.event_name == 'schedule' || github.event_name == 'pull_request' || inputs.continue_on_error) && '--continue-on-error' || '' }}
     secrets: inherit
 
   # Stage C GPU: e2e suite on the 5gpu runner.
@@ -445,5 +445,5 @@ jobs:
         --labels ${{ join(github.event.pull_request.labels.*.name, ' ') }}
         ${{ (github.event_name == 'workflow_dispatch' || github.event_name == 'schedule' || contains(github.event.pull_request.labels.*.name, 'run-ci-image') || contains(github.event.pull_request.labels.*.name, 'run-ci-all')) && '--match-all-labels' || '' }}
         ${{ (github.event_name == 'schedule' || inputs.nightly) && '--nightly' || '' }}
-        ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }}
+        ${{ (github.event_name == 'schedule' || github.event_name == 'pull_request' || inputs.continue_on_error) && '--continue-on-error' || '' }}
     secrets: inherit

docs/advanced/streaming-reward.md

@@ -21,9 +21,8 @@ flowchart LR
 
 
 The rollout pipeline decomposes into `generate` / `deserialize` / `reward` components. With `--use-output-store`, the
-output is read back from the store and deleted there between `deserialize` and `reward`, still inside the generation
-slot. Because every microgroup is its own `asyncio` task, the engine starts the next request while previous responses
-are still being unpacked and scored.
+reward actors read the output from the store, and it is deleted there once scored. Because every microgroup is its
+own `asyncio` task, the engine starts the next request while previous responses are still being unpacked and scored.
 Nothing waits for the full batch.
 
 ## 2. Deserialization: msgpack + parser-actor pool

docs/user-guide/customization.md

@@ -57,6 +57,10 @@ def generate_rollout(
 Legacy list/dict returns are still wrapped into the corresponding output dataclass by
 `miles.rollout.base_types.call_rollout_fn`.
 
+With `--use-output-store`, a rollout function that does not go through `generate_rollout_async` sets
+`GenerateState(args).media_outputs_left` to the number of outputs to keep for its media logging at the start of each
+rollout, and calls `release_generated_outputs` after a reward it runs itself.
+
 ### `--eval-function-path`
 
 Same shape as the rollout function. Defaults to `--rollout-function-path` when
@@ -92,10 +96,8 @@ async def custom_generate(
 
 `evaluation=` is optional — if your signature omits it, the caller skips that kwarg.
 
-With `--use-output-store`, `generate_microgroup` leaves `sample.generated_output` as a `StoredTensor`.
-`generate_and_rm_microgroup` reads it into the tensor after your function returns and before the reward. Code that
-reads the output earlier, such as a custom generate function that refines its own first pass, calls
-`await load_generated_outputs(GenerateState(args), microgroup)` first.
+With `--use-output-store`, `generate_microgroup` leaves `sample.generated_output` as a `StoredTensor`: the output
+stays in the store for the reward actors, and a custom generate function cannot read it.
 
 ### `--diffusion-step-strategy-path`
 
@@ -134,6 +136,13 @@ Built-in scorers (`--rm-type pickscore` / `hps` / `ocr` / `custom_api` / `openai
 [Rewards](rewards.md). `--custom-rm-path` replaces that dispatch entirely; `--custom-reward-post-process-path`
 replaces the normalization that runs after scoring.
 
+With `--use-output-store`, the rollout process does not read a training rollout's output. There,
+`sample.generated_output` is a `StoredTensor`, and the actors of a reward pool read it when they score it. A
+`--custom-rm-path` function therefore scores through a pool, a built-in `*_rm` or its own `AsyncRewardActorPool`, and
+does not read the tensor itself. Once a sample is scored its output is `None`; the first `--wandb-log-num-images`
+scored samples of a rollout that logs media keep it as a tensor, so a rollout whose first groups a filter drops logs
+fewer. Every hook that runs after scoring sees that.
+
 ### `--custom-rm-path`
 
 ```python
@@ -183,6 +192,9 @@ async def api_rm(args, samples: list[Sample], **kwargs) -> list[float]:
 --rm-url http://localhost:8000/score
 ```
 
+With `--use-output-store` this function cannot read the output. Send it through `--custom-api-rm-config`, whose
+`ApiRewardActor` reads it in the actor, or through your own `AsyncRewardActorPool(args=args, ...)`.
+
 ### `--custom-reward-post-process-path`
 
 Replace GRPO advantage normalization in

miles/dashboard/events.py

@@ -33,16 +33,13 @@ class TrajectoryEventKind(StrEnum):
     DESER_END = "deser_end"
     REWARD_START = "reward_start"
     REWARD_END = "reward_end"
-    STORE_READ_START = "store_read_start"
-    STORE_READ_END = "store_read_end"
 
 
 # stage name (StageTimer) -> (start kind, end kind)
 STAGE_KINDS = {
     "generate": (TrajectoryEventKind.GEN_START, TrajectoryEventKind.GEN_END),
     "deserialize": (TrajectoryEventKind.DESER_START, TrajectoryEventKind.DESER_END),
     "reward": (TrajectoryEventKind.REWARD_START, TrajectoryEventKind.REWARD_END),
-    "store_read": (TrajectoryEventKind.STORE_READ_START, TrajectoryEventKind.STORE_READ_END),
 }
 
 # kind -> (base span, is_start), for the reader to fold start/end back into spans
@@ -53,8 +50,6 @@ class TrajectoryEventKind(StrEnum):
     "deser_end": ("deser", False),
     "reward_start": ("reward", True),
     "reward_end": ("reward", False),
-    "store_read_start": ("store_read", True),
-    "store_read_end": ("store_read", False),
 }
 
 

miles/ray/rollout.py

@@ -449,12 +449,15 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl
         reward_stats["rollout/step"] = compute_rollout_step(self.args, self.rollout_id)
         tracking_utils.log(self.args, reward_stats, step_key="rollout/step")
 
-        max_images = self.args.wandb_log_num_images
-        interval = self.args.wandb_log_image_interval
-        if max_images > 0 and self.rollout_id % interval == 0:
+        if tracking_utils.logs_rollout_media(self.args, self.rollout_id):
+            # with the output store on, only the samples the rollout kept for media hold a tensor
             self._log_images_async(
-                {"rollout_media/sample_images": samples},
-                max_images=max_images,
+                {
+                    "rollout_media/sample_images": [
+                        sample for sample in samples if isinstance(sample.generated_output, torch.Tensor)
+                    ]
+                },
+                max_images=self.args.wandb_log_num_images,
                 step_key="rollout/step",
                 step_value=compute_rollout_step(self.args, self.rollout_id),
                 reward_key=self.args.reward_key,

miles/rollout/rm_hub/api.py

@@ -66,6 +66,7 @@ def __init__(self, args) -> None:
         if not isinstance(actor_cls, type) or not issubclass(actor_cls, self.actor_base_cls):
             raise TypeError(f"API reward actor_class must be an {self.actor_base_cls.__name__} subclass")
         super().__init__(
+            args=args,
             actor_cls=actor_cls,
             actor_kwargs=config.actor_kwargs,
             num_workers=1,

miles/rollout/rm_hub/core.py

@@ -9,6 +9,8 @@
 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
 
 from miles.ray.utils import COLOCATED_REWARD_GPU
+from miles.utils.diffusion_rollout_response import read_generated_outputs
+from miles.utils.output_store import RegisteredTensorReader, StoredTensor, connect_output_store
 
 logger = logging.getLogger(__name__)
 
@@ -79,10 +81,27 @@ def record_reward_queue_depth(samples, name: str, max_queue_depth: int) -> None:
         sample.reward_max_queue_depth = depths
 
 
+class _StoredOutputRewardActor:
+    """A reward actor that reads the stored outputs it is given from the output store, then scores them."""
+
+    def __init__(self, *, scorer_cls, scorer_kwargs: dict, args) -> None:
+        self._scorer = scorer_cls(**scorer_kwargs)
+        self._reader = RegisteredTensorReader(connect_output_store(args))
+
+    def score_batch(self, outputs: list, prompts: list[str]) -> list[float]:
+        # eval outputs arrive inline
+        stored_outputs = [output for output in outputs if isinstance(output, StoredTensor)]
+        if stored_outputs:
+            read_outputs = iter(read_generated_outputs(self._reader, stored_outputs))
+            outputs = [next(read_outputs) if isinstance(output, StoredTensor) else output for output in outputs]
+        return self._scorer.score_batch(outputs, prompts)
+
+
 class AsyncRewardActorPool:
     """Round-robin pool for Ray reward actors exposing ``score_batch(outputs, prompts)``.
 
-    Actors receive the float ``generated_output`` tensors and quantise to uint8 themselves.
+    Actors receive the float ``generated_output`` tensors and quantise to uint8 themselves. With the output store
+    on they are given stored outputs and read them first.
     Colocated pools take one slot per rollout bundle from the manager's
     ``slots``; standalone pools are default-scheduled at ``num_gpus_per_worker``, which
     only lands on GPUs outside every placement group.
@@ -91,6 +110,7 @@ class AsyncRewardActorPool:
     def __init__(
         self,
         *,
+        args,
         actor_cls,
         actor_kwargs: dict,
         num_workers: int,
@@ -113,6 +133,9 @@ def __init__(
             num_gpus_per_worker = COLOCATED_REWARD_GPU
         else:
             strategies = ["DEFAULT"] * num_workers
+        if args.use_output_store:
+            actor_kwargs = {"scorer_cls": actor_cls, "scorer_kwargs": actor_kwargs, "args": args}
+            actor_cls = _StoredOutputRewardActor
 
         self._actors = [
             ray.remote(actor_cls)

miles/rollout/rm_hub/hps.py

@@ -128,6 +128,7 @@ class AsyncHPSPool(AsyncRewardActorPool, metaclass=SingletonMeta):
 
     def __init__(self, args, placement_group=None, slots=None) -> None:
         super().__init__(
+            args=args,
             actor_cls=HPSRewardActor,
             actor_kwargs={
                 "hps_version": args.hps_version,

miles/rollout/rm_hub/hpsv3.py

@@ -233,6 +233,7 @@ class AsyncHPSv3Pool(AsyncRewardActorPool, metaclass=SingletonMeta):
 
     def __init__(self, args, placement_group=None, slots=None) -> None:
         super().__init__(
+            args=args,
             actor_cls=HPSv3RewardActor,
             actor_kwargs={"checkpoint_path": args.hpsv3_checkpoint_path},
             num_workers=args.hpsv3_num_workers,

miles/rollout/rm_hub/ocr.py

@@ -99,6 +99,7 @@ class AsyncOcrPool(AsyncRewardActorPool, metaclass=SingletonMeta):
 
     def __init__(self, args) -> None:
         super().__init__(
+            args=args,
             actor_cls=OcrRewardActor,
             actor_kwargs={"use_gpu": False},
             num_workers=args.ocr_num_workers,

miles/rollout/rm_hub/pickscore.py

@@ -7,6 +7,7 @@
 from PIL import Image
 
 from miles.utils.misc import SingletonMeta
+from miles.utils.output_store import StoredTensor
 from miles.utils.processing_utils import generated_output_to_rgb_hwc_uint8_frames, sample_frame_indices
 from miles.utils.types import Sample
 
@@ -74,8 +75,10 @@ def __init__(
         processor_path: str,
         model_path: str,
         frames_per_forward: int,
+        num_frames: int | None,
     ) -> None:
         self.frames_per_forward = frames_per_forward
+        self.num_frames = num_frames
         gpu_ids = ray.get_gpu_ids()
         use_cuda = bool(gpu_ids) and torch.cuda.is_available()
         if use_cuda:
@@ -88,10 +91,10 @@ def __init__(
         )
 
     def score_batch(self, outputs: list[torch.Tensor], prompts: list[str]) -> list[float]:
-        # a video sample is scored on every frame it arrives with and gets their mean
+        # a video is scored on --pickscore-num-frames of its frames, averaged; already picked frames pass unchanged
         images, frame_counts = [], []
         for output in outputs:
-            frames = generated_output_to_rgb_hwc_uint8_frames(output, None)
+            frames = generated_output_to_rgb_hwc_uint8_frames(output, self.num_frames)
             images.extend(Image.fromarray(frame) for frame in frames)
             frame_counts.append(len(frames))
         flat_prompts = [p for p, n in zip(prompts, frame_counts, strict=True) for _ in range(n)]
@@ -112,11 +115,13 @@ class AsyncPickScorePool(AsyncRewardActorPool, metaclass=SingletonMeta):
 
     def __init__(self, args, placement_group=None, slots=None) -> None:
         super().__init__(
+            args=args,
             actor_cls=PickScoreRewardActor,
             actor_kwargs={
                 "processor_path": args.pickscore_processor_path,
                 "model_path": args.pickscore_model_path,
                 "frames_per_forward": args.pickscore_batch_size,
+                "num_frames": args.pickscore_num_frames,
             },
             num_workers=args.pickscore_num_workers,
             batch_size=args.pickscore_batch_size,
@@ -130,10 +135,15 @@ def __init__(self, args, placement_group=None, slots=None) -> None:
 
 async def pickscore_rm(args, samples: Sequence[Sample]) -> list[float]:
     pool = AsyncPickScorePool(args)
-    # pick --pickscore-num-frames here so only the scored frames cross the object store
+    # pick --pickscore-num-frames here so only the scored frames cross the object store; a stored output is not
+    # here to pick from, so its actor picks them
     outputs = [
-        s.generated_output[:, sample_frame_indices(s.generated_output.shape[1], args.pickscore_num_frames)]
-        for s in samples
+        (
+            output
+            if isinstance(output, StoredTensor)
+            else output[:, sample_frame_indices(output.shape[1], args.pickscore_num_frames)]
+        )
+        for output in (s.generated_output for s in samples)
     ]
     scores, max_queue_depth = await pool.score(outputs, [s.prompt for s in samples])
     record_reward_queue_depth(samples, "pickscore", max_queue_depth)