NVIDIA PyTorch/CUDA Job #32
This file contains 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: NVIDIA PyTorch/CUDA Job | |
| on: | |
| workflow_dispatch: | |
| inputs: | |
| script_content: | |
| description: 'Content of script (Python or CUDA)' | |
| required: true | |
| type: string | |
| filename: | |
| description: 'Name of script (supports .py and .cu files)' | |
| required: true | |
| type: string | |
| jobs: | |
| build-and-train: | |
| runs-on: [gpumode-nvidia-arc] | |
| timeout-minutes: 10 | |
| container: | |
| image: nvidia/cuda:12.4.0-devel-ubuntu22.04 | |
| steps: | |
| - name: Create script | |
| shell: bash | |
| run: | | |
| echo '${{ github.event.inputs.script_content }}' > '${{ github.event.inputs.filename }}' | |
| - name: Setup Python | |
| uses: actions/setup-python@v4 | |
| with: | |
| python-version: '3.10' | |
| - name: Install uv if Python | |
| if: endsWith(github.event.inputs.filename, '.py') | |
| uses: astral-sh/setup-uv@v3 | |
| with: | |
| version: "latest" | |
| - name: Setup Python environment if needed | |
| if: endsWith(github.event.inputs.filename, '.py') | |
| run: | | |
| uv venv .venv | |
| echo "VIRTUAL_ENV=$PWD/.venv" >> $GITHUB_ENV | |
| echo "$PWD/.venv/bin" >> $GITHUB_PATH | |
| - name: Install Python dependencies if needed | |
| if: endsWith(github.event.inputs.filename, '.py') | |
| run: | | |
| # Check if 'import torch' is in any Python file | |
| if grep -rE "(import torch|from torch)" "${{ github.event.inputs.filename }}"; then | |
| echo "PyTorch detected, installing torch" | |
| uv pip install numpy torch | |
| fi | |
| # Check if 'import triton' is in any Python file | |
| if grep -rE "(import triton|from triton)" "${{ github.event.inputs.filename }}"; then | |
| echo "Triton detected, installing triton" | |
| uv pip install triton | |
| fi | |
| - name: Compile CUDA if needed | |
| if: endsWith(github.event.inputs.filename, '.cu') | |
| run: | | |
| echo "Compiling CUDA file..." | |
| nvcc "${{ github.event.inputs.filename }}" -o cuda_program | |
| - name: Run script | |
| run: | | |
| if [[ "${{ github.event.inputs.filename }}" == *.cu ]]; then | |
| ./cuda_program > output.log 2>&1 | |
| else | |
| python "${{ github.event.inputs.filename }}" > output.log 2>&1 | |
| fi | |
| - name: Upload artifacts | |
| uses: actions/upload-artifact@v4 | |
| if: always() | |
| with: | |
| name: run-artifacts | |
| path: | | |
| output.log | |
| ${{ github.event.inputs.filename }} | |
| env: | |
| CUDA_VISIBLE_DEVICES: 0 |