ci: cache the embed model for the k8s smoke tests

.github/workflows/ci.yaml

@@ -149,6 +149,7 @@ jobs:
           image_registry="ghcr.io/${GITHUB_REPOSITORY,,}"
           source_sha="${SOURCE_SHA:-$GITHUB_SHA}"
           bake_tag="${INPUT_IMAGE_TAG:-$source_sha}"
+          fastembed_cache_dir="${RUNNER_TEMP}/fastembed-cache"
           publish_images="false"
 
           if [ "$GITHUB_EVENT_NAME" = "push" ] && [ "$GITHUB_REF" = "refs/heads/main" ]; then
@@ -173,6 +174,10 @@ jobs:
             printf 'CACHE_REGISTRY=%s\n' "$image_registry"
             printf 'BAKE_TAG=%s\n' "$bake_tag"
             printf 'CI_COMMIT_SHA=%s\n' "$source_sha"
+            printf 'FASTEMBED_CACHE_CONTEXT=%s\n' "$fastembed_cache_dir"
+            printf 'FASTEMBED_CACHE_DIR=%s\n' "$fastembed_cache_dir"
+            printf 'FASTEMBED_MODEL_REPO=%s\n' "qdrant/all-MiniLM-L6-v2-onnx"
+            printf 'FASTEMBED_MODEL_REVISION=%s\n' "main"
             printf 'PUBLISH_IMAGES=%s\n' "$publish_images"
           } >> "$GITHUB_ENV"
           {
@@ -181,6 +186,39 @@ jobs:
             printf 'publish_images=%s\n' "$publish_images"
           } >> "$GITHUB_OUTPUT"
 
+      - name: Prepare fastembed model cache directory
+        shell: bash
+        run: |
+          set -euo pipefail
+          mkdir -p "${FASTEMBED_CACHE_DIR}"
+          touch "${FASTEMBED_CACHE_DIR}/.cache-context"
+
+      - name: Restore fastembed model cache
+        id: fastembed-cache
+        uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
+        with:
+          path: ${{ runner.temp }}/fastembed-cache
+          key: fastembed-qdrant-all-MiniLM-L6-v2-onnx-main-v1
+
+      - name: Install uv for fastembed model cache
+        if: steps.fastembed-cache.outputs.cache-hit != 'true'
+        uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
+        with:
+          version-file: pyproject.toml
+
+      - name: Populate fastembed model cache
+        if: steps.fastembed-cache.outputs.cache-hit != 'true'
+        shell: bash
+        run: |
+          set -euo pipefail
+          HF_HUB_DISABLE_UPDATE_CHECK=1 uv run --no-project --with 'huggingface-hub>=1.0.1,<2.0.0' hf download \
+            "${FASTEMBED_MODEL_REPO}" \
+            --revision "${FASTEMBED_MODEL_REVISION}" \
+            --cache-dir "${FASTEMBED_CACHE_DIR}" \
+            --format quiet
+          find "${FASTEMBED_CACHE_DIR}/models--qdrant--all-MiniLM-L6-v2-onnx/snapshots" \
+            -mindepth 1 -maxdepth 1 -type d -print -quit | grep -q .
+
       - name: Log in to GHCR
         if: env.PUBLISH_IMAGES == 'true'
         shell: bash
@@ -197,17 +235,21 @@ jobs:
       - name: Build CPU images
         if: env.PUBLISH_IMAGES != 'true'
         shell: bash
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: make docker-load TARGET=docker-cpu
 
       - name: Build and publish CPU images
         if: env.PUBLISH_IMAGES == 'true'
         shell: bash
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: make docker-push TARGET=docker-cpu
 
+      - name: Save fastembed model cache
+        if: steps.fastembed-cache.outputs.cache-hit != 'true'
+        uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
+        continue-on-error: true
+        with:
+          path: ${{ runner.temp }}/fastembed-cache
+          key: fastembed-qdrant-all-MiniLM-L6-v2-onnx-main-v1
+
   kind-cpu-smoke:
     name: Kind CPU smoke test
     needs: [changes, build-cpu-smoke-images]
@@ -308,7 +350,6 @@ jobs:
       - name: Start kind cluster
         shell: bash
         env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN }}
           NGC_API_KEY: not-used-for-ghcr-cpu-smoke
         run: bash "${K8S_E2E_SCRIPTS}/setup_local_kind_cpu.sh"
 

docker-bake.hcl

@@ -103,8 +103,8 @@ variable "BUILD_ARCH" {
   default = ""
 }
 
-variable "HF_TOKEN" {
-  default = ""
+variable "FASTEMBED_CACHE_CONTEXT" {
+  default = "docker/fastembed-cache-empty"
 }
 
 variable "CUDA_VERSION" {
@@ -203,16 +203,6 @@ function "image_output" {
   result = ["type=image,compression=zstd,force-compression=true"]
 }
 
-function "maybe_hf_token_secret" {
-  params = []
-  result = notequal(HF_TOKEN, "") ? [
-    {
-      type = "env"
-      id   = "HF_TOKEN"
-    }
-  ] : []
-}
-
 function "maybe_registry_cache_from" {
   params = [name]
   result = [
@@ -463,6 +453,7 @@ target "nmp-api-docker" {
     nmp-jobs-launcher         = "target:nmp-jobs-launcher"
     nmp-studio-ui             = "target:nmp-studio-ui"
     policy-wasm-artifacts     = "target:root-policy-wasm-artifacts"
+    fastembed-cache           = FASTEMBED_CACHE_CONTEXT
   }
   args = {
     NMP_PLATFORM_VERSION = notequal(BAKE_TAG, "") ? BAKE_TAG : "dev"
@@ -472,7 +463,6 @@ target "nmp-api-docker" {
   cache-from = maybe_registry_cache_from("nmp-api")
   tags       = sha_and_maybe_latest_tags("nmp-api")
   output     = image_output()
-  secret     = maybe_hf_token_secret()
   platforms  = get_platforms()
 }
 

docker/Dockerfile.nmp-api

@@ -30,9 +30,12 @@ RUN sh /app/script/install_duckdb_extensions.sh
 # Download embedding model used by `nemoguardrails` library.
 # Uses snapshot_download directly to avoid importing onnxruntime, which crashes under QEMU (ARM64 cross-compilation).
 # Stored in /tmp/fastembed_cache (fastembed's default cache location) and copied to the runtime image below.
-RUN --mount=type=secret,id=HF_TOKEN,required=false \
-    if [ -s /run/secrets/HF_TOKEN ]; then export HF_TOKEN="$(cat /run/secrets/HF_TOKEN)"; fi; \
-    uv run --no-sync python -c 'import os, tempfile; from huggingface_hub import snapshot_download; snapshot_download(repo_id="qdrant/all-MiniLM-L6-v2-onnx", cache_dir=os.path.join(tempfile.gettempdir(), "fastembed_cache"))'
+COPY --from=fastembed-cache . /tmp/fastembed_cache/
+RUN if find /tmp/fastembed_cache/models--qdrant--all-MiniLM-L6-v2-onnx/snapshots -mindepth 1 -maxdepth 1 -type d -print -quit 2>/dev/null | grep -q .; then \
+        echo "Using cached qdrant/all-MiniLM-L6-v2-onnx fastembed model"; \
+    else \
+        uv run --no-sync python -c 'import os, tempfile; from huggingface_hub import snapshot_download; snapshot_download(repo_id="qdrant/all-MiniLM-L6-v2-onnx", cache_dir=os.path.join(tempfile.gettempdir(), "fastembed_cache"))'; \
+    fi
 
 FROM ${NMP_PYTHON_BASE} AS runtime
 ARG USERNAME=nvs