From dbe261a6b482352dc0564ff5ef9fd6bd653a1ec2 Mon Sep 17 00:00:00 2001
From: Morgan Funtowicz <funtowiczmo@gmail.com>
Date: Tue, 25 Jun 2024 15:31:47 +0200
Subject: [PATCH] feat(tests): Move to k8s cluster for CI/CD

---
 .github/workflows/pr_functional_tests.yml     | 69 -------------------
 ...{pr_integration_tests.yml => pr_tests.yml} |  6 +-
 src/optimum/nvidia/__init__.py                |  2 +-
 3 files changed, 4 insertions(+), 73 deletions(-)
 delete mode 100644 .github/workflows/pr_functional_tests.yml
 rename .github/workflows/{pr_integration_tests.yml => pr_tests.yml} (92%)

diff --git a/.github/workflows/pr_functional_tests.yml b/.github/workflows/pr_functional_tests.yml
deleted file mode 100644
index f13e77ad..00000000
--- a/.github/workflows/pr_functional_tests.yml
+++ /dev/null
@@ -1,69 +0,0 @@
-name: GPU-Enabled Functional Test on PRs
-
-on:
-  pull_request:
-    branches:
-      - main
-  push:
-    branches:
-      - ci-*
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
-  cancel-in-progress: true
-
-env:
-  OPTIMUM_NVIDIA_IS_CI: ON
-  RUN_CPU_ONLY: ON
-
-jobs:
-  run_fast_tests:
-    strategy:
-      fail-fast: false
-      matrix:
-        config:
-          - name: GPU-Enabled Optimum-Nvidia Functional Test Suite
-            runner: [ci, nvidia-gpu]
-            image: nvidia/cuda:12.1.0-devel-ubuntu22.04
-            report: dev
-        cuda_arch: [sm_86, sm_89]
-
-    name: ${{ matrix.config.name }}
-    runs-on: [ci, nvidia-gpu, multi-gpu, "${{ matrix.cuda_arch }}"]
-
-    container:
-      image: ${{ matrix.config.image }}
-      options: --mount type=tmpfs,destination=/tmp --shm-size 64gb --gpus all --ipc host -v /mnt/hf_cache:/mnt/cache/
-      env:
-        HF_TOKEN: ${{ secrets.OPTIMUM_NVIDIA_HUB_READ_TOKEN }}
-
-    defaults:
-      run:
-        shell: bash
-
-    steps:
-      - uses: actions/setup-python@v5
-        with:
-            python-version: '3.10'
-
-      - name: Checkout optimum-nvidia
-        uses: actions/checkout@v4
-        with:
-          fetch-depth: 1
-
-      - name: Install dependencies
-        run: |
-          apt update && apt install -y openmpi-bin libopenmpi-dev git
-          python3 -m pip install --upgrade -e .[quality,tests] --pre --extra-index-url https://pypi.nvidia.com 
-
-      - name: Run nvidia-smi
-        run: |
-          nvidia-smi
-
-      - name: Print TensorRT-LLM version
-        run: |
-          python -c "from tensorrt_llm import __version__; print(__version__)"
-
-      - name: Run optimum-nvidia functional test-suite
-        run: |
-          pytest -n 4 -s -vvvvv -p no:warnings -o log_cli=true --ignore=tests/integration/ tests/
diff --git a/.github/workflows/pr_integration_tests.yml b/.github/workflows/pr_tests.yml
similarity index 92%
rename from .github/workflows/pr_integration_tests.yml
rename to .github/workflows/pr_tests.yml
index 702f5051..de19538a 100644
--- a/.github/workflows/pr_integration_tests.yml
+++ b/.github/workflows/pr_tests.yml
@@ -14,7 +14,6 @@ concurrency:
 
 env:
   OPTIMUM_NVIDIA_IS_CI: ON
-  RUN_CPU_ONLY: OFF
   RUN_NIGHTLY: OFF
   RUN_SLOW: ON
   HF_HUB_ENABLE_HF_TRANSFER: ON
@@ -27,10 +26,11 @@ jobs:
         config:
           - name: GPU-enabled Optimum-Nvidia Integration Test Suite
             image: nvidia/cuda:12.1.0-devel-ubuntu22.04
-        cuda_arch: [sm_86, sm_89]
+        gpu_target: ["nvidia-multi-gpu-l4-runners", "nvidia-multi-gpu-a10-runners"]
 
     name: ${{ matrix.config.name }}
-    runs-on: [ci, nvidia-gpu, multi-gpu, "${{ matrix.cuda_arch }}"]
+    runs-on:
+      group: "${{matrix.gpu_target}}"
 
     container:
       image: ${{ matrix.config.image }}
diff --git a/src/optimum/nvidia/__init__.py b/src/optimum/nvidia/__init__.py
index ef751b18..2357843b 100644
--- a/src/optimum/nvidia/__init__.py
+++ b/src/optimum/nvidia/__init__.py
@@ -21,5 +21,5 @@
 from .models import AutoModelForCausalLM
 from .optimizations import IntoModelOptQuantizeConfig
 
-# from .pipelines import pipeline
+from .pipelines import pipeline
 from .version import VERSION, __version__