From 228b3d6f7c7e23024917cff525378c03ce13a517 Mon Sep 17 00:00:00 2001
From: Morgan Funtowicz <funtowiczmo@gmail.com>
Date: Tue, 7 Jan 2025 11:48:42 +0100
Subject: [PATCH] misc(ci): make the wf callable for reuse

---
 .github/workflows/build_trtllm.yaml | 56 +++++++++--------------------
 .github/workflows/trtllm_tests.yaml | 43 ++++++++++++++++++++++
 2 files changed, 60 insertions(+), 39 deletions(-)
 create mode 100644 .github/workflows/trtllm_tests.yaml

diff --git a/.github/workflows/build_trtllm.yaml b/.github/workflows/build_trtllm.yaml
index ff8d06b6..1c96ea3e 100644
--- a/.github/workflows/build_trtllm.yaml
+++ b/.github/workflows/build_trtllm.yaml
@@ -1,24 +1,19 @@
 name: Build TensorRT-LLM
-
 on:
-  push:
-    branches:
-      - 'main'
-      - 'trtllm/ci'
-    tags:
-      - 'v*'
-  pull_request:
-    paths:
-      - "backends/trtllm"
-      - "server/**"
-      - "proto/**"
-      - "router/**"
-      - "launcher/**"
-      - "Cargo.lock"
-      - "rust-toolchain.toml"
-      - ".github/workflows/build_trtllm.yaml"
-    branches:
-      - "main"
+  workflow_call:
+    inputs:
+      runs-on:
+        type: string
+        description: "Which instance type to use to run the workflow"
+        required: true
+
+    outputs:
+      docker_image:
+        description: "Reference to the Docker Image build by this workflow"
+        value: ${{ jobs.build-and-push.outputs.docker_image }}
+      label:
+        description: "Label generated for this build"
+        value: ${{ jobs.build-and-push.outputs.label }}
 
 permissions:
   contents: read # Required to check out repository.
@@ -26,14 +21,13 @@ permissions:
 
 jobs:
   build-and-push:
+    concurrency:
+      group: ${{ github.workflow }}-${{ github.job }}-${{ github.head_ref || github.run_id }}
     outputs:
       docker_image: ${{ steps.final.outputs.docker_image }}
       label: ${ steps.final.outputs.label }}
-    concurrency:
-      group: ${{ github.workflow }}-${{ github.job }}-build-and-push-${{ github.head_ref || github.run_id }}
-      cancel-in-progress: true
     runs-on:
-      group: aws-highmemory-64-plus-priv
+      group: ${{ inputs.runs-on }}
     steps:
       - uses: actions/checkout@v4
 
@@ -134,21 +128,5 @@ jobs:
         run: |
           echo "docker_image=registry.internal.huggingface.tech/api-inference/community/text-generation-inference/tensorrt-llm:sha-${{ env.GITHUB_SHA_SHORT }}${{ env.LABEL }}" >> "$GITHUB_OUTPUT"
           echo "label=${{ env.LABEL }}" >> "$GITHUB_OUTPUT"
-  tests:
-    needs: build-and-push
-    concurrency:
-      group: ${{ github.workflow }}-${{ github.job }}-trtllm-${{ github.head_ref || github.run_id }}
-      cancel-in-progress: true
-    runs-on:
-      group: aws-g6-12xl-plus-priv-cache
-    container:
-      image: ${{ needs.build-and-push.outputs.docker_image }}
-      credentials:
-        username: ${{ secrets.REGISTRY_USERNAME }}
-        password: ${{ secrets.REGISTRY_PASSWORD }}
-      options: --gpus all
 
-    steps:
-      - name: Run C++/CUDA tests
-        run: /usr/local/tgi/bin/tgi_trtllm_backend_tests
 
diff --git a/.github/workflows/trtllm_tests.yaml b/.github/workflows/trtllm_tests.yaml
new file mode 100644
index 00000000..a49695ad
--- /dev/null
+++ b/.github/workflows/trtllm_tests.yaml
@@ -0,0 +1,43 @@
+on:
+  push:
+    branches:
+      - 'main'
+      - 'trtllm/ci'
+    tags:
+      - 'v*'
+  pull_request:
+    paths:
+      - "backends/trtllm"
+      - "server/**"
+      - "proto/**"
+      - "router/**"
+      - "launcher/**"
+      - "Cargo.lock"
+      - "rust-toolchain.toml"
+      - ".github/workflows/build_trtllm.yaml"
+      - ".github/workflows/trtllm_tests.yaml"
+    branches:
+      - "main"
+
+jobs:
+  build-and-push:
+    permissions:
+      contents: write
+      packages: write
+    uses: ./.github/workflows/build_trtllm.yaml
+    with:
+      runs-on: aws-highmemory-64-plus-priv
+
+  run-tests:
+    needs: build-and-push
+    runs-on:
+      container:
+        image: ${{ needs.build-and-push.outputs.docker_image }}
+        credentials:
+          username: ${{ secrets.REGISTRY_USERNAME }}
+          password: ${{ secrets.REGISTRY_PASSWORD }}
+        options: --gpus all
+
+    steps:
+      - name: Run C++/CUDA tests
+        run: /usr/local/tgi/bin/tgi_trtllm_backend_tests
\ No newline at end of file