nvidia_workflow.yml

name: NVIDIA PyTorch/CUDA Job
on:
  workflow_dispatch:
    inputs:
      payload:
        description: 'Content of the user submission, as json string'
        required: true
        type: string
      requirements:
        description: 'Contents for a requirements.txt file'
        required: false
        type: string

jobs:
  run:
    runs-on: [gpumode-nvidia-arc]
    timeout-minutes: 10
    container:
      image: nvidia/cuda:12.4.0-devel-ubuntu22.04
    steps:
    - uses: actions/checkout@v4

    - name: Setup Python
      uses: actions/setup-python@v5
      with:
        python-version: '3.10'
    
    - name: Create input files
      uses: nick-fields/retry@v3
      with:
        timeout_minutes: 2
        max_attempts: 5
        shell: bash
        command: |
          # install jq
          apt update
          apt install -y jq
          # Extract the payload content without printing it
          PAYLOAD=$(jq -r '.inputs.payload' $GITHUB_EVENT_PATH)
          
          # Apply mask to the extracted content
          echo "::add-mask::$PAYLOAD"
          
          # Now write to file (won't be logged since it's masked)
          echo "$PAYLOAD" > payload.json

    - name: Install uv
      uses: astral-sh/setup-uv@v3
      with:
        version: "latest"

    - name: Setup Python environment
      shell: bash
      run: |
        uv venv .venv
        echo "VIRTUAL_ENV=$PWD/.venv" >> $GITHUB_ENV
        echo "$PWD/.venv/bin" >> $GITHUB_PATH

        if [[ -n "${{ github.event.inputs.requirements }}" ]]; then
          cat > "requirements.txt" <<'EOL'
          ${{ github.event.inputs.requirements }}
        EOL
        uv pip install -r "requirements.txt"
        fi

    - name: Run script
      shell: bash
      run: |
        python .github/workflows/runner.py
        cat result.json  # Debug: show output

    - name: Upload training artifacts
      uses: actions/upload-artifact@v4
      if: always()
      with:
        name: run-result
        path: |
          result.json

    env:
      CUDA_VISIBLE_DEVICES: 0