# Copyright 2024 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

name: PkgCI Test PJRT plugin
on:
  workflow_call:
    inputs:
      artifact_run_id:
        type: string
        default: ""
      write-caches:
        required: true
        type: string
  workflow_dispatch:
    inputs:
      artifact_run_id:
        type: string
        default: ""
      write-caches:
        required: true
        type: string

jobs:
  build_and_test:
    strategy:
      matrix:
        include:
          - runner: ubuntu-24.04
            pjrt_platform: cpu
          - runner: ubuntu-24.04
            pjrt_platform: cuda
          # TODO: enable these AMD runners
          # - runner: nodai-amdgpu-w7900-x86-64
          #   pjrt_platform: rocm
          # - runner: nodai-amdgpu-w7900-x86-64
          #   pjrt_platform: vulkan
    name: Build and test
    runs-on: ${{ matrix.runner }}
    env:
      PACKAGE_DOWNLOAD_DIR: ${{ github.workspace }}/.packages
      VENV_DIR: ${{ github.workspace }}/.venv
      GH_TOKEN: ${{ github.token }}
    steps:
      - name: Checking out repository
        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
        with:
          submodules: true
      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
        with:
          # Must match the subset of versions built in pkgci_build_packages.
          python-version: "3.11"
      - uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0 # v5.0.0
        if: ${{ inputs.artifact_run_id == '' }}
        with:
          name: linux_x86_64_release_packages
          path: ${{ env.PACKAGE_DOWNLOAD_DIR }}
      - name: Setup base venv
        run: |
          ./build_tools/pkgci/setup_venv.py ${VENV_DIR} \
            --artifact-path=${PACKAGE_DOWNLOAD_DIR} \
            --fetch-gh-workflow=${{ inputs.artifact_run_id }}
      - name: ccache
        uses: hendrikmuhs/ccache-action@63069e3931dedbf3b63792097479563182fe70d1 # v1.2.18
        with:
          key: ${{ github.job }}-${{ matrix.pjrt_platform }}
          save: ${{ inputs.write-caches == 1 }}
      - name: Build and install
        run: |
          # install editable into venv
          source ${VENV_DIR}/bin/activate
          export CMAKE_C_COMPILER_LAUNCHER=ccache
          export CMAKE_CXX_COMPILER_LAUNCHER=ccache
          python -m pip install -v --no-deps -e integrations/pjrt/python_packages/iree_${{ matrix.pjrt_platform }}_plugin
          # install
          python -m pip install jax==0.6.1
      # TODO: We can only run on CPU for now, as the CUDA runner
      # is not available yet, refer to #18814.
      - name: Run tests
        if: ${{ matrix.pjrt_platform == 'cpu' }}
        run: |
          source ${VENV_DIR}/bin/activate
          build_tools/testing/run_jax_tests.sh ${{ matrix.pjrt_platform }}