nvidia_workflow.yml

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 if needed
      if: endsWith(github.event.inputs.filename, '.py')
      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