trigger:
  tags:
    include: ["*"]
  branches:
    include: ["main"]
  paths:
    include:
      - ".azure/docker-build.yml"
      - "dockers/**"
      - "requirements.txt"
      - "requirements/*.txt"
      - "setup.py"
    exclude:
      - "*.md"
      - "**/*.md"

pr:
  branches:
    include: ["*"]
  paths:
    include:
      - ".azure/docker-build.yml"
      - "dockers/**"
      - "requirements.txt"
      - "requirements/*.txt"
      - "setup.py"
    exclude:
      - "*.md"
      - "**/*.md"

schedules:
  - cron: "0 */2 * * *"
    displayName: rebuild dockers for CI every 2 hours
    branches:
      include: ["main"]

jobs:
  - job: build_push
    strategy:
      matrix:
        "cuda 12.8 | torch 2.8.0 | cudnn FE v1.15.0":
          { CUDA_VERSION: "12.8.1", TORCH_VERSION: "2.8.0", TRITON_VERSION: "3.4.0", CUDNN_FRONTEND_VERSION: "1.15.0" }
        "cuda 12.8 | torch nightly | cudnn FE v1.15.0":
          { CUDA_VERSION: "12.8.1", TORCH_VERSION: "main", TORCH_INSTALL: "source", CUDNN_FRONTEND_VERSION: "1.15.0" }
        #'cuda 12.1': # this version - '8.9.5.29-1+cuda12.1' for 'libcudnn8' was not found
    # how much time to give 'run always even if cancelled tasks' before stopping them
    cancelTimeoutInMinutes: "2"
    timeoutInMinutes: "95"
    variables:
      UBUNTU_VERSION: "24.04"
      PYTHON_VERSION: "3.12"
      imageRepository: "pytorchlightning/lightning-thunder"
      imageTag: "ubuntu$(UBUNTU_VERSION)-cuda$(CUDA_VERSION)-cudnn-fe$(CUDNN_FRONTEND_VERSION)-py$(PYTHON_VERSION)-pt_${TORCH_VERSION/v/}"
    pool: "lit-rtx-3090"
    workspace:
      clean: all
    steps:
      - bash: |
          set -e
          echo $imageTag
          nvidia-smi
          docker image build \
            -t $(imageRepository):$(imageTag) \
            -f "dockers/ubuntu-cuda/Dockerfile" \
            --build-arg UBUNTU_VERSION="$(UBUNTU_VERSION)" \
            --build-arg CUDA_VERSION="$(CUDA_VERSION)" \
            --build-arg CUDNN_FRONTEND_VERSION="v$(CUDNN_FRONTEND_VERSION)" \
            --build-arg PYTHON_VERSION="$(PYTHON_VERSION)" \
            --build-arg TORCH_VERSION="$(TORCH_VERSION)" \
            --build-arg TORCH_INSTALL="$(TORCH_INSTALL)" \
            --build-arg TRITON_VERSION="$(TRITON_VERSION)" \
            . --no-cache
        timeoutInMinutes: "95"
        displayName: "Build base image"

      - bash: |
          docker image build \
            -t $(imageRepository):$(imageTag)-apex \
            -f "dockers/with-apex/Dockerfile" \
            --build-arg BASE_IMAGE_TAG="$(imageTag)" \
            . --no-cache
        timeoutInMinutes: "25"
        displayName: "Build Apex image"

      - bash: |
          docker image build \
            -t $(imageRepository):$(imageTag)-dev \
            -f "dockers/with-dev/Dockerfile" \
            --build-arg BASE_IMAGE_TAG="$(imageTag)-apex" \
            . --no-cache
        timeoutInMinutes: "25"
        displayName: "Build Dev image"

      - bash: |
          docker image ls | grep $(imageRepository)
          # drop pt from requirements so not to interfere with the existing one
          bash scripts/remove-torch-lines.sh requirements/base.txt
          docker run --rm --gpus=all -v .:/workspace $(imageRepository):$(imageTag)-dev \
            bash -c  "cd /workspace && ls -lh . && \
                      pip install -q . && \
                      bash scripts/sanity-check.sh"
        timeoutInMinutes: "5"
        displayName: "Sanity check"

      - bash: |
          set -e
          echo $(imageRepository):$(imageTag)
          echo $(DOCKERHUB_PAT) | docker login --username $(DOCKERHUB_USER) --password-stdin
          docker push $(imageRepository):$(imageTag)-dev
        condition: ne(variables['Build.Reason'], 'PullRequest')
        timeoutInMinutes: "35"
        displayName: "Push base image"
