diff --git a/tests/functional_tests/python_test_utils/get_test_results_from_tensorboard_logs.py b/tests/functional_tests/python_test_utils/get_test_results_from_tensorboard_logs.py index 7b74a6879ad..091623b9b84 100644 --- a/tests/functional_tests/python_test_utils/get_test_results_from_tensorboard_logs.py +++ b/tests/functional_tests/python_test_utils/get_test_results_from_tensorboard_logs.py @@ -1,3 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import os os.environ["OPENBLAS_NUM_THREADS"] = "1" @@ -22,27 +23,13 @@ help="Use first or all tensorboard logs", default=False, ) -@click.option( - "--is-second-run/--is-not-second-run", - type=bool, - help="Use second run of tensorboard logs", - default=False, -) @click.option("--step-size", required=False, default=5, type=int, help="Step size of sampling") def collect_train_test_metrics( - logs_dir: str, - train_iters: str, - output_path: str, - is_convergence_test: bool, - is_second_run: bool, - step_size: int, + logs_dir: str, train_iters: str, output_path: str, is_convergence_test: bool, step_size: int ): - if is_convergence_test and is_second_run: - raise ValueError("Convergence test cannot be run on second run of tensorboard logs") - summaries = common.read_tb_logs_as_list( logs_dir, - index=(-1 if is_convergence_test else (1 if is_second_run else 0)), + index=(-1 if is_convergence_test else 0), train_iters=train_iters, start_idx=1, step_size=step_size, diff --git a/tests/functional_tests/shell_test_utils/run_ci_test.sh b/tests/functional_tests/shell_test_utils/run_ci_test.sh index 0b360530b7d..efd4f08f2b3 100644 --- a/tests/functional_tests/shell_test_utils/run_ci_test.sh +++ b/tests/functional_tests/shell_test_utils/run_ci_test.sh @@ -226,6 +226,11 @@ for i in $(seq 1 $N_REPEAT); do FILE=$(basename "$_TENSORBOARD_PATH") export TENSORBOARD_PATH=$DIR/$i/$FILE mkdir -p $(dirname $TENSORBOARD_PATH) + # Per-repeat base; for ckpt-resume / frozen-resume tests each training + # phase writes to its own ${_REPEAT_TENSORBOARD_PATH}/run_N subdir so the + # analyzer can read by explicit path rather than guessing phase from + # mtime-sorted glob order. + _REPEAT_TENSORBOARD_PATH=$TENSORBOARD_PATH export REPEAT=$i wait_for_all_nodes "repeat_${REPEAT}_start" @@ -294,6 +299,10 @@ for i in $(seq 1 $N_REPEAT); do done else # The standard single-run test that otherwise runs + if [[ "$TEST_TYPE" == "ckpt-resume" || "$TEST_TYPE" == "frozen-resume" ]]; then + export TENSORBOARD_PATH="$_REPEAT_TENSORBOARD_PATH/run_${RUN_NUMBER}" + mkdir -p "$TENSORBOARD_PATH" + fi run_training_phase "repeat_${REPEAT}_run_${RUN_NUMBER}" fi @@ -316,6 +325,8 @@ for i in $(seq 1 $N_REPEAT); do echo $((TRAIN_ITERS / 2)) >$CHECKPOINT_LOAD_PATH/latest_checkpointed_iteration.txt export RUN_NUMBER=2 + export TENSORBOARD_PATH="$_REPEAT_TENSORBOARD_PATH/run_${RUN_NUMBER}" + mkdir -p "$TENSORBOARD_PATH" run_training_phase "repeat_${REPEAT}_run_${RUN_NUMBER}" fi @@ -326,6 +337,8 @@ for i in $(seq 1 $N_REPEAT); do export CHECKPOINT_SAVE_PATH=/tmp/checkpoints/ export RUN_NUMBER=2 + export TENSORBOARD_PATH="$_REPEAT_TENSORBOARD_PATH/run_${RUN_NUMBER}" + mkdir -p "$TENSORBOARD_PATH" run_training_phase "repeat_${REPEAT}_run_${RUN_NUMBER}" export CHECKPOINT_SAVE_PATH=$_CHECKPOINT_SAVE_PATH @@ -371,8 +384,13 @@ for i in $(seq 1 $N_REPEAT); do # Read test values from Tensorboard for non-inference tests. # Inference tests will load from JSON instead. if [[ "$MODE" == "pretraining" ]]; then + if [[ "$TEST_TYPE" == "ckpt-resume" || "$TEST_TYPE" == "frozen-resume" ]]; then + FIRST_RUN_TENSORBOARD_PATH="$_REPEAT_TENSORBOARD_PATH/run_1" + else + FIRST_RUN_TENSORBOARD_PATH="$TENSORBOARD_PATH" + fi uv run --no-sync python $ROOT_DIR/tests/functional_tests/python_test_utils/get_test_results_from_tensorboard_logs.py \ - --logs-dir $TENSORBOARD_PATH \ + --logs-dir $FIRST_RUN_TENSORBOARD_PATH \ --train-iters $TRAIN_ITERS \ --output-path ${OUTPUT_PATH}/$(basename $GOLDEN_VALUES_PATH) \ "${EXTRACT_ARGS[@]}" @@ -424,10 +442,9 @@ for i in $(seq 1 $N_REPEAT); do if [[ "$TEST_TYPE" == "ckpt-resume" || "$TEST_TYPE" == "frozen-resume" ]]; then uv run --no-sync python $ROOT_DIR/tests/functional_tests/python_test_utils/get_test_results_from_tensorboard_logs.py \ - --logs-dir $TENSORBOARD_PATH \ + --logs-dir "$_REPEAT_TENSORBOARD_PATH/run_2" \ --train-iters $TRAIN_ITERS \ --output-path "${OUTPUT_PATH}/$(basename $GOLDEN_VALUES_PATH .json)_2nd.json" \ - --is-second-run \ "${EXTRACT_ARGS[@]}" echo "Running pytest 1st vs 2nd run comparison"