Fix off-by-one in GRPO rollout chunk-slicing that produces empty chunks

nvflow/lib/rl/rollout.py

@@ -580,6 +580,22 @@ def _build_chunk_slice_segment(max_num_samples: int) -> str:
     No physical pre-splitting on the launcher — keeps the launcher
     lightweight and filesystem-agnostic.  The ``if [ $NUM_CHUNKS -gt 1 ]``
     runtime guard suppresses slicing in the single-chunk case.
+
+    Chunk boundaries use balanced-remainder division (the first
+    ``EFFECTIVE % NUM_CHUNKS`` chunks get one extra line) rather than
+    ceiling division.  Ceiling division overshoots ``EFFECTIVE`` by up
+    to a full chunk once ``NUM_CHUNKS * ceil(EFFECTIVE / NUM_CHUNKS)``
+    exceeds it, which silently produces an empty trailing chunk (or
+    several) any time ``EFFECTIVE`` doesn't divide evenly by
+    ``NUM_CHUNKS`` -- a legitimately non-empty chunk range then reads
+    as a merge failure downstream, since the merge step correctly
+    treats a missing/empty chunk file as an upstream error.  Balanced
+    division covers exactly ``1..EFFECTIVE`` with no gaps or overlaps
+    for every ``CHUNK_ID`` in ``[0, NUM_CHUNKS)`` whenever
+    ``NUM_CHUNKS <= EFFECTIVE``; a chunk is only empty when
+    ``CHUNK_ID >= EFFECTIVE`` (more chunks requested than samples
+    exist), which is a real misconfiguration the merge check should
+    still catch.
     """
     return (
         'CHUNK_INPUT=""\n'
@@ -592,10 +608,16 @@ def _build_chunk_slice_segment(max_num_samples: int) -> str:
         "    if [ $MAX_SAMPLES -gt 0 ] && [ $MAX_SAMPLES -lt $TOTAL_LINES ]; then\n"
         "        EFFECTIVE=$MAX_SAMPLES\n"
         "    fi\n"
-        "    CHUNK_SIZE=$(( (EFFECTIVE + NUM_CHUNKS - 1) / NUM_CHUNKS ))\n"
-        "    START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 ))\n"
-        "    END_LINE=$(( (CHUNK_ID + 1) * CHUNK_SIZE ))\n"
-        "    [ $END_LINE -gt $EFFECTIVE ] && END_LINE=$EFFECTIVE\n"
+        "    BASE_SIZE=$(( EFFECTIVE / NUM_CHUNKS ))\n"
+        "    REMAINDER=$(( EFFECTIVE % NUM_CHUNKS ))\n"
+        "    if [ $CHUNK_ID -lt $REMAINDER ]; then\n"
+        "        CHUNK_SIZE=$(( BASE_SIZE + 1 ))\n"
+        "        START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 ))\n"
+        "    else\n"
+        "        CHUNK_SIZE=$BASE_SIZE\n"
+        "        START_LINE=$(( REMAINDER * (BASE_SIZE + 1) + (CHUNK_ID - REMAINDER) * BASE_SIZE + 1 ))\n"
+        "    fi\n"
+        "    END_LINE=$(( START_LINE + CHUNK_SIZE - 1 ))\n"
         '    CHUNK_INPUT="$OUTPUT_DIR/chunk_input_chunk$CHUNK_ID.jsonl"\n'
         '    head -n $END_LINE "$INPUT_DATA" | tail -n +$START_LINE > "$CHUNK_INPUT"\n'
         '    echo "  Chunk $CHUNK_ID/$NUM_CHUNKS: lines $START_LINE-$END_LINE ($((END_LINE - START_LINE + 1)) samples)"\n'

tests/fixtures/rollout/client_cmd_dual_server.txt

@@ -103,10 +103,16 @@ if [ $NUM_CHUNKS -gt 1 ]; then
     if [ $MAX_SAMPLES -gt 0 ] && [ $MAX_SAMPLES -lt $TOTAL_LINES ]; then
         EFFECTIVE=$MAX_SAMPLES
     fi
-    CHUNK_SIZE=$(( (EFFECTIVE + NUM_CHUNKS - 1) / NUM_CHUNKS ))
-    START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 ))
-    END_LINE=$(( (CHUNK_ID + 1) * CHUNK_SIZE ))
-    [ $END_LINE -gt $EFFECTIVE ] && END_LINE=$EFFECTIVE
+    BASE_SIZE=$(( EFFECTIVE / NUM_CHUNKS ))
+    REMAINDER=$(( EFFECTIVE % NUM_CHUNKS ))
+    if [ $CHUNK_ID -lt $REMAINDER ]; then
+        CHUNK_SIZE=$(( BASE_SIZE + 1 ))
+        START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 ))
+    else
+        CHUNK_SIZE=$BASE_SIZE
+        START_LINE=$(( REMAINDER * (BASE_SIZE + 1) + (CHUNK_ID - REMAINDER) * BASE_SIZE + 1 ))
+    fi
+    END_LINE=$(( START_LINE + CHUNK_SIZE - 1 ))
     CHUNK_INPUT="$OUTPUT_DIR/chunk_input_chunk$CHUNK_ID.jsonl"
     head -n $END_LINE "$INPUT_DATA" | tail -n +$START_LINE > "$CHUNK_INPUT"
     echo "  Chunk $CHUNK_ID/$NUM_CHUNKS: lines $START_LINE-$END_LINE ($((END_LINE - START_LINE + 1)) samples)"

tests/fixtures/rollout/client_cmd_max_samples.txt

@@ -103,10 +103,16 @@ if [ $NUM_CHUNKS -gt 1 ]; then
     if [ $MAX_SAMPLES -gt 0 ] && [ $MAX_SAMPLES -lt $TOTAL_LINES ]; then
         EFFECTIVE=$MAX_SAMPLES
     fi
-    CHUNK_SIZE=$(( (EFFECTIVE + NUM_CHUNKS - 1) / NUM_CHUNKS ))
-    START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 ))
-    END_LINE=$(( (CHUNK_ID + 1) * CHUNK_SIZE ))
-    [ $END_LINE -gt $EFFECTIVE ] && END_LINE=$EFFECTIVE
+    BASE_SIZE=$(( EFFECTIVE / NUM_CHUNKS ))
+    REMAINDER=$(( EFFECTIVE % NUM_CHUNKS ))
+    if [ $CHUNK_ID -lt $REMAINDER ]; then
+        CHUNK_SIZE=$(( BASE_SIZE + 1 ))
+        START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 ))
+    else
+        CHUNK_SIZE=$BASE_SIZE
+        START_LINE=$(( REMAINDER * (BASE_SIZE + 1) + (CHUNK_ID - REMAINDER) * BASE_SIZE + 1 ))
+    fi
+    END_LINE=$(( START_LINE + CHUNK_SIZE - 1 ))
     CHUNK_INPUT="$OUTPUT_DIR/chunk_input_chunk$CHUNK_ID.jsonl"
     head -n $END_LINE "$INPUT_DATA" | tail -n +$START_LINE > "$CHUNK_INPUT"
     echo "  Chunk $CHUNK_ID/$NUM_CHUNKS: lines $START_LINE-$END_LINE ($((END_LINE - START_LINE + 1)) samples)"

tests/fixtures/rollout/client_cmd_no_chunk.txt

@@ -103,10 +103,16 @@ if [ $NUM_CHUNKS -gt 1 ]; then
     if [ $MAX_SAMPLES -gt 0 ] && [ $MAX_SAMPLES -lt $TOTAL_LINES ]; then
         EFFECTIVE=$MAX_SAMPLES
     fi
-    CHUNK_SIZE=$(( (EFFECTIVE + NUM_CHUNKS - 1) / NUM_CHUNKS ))
-    START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 ))
-    END_LINE=$(( (CHUNK_ID + 1) * CHUNK_SIZE ))
-    [ $END_LINE -gt $EFFECTIVE ] && END_LINE=$EFFECTIVE
+    BASE_SIZE=$(( EFFECTIVE / NUM_CHUNKS ))
+    REMAINDER=$(( EFFECTIVE % NUM_CHUNKS ))
+    if [ $CHUNK_ID -lt $REMAINDER ]; then
+        CHUNK_SIZE=$(( BASE_SIZE + 1 ))
+        START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 ))
+    else
+        CHUNK_SIZE=$BASE_SIZE
+        START_LINE=$(( REMAINDER * (BASE_SIZE + 1) + (CHUNK_ID - REMAINDER) * BASE_SIZE + 1 ))
+    fi
+    END_LINE=$(( START_LINE + CHUNK_SIZE - 1 ))
     CHUNK_INPUT="$OUTPUT_DIR/chunk_input_chunk$CHUNK_ID.jsonl"
     head -n $END_LINE "$INPUT_DATA" | tail -n +$START_LINE > "$CHUNK_INPUT"
     echo "  Chunk $CHUNK_ID/$NUM_CHUNKS: lines $START_LINE-$END_LINE ($((END_LINE - START_LINE + 1)) samples)"

tests/fixtures/rollout/client_cmd_no_rcp.txt

@@ -103,10 +103,16 @@ if [ $NUM_CHUNKS -gt 1 ]; then
     if [ $MAX_SAMPLES -gt 0 ] && [ $MAX_SAMPLES -lt $TOTAL_LINES ]; then
         EFFECTIVE=$MAX_SAMPLES
     fi
-    CHUNK_SIZE=$(( (EFFECTIVE + NUM_CHUNKS - 1) / NUM_CHUNKS ))
-    START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 ))
-    END_LINE=$(( (CHUNK_ID + 1) * CHUNK_SIZE ))
-    [ $END_LINE -gt $EFFECTIVE ] && END_LINE=$EFFECTIVE
+    BASE_SIZE=$(( EFFECTIVE / NUM_CHUNKS ))
+    REMAINDER=$(( EFFECTIVE % NUM_CHUNKS ))
+    if [ $CHUNK_ID -lt $REMAINDER ]; then
+        CHUNK_SIZE=$(( BASE_SIZE + 1 ))
+        START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 ))
+    else
+        CHUNK_SIZE=$BASE_SIZE
+        START_LINE=$(( REMAINDER * (BASE_SIZE + 1) + (CHUNK_ID - REMAINDER) * BASE_SIZE + 1 ))
+    fi
+    END_LINE=$(( START_LINE + CHUNK_SIZE - 1 ))
     CHUNK_INPUT="$OUTPUT_DIR/chunk_input_chunk$CHUNK_ID.jsonl"
     head -n $END_LINE "$INPUT_DATA" | tail -n +$START_LINE > "$CHUNK_INPUT"
     echo "  Chunk $CHUNK_ID/$NUM_CHUNKS: lines $START_LINE-$END_LINE ($((END_LINE - START_LINE + 1)) samples)"

tests/fixtures/rollout/client_cmd_policy_only.txt

@@ -102,10 +102,16 @@ if [ $NUM_CHUNKS -gt 1 ]; then
     if [ $MAX_SAMPLES -gt 0 ] && [ $MAX_SAMPLES -lt $TOTAL_LINES ]; then
         EFFECTIVE=$MAX_SAMPLES
     fi
-    CHUNK_SIZE=$(( (EFFECTIVE + NUM_CHUNKS - 1) / NUM_CHUNKS ))
-    START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 ))
-    END_LINE=$(( (CHUNK_ID + 1) * CHUNK_SIZE ))
-    [ $END_LINE -gt $EFFECTIVE ] && END_LINE=$EFFECTIVE
+    BASE_SIZE=$(( EFFECTIVE / NUM_CHUNKS ))
+    REMAINDER=$(( EFFECTIVE % NUM_CHUNKS ))
+    if [ $CHUNK_ID -lt $REMAINDER ]; then
+        CHUNK_SIZE=$(( BASE_SIZE + 1 ))
+        START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 ))
+    else
+        CHUNK_SIZE=$BASE_SIZE
+        START_LINE=$(( REMAINDER * (BASE_SIZE + 1) + (CHUNK_ID - REMAINDER) * BASE_SIZE + 1 ))
+    fi
+    END_LINE=$(( START_LINE + CHUNK_SIZE - 1 ))
     CHUNK_INPUT="$OUTPUT_DIR/chunk_input_chunk$CHUNK_ID.jsonl"
     head -n $END_LINE "$INPUT_DATA" | tail -n +$START_LINE > "$CHUNK_INPUT"
     echo "  Chunk $CHUNK_ID/$NUM_CHUNKS: lines $START_LINE-$END_LINE ($((END_LINE - START_LINE + 1)) samples)"

tests/test_rollout.py

@@ -38,6 +38,7 @@
 
 from __future__ import annotations
 
+import subprocess
 from pathlib import Path
 
 import pytest
@@ -46,6 +47,7 @@
 
 from nvflow.lib.rl.helpers import LauncherFS
 from nvflow.lib.rl.rollout import (
+    _build_chunk_slice_segment,
     _build_client_cmd,
     _build_merge_cmd,
     _build_port_read_preamble,
@@ -832,3 +834,114 @@ def test_get_remaining_jobs_mixed_state(tmp_path: Path) -> None:
 
     remaining = _get_remaining_jobs(LauncherFS({}), str(tmp_path), seeds, chunks, rerun_done=False)
     assert remaining == [(1, 1)]
+
+
+# ===========================================================================
+# _build_chunk_slice_segment — runtime chunk-boundary arithmetic
+# ===========================================================================
+#
+# These tests actually execute the rendered bash against real files rather
+# than snapshotting the text, because the bug class here is arithmetic, not
+# formatting.  The previous ceiling-division formula gave
+# EFFECTIVE=4/NUM_CHUNKS=3 a last chunk with START_LINE(5) > END_LINE(4) --
+# a silently empty chunk file -- which then tripped the merge step's own
+# "chunk .done exists but file missing/empty" safety check (see the
+# _build_merge_cmd tests above) as if rollout collection had failed
+# upstream, when nothing had actually gone wrong.  A text snapshot would
+# not catch a wrong split point; only running the arithmetic does.
+
+
+def _run_chunk_slice(
+    tmp_path: Path,
+    total_lines: int,
+    num_chunks: int,
+    chunk_id: int,
+    max_num_samples: int = 0,
+) -> list[str]:
+    """Execute the rendered chunk-slice segment against a real input file.
+
+    Returns the lines (no trailing newline) written to this chunk's
+    output file, or ``[]`` if no chunk file was produced (either the
+    ``NUM_CHUNKS <= 1`` no-op path, or -- pre-fix -- an empty chunk).
+    """
+    tmp_path.mkdir(parents=True, exist_ok=True)
+    input_data = tmp_path / "input.jsonl"
+    input_data.write_text("".join(f"line{i}\n" for i in range(1, total_lines + 1)))
+    output_dir = tmp_path / "out"
+    output_dir.mkdir()
+
+    script = (
+        "set -eu\n"
+        f'INPUT_DATA="{input_data}"\n'
+        f'OUTPUT_DIR="{output_dir}"\n'
+        f"NUM_CHUNKS={num_chunks}\n"
+        f"CHUNK_ID={chunk_id}\n"
+    ) + _build_chunk_slice_segment(max_num_samples)
+
+    result = subprocess.run(["bash", "-c", script], capture_output=True, text=True)
+    assert result.returncode == 0, f"chunk-slice script failed: {result.stderr}"
+
+    chunk_file = output_dir / f"chunk_input_chunk{chunk_id}.jsonl"
+    if not chunk_file.exists():
+        return []
+    return chunk_file.read_text().splitlines()
+
+
+def test_chunk_slice_no_gaps_or_overlaps_across_shapes(tmp_path: Path) -> None:
+    """For every (total_lines, num_chunks) with 2 <= num_chunks <= total_lines,
+    concatenating all chunks in order must reproduce the input exactly --
+    every line appears, in order, in exactly one chunk.
+    """
+    for total_lines in range(2, 13):
+        for num_chunks in range(2, total_lines + 1):
+            all_lines: list[str] = []
+            for chunk_id in range(num_chunks):
+                case_dir = tmp_path / f"n{total_lines}_k{num_chunks}_c{chunk_id}"
+                all_lines.extend(_run_chunk_slice(case_dir, total_lines, num_chunks, chunk_id))
+            assert all_lines == [f"line{i}" for i in range(1, total_lines + 1)], (
+                f"total_lines={total_lines}, num_chunks={num_chunks}"
+            )
+
+
+def test_chunk_slice_no_chunk_is_empty_when_num_chunks_le_total_lines(tmp_path: Path) -> None:
+    """Regression test for the ceiling-division bug: whenever there are at
+    least as many samples as chunks, every chunk must get at least one
+    line.  Previously e.g. total_lines=4/num_chunks=3 left the last
+    chunk empty, and that empty chunk was then reported by the merge
+    step as an upstream rollout failure.
+    """
+    for total_lines in range(2, 13):
+        for num_chunks in range(2, total_lines + 1):
+            for chunk_id in range(num_chunks):
+                case_dir = tmp_path / f"n{total_lines}_k{num_chunks}_c{chunk_id}"
+                lines = _run_chunk_slice(case_dir, total_lines, num_chunks, chunk_id)
+                assert lines, (
+                    f"chunk {chunk_id}/{num_chunks} was empty for total_lines={total_lines} "
+                    "-- this is the previously-fixed ceiling-division bug"
+                )
+
+
+def test_chunk_slice_previously_buggy_case_four_lines_three_chunks(tmp_path: Path) -> None:
+    """Exact repro of the bug found in review: EFFECTIVE=4, NUM_CHUNKS=3
+    used to give chunk 2 START_LINE(5) > END_LINE(4), producing a 0-byte
+    chunk file.  Balanced-remainder division instead gives chunk sizes
+    [2, 1, 1] (the first EFFECTIVE % NUM_CHUNKS chunks get one extra line).
+    """
+    assert _run_chunk_slice(tmp_path / "c0", 4, 3, 0) == ["line1", "line2"]
+    assert _run_chunk_slice(tmp_path / "c1", 4, 3, 1) == ["line3"]
+    assert _run_chunk_slice(tmp_path / "c2", 4, 3, 2) == ["line4"]
+
+
+def test_chunk_slice_single_chunk_is_a_noop(tmp_path: Path) -> None:
+    """``NUM_CHUNKS <= 1`` must skip slicing entirely (no chunk file is
+    created) -- this behavior is unchanged by the arithmetic fix.
+    """
+    assert _run_chunk_slice(tmp_path, 5, 1, 0) == []
+
+
+def test_chunk_slice_respects_max_num_samples_clamp(tmp_path: Path) -> None:
+    """``max_num_samples`` still clamps EFFECTIVE before chunking -- lines
+    beyond the clamp must not appear in any chunk.
+    """
+    assert _run_chunk_slice(tmp_path / "c0", 10, 2, 0, max_num_samples=4) == ["line1", "line2"]
+    assert _run_chunk_slice(tmp_path / "c1", 10, 2, 1, max_num_samples=4) == ["line3", "line4"]