lightning/.azure/gpu-benchmarks.yml

# Python package
# Create and test a Python package on multiple Python versions.
# Add steps that analyze code, save the dist with the build record, publish to a PyPI-compatible index, and more:
# https://docs.microsoft.com/azure/devops/pipelines/languages/python

trigger:
  tags:
    include: ["*"]
  branches:
    include:
      - "master"
      - "release/*"
      - "refs/tags/*"

pr:
  branches:
    include:
      - "master"
      - "release/*"
  paths:
    include:
      - ".azure/gpu-benchmarks.yml"
      - "requirements/fabric/**"
      - "requirements/pytorch/**"
      - "src/lightning/fabric/**"
      - "src/lightning/pytorch/**"
      - "tests/parity_fabric/**"
      - "tests/parity_pytorch/**"
    exclude:
      - "requirements/*/docs.txt"
      - "*.md"
      - "**/*.md"

schedules:
  - cron: "0 0 * * *" # At the end of every day
    displayName: Daily midnight benchmark
    branches:
      include:
        - "master"

jobs:
  - job: benchmarks
    timeoutInMinutes: "90"
    cancelTimeoutInMinutes: "2"
    pool: lit-rtx-3090
    variables:
      DEVICES: $( python -c 'print("$(Agent.Name)".split("_")[-1])' )
    container:
      # TODO: Upgrade to Python 3.11
      image: "pytorchlightning/pytorch_lightning:base-cuda-py3.10-torch2.2-cuda12.1.0"
      options: "--gpus=all --shm-size=32g"
    strategy:
      matrix:
        "pkg: Fabric":
          PACKAGE_NAME: "fabric"
        "pkg: Pytorch":
          PACKAGE_NAME: "pytorch"
    workspace:
      clean: all

    steps:
      - bash: |
          echo "##vso[task.setvariable variable=CUDA_VISIBLE_DEVICES]$(DEVICES)"
          cuda_ver=$(python -c "import torch ; print(''.join(map(str, torch.version.cuda.split('.')[:2])))")
          echo "##vso[task.setvariable variable=TORCH_URL]https://download.pytorch.org/whl/cu${cuda_ver}/torch_stable.html"
        displayName: "set env. vars"

      - bash: |
          echo $CUDA_VISIBLE_DEVICES
          echo $TORCH_URL
          whereis nvidia
          nvidia-smi
          which python && which pip
          python --version
          pip --version
          pip list
        displayName: "Image info & NVIDIA"

      - bash: pip install -e .[dev] --find-links ${TORCH_URL}
        env:
          FREEZE_REQUIREMENTS: "1"
        displayName: "Install package"

      - bash: |
          set -e
          python requirements/collect_env_details.py
          python -c "import torch ; mgpu = torch.cuda.device_count() ; assert mgpu == 2, f'GPU: {mgpu}'"
        displayName: "Env details"

      - bash: |
          pip install -q -r .actions/requirements.txt
          python .actions/assistant.py copy_replace_imports --source_dir="./tests" \
            --source_import="lightning.fabric,lightning.pytorch" \
            --target_import="lightning_fabric,pytorch_lightning"
        displayName: "Adjust tests"

      - bash: python -m pytest parity_$(PACKAGE_NAME) -v --durations=0
        env:
          PL_RUNNING_BENCHMARKS: "1"
          PL_RUN_CUDA_TESTS: "1"
        workingDirectory: tests/
        displayName: "Testing: benchmarks"

      - bash: bash run_standalone_tasks.sh
        workingDirectory: tests/parity_fabric
        # without succeeded this could run even if the job has already failed
        condition: and(succeeded(), eq(variables['PACKAGE_NAME'], 'fabric'))
        env:
          PL_RUN_CUDA_TESTS: "1"
        displayName: "Testing: fabric standalone tasks"
        timeoutInMinutes: "10"