diff --git a/nvflow/lib/rl/rollout.py b/nvflow/lib/rl/rollout.py index 37b50aa..ae84503 100644 --- a/nvflow/lib/rl/rollout.py +++ b/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' diff --git a/tests/fixtures/rollout/client_cmd_dual_server.txt b/tests/fixtures/rollout/client_cmd_dual_server.txt index 9c54ac9..4ff49c3 100644 --- a/tests/fixtures/rollout/client_cmd_dual_server.txt +++ b/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)" diff --git a/tests/fixtures/rollout/client_cmd_max_samples.txt b/tests/fixtures/rollout/client_cmd_max_samples.txt index 83c5902..6a11de8 100644 --- a/tests/fixtures/rollout/client_cmd_max_samples.txt +++ b/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)" diff --git a/tests/fixtures/rollout/client_cmd_no_chunk.txt b/tests/fixtures/rollout/client_cmd_no_chunk.txt index 4a163bc..1478f56 100644 --- a/tests/fixtures/rollout/client_cmd_no_chunk.txt +++ b/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)" diff --git a/tests/fixtures/rollout/client_cmd_no_rcp.txt b/tests/fixtures/rollout/client_cmd_no_rcp.txt index e6e7937..ed7ec6c 100644 --- a/tests/fixtures/rollout/client_cmd_no_rcp.txt +++ b/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)" diff --git a/tests/fixtures/rollout/client_cmd_policy_only.txt b/tests/fixtures/rollout/client_cmd_policy_only.txt index 69062a3..d44c815 100644 --- a/tests/fixtures/rollout/client_cmd_policy_only.txt +++ b/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)" diff --git a/tests/test_rollout.py b/tests/test_rollout.py index be6b327..f267fd1 100644 --- a/tests/test_rollout.py +++ b/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"]