[train][CI] Add regression thresholds for E2E CI runs

[SumanthRH]

What does this PR do?

Adds regression thresholds for E2E CI runs.

Methodology

Currently, I've chosen the following metrics:

1. Evaluation reward (eval/all/avg_score)
2. Train reward (loss/avg_final_rewards)
3. Avg num tokens generated (training) (generate/avg_num_tokens)
4. Rollout logprobs train mean difference (policy/rollout_train_logprobs_abs_diff_mean)

For each metric, I found the min/ max values over the previous 10 CI runs and used a threshold with a 5 % allowance. I found this to be reasonable for all the metrics chosen.

---

[gemini-code-assist]

Code Review

This pull request introduces regression testing for E2E CI runs by adding a Python script to check metrics from wandb and updating CI shell scripts to use it. The changes are logical and well-structured. I've provided a few suggestions to improve the implementation, mainly focusing on making the run identification more robust, improving code readability, and following best practices.

[gemini-code-assist]

For better performance and cleaner code, you can use the wandb API's filtering capabilities to fetch the specific run directly, instead of fetching a list of runs and iterating through them. You can replace this block with a more direct approach like the following:

# get latest run with the run name
# get all runs in the project in the order of latest to oldest
runs = api.runs(f"{args.project_name}", filters={"display_name": args.run_name}, order="-created_at")
matched_run = next(iter(runs), None)

[gemini-code-assist]

This line is quite long and can be hard to read. For better readability and maintainability, you can break it into multiple lines using backslashes, like this:

uv run --isolated --extra fsdp $SCRIPT_DIR/get_summary.py \
    --run_name $RUN_NAME \
    --project_name "gsm8k_ci" \
    --asserts "eval/all/avg_score >= $EVAL_ACC_MIN_VALUE" \
              "loss/avg_final_rewards >= $TRAIN_ACC_MIN_VALUE" \
              "generate/avg_num_tokens <= $NUM_TOKENS_MAX_VALUE" \
              "policy/rollout_train_logprobs_abs_diff_mean <= $LOGPROBS_DIFF_MAX_VALUE"

[gemini-code-assist]

The current RUN_NAME format run_$(date +%Y%m%d%H) might not be unique if multiple CI runs are triggered within the same hour. This could lead to the test script checking the wrong run. To ensure uniqueness, consider adding minutes and seconds to the timestamp.

Suggested change

	RUN_NAME="run_$(date +%Y%m%d%H)"
	RUN_NAME="run_$(date +%Y%m%d%H%M%S)"

[gemini-code-assist]

This line is quite long and can be hard to read. For better readability and maintainability, you can break it into multiple lines using backslashes, like this:

uv run --isolated --extra fsdp $SCRIPT_DIR/get_summary.py \
    --run_name $RUN_NAME \
    --project_name $PROJECT_NAME \
    --asserts "eval/all/avg_score >= $EVAL_ACC_MIN_VALUE" \
              "loss/avg_final_rewards >= $TRAIN_ACC_MIN_VALUE" \
              "generate/avg_num_tokens <= $NUM_TOKENS_MAX_VALUE" \
              "policy/rollout_train_logprobs_abs_diff_mean <= $LOGPROBS_DIFF_MAX_VALUE"

[gemini-code-assist]

This line is quite long and can be hard to read. For better readability and maintainability, you can break it into multiple lines using backslashes, like this:

uv run --isolated --extra fsdp $SCRIPT_DIR/get_summary.py \
    --run_name $RUN_NAME \
    --project_name "gsm8k_fully_async_ci" \
    --asserts "eval/all/avg_score >= $EVAL_ACC_MIN_VALUE" \
              "loss/avg_final_rewards >= $TRAIN_ACC_MIN_VALUE" \
              "generate/avg_num_tokens <= $AVG_NUM_TOKENS_MAX_VALUE" \
              "policy/rollout_train_logprobs_abs_diff_mean <= $LOGPROBS_DIFF_MAX_VALUE"

[CharlieFRuan]

i've always thought that our CI runs will just end after we hit the timelimit. if so, will this get summary run at all?

[SumanthRH]

The CI workflow will fail when it hits timeout , yes.

The get_summary script will not run then - and that seems fine by design - because the Github workflow will fail.