notebook-run.yml

name: Simple Notebook Run
on:
  workflow_dispatch:
    inputs:
      notebookFile:
        description: 'Notebook to run (inside `notebook` dir)'
        required: true
        default: 'trainer-v4-unit-test/model-init.ipynb'
      gpuTarget:
        description: 'GPU Target (use either 8x12vgb / 1x24vgb)'
        # required: true
        default: '1x24vgb'
      lanePrefence:
        description: 'Lane preferenc (lane2,laneX)'
        default: "lane2"
      cudaVersion:
        description: 'cuda version'
        # required: true
        default: '11-8'
      huggingFaceRepo:
        description: 'Hugging Face repo to upload to'
        default: 'rwkv-x-dev/rwkv-x-playground'

env:
  # Notebook runner settings
  gpuTarget: ${{github.event.inputs.gpuTarget}}
  cudaVersion: ${{github.event.inputs.cudaVersion}}
  
  # HF repo to sync to
  HF_REPO_SYNC: ${{github.event.inputs.huggingFaceRepo}}

  # Cache dir size limit
  RUNNER_CACHE_SIZE_LIMIT: 100G

  # Secrets
  HUGGING_FACE_HUB_TOKEN: ${{secrets.HUGGING_FACE_HUB_TOKEN}}
  WANDB_API_KEY: ${{secrets.WANDB_API_KEY}}

jobs:
  notebook-run:
    # Due to github worker hard limitation, of 24 hours
    # we apply a timeout of 23 hours instead.
    timeout-minutes: 1380
    name: ${{github.event.inputs.notebookFile}} - ${{github.event.inputs.gpuTarget}} / cuda-${{github.event.inputs.cudaVersion}}
    runs-on: 
      - cuda-${{github.event.inputs.cudaVersion}}
      - ${{github.event.inputs.gpuTarget}}
      - ${{github.event.inputs.lanePrefence}}
    steps:
      - name: Checkout repository
        uses: actions/checkout@v3
      - name: Update external services
        run: |
          pip install --upgrade huggingface_hub
          pip install --upgrade wandb
      - name: Setup HF cache
        run: |
          cd ./notebook/github-runner
          chmod +x ./hf-cache-setup.sh
          ./hf-cache-setup.sh
      - name: Run the notebook, with harness
        run: |
          cd ./notebook/github-runner
          chmod +x ./github-runner.sh
          ./github-runner.sh "${{github.event.inputs.notebookFile}}" 
      - name: Save output notebook
        if: always()
        uses: actions/upload-artifact@v3
        with:
          name: output-notebook
          path: |
            output
      - name: Upload outputs to HF
        if: always()
        run: |
          cd ./notebook/github-runner
          chmod +x ./hf-upload-runner.sh 
          ./hf-upload-runner.sh "${{github.event.inputs.notebookFile}}"
      # Note that this is meant to be a contigency measure, in case the HF upload failed
      - name: Save output models
        uses: actions/upload-artifact@v3
        # if: failure()
        if: always()
        with:
          name: output-models
          path: |
            model/*.pth
          retention-days: 14