group: aarch64 tests
sort_key: "~aarch64"
steps:
  # block on premerge and microcheck
  - block: "Run Linux aarch64 tests"
    if: build.env("BUILDKITE_PIPELINE_ID") == "0189942e-0876-4b8f-80a4-617f988ec59b" || build.env("BUILDKITE_PIPELINE_ID") == "018f4f1e-1b73-4906-9802-92422e3badaa"

  - name: forge-aarch64
    wanda: ci/docker/forge.aarch64.wanda.yaml
    instance_type: builder-arm64

  - name: manylinux-aarch64
    wanda: ci/docker/manylinux.aarch64.wanda.yaml
    instance_type: builder-arm64

  - name: raycudabase-aarch64
    label: "wanda: ray.py{{matrix.python}}.cu{{matrix.cuda}}.base (aarch64)"
    tags:
      - python_dependencies
      - docker
      - core_cpp
    wanda: ci/docker/ray.cuda.base.aarch64.wanda.yaml
    matrix:
      setup:
        python:
          - "3.9"
          - "3.10"
          - "3.11"
          - "3.12"
        cuda:
          - "11.7.1-cudnn8"
          - "11.8.0-cudnn8"
          - "12.1.1-cudnn8"
          - "12.3.2-cudnn9"
    instance_type: builder-arm64
    env:
      PYTHON_VERSION: "{{matrix.python}}"
      CUDA_VERSION: "{{matrix.cuda}}"

  - name: raycpubase-aarch64
    label: "wanda: ray.py{{matrix}}.cpu.base (aarch64)"
    tags:
      - python_dependencies
      - docker
      - core_cpp
    wanda: ci/docker/ray.cpu.base.aarch64.wanda.yaml
    matrix:
      - "3.9"
      - "3.10"
      - "3.11"
      - "3.12"
    instance_type: builder-arm64
    env:
      PYTHON_VERSION: "{{matrix}}"

  - label: ":tapioca: build: wheel {{matrix}} (aarch64)"
    tags:
      - linux_wheels
      - oss
    instance_type: medium-arm64
    commands:
      - bazel run //ci/ray_ci:build_in_docker -- wheel --python-version {{matrix}} --architecture aarch64 --upload
    matrix:
      - "3.9"
      - "3.10"
      - "3.11"
      - "3.12"
    depends_on:
      - manylinux-aarch64
      - forge-aarch64
    job_env: forge-aarch64

  - label: ":tapioca: build: ray py{{matrix}} docker (aarch64)"
    tags:
      - python_dependencies
      - docker
      - core_cpp
      - oss
    instance_type: medium-arm64
    commands:
      - bazel run //ci/ray_ci:build_in_docker -- docker --python-version {{matrix}}
        --platform cu11.7.1-cudnn8 --platform cu11.8.0-cudnn8
        --platform cu12.1.1-cudnn8 --platform cu12.3.2-cudnn9
        --platform cpu
        --image-type ray --architecture aarch64 --upload
    depends_on:
      - manylinux-aarch64
      - forge-aarch64
      - raycudabase-aarch64
      - raycpubase-aarch64
    job_env: forge-aarch64
    matrix:
      - "3.9"
      - "3.10"
      - "3.11"
      - "3.12"

  - label: ":ray: core: wheel-aarch64 tests"
    tags: linux_wheels
    instance_type: medium-arm64
    commands:
      - bazel run //ci/ray_ci:test_in_docker -- //python/ray/tests/... //doc/... core 
        --build-type wheel-aarch64
        --build-name oss-ci-base_build-aarch64
        --parallelism-per-worker 3
        --only-tags post_wheel_build
        --test-env=RAY_CI_POST_WHEEL_TESTS=True
    depends_on:
      - manylinux-aarch64
      - oss-ci-base_build-aarch64
      - forge-aarch64
    job_env: forge-aarch64

  - label: ":ray-serve: serve: wheel-aarch64 tests"
    tags: linux_wheels
    instance_type: medium-arm64
    commands:
      - bazel run //ci/ray_ci:test_in_docker -- //python/ray/serve/... //doc/... serve 
        --build-type wheel-aarch64
        --build-name oss-ci-base_build-aarch64
        --parallelism-per-worker 3
        --only-tags post_wheel_build
        --test-env=RAY_CI_POST_WHEEL_TESTS=True
    depends_on:
      - manylinux-aarch64
      - oss-ci-base_build-aarch64
      - forge-aarch64
    job_env: forge-aarch64