nvidia_workflow.yml

name: NVIDIA PyTorch Job

on:
  workflow_dispatch:
    inputs:
      script_content:
        description: 'Content of Python script'
        required: true
        type: string
      filename:
        description: 'Name of Python script'
        required: true
        type: string

jobs:
  train:
    runs-on: [gpumode-nvidia-arc]
    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: Create script
      shell: python
      run: |
        with open('${{ github.event.inputs.filename }}', 'w') as f:
            f.write('''${{ github.event.inputs.script_content }}''')

    - name: Install dependencies
      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"
          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"
          pip install triton
        fi

    - name: Run script with profiler
      run: |
        # Run the script with NSight Compute profiler and save to CSV
        # ncu --csv python "${{ github.event.inputs.filename }}" > profile_results.csv 2>&1
        # Also keep regular output in training.log
        python "${{ github.event.inputs.filename }}" > training.log 2>&1

    - name: Upload training artifacts
      uses: actions/upload-artifact@v3
      if: always()
      with:
        name: training-artifacts
        path: |
          training.log
          # profile_results.csv
          ${{ github.event.inputs.filename }}
      env:
        CUDA_VISIBLE_DEVICES: 0 # Make sure only one GPU is used for testing