Fix GPU pipeline issues (#17)

Author: ultmaster

.github/workflows/examples.yml

@@ -39,7 +39,7 @@ jobs:
       - name: Upload dependencies artifact
         uses: actions/upload-artifact@v4
         with:
-          name: dependencies-python
+          name: dependencies-${{ matrix.setup }}
           path: requirements-freeze.txt
           compression-level: 0
       - name: Prepare Spider dataset
@@ -97,6 +97,7 @@ jobs:
           source .venv/bin/activate
           cd examples/calc_x
           ../../scripts/restart_ray.sh
+          sleep 5
           PYTHONUNBUFFERED=1 python calc_agent.py &
           bash train_ci.sh
           pkill -f calc_agent.py && echo "SIGTERM sent to calc_agent.py" || echo "No calc_agent.py process found"
@@ -105,15 +106,27 @@ jobs:
             sleep 5
           done
           echo "calc_agent.py has finished."
+          sleep 10
         shell: bash
         env:
           WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
+        id: calc_x_train
+
+      - name: Validate Calc-X training
+        run: |
+          set -ex
+          . .venv/bin/activate
+          python scripts/validate_example_wandb.py ${{ steps.calc_x_train.outputs.project_name }} ${{ steps.calc_x_train.outputs.run_name }}
+        env:
+          WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
+
       - name: Spider training
         run: |
           set -ex
           source .venv/bin/activate
           cd examples/spider
           ../../scripts/restart_ray.sh
+          sleep 5
           PYTHONUNBUFFERED=1 python sql_agent.py --trainer.n-workers 10 &
           bash train_ci.sh
           pkill -f sql_agent.py && echo "SIGTERM sent to sql_agent.py" || echo "No sql_agent.py process found"
@@ -122,11 +135,22 @@ jobs:
             sleep 5
           done
           echo "sql_agent.py has finished."
+          sleep 10
         shell: bash
         env:
           VERL_API_BASE: http://localhost:9991/
           WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
+        id: spider_train
         if: success() || failure()
+
+      - name: Validate Spider training
+        run: |
+          set -ex
+          . .venv/bin/activate
+          python scripts/validate_example_wandb.py ${{ steps.spider_train.outputs.project_name }} ${{ steps.spider_train.outputs.run_name }}
+        env:
+          WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
+
       - name: Cleanup
         run: ./scripts/cleanup.sh
         if: success() || failure()

agentlightning/verl/trainer.py

@@ -336,6 +336,11 @@ def fit(self):
                 if is_last_step:
                     pprint(f"Final validation metrics: {last_val_metrics}")
                     progress_bar.close()
+
+                    # This exit logic is to ensure a robust CI.
+                    pprint(f"Flush the logger...")
+                    del logger  # Make sure the loggers are flushed and closed properly
+                    pprint(f"Training finished at step {self.global_steps}.")
                     return
 
                 progress_bar.update(1)

examples/calc_x/train_ci.sh

@@ -1,14 +1,15 @@
-
 #!/bin/bash
 
-set -e
+set -ex
 
 export N_GPUS=1
 export BASE_MODEL=Qwen/Qwen2.5-1.5B-Instruct
 export DATA_DIR=data
 export ROLLOUT_TP_SIZE=1
 export EXPERIMENT_NAME="calc_x_$(date +%Y%m%d%H%M%S)"
 export PROJECT_NAME=AgentLightningCI
+echo "project_name=${PROJECT_NAME}" >> $GITHUB_OUTPUT
+echo "run_name=${EXPERIMENT_NAME}" >> $GITHUB_OUTPUT
 
 PYTHONUNBUFFERED=1 python -m agentlightning.verl \
     algorithm.adv_estimator=grpo \
@@ -48,6 +49,6 @@ PYTHONUNBUFFERED=1 python -m agentlightning.verl \
     trainer.experiment_name=${EXPERIMENT_NAME} \
     trainer.nnodes=1 \
     trainer.save_freq=256 \
-    trainer.test_freq=3 \
+    trainer.test_freq=6 \
     trainer.total_epochs=1 \
-    trainer.total_training_steps=3 $@
+    trainer.total_training_steps=6 $@

examples/spider/train_ci.sh

@@ -1,17 +1,17 @@
 #!/bin/bash
 
-set -e
+set -ex
 
 export N_GPUS=1
 export BASE_MODEL=Qwen/Qwen2.5-Coder-1.5B-Instruct
 export DATA_DIR=data
 export ROLLOUT_TP_SIZE=1
 export EXPERIMENT_NAME="spider_$(date +%Y%m%d%H%M%S)"
-export PROJECT_NAME=AgentLightning
+export PROJECT_NAME=AgentLightningCI
+echo "project_name=${PROJECT_NAME}" >> $GITHUB_OUTPUT
+echo "run_name=${EXPERIMENT_NAME}" >> $GITHUB_OUTPUT
 
-echo "Starting training script..."
-
-python -m agentlightning.verl \
+PYTHONUNBUFFERED=1 python -m agentlightning.verl \
     agentlightning.port=9991 \
     algorithm.adv_estimator=grpo \
     data.train_files=${DATA_DIR}/train_spider.parquet \

scripts/restart_ray.sh

@@ -1,4 +1,4 @@
 #!/bin/bash
 
-ray stop
+ray stop --force
 env RAY_DEBUG=legacy HYDRA_FULL_ERROR=1 VLLM_USE_V1=1 ray start --head --dashboard-host=0.0.0.0

scripts/validate_example_wandb.py

@@ -0,0 +1,38 @@
+import wandb
+import sys
+
+if len(sys.argv) != 3:
+    print("Usage: python validate_example_wandb.py <project> <run_name>")
+
+project = sys.argv[1]
+run_name = sys.argv[2]
+api = wandb.Api()
+entity_name = api.default_entity
+print("Default entity:", entity_name)
+print("Project:", project)
+print("Run name:", run_name)
+
+runs = api.runs(f"{entity_name}/{project}", filters={"displayName": run_name})
+for run in runs:
+    print(f"Found run: {run.name} (ID: {run.id})")
+    if run.name == run_name:
+        break
+else:
+    print(f"::error::Run with name '{run_name}' not found in project '{project}'.")
+    sys.exit(1)
+
+hist = run.history(keys=["val/reward"], pandas=True)
+print("History:", hist)
+if hist.empty:
+    print("::error::No history found for the run.")
+    sys.exit(1)
+else:
+    first, last = hist["val/reward"].iloc[0], hist["val/reward"].iloc[-1]
+    if last <= first:
+        print(
+            f"::warning title=Training no improvement::No improvement (run_name={run_name} start={first:.4f}, end={last:.4f})"
+        )
+    else:
+        print(
+            f"::notice title=Training completed::Run has improved (run_name={run_name} start={first:.4f}, end={last:.4f})"
+        )