lightning/.azure/gpu-benchmark.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-benchmark.yml"

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:
      image: "pytorchlightning/pytorch_lightning:base-cuda-py3.9-torch1.12-cuda11.6.1"
      options: "--gpus=all --shm-size=32g"
    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=CUDA_VERSION_MM]$cuda_ver"
        echo "##vso[task.setvariable variable=TORCH_URL]https://download.pytorch.org/whl/cu${cuda_ver}/torch_stable.html"
      displayName: 'set env. vars'

    - bash: |
        pip install -e .[strategies] --find-links ${TORCH_URL}
        pip list
      env:
        PACKAGE_NAME: pytorch
        FREEZE_REQUIREMENTS: 1
      displayName: 'Install package'

    - bash: python -m pytest benchmarks -v --durations=0
      env:
        PL_RUNNING_BENCHMARKS: "1"
        PL_RUN_CUDA_TESTS: "1"
      workingDirectory: tests/tests_pytorch
      displayName: 'Testing: PyTorch benchmarks'