Add DeepSpeed CI regression tests for QLoRA and GPT-103B #9
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: DeepSpeed Regression Test | |
| on: | |
| schedule: | |
| - cron: '0 14 * * 6' | |
| push: | |
| branches: [main] | |
| paths: | |
| - '3.test_cases/pytorch/deepspeed/**' | |
| pull_request: | |
| paths: | |
| - '3.test_cases/pytorch/deepspeed/**' | |
| workflow_dispatch: | |
| env: | |
| AWS_REGION: us-east-1 | |
| SLURM_HOST: p5en.smml.aiml.aws.dev | |
| SLURM_USER: ghactions | |
| AWS_ROLE_ARN: arn:aws:iam::159553542841:role/awslabs-AOSH-GitHubActionsRole | |
| BASE_PATH: /fsx/agents/pr-reviews | |
| HOME_PATH: /home/ghactions | |
| permissions: | |
| id-token: write | |
| contents: read | |
| jobs: | |
| qlora: | |
| strategy: | |
| fail-fast: false | |
| max-parallel: 1 | |
| matrix: | |
| zero_config: [zero2, zero3] | |
| runs-on: ubuntu-latest | |
| concurrency: | |
| group: ${{ github.workflow }}-qlora-${{ matrix.zero_config }} | |
| cancel-in-progress: false | |
| timeout-minutes: 180 | |
| steps: | |
| - name: Checkout code | |
| uses: actions/checkout@v4 | |
| with: | |
| path: source-code | |
| - name: Configure AWS Credentials | |
| uses: aws-actions/configure-aws-credentials@v4 | |
| with: | |
| role-to-assume: ${{ env.AWS_ROLE_ARN }} | |
| aws-region: ${{ env.AWS_REGION }} | |
| - name: Setup SSH Key | |
| run: | | |
| mkdir -p ~/.ssh | |
| echo "${{ secrets.SLURM_SSH_KEY }}" > ~/.ssh/slurm_key | |
| chmod 600 ~/.ssh/slurm_key | |
| for i in {1..5}; do | |
| if ssh-keyscan -H ${{ env.SLURM_HOST }} >> ~/.ssh/known_hosts 2>/dev/null; then | |
| echo "SSH keyscan successful" | |
| break | |
| fi | |
| echo "SSH keyscan attempt $i failed, retrying..." | |
| sleep 5 | |
| done | |
| - name: Setup Environment Variables | |
| id: setup | |
| run: | | |
| BUILD_ID="${{ github.run_id }}" | |
| REMOTE_TEST_PATH="${{ env.BASE_PATH }}/deepspeed-tests/qlora-${BUILD_ID}-${{ matrix.zero_config }}" | |
| LOG_DIR="${{ env.HOME_PATH }}/ds-qlora-logs-${BUILD_ID}-${{ matrix.zero_config }}" | |
| echo "remote_test_path=$REMOTE_TEST_PATH" >> $GITHUB_OUTPUT | |
| echo "log_dir=$LOG_DIR" >> $GITHUB_OUTPUT | |
| echo "REMOTE_TEST_PATH=$REMOTE_TEST_PATH" >> $GITHUB_ENV | |
| echo "LOG_DIR=$LOG_DIR" >> $GITHUB_ENV | |
| echo "BUILD_ID=$BUILD_ID" >> $GITHUB_ENV | |
| - name: Create Remote Directories | |
| run: | | |
| ssh -i ~/.ssh/slurm_key -o StrictHostKeyChecking=no -o ConnectTimeout=30 \ | |
| ${{ env.SLURM_USER }}@${{ env.SLURM_HOST }} \ | |
| "mkdir -p ${{ env.REMOTE_TEST_PATH }} ${{ env.LOG_DIR }}" | |
| - name: Transfer Code to Cluster | |
| run: | | |
| for i in {1..3}; do | |
| if scp -i ~/.ssh/slurm_key -o StrictHostKeyChecking=no -o ConnectTimeout=30 -r \ | |
| source-code/* ${{ env.SLURM_USER }}@${{ env.SLURM_HOST }}:${{ env.REMOTE_TEST_PATH }}/; then | |
| echo "Code transfer successful" | |
| break | |
| fi | |
| echo "Transfer attempt $i failed, retrying..." | |
| sleep 10 | |
| done | |
| - name: Build QLoRA Container | |
| run: | | |
| QLORA_DIR="${{ env.REMOTE_TEST_PATH }}/3.test_cases/pytorch/deepspeed/qlora" | |
| SQSH_PATH="/fsx/apps/qwen3-qlora-${{ env.BUILD_ID }}-${{ matrix.zero_config }}.sqsh" | |
| echo "SQSH_PATH=$SQSH_PATH" >> $GITHUB_ENV | |
| ssh -i ~/.ssh/slurm_key -o StrictHostKeyChecking=no -o ConnectTimeout=30 \ | |
| -o ServerAliveInterval=60 -o ServerAliveCountMax=30 \ | |
| ${{ env.SLURM_USER }}@${{ env.SLURM_HOST }} << EOF | |
| set -ex | |
| # Free disk space before build | |
| docker system prune -af || true | |
| # Use /fsx for temp storage (login node root disk is too small) | |
| export TMPDIR=/fsx/tmp ENROOT_TEMP_PATH=/fsx/tmp | |
| mkdir -p /fsx/tmp | |
| cd $QLORA_DIR | |
| echo "Building Docker image..." | |
| docker build -t qwen3-qlora -f Dockerfile . | |
| echo "Converting to enroot squash file..." | |
| mkdir -p /fsx/apps | |
| enroot import -o $SQSH_PATH dockerd://qwen3-qlora:latest | |
| echo "Cleaning up Docker image..." | |
| docker rmi qwen3-qlora || true | |
| rm -rf /fsx/tmp/* | |
| echo "Container build complete: $SQSH_PATH" | |
| EOF | |
| - name: Pre-cache Model and Dataset | |
| env: | |
| HF_TOKEN: ${{ secrets.HF_TOKEN }} | |
| run: | | |
| ssh -i ~/.ssh/slurm_key -o StrictHostKeyChecking=no -o ConnectTimeout=30 \ | |
| ${{ env.SLURM_USER }}@${{ env.SLURM_HOST }} << EOF | |
| set -ex | |
| export HF_HOME=/fsx/.cache/huggingface | |
| export HF_TOKEN=${HF_TOKEN} | |
| mkdir -p \$HF_HOME | |
| python3 -c " | |
| from huggingface_hub import snapshot_download | |
| from datasets import load_dataset | |
| print('Downloading Qwen/Qwen3-8B...') | |
| snapshot_download('Qwen/Qwen3-8B', token='${HF_TOKEN}') | |
| print('Downloading medical-o1-reasoning-SFT dataset...') | |
| load_dataset('FreedomIntelligence/medical-o1-reasoning-SFT', 'en', split='train') | |
| print('Pre-cache complete.') | |
| " | |
| EOF | |
| - name: Prepare and Submit Slurm Job | |
| id: submit_job | |
| env: | |
| HF_TOKEN: ${{ secrets.HF_TOKEN }} | |
| run: | | |
| QLORA_SLURM_DIR="${{ env.REMOTE_TEST_PATH }}/3.test_cases/pytorch/deepspeed/qlora/slurm" | |
| SBATCH_FILE="qwen3_8b-qlora-${{ matrix.zero_config }}.sbatch" | |
| ssh -i ~/.ssh/slurm_key -o StrictHostKeyChecking=no -o ConnectTimeout=30 \ | |
| ${{ env.SLURM_USER }}@${{ env.SLURM_HOST }} << EOF | |
| set -ex | |
| export PATH=/opt/slurm/bin:\$PATH | |
| cd $QLORA_SLURM_DIR | |
| cp "$SBATCH_FILE" regression_test.sbatch | |
| # Patch output/error paths | |
| sed -i "s|#SBATCH --output=.*|#SBATCH --output=${{ env.LOG_DIR }}/%x_%j.out|" regression_test.sbatch | |
| sed -i "s|#SBATCH --error=.*|#SBATCH --error=${{ env.LOG_DIR }}/%x_%j.err|" regression_test.sbatch | |
| # Set container image | |
| sed -i "s|^CONTAINER_IMAGE=\${CONTAINER_IMAGE:-}|CONTAINER_IMAGE=${{ env.SQSH_PATH }}|" regression_test.sbatch | |
| # Override output dir for CI | |
| sed -i "s|OUTPUT_DIR=.*|OUTPUT_DIR=${{ env.LOG_DIR }}/training_output|" regression_test.sbatch | |
| # Inject --max_samples 100 for quick CI run | |
| sed -i 's|--resume_from_checkpoint auto|--resume_from_checkpoint auto --max_samples 100|' regression_test.sbatch | |
| # Set HF_HOME and inject HF_TOKEN | |
| sed -i "s|export HF_HOME=.*|export HF_HOME=/fsx/.cache/huggingface|" regression_test.sbatch | |
| sed -i "/export HF_HOME/a export HF_TOKEN=${HF_TOKEN}" regression_test.sbatch | |
| # Submit job | |
| JOB_ID=\$(sbatch --parsable --partition=p5en regression_test.sbatch) | |
| echo "\$JOB_ID" > ${{ env.REMOTE_TEST_PATH }}/job_id.txt | |
| echo "Submitted Slurm job: \$JOB_ID" | |
| EOF | |
| sleep 2 | |
| JOB_ID=$(ssh -i ~/.ssh/slurm_key -o StrictHostKeyChecking=no \ | |
| ${{ env.SLURM_USER }}@${{ env.SLURM_HOST }} \ | |
| "cat ${{ env.REMOTE_TEST_PATH }}/job_id.txt") | |
| echo "job_id=$JOB_ID" >> $GITHUB_OUTPUT | |
| echo "JOB_ID=$JOB_ID" >> $GITHUB_ENV | |
| echo "Submitted Slurm job: $JOB_ID" | |
| - name: Monitor Slurm Job | |
| run: | | |
| echo "Monitoring job ${{ env.JOB_ID }}..." | |
| START_TIME=$(date +%s) | |
| TIMEOUT=7200 # 2 hours | |
| ssh_cmd() { | |
| for attempt in 1 2 3; do | |
| if ssh -i ~/.ssh/slurm_key -o StrictHostKeyChecking=no -o ConnectTimeout=30 \ | |
| ${{ env.SLURM_USER }}@${{ env.SLURM_HOST }} "$@"; then | |
| return 0 | |
| fi | |
| echo "SSH attempt $attempt failed, retrying..." | |
| sleep 5 | |
| done | |
| return 1 | |
| } | |
| while true; do | |
| CURRENT_TIME=$(date +%s) | |
| ELAPSED=$((CURRENT_TIME - START_TIME)) | |
| if [ $ELAPSED -gt $TIMEOUT ]; then | |
| echo "Timeout reached after $((TIMEOUT / 60)) minutes" | |
| exit 1 | |
| fi | |
| JOB_STATUS=$(ssh_cmd "export PATH=/opt/slurm/bin:\$PATH; squeue -j ${{ env.JOB_ID }} -h -o %T 2>/dev/null || true") | |
| if [ -z "$JOB_STATUS" ]; then | |
| # Job no longer in queue — check final state via sacct | |
| FINAL_STATE=$(ssh_cmd "export PATH=/opt/slurm/bin:\$PATH; sacct -j ${{ env.JOB_ID }} --format=State --noheader -P | head -1" | tr -d '[:space:]') | |
| echo "Job finished with state: $FINAL_STATE" | |
| if [ "$FINAL_STATE" = "COMPLETED" ]; then | |
| echo "Job completed successfully" | |
| break | |
| else | |
| echo "Job failed with state: $FINAL_STATE" | |
| echo "=== Last 200 lines of stdout ===" | |
| ssh_cmd "tail -200 ${{ env.LOG_DIR }}/*_${{ env.JOB_ID }}.out 2>/dev/null || true" | |
| echo "=== Last 200 lines of stderr ===" | |
| ssh_cmd "tail -200 ${{ env.LOG_DIR }}/*_${{ env.JOB_ID }}.err 2>/dev/null || true" | |
| exit 1 | |
| fi | |
| fi | |
| echo "--- Job status: $JOB_STATUS (elapsed: $((ELAPSED / 60)) min) ---" | |
| sleep 30 | |
| done | |
| - name: Retrieve Logs | |
| if: always() | |
| run: | | |
| mkdir -p ./logs | |
| for i in {1..3}; do | |
| if scp -i ~/.ssh/slurm_key -o StrictHostKeyChecking=no -o ConnectTimeout=30 -r \ | |
| ${{ env.SLURM_USER }}@${{ env.SLURM_HOST }}:${{ env.LOG_DIR }}/* ./logs/ 2>/dev/null; then | |
| echo "Logs retrieved successfully" | |
| break | |
| fi | |
| echo "Log retrieval attempt $i failed, retrying..." | |
| sleep 10 | |
| done | |
| - name: Upload Logs as Artifacts | |
| if: always() | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: ds-qlora-${{ matrix.zero_config }}-logs-${{ github.run_id }} | |
| path: ./logs | |
| retention-days: 60 | |
| - name: Cleanup | |
| if: always() | |
| run: | | |
| # Cancel job if still running | |
| if [ -n "${{ env.JOB_ID }}" ]; then | |
| ssh -i ~/.ssh/slurm_key -o StrictHostKeyChecking=no \ | |
| ${{ env.SLURM_USER }}@${{ env.SLURM_HOST }} \ | |
| "export PATH=/opt/slurm/bin:\$PATH; scancel ${{ env.JOB_ID }} 2>/dev/null || true" | |
| fi | |
| # Clean up remote directories and container image | |
| ssh -i ~/.ssh/slurm_key -o StrictHostKeyChecking=no \ | |
| ${{ env.SLURM_USER }}@${{ env.SLURM_HOST }} << EOF | |
| rm -rf ${{ env.REMOTE_TEST_PATH }} | |
| rm -rf ${{ env.LOG_DIR }} | |
| rm -f ${{ env.SQSH_PATH }} | |
| EOF | |
| rm -rf ./logs | |
| echo "Cleanup completed!" | |
| gpt-103b: | |
| runs-on: ubuntu-latest | |
| concurrency: | |
| group: ${{ github.workflow }}-gpt-103b | |
| cancel-in-progress: false | |
| timeout-minutes: 180 | |
| steps: | |
| - name: Checkout code | |
| uses: actions/checkout@v4 | |
| with: | |
| path: source-code | |
| - name: Configure AWS Credentials | |
| uses: aws-actions/configure-aws-credentials@v4 | |
| with: | |
| role-to-assume: ${{ env.AWS_ROLE_ARN }} | |
| aws-region: ${{ env.AWS_REGION }} | |
| - name: Setup SSH Key | |
| run: | | |
| mkdir -p ~/.ssh | |
| echo "${{ secrets.SLURM_SSH_KEY }}" > ~/.ssh/slurm_key | |
| chmod 600 ~/.ssh/slurm_key | |
| for i in {1..5}; do | |
| if ssh-keyscan -H ${{ env.SLURM_HOST }} >> ~/.ssh/known_hosts 2>/dev/null; then | |
| echo "SSH keyscan successful" | |
| break | |
| fi | |
| echo "SSH keyscan attempt $i failed, retrying..." | |
| sleep 5 | |
| done | |
| - name: Setup Environment Variables | |
| id: setup | |
| run: | | |
| BUILD_ID="${{ github.run_id }}" | |
| REMOTE_TEST_PATH="${{ env.BASE_PATH }}/deepspeed-tests/gpt103b-${BUILD_ID}" | |
| LOG_DIR="${{ env.HOME_PATH }}/ds-gpt103b-logs-${BUILD_ID}" | |
| echo "remote_test_path=$REMOTE_TEST_PATH" >> $GITHUB_OUTPUT | |
| echo "log_dir=$LOG_DIR" >> $GITHUB_OUTPUT | |
| echo "REMOTE_TEST_PATH=$REMOTE_TEST_PATH" >> $GITHUB_ENV | |
| echo "LOG_DIR=$LOG_DIR" >> $GITHUB_ENV | |
| echo "BUILD_ID=$BUILD_ID" >> $GITHUB_ENV | |
| - name: Create Remote Directories | |
| run: | | |
| ssh -i ~/.ssh/slurm_key -o StrictHostKeyChecking=no -o ConnectTimeout=30 \ | |
| ${{ env.SLURM_USER }}@${{ env.SLURM_HOST }} \ | |
| "mkdir -p ${{ env.REMOTE_TEST_PATH }} ${{ env.LOG_DIR }}" | |
| - name: Transfer Code to Cluster | |
| run: | | |
| for i in {1..3}; do | |
| if scp -i ~/.ssh/slurm_key -o StrictHostKeyChecking=no -o ConnectTimeout=30 -r \ | |
| source-code/* ${{ env.SLURM_USER }}@${{ env.SLURM_HOST }}:${{ env.REMOTE_TEST_PATH }}/; then | |
| echo "Code transfer successful" | |
| break | |
| fi | |
| echo "Transfer attempt $i failed, retrying..." | |
| sleep 10 | |
| done | |
| - name: Build DeepSpeed Container | |
| run: | | |
| DS_DIR="${{ env.REMOTE_TEST_PATH }}/3.test_cases/pytorch/deepspeed" | |
| SQSH_PATH="/fsx/apps/deepspeed-${{ env.BUILD_ID }}.sqsh" | |
| echo "SQSH_PATH=$SQSH_PATH" >> $GITHUB_ENV | |
| ssh -i ~/.ssh/slurm_key -o StrictHostKeyChecking=no -o ConnectTimeout=30 \ | |
| -o ServerAliveInterval=60 -o ServerAliveCountMax=30 \ | |
| ${{ env.SLURM_USER }}@${{ env.SLURM_HOST }} << EOF | |
| set -ex | |
| # Free disk space before build | |
| docker system prune -af || true | |
| # Use /fsx for temp storage (login node root disk is too small) | |
| export TMPDIR=/fsx/tmp ENROOT_TEMP_PATH=/fsx/tmp | |
| mkdir -p /fsx/tmp | |
| cd $DS_DIR | |
| echo "Building Docker image..." | |
| docker build -t deepspeed -f 0.deepspeed.dockerfile . | |
| echo "Converting to enroot squash file..." | |
| mkdir -p /fsx/apps | |
| enroot import -o $SQSH_PATH dockerd://deepspeed:latest | |
| echo "Cleaning up Docker image..." | |
| docker rmi deepspeed || true | |
| rm -rf /fsx/tmp/* | |
| echo "Container build complete: $SQSH_PATH" | |
| EOF | |
| - name: Setup Megatron-DeepSpeed and Data | |
| run: | | |
| ssh -i ~/.ssh/slurm_key -o StrictHostKeyChecking=no -o ConnectTimeout=30 \ | |
| ${{ env.SLURM_USER }}@${{ env.SLURM_HOST }} << EOF | |
| set -ex | |
| # Clone Megatron-DeepSpeed if not present | |
| if [ ! -d /fsx/deepspeed/Megatron-DeepSpeed ]; then | |
| mkdir -p /fsx/deepspeed | |
| git clone https://github.com/microsoft/Megatron-DeepSpeed.git /fsx/deepspeed/Megatron-DeepSpeed | |
| fi | |
| # Download GPT-2 tokenizer files if not present | |
| mkdir -p /fsx/deepspeed/data | |
| if [ ! -f /fsx/deepspeed/data/gpt2-vocab.json ]; then | |
| wget -q -O /fsx/deepspeed/data/gpt2-vocab.json \ | |
| https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json | |
| fi | |
| if [ ! -f /fsx/deepspeed/data/gpt2-merges.txt ]; then | |
| wget -q -O /fsx/deepspeed/data/gpt2-merges.txt \ | |
| https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt | |
| fi | |
| # Generate synthetic BookCorpus data if not present | |
| if [ ! -f /fsx/deepspeed/data/BookCorpusDataset_text_document.idx ]; then | |
| echo "Generating synthetic training data..." | |
| # Create synthetic corpus | |
| python3 -c " | |
| import json, os | |
| corpus_path = '/fsx/deepspeed/data/synthetic_corpus.json' | |
| if not os.path.exists(corpus_path): | |
| samples = [] | |
| for i in range(10000): | |
| samples.append({'text': f'This is synthetic training sample number {i}. ' * 20}) | |
| with open(corpus_path, 'w') as f: | |
| for s in samples: | |
| f.write(json.dumps(s) + chr(10)) | |
| print(f'Created {len(samples)} synthetic samples') | |
| " | |
| # Preprocess data inside the container | |
| enroot create --name deepspeed-preprocess ${{ env.SQSH_PATH }} | |
| enroot start --rw \ | |
| --mount /fsx:/fsx \ | |
| deepspeed-preprocess \ | |
| bash -c "pip install nltk && python3 -c 'import nltk; nltk.download(\"punkt\", quiet=True)' && python3 /fsx/deepspeed/Megatron-DeepSpeed/tools/preprocess_data.py \ | |
| --input /fsx/deepspeed/data/synthetic_corpus.json \ | |
| --output-prefix /fsx/deepspeed/data/BookCorpusDataset \ | |
| --vocab-file /fsx/deepspeed/data/gpt2-vocab.json \ | |
| --merge-file /fsx/deepspeed/data/gpt2-merges.txt \ | |
| --tokenizer-type GPT2BPETokenizer \ | |
| --append-eod \ | |
| --workers 8" | |
| enroot remove -f deepspeed-preprocess || true | |
| echo "Data preprocessing complete." | |
| fi | |
| echo "Megatron-DeepSpeed and data setup complete." | |
| EOF | |
| - name: Prepare and Submit Slurm Job | |
| id: submit_job | |
| run: | | |
| GPT_SLURM_DIR="${{ env.REMOTE_TEST_PATH }}/3.test_cases/pytorch/deepspeed/gpt/slurm" | |
| ssh -i ~/.ssh/slurm_key -o StrictHostKeyChecking=no -o ConnectTimeout=30 \ | |
| ${{ env.SLURM_USER }}@${{ env.SLURM_HOST }} << EOF | |
| set -ex | |
| export PATH=/opt/slurm/bin:\$PATH | |
| cd $GPT_SLURM_DIR | |
| cp pretrain_gpt_103b.sbatch regression_test.sbatch | |
| # Patch output/error paths | |
| sed -i "s|#SBATCH --output=.*|#SBATCH --output=${{ env.LOG_DIR }}/%x_%j.out|" regression_test.sbatch | |
| sed -i "s|#SBATCH --error=.*|#SBATCH --error=${{ env.LOG_DIR }}/%x_%j.err|" regression_test.sbatch | |
| # Submit with CI overrides: 4 nodes, TP=8, PP=4, ZeRO=0, 10 iters, fusions on | |
| JOB_ID=\$(sbatch --parsable --partition=p5en --nodes=4 \ | |
| --export=ALL,TP=8,PP=4,ZERO_STAGE=0,TRAIN_ITERS=10,ENABLE_FUSIONS=1,IMAGE=${{ env.SQSH_PATH }} \ | |
| regression_test.sbatch) | |
| echo "\$JOB_ID" > ${{ env.REMOTE_TEST_PATH }}/job_id.txt | |
| echo "Submitted Slurm job: \$JOB_ID" | |
| EOF | |
| sleep 2 | |
| JOB_ID=$(ssh -i ~/.ssh/slurm_key -o StrictHostKeyChecking=no \ | |
| ${{ env.SLURM_USER }}@${{ env.SLURM_HOST }} \ | |
| "cat ${{ env.REMOTE_TEST_PATH }}/job_id.txt") | |
| echo "job_id=$JOB_ID" >> $GITHUB_OUTPUT | |
| echo "JOB_ID=$JOB_ID" >> $GITHUB_ENV | |
| echo "Submitted Slurm job: $JOB_ID" | |
| - name: Monitor Slurm Job | |
| run: | | |
| echo "Monitoring job ${{ env.JOB_ID }}..." | |
| START_TIME=$(date +%s) | |
| TIMEOUT=7200 # 2 hours | |
| ssh_cmd() { | |
| for attempt in 1 2 3; do | |
| if ssh -i ~/.ssh/slurm_key -o StrictHostKeyChecking=no -o ConnectTimeout=30 \ | |
| ${{ env.SLURM_USER }}@${{ env.SLURM_HOST }} "$@"; then | |
| return 0 | |
| fi | |
| echo "SSH attempt $attempt failed, retrying..." | |
| sleep 5 | |
| done | |
| return 1 | |
| } | |
| while true; do | |
| CURRENT_TIME=$(date +%s) | |
| ELAPSED=$((CURRENT_TIME - START_TIME)) | |
| if [ $ELAPSED -gt $TIMEOUT ]; then | |
| echo "Timeout reached after $((TIMEOUT / 60)) minutes" | |
| exit 1 | |
| fi | |
| JOB_STATUS=$(ssh_cmd "export PATH=/opt/slurm/bin:\$PATH; squeue -j ${{ env.JOB_ID }} -h -o %T 2>/dev/null || true") | |
| if [ -z "$JOB_STATUS" ]; then | |
| FINAL_STATE=$(ssh_cmd "export PATH=/opt/slurm/bin:\$PATH; sacct -j ${{ env.JOB_ID }} --format=State --noheader -P | head -1" | tr -d '[:space:]') | |
| echo "Job finished with state: $FINAL_STATE" | |
| if [ "$FINAL_STATE" = "COMPLETED" ]; then | |
| echo "Job completed successfully" | |
| break | |
| else | |
| echo "Job failed with state: $FINAL_STATE" | |
| echo "=== Last 200 lines of stdout ===" | |
| ssh_cmd "tail -200 ${{ env.LOG_DIR }}/*_${{ env.JOB_ID }}.out 2>/dev/null || true" | |
| echo "=== Last 200 lines of stderr ===" | |
| ssh_cmd "tail -200 ${{ env.LOG_DIR }}/*_${{ env.JOB_ID }}.err 2>/dev/null || true" | |
| exit 1 | |
| fi | |
| fi | |
| echo "--- Job status: $JOB_STATUS (elapsed: $((ELAPSED / 60)) min) ---" | |
| sleep 30 | |
| done | |
| - name: Retrieve Logs | |
| if: always() | |
| run: | | |
| mkdir -p ./logs | |
| for i in {1..3}; do | |
| if scp -i ~/.ssh/slurm_key -o StrictHostKeyChecking=no -o ConnectTimeout=30 -r \ | |
| ${{ env.SLURM_USER }}@${{ env.SLURM_HOST }}:${{ env.LOG_DIR }}/* ./logs/ 2>/dev/null; then | |
| echo "Logs retrieved successfully" | |
| break | |
| fi | |
| echo "Log retrieval attempt $i failed, retrying..." | |
| sleep 10 | |
| done | |
| - name: Upload Logs as Artifacts | |
| if: always() | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: ds-gpt103b-logs-${{ github.run_id }} | |
| path: ./logs | |
| retention-days: 60 | |
| - name: Cleanup | |
| if: always() | |
| run: | | |
| # Cancel job if still running | |
| if [ -n "${{ env.JOB_ID }}" ]; then | |
| ssh -i ~/.ssh/slurm_key -o StrictHostKeyChecking=no \ | |
| ${{ env.SLURM_USER }}@${{ env.SLURM_HOST }} \ | |
| "export PATH=/opt/slurm/bin:\$PATH; scancel ${{ env.JOB_ID }} 2>/dev/null || true" | |
| fi | |
| # Clean up remote directories and container image | |
| ssh -i ~/.ssh/slurm_key -o StrictHostKeyChecking=no \ | |
| ${{ env.SLURM_USER }}@${{ env.SLURM_HOST }} << EOF | |
| rm -rf ${{ env.REMOTE_TEST_PATH }} | |
| rm -rf ${{ env.LOG_DIR }} | |
| rm -f ${{ env.SQSH_PATH }} | |
| EOF | |
| rm -rf ./logs | |
| echo "Cleanup completed!" |