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

pr:
  branches:
    include: ["*"]

jobs:
  - job: testing
    strategy:
      matrix:
        "main w/ torch 2.8.0":
          docker-image: "ubuntu24.04-cuda12.8.1-cudnn-fe1.15.0-py3.12-pt_2.8.0-dev"
          testing: "main"
        "ops w/ torch 2.8.0":
          docker-image: "ubuntu24.04-cuda12.8.1-cudnn-fe1.15.0-py3.12-pt_2.8.0-dev"
          testing: "ops"
        "grads w/ torch 2.8.0":
          docker-image: "ubuntu24.04-cuda12.8.1-cudnn-fe1.15.0-py3.12-pt_2.8.0-dev"
          testing: "grads"
        "distributed w/ torch 2.8.0":
          docker-image: "ubuntu24.04-cuda12.8.1-cudnn-fe1.15.0-py3.12-pt_2.8.0-dev"
          testing: "distributed"
        "main w/ torch-nightly":
          docker-image: "ubuntu24.04-cuda12.8.1-cudnn-fe1.15.0-py3.12-pt_main-dev"
          testing: "main"
        "ops w/ torch-nightly":
          docker-image: "ubuntu24.04-cuda12.8.1-cudnn-fe1.15.0-py3.12-pt_main-dev"
          testing: "ops"
        "grads w/ torch-nightly":
          docker-image: "ubuntu24.04-cuda12.8.1-cudnn-fe1.15.0-py3.12-pt_main-dev"
          testing: "grads"
        "distributed w/ torch-nightly":
          docker-image: "ubuntu24.04-cuda12.8.1-cudnn-fe1.15.0-py3.12-pt_main-dev"
          testing: "distributed"
    # how much time to give 'run always even if cancelled tasks' before stopping them
    cancelTimeoutInMinutes: "2"
    pool: "lit-rtx-3090"
    variables:
      DEVICES: $( python -c 'name = "$(Agent.Name)" ; gpus = name.split("_")[-1] if "_" in name else "0"; print(gpus)' )
      TORCH_HOME: "/var/tmp/torch"
      PIP_CACHE_DIR: "/var/tmp/pip"
      PYTHONHASHSEED: "0"
      NCCL_DEBUG: "INFO"
      CI: "true"
    container:
      image: "pytorchlightning/lightning-thunder:$(docker-image)"
      options: "--gpus=all --shm-size=16g -v /var/tmp:/var/tmp"
    workspace:
      clean: all
    steps:
      - bash: |
          echo $(DEVICES)
          lspci | egrep 'VGA|3D'
          dpkg-query -W -f='${Package} ${Version}\n' libnccl2 libnccl-dev
          whereis nvidia
          nvidia-smi
          which python && which pip
          python --version
          pip --version
          pip list
          echo "##vso[task.setvariable variable=CUDA_VISIBLE_DEVICES]$(DEVICES)"
        displayName: "Image info & NVIDIA"

      - bash: |
          set -ex
          # drop pt from requirements so not to interfere with the existing one
          bash scripts/remove-torch-lines.sh requirements/base.txt
          cat requirements/base.txt

          # double check on test requirements
          pip install -U -r requirements/base.txt -r requirements/test.txt

          # https://docs.codecov.com/docs/codecov-uploader
          curl -Os https://uploader.codecov.io/latest/linux/codecov
          chmod +x codecov

          # install this package
          pip install -e .
        displayName: "Install package & ..."

      - bash: bash scripts/sanity-check.sh
        displayName: "Sanity check / details"

      - bash: |
          set -ex
          export CUDA_LAUNCH_BLOCKING=1
          coverage run --source thunder -m \
            pytest thunder/tests/ \
              -m "not standalone" \
              -v --datefmt="%Y%m%d-%H:%M:%S.%f" \
              --random-order-seed=42 \
              --durations=250 \
              --timeout=240 \
              --numprocesses=9 \
              --gpu-mem-limit=2 \
              --ignore=thunder/tests/distributed --ignore=thunder/tests/test_networks.py \
              --ignore=thunder/tests/test_ops.py --ignore=thunder/tests/test_grad.py
          # compile coverage results
          python -m coverage report
          python -m coverage xml
          # upload to codecov
          ./codecov --token=$(CODECOV_TOKEN) --commit=$(Build.SourceVersion) \
            --flags=gpu,pytest,regular --name="GPU-coverage" --env=linux,azure
        condition: eq(variables['testing'], 'main')
        timeoutInMinutes: "40"
        displayName: "Testing: main"

      - bash: |
          set -ex
          # these test need to run in single thread as they occurs with CUDA OOM
          coverage run --source thunder -m \
             pytest \
               thunder/tests/test_networks.py \
               -m "not standalone" \
               -v --durations=0 \
               --random-order-seed=42 \
               --gpu-mem-limit=6 \
               --numprocesses=3
          # compile coverage results
          python -m coverage report
          python -m coverage xml
          # upload to codecov
          ./codecov --token=$(CODECOV_TOKEN) --commit=$(Build.SourceVersion) \
            --flags=gpu,pytest,networks --name="GPU-coverage" --env=linux,azure
        condition: eq(variables['testing'], 'main')
        timeoutInMinutes: "15"
        displayName: "Testing: networks"

      - bash: |
          set -ex
          coverage run --source thunder -m \
            pytest thunder/tests/test_ops.py \
              -m "not standalone" \
              -v --datefmt="%Y%m%d-%H:%M:%S.%f" \
              --random-order-seed=42 \
              --durations=250 \
              --timeout=240 \
              --gpu-mem-limit=2 \
              --numprocesses=9
          # compile coverage results
          python -m coverage report
          python -m coverage xml
          # upload to codecov
          ./codecov --token=$(CODECOV_TOKEN) --commit=$(Build.SourceVersion) \
            --flags=gpu,pytest,regular --name="GPU-coverage" --env=linux,azure
        condition: eq(variables['testing'], 'ops')
        env:
          CUDA_LAUNCH_BLOCKING: "1"
        timeoutInMinutes: "40"
        displayName: "Testing: ops"

      - bash: |
          set -ex
          coverage run --source thunder -m \
            pytest thunder/tests/test_grad.py \
              -m "not standalone" \
              -v --datefmt="%Y%m%d-%H:%M:%S.%f" \
              --random-order-seed=42 \
              --durations=250 \
              --timeout=240 \
              --gpu-mem-limit=2 \
              --numprocesses=9
          # compile coverage results
          python -m coverage report
          python -m coverage xml
          # upload to codecov
          ./codecov --token=$(CODECOV_TOKEN) --commit=$(Build.SourceVersion) \
            --flags=gpu,pytest,regular --name="GPU-coverage" --env=linux,azure
        condition: eq(variables['testing'], 'grads')
        env:
          CUDA_LAUNCH_BLOCKING: "1"
        timeoutInMinutes: "35"
        displayName: "Testing: grads"

      - bash: |
          set -ex
          # run all found tests in given past as standalone
          pytest \
             thunder/tests/distributed \
             -v --durations=0 \
             --random-order-seed=42
          # compile coverage results
          # TODO: collect and merge reports
          #  python -m coverage report
          #  python -m coverage xml
          #  # upload to codecov
          #  ./codecov --token=$(CODECOV_TOKEN) --commit=$(Build.SourceVersion) \
          #    --flags=gpu,pytest,distributed --name="GPU-coverage" --env=linux,azure
        condition: eq(variables['testing'], 'distributed')
        timeoutInMinutes: "30"
        displayName: "Testing: distributed"
