nvidia_workflow.yml

name: NVIDIA PyTorch/CUDA Job
on:
  workflow_dispatch:
    inputs:
      script_content:
        description: 'Content of Python/CUDA script (.py or .cu file)'
        required: true
        type: string
      filename:
        description: 'Name of script (supports .py or .cu)'
        required: true
        type: string
jobs:
  train:
    runs-on: [gpumode-nvidia-arc]
    timeout-minutes: 10
    container:
      image: nvidia/cuda:12.4.0-devel-ubuntu22.04
    steps:
      - name: Setup Python
        uses: actions/setup-python@v4
        with:
          python-version: '3.10'
      
      - name: Install uv
        uses: astral-sh/setup-uv@v3
        with:
          version: "latest"
      
      - name: Setup Python environment
        run: |
          uv venv .venv
          echo "VIRTUAL_ENV=$PWD/.venv" >> $GITHUB_ENV
          echo "$PWD/.venv/bin" >> $GITHUB_PATH
      
      - name: Create script
        shell: python
        run: |
          with open('${{ github.event.inputs.filename }}', 'w') as f:
            f.write('''${{ github.event.inputs.script_content }}''')
      
      - name: Install dependencies and setup NVCC
        run: |          
          # Check for PyTorch
          if grep -rE "(import torch|from torch)" "${{ github.event.inputs.filename }}"; then
            echo "PyTorch detected, installing torch"
            uv pip install numpy torch
          fi
          
          # Check for Triton
          if grep -rE "(import triton|from triton)" "${{ github.event.inputs.filename }}"; then
            echo "Triton detected, installing triton"
            uv pip install triton
          fi
      
      - name: Run script with profiler
        run: |
          if [[ "${{ github.event.inputs.filename }}" == *.cu ]]; then
            # Compile and run CUDA file
            nvcc "${{ github.event.inputs.filename }}" -o cuda_program
            ./cuda_program > training.log 2>&1
          else
            # Run Python file
            python "${{ github.event.inputs.filename }}" > training.log 2>&1
          fi
      
      - name: Upload training artifacts
        uses: actions/upload-artifact@v4
        if: always()
        with:
          name: training-artifacts
          path: |
            training.log
            # profile_results.csv
            ${{ github.event.inputs.filename }}
            cuda_program*
        env:
          CUDA_VISIBLE_DEVICES: 0  # Make sure only one GPU is used for testing