Fix paibench_c shard task counting and enforce evaluation coverage

evaluation/cosmos3/generator/paibench_c/README.md

@@ -52,6 +52,8 @@ SKIP_EVAL=1 bash run_paibench_c.sh
 SKIP_GEN=1 bash run_paibench_c.sh
 ```
 
+Generation is resume-safe: tasks with a valid video are skipped on re-run, and evaluation refuses to score a tree whose video count does not match the requested task set, so a partial run can never be reported as a complete benchmark result.
+
 Alternatively, open `run_with_cosmos_framework.ipynb` for an interactive step-by-step walkthrough
 with inline video previews and metric display.
 

evaluation/cosmos3/generator/paibench_c/run_paibench_c.sh

@@ -540,7 +540,9 @@ for tid in task_ids:
 
 dst = output_dir / "input.jsonl"
 dst.parent.mkdir(parents=True, exist_ok=True)
-dst.write_text("\n".join(json.dumps(r) for r in rows))
+# Trailing newline: wc -l and line-oriented tools count the final row only
+# when the file ends with one.
+dst.write_text("".join(json.dumps(r) + "\n" for r in rows))
 print(f"Wrote {len(rows)} rows → {dst}")
 PY
 
@@ -609,7 +611,8 @@ for row in shard_rows:
         vision_mp4.unlink()
     pending.append(row)
 
-shard_jsonl.write_text("\n".join(json.dumps(r) for r in pending))
+# Trailing newline so line counting (wc -l) sees every pending row.
+shard_jsonl.write_text("".join(json.dumps(r) + "\n" for r in pending))
 skipped = len(shard_rows) - len(pending)
 print(f"Shard {shard_idx}/{num_shards}: {len(pending)} to generate, {skipped} already done")
 PY
@@ -620,7 +623,9 @@ PY
     for (( gi=0; gi<_num_gpus; gi++ )); do
         local _shard_jsonl="${_shard_jsonls[$gi]}"
         local _pending_count
-        _pending_count="$(wc -l < "$_shard_jsonl" | tr -d ' ')"
+        # NR counts the last row even without a trailing newline (and unlike
+        # wc -l, awk exits 0 on an empty file under set -o pipefail).
+        _pending_count="$(awk 'END{print NR}' "$_shard_jsonl")"
         local _gpu_id="${_gpu_ids[$gi]}"
 
         if [[ "$_pending_count" -eq 0 ]]; then
@@ -645,9 +650,11 @@ PY
         _worker_pids+=("$!")
     done
 
-    # Wait for all GPU workers to finish.
+    # Wait for all GPU workers to finish.  The ${var+...} expansion keeps the
+    # loop a no-op (instead of an "unbound variable" abort under set -u on
+    # bash < 4.4) when every shard was already complete.
     local _fail=0
-    for _pid in "${_worker_pids[@]}"; do
+    for _pid in ${_worker_pids[@]+"${_worker_pids[@]}"}; do
         wait "$_pid" || _fail=1
     done
     [[ "$_fail" -eq 0 ]] || die "One or more GPU workers failed for modality=$MODALITY"
@@ -715,10 +722,12 @@ _py_for_summary() {
     fi
 }
 
-for _mod in $PAIBENCH_C_MODALITIES; do
-    _videos_parent="$PAIBENCH_C_OUTPUT_ROOT/${_ckpt_slug}/$_mod"
-    _videos_dir="$_videos_parent/videos"
-    _metrics_out="$_videos_parent/metrics.json"
+evaluate_modality() {
+    local _mod="$1"
+    local _videos_parent="$PAIBENCH_C_OUTPUT_ROOT/${_ckpt_slug}/$_mod"
+    local _videos_dir="$_videos_parent/videos"
+    local _metrics_out="$_videos_parent/metrics.json"
+    local _eval_ngpu
     _eval_ngpu="$(int_min "$PAIBENCH_C_NUM_SAMPLES" "$COSMOS3_NUM_GPUS")"
 
     # Verify videos exist before launching torchrun to surface errors early.
@@ -727,9 +736,18 @@ for _mod in $PAIBENCH_C_MODALITIES; do
   Run without SKIP_GEN=1, or point PAIBENCH_C_OUTPUT_ROOT at an existing output tree."
     fi
     _nvideos="$(find "$_videos_dir" -name '*.mp4' | wc -l)"
-    if [[ "$_nvideos" -eq 0 ]]; then
-        die "No .mp4 files found in $_videos_dir
-  Re-run generation (remove SKIP_GEN=1) to produce videos."
+
+    # Evaluation must cover exactly the requested task set — scoring a subset
+    # would silently report the subset metric as the benchmark result.
+    local _expected_tasks
+    if [[ -f "$_videos_parent/input.jsonl" ]]; then
+        _expected_tasks="$(awk 'END{print NR}' "$_videos_parent/input.jsonl")"
+    else
+        _expected_tasks="$PAIBENCH_C_NUM_SAMPLES"
+    fi
+    if [[ "$_nvideos" -ne "$_expected_tasks" ]]; then
+        die "Coverage mismatch for modality=$_mod: found $_nvideos video(s), expected $_expected_tasks.
+  Evaluation refuses to score a partial benchmark. Re-run generation (remove SKIP_GEN=1) to produce the missing videos, or point PAIBENCH_C_OUTPUT_ROOT at the output tree that matches this input set."
     fi
     log "  Found $_nvideos video(s) in $_videos_dir"
 
@@ -765,6 +783,10 @@ m = json.loads(pathlib.Path(os.environ["_mout"]).read_text())
 for k, v in m["global"].items():
     print(f"    {k}: {v:.4f}")
 PY
+}
+
+for _mod in $PAIBENCH_C_MODALITIES; do
+    evaluate_modality "$_mod"
 done
 
 log "=== PAI-Bench-C complete ==="

evaluation/cosmos3/generator/paibench_c/test_run_paibench_c.sh

@@ -0,0 +1,268 @@
+#!/usr/bin/env bash
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+#
+# Regression tests for run_paibench_c.sh generation-shard bookkeeping.
+#
+# The tests extract the real functions from run_paibench_c.sh (no logic is
+# duplicated here) and drive them with a stub python that fakes the
+# cosmos-framework inference entrypoint. No GPU, network, or venv is needed.
+#
+# Run:  bash test_run_paibench_c.sh
+set -euo pipefail
+
+HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+SCRIPT_UNDER_TEST="$HERE/run_paibench_c.sh"
+[[ -f "$SCRIPT_UNDER_TEST" ]] || { echo "FAIL: $SCRIPT_UNDER_TEST not found" >&2; exit 1; }
+
+PASS=0
+FAIL=0
+
+assert_eq() { # assert_eq <actual> <expected> <message>
+    if [[ "$1" == "$2" ]]; then
+        PASS=$((PASS + 1))
+    else
+        FAIL=$((FAIL + 1))
+        echo "FAIL: $3"
+        echo "  expected: $2"
+        echo "  actual:   $1"
+    fi
+}
+
+assert_contains() { # assert_contains <haystack> <needle> <message>
+    if [[ "$1" == *"$2"* ]]; then
+        PASS=$((PASS + 1))
+    else
+        FAIL=$((FAIL + 1))
+        echo "FAIL: $3"
+        echo "  haystack does not contain: $2"
+    fi
+}
+
+# Extract one function definition verbatim from the script under test and
+# evaluate it in this shell so the tests run the shipped code, not a copy.
+# Handles both one-line definitions ( ... ; }) and multi-line ones closed by
+# a lone "}" at column zero.
+load_fn() {
+    local _fn="$1"
+    local _src
+    _src="$(awk -v fn="^${_fn}[(]" '
+        $0 ~ fn {
+            print
+            if ($0 ~ /}[[:space:]]*$/) { exit }
+            in_body = 1
+            next
+        }
+        in_body {
+            print
+            if ($0 == "}") { exit }
+        }
+    ' "$SCRIPT_UNDER_TEST")"
+    [[ -n "$_src" && "$_src" == *}* ]] \
+        || { echo "FAIL: function ${_fn}() not found in $SCRIPT_UNDER_TEST" >&2; exit 1; }
+    eval "$_src"
+}
+
+log() { echo "[test-log] $*"; }
+die() { echo "[test-die] $*" >&2; exit 9; }
+
+# Stub venv: a bin dir whose python passes heredoc scripts through to the real
+# python3 and fakes the framework inference entrypoint by materializing one
+# valid (>50 KB) vision.mp4 per JSONL row.
+make_stub_env() {
+    local _env="$1"
+    mkdir -p "$_env/bin"
+    cat > "$_env/bin/python" <<'STUB'
+#!/usr/bin/env bash
+if [[ "${1:-}" == "-" ]]; then
+    exec python3 -
+fi
+if [[ "${1:-}" == "-m" ]]; then
+    stub_shard=""; stub_raw=""
+    while [[ $# -gt 0 ]]; do
+        case "$1" in
+            -i) stub_shard="$2"; shift 2 ;;
+            -o) stub_raw="$2"; shift 2 ;;
+            *) shift ;;
+        esac
+    done
+    export stub_shard stub_raw
+    python3 - <<'FAKE'
+import json, os, pathlib
+shard, raw = pathlib.Path(os.environ["stub_shard"]), pathlib.Path(os.environ["stub_raw"])
+for line in shard.read_text().splitlines():
+    if not line.strip():
+        continue
+    name = json.loads(line)["name"]
+    d = raw / name
+    d.mkdir(parents=True, exist_ok=True)
+    (d / "vision.mp4").write_bytes(b"\x00" * 52000)
+FAKE
+    exit 0
+fi
+exec python3 "$@"
+STUB
+    chmod +x "$_env/bin/python"
+}
+
+# Shared fixture: 4 tasks on 4 GPUs through the real build_jsonl_and_run().
+run_generation() { # run_generation <output_root> <logfile>
+    local _out="$1" _logfile="$2"
+    PAIBENCH_C_ROOT="$HERE" \
+    PAIBENCH_C_DATASET_ROOT="$PAIBENCH_C_OUTPUT_ROOT_FIXTURE_DATASET" \
+    PAIBENCH_C_OUTPUT_ROOT="$_out" \
+    PAIBENCH_C_CHECKPOINT="Cosmos3-Nano" \
+    PAIBENCH_C_NUM_SAMPLES="4" \
+    COSMOS3_UV_ENV="$PAIBENCH_C_OUTPUT_ROOT_FIXTURE_STUB" \
+    COSMOS3_REPO="$PAIBENCH_C_OUTPUT_ROOT_FIXTURE_REPO" \
+    TRITON_CACHE_DIR="$_out/triton" \
+    CUDA_VISIBLE_DEVICES="0,1,2,3" \
+        build_jsonl_and_run edge >"$_logfile" 2>&1
+}
+
+# ---------------------------------------------------------------------------
+# Scenario 1: fresh 4-task / 4-GPU run — every shard must report 1 pending
+# task (the wc -l bug counted 0 for a 1-row file without a trailing newline)
+# and all four videos must be generated.
+# ---------------------------------------------------------------------------
+test_fresh_run_counts_single_row_shards() {
+    local _tmp; _tmp="$(mktemp -d)"
+    make_stub_env "$_tmp/stub-venv"
+    mkdir -p "$_tmp/dataset" "$_tmp/repo"
+    PAIBENCH_C_OUTPUT_ROOT_FIXTURE_DATASET="$_tmp/dataset"
+    PAIBENCH_C_OUTPUT_ROOT_FIXTURE_STUB="$_tmp/stub-venv"
+    PAIBENCH_C_OUTPUT_ROOT_FIXTURE_REPO="$_tmp/repo"
+
+    load_fn build_jsonl_and_run
+
+    local _rc=0
+    run_generation "$_tmp/out" "$_tmp/run.log" || _rc=$?
+
+    assert_eq "$_rc" "0" "fresh run exits 0"
+    assert_eq "$(wc -l < "$_tmp/out/Cosmos3-Nano/edge/input.jsonl" | tr -d ' ')" "4" \
+        "input.jsonl has a trailing newline so wc -l counts all 4 rows"
+    for gpu in 0 1 2 3; do
+        assert_contains "$(cat "$_tmp/run.log")" "GPU $gpu: 1 tasks" \
+            "GPU $gpu shard with exactly 1 pending row is counted as 1 task"
+    done
+    assert_eq "$(find "$_tmp/out/Cosmos3-Nano/edge/videos" -name '*.mp4' | wc -l)" "4" \
+        "all 4 shard tasks were generated (none skipped as already done)"
+    assert_eq "$(find "$_tmp/out/Cosmos3-Nano/edge" -name 'shard_*.jsonl' | wc -l)" "0" \
+        "worker shard files are cleaned up"
+
+    rm -rf "$_tmp"
+}
+
+# ---------------------------------------------------------------------------
+# Scenario 2: fully-complete resume — every shard reports 0 pending, the
+# worker-pids array stays empty, and the wait loop must not trip set -u
+# (unbound variable) semantics on bash < 4.4.
+# ---------------------------------------------------------------------------
+test_resume_with_no_pending_tasks() {
+    local _tmp; _tmp="$(mktemp -d)"
+    make_stub_env "$_tmp/stub-venv"
+    mkdir -p "$_tmp/dataset" "$_tmp/repo"
+    PAIBENCH_C_OUTPUT_ROOT_FIXTURE_DATASET="$_tmp/dataset"
+    PAIBENCH_C_OUTPUT_ROOT_FIXTURE_STUB="$_tmp/stub-venv"
+    PAIBENCH_C_OUTPUT_ROOT_FIXTURE_REPO="$_tmp/repo"
+
+    load_fn build_jsonl_and_run
+
+    run_generation "$_tmp/out" "$_tmp/run1.log"
+    local _rc=0
+    run_generation "$_tmp/out" "$_tmp/run2.log" || _rc=$?
+
+    assert_eq "$_rc" "0" "fully-complete resume exits 0 under set -euo pipefail"
+    local _skipped
+    _skipped="$(grep -c 'all tasks done, skipping' "$_tmp/run2.log" || true)"
+    assert_eq "$_skipped" "4" "all four GPUs report nothing pending on resume"
+    assert_eq "$(find "$_tmp/out/Cosmos3-Nano/edge/videos" -name '*.mp4' | wc -l)" "4" \
+        "resume regenerates nothing and keeps the four videos"
+
+    rm -rf "$_tmp"
+}
+
+# ---------------------------------------------------------------------------
+# Scenario 3: evaluation coverage gate — metrics are only computed when the
+# collected video count equals the requested task count in input.jsonl; a
+# partial tree must be rejected instead of silently scored as the benchmark
+# result.
+# ---------------------------------------------------------------------------
+test_evaluation_coverage_gate() {
+    local _tmp; _tmp="$(mktemp -d)"
+    make_stub_env "$_tmp/stub-venv"
+    # Stub eval venv: torchrun writes a metrics.json and succeeds; the real
+    # call never happens when the gate rejects the tree.
+    mkdir -p "$_tmp/eval-venv/bin" "$_tmp/eval-root"
+    cat > "$_tmp/eval-venv/bin/torchrun" <<'STUB'
+#!/usr/bin/env bash
+out=""
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --output_path) out="$2"; shift 2 ;;
+        *) shift ;;
+    esac
+done
+printf '{"global": {"dummy_metric": 1.0}}\n' > "$out"
+STUB
+    chmod +x "$_tmp/eval-venv/bin/torchrun"
+
+    local _tree="$_tmp/out/Cosmos3-Nano/edge"
+    mkdir -p "$_tree/videos"
+    for i in 0 1 2 3; do
+        printf 'x' > "$_tree/videos/task_000$i.mp4"
+        printf '{"name": "task_000%i"}\n' "$i" >> "$_tree/input.jsonl"
+    done
+
+    load_fn int_min
+    load_fn _py_for_summary
+    load_fn evaluate_modality
+
+    local _rc
+    _run_eval() {
+        PAIBENCH_C_OUTPUT_ROOT="$_tmp/out" \
+        PAIBENCH_C_CHECKPOINT="Cosmos3-Nano" \
+        _ckpt_slug="Cosmos3-Nano" \
+        PAIBENCH_C_NUM_SAMPLES="4" \
+        COSMOS3_NUM_GPUS="2" \
+        COSMOS3_UV_ENV="$_tmp/stub-venv" \
+        PAIBENCH_UV_ENV="$_tmp/eval-venv" \
+        PAIBENCH_EVAL_ROOT="$_tmp/eval-root" \
+        PAIBENCH_C_DATASET_ROOT="$_tmp/dataset" \
+        CUDA_VISIBLE_DEVICES="0,1" \
+        evaluate_modality edge
+    }
+
+    # Case A: full coverage — evaluation runs and reports the metrics.
+    local _rc=0
+    _run_eval >"$_tmp/eval.log" 2>&1 || _rc=$?
+    assert_eq "$_rc" "0" "evaluation with full coverage exits 0"
+    assert_contains "$(cat "$_tmp/eval.log")" "dummy_metric: 1.0000" \
+        "full-coverage evaluation reports the computed metrics"
+
+    # Case B: one video missing — evaluation must refuse to score.
+    rm -f "$_tree/videos/task_0002.mp4" "$_tree/metrics.json"
+    _rc=0
+    ( _run_eval ) >"$_tmp/eval2.log" 2>&1 || _rc=$?
+    assert_eq "$_rc" "9" "partial-coverage evaluation is rejected"
+    assert_contains "$(cat "$_tmp/eval2.log")" "found 3 video(s), expected 4" \
+        "rejection message states the missing coverage"
+    [[ ! -f "$_tmp/out/Cosmos3-Nano/edge/metrics.json" ]] && PASS=$((PASS + 1)) || {
+        FAIL=$((FAIL + 1))
+        echo "FAIL: partial-coverage evaluation must not produce metrics.json"
+    }
+
+    rm -rf "$_tmp"
+}
+
+# ---------------------------------------------------------------------------
+run_tests() {
+    test_fresh_run_counts_single_row_shards
+    test_resume_with_no_pending_tasks
+    test_evaluation_coverage_gate
+    echo
+    echo "passed: $PASS  failed: $FAIL"
+    [[ "$FAIL" -eq 0 ]]
+}
+
+run_tests