Demo – toy GPU workload, optional sleep (for debugging) #2
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: Demo – toy GPU workload, optional sleep (for debugging) | |
| on: | |
| workflow_dispatch: | |
| inputs: | |
| sleep: | |
| description: "Sleep for this many seconds before completing job (optional, helps keep instance alive for SSH access and debugging)" | |
| required: false | |
| type: number | |
| default: 0 | |
| ssh_pubkey: | |
| description: "Add this SSH public key to instance's `~/.ssh/authorized_keys` (optional, for debugging)" | |
| required: false | |
| type: string | |
| workflow_call: # Tested by `demos.yml` | |
| inputs: | |
| sleep: | |
| required: false | |
| type: number | |
| default: 0 | |
| ssh_pubkey: | |
| required: false | |
| type: string | |
| permissions: | |
| id-token: write # Required for AWS OIDC authentication | |
| contents: read # Required for actions/checkout | |
| jobs: | |
| ec2: | |
| uses: ./.github/workflows/runner.yml | |
| secrets: inherit | |
| with: | |
| ec2_instance_type: g4dn.xlarge | |
| ec2_image_id: ami-00096836009b16a22 # Deep Learning OSS Nvidia Driver AMI GPU PyTorch | |
| ssh_pubkey: ${{ inputs.ssh_pubkey }} | |
| gpu-workload: | |
| needs: ec2 | |
| runs-on: ${{ needs.ec2.outputs.id }} | |
| steps: | |
| - uses: actions/checkout@v4 | |
| - name: System info | |
| run: | | |
| echo "=== System Info ===" | |
| uname -a | |
| lscpu | grep "Model name" || true | |
| free -h | |
| echo -e "\n=== GPU Info ===" | |
| nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv | |
| - name: Install PyTorch with CUDA support | |
| run: | | |
| echo "=== Installing PyTorch with CUDA support ===" | |
| # Install python3-venv if not available | |
| sudo apt-get update && sudo apt-get install -y python3-venv | |
| # Use system Python3 and create a virtual environment to avoid pip root warning | |
| python3 --version | |
| python3 -m venv /tmp/venv | |
| source /tmp/venv/bin/activate | |
| pip install torch --index-url https://download.pytorch.org/whl/cu121 | |
| echo "PyTorch installed successfully" | |
| - name: Run PyTorch GPU benchmark | |
| run: | | |
| echo "=== PyTorch GPU Benchmark ===" | |
| source /tmp/venv/bin/activate | |
| python .github/test-scripts/gpu-benchmark.py | |
| - name: Sleep (${{ inputs.sleep }}s) | |
| if: ${{ inputs.sleep > 0 }} | |
| run: | | |
| echo "Sleeping for ${{ inputs.sleep }} seconds (useful for SSH debugging)..." | |
| echo "Instance will remain available at the IP shown in the GitHub Actions log" | |
| sleep ${{ inputs.sleep }} |