From 5429a116d413006cb8af4e18072d434d7087dac5 Mon Sep 17 00:00:00 2001
From: Morgan Funtowicz <funtowiczmo@gmail.com>
Date: Mon, 6 Jan 2025 15:23:11 +0100
Subject: [PATCH] misc(ci): gogogo

---
 .github/workflows/build_trtllm.yaml | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/.github/workflows/build_trtllm.yaml b/.github/workflows/build_trtllm.yaml
index eeebe85f..ebfe068c 100644
--- a/.github/workflows/build_trtllm.yaml
+++ b/.github/workflows/build_trtllm.yaml
@@ -126,20 +126,16 @@ jobs:
       - name: Final
         id: final
         run: |
-          echo "docker_image=registry.internal.huggingface.tech/api-inference/community/text-generation-inference:sha-${{ env.GITHUB_SHA_SHORT}}${{ env.LABEL }}" >> "$GITHUB_OUTPUT"
-          echo "docker_devices=${{ env.DOCKER_DEVICES }}" >> "$GITHUB_OUTPUT"
-          echo "docker_volume=${{ env.DOCKER_VOLUME }}" >> "$GITHUB_OUTPUT"
-          echo "runs_on=${{ env.RUNS_ON }}" >> "$GITHUB_OUTPUT"
-          echo "label=${{ env.LABEL }}" >> "$GITHUB_OUTPUT"
-          echo "extra_pytest=${{ env.EXTRA_PYTEST }}" >> "$GITHUB_OUTPUT"
+          echo "docker_image=registry.internal.huggingface.tech/api-inference/community/text-generation-inference:sha-${{ env.GITHUB_SHA_SHORT }}${{ env.LABEL }}" >> "$GITHUB_OUTPUT"
 
   tests:
     needs: build-and-push
     runs-on:
       group: aws-g6-12xl-plus-priv-cache
+    env:
+      DOCKER_IMAGE: ${{ needs.build-and-push.outputs.docker_image }}
     container:
-      image:
-        registry.internal.huggingface.tech/api-inference/community/text-generation-inference/tensorrt-llm:${{ needs.build-and-push.outputs.docker_image }}
+      image: ${DOCKER_IMAGE}
       credentials:
         username: ${{ secrets.REGISTRY_USERNAME }}
         password: ${{ secrets.REGISTRY_PASSWORD }}