Add GPU embedding preflight and tuning

docker/README.md

@@ -83,6 +83,12 @@ CPU and GPU services are separate so they can be benchmarked against the same
 model and workload. Model files are cached in persistent, model-specific
 volumes.
 
+The CPU services are defined in `docker-compose-local-embeddings.yaml`; the
+GPU services are defined in the separately included
+`docker-compose-local-embeddings-gpu.yaml`. This keeps the GPU runtime setup
+isolated while the root Compose file still exposes one consistent service
+matrix.
+
 | Model | Vector dimensions | CPU service / host port | GPU service / host port |
 |-------|-------------------:|-------------------------|-------------------------|
 | Qwen3-Embedding-0.6B | 1024 | qwen3-embedding-06b-cpu / 8101 | qwen3-embedding-06b-gpu / 8201 |
@@ -108,6 +114,13 @@ VERSION=dev docker compose -f docker-compose.yaml --profile embeddings-both up -
   qwen3-embedding-06b-gpu qwen3-embedding-4b-gpu qwen3-embedding-8b-gpu
 ```
 
+Before starting either GPU profile, verify that the host can see an NVIDIA
+device and driver:
+
+```bash
+python3 docker/benchmarks/gpu_preflight.py --json
+```
+
 GPU services require the NVIDIA Container Toolkit and a compatible NVIDIA
 driver. The default CUDA image targets the TEI CUDA 1.9 runtime; set
 QWEN3_TEI_GPU_IMAGE when an architecture-specific image is needed.
@@ -117,6 +130,12 @@ service in the matrix and may exceed available GPU memory if all six are
 launched together. For a fair comparison, start one profile at a time or
 benchmark endpoints sequentially.
 
+GPU services have independent tuning variables so a GPU run does not change
+the CPU run's request limits: `QWEN3_GPU_MAX_BATCH_TOKENS` (default `8192`),
+`QWEN3_GPU_MAX_CLIENT_BATCH_SIZE` (default `32`),
+`QWEN3_GPU_MAX_CONCURRENT_REQUESTS` (default `4`), and
+`QWEN3_GPU_TOKENIZATION_WORKERS` (default `4`).
+
 Each service exposes an OpenAI-compatible endpoint. From an Unstract
 container, use the internal URL; from the host, use the localhost URL:
 

docker/benchmarks/gpu_preflight.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python3
+"""Check host NVIDIA visibility before launching GPU embedding services."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import shutil
+import subprocess
+import sys
+from typing import Any
+
+GPU_QUERY = "index,name,memory.total,driver_version"
+GPU_FIELDS = ("index", "name", "memory_total_mib", "driver_version")
+
+
+def query_gpus(nvidia_smi: str) -> tuple[list[dict[str, str]], str | None]:
+    """Return visible GPUs or a human-readable command error."""
+    try:
+        completed = subprocess.run(
+            [
+                nvidia_smi,
+                f"--query-gpu={GPU_QUERY}",
+                "--format=csv,noheader,nounits",
+            ],
+            capture_output=True,
+            check=False,
+            text=True,
+        )
+    except OSError as error:
+        return [], f"could not execute {nvidia_smi}: {error}"
+
+    if completed.returncode:
+        detail = completed.stderr.strip() or "no diagnostic was returned"
+        return [], f"{nvidia_smi} failed with exit code {completed.returncode}: {detail}"
+
+    gpus: list[dict[str, str]] = []
+    for row in csv.reader(line for line in completed.stdout.splitlines() if line.strip()):
+        if len(row) != len(GPU_FIELDS):
+            return [], f"unexpected {nvidia_smi} output row: {row!r}"
+        gpus.append(
+            {field: value.strip() for field, value in zip(GPU_FIELDS, row, strict=True)}
+        )
+    return gpus, None
+
+
+def build_report(
+    minimum_gpus: int,
+    minimum_memory_mib: int,
+) -> dict[str, Any]:
+    """Build a JSON-serializable preflight report."""
+    nvidia_smi = shutil.which("nvidia-smi")
+    report: dict[str, Any] = {
+        "available": False,
+        "minimum_gpus": minimum_gpus,
+        "minimum_memory_mib": minimum_memory_mib,
+        "nvidia_smi": nvidia_smi,
+        "gpus": [],
+    }
+    if nvidia_smi is None:
+        report["error"] = "nvidia-smi was not found on PATH"
+        return report
+
+    gpus, error = query_gpus(nvidia_smi)
+    report["gpus"] = gpus
+    if error:
+        report["error"] = error
+        return report
+    if len(gpus) < minimum_gpus:
+        report["error"] = (
+            f"found {len(gpus)} visible GPU(s), need at least {minimum_gpus}"
+        )
+        return report
+
+    if minimum_memory_mib:
+        low_memory = [
+            gpu["index"]
+            for gpu in gpus
+            if int(gpu["memory_total_mib"]) < minimum_memory_mib
+        ]
+        if low_memory:
+            report["error"] = (
+                f"GPU(s) {', '.join(low_memory)} have less than "
+                f"{minimum_memory_mib} MiB of memory"
+            )
+            return report
+
+    report["available"] = True
+    return report
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Check NVIDIA GPU availability for local TEI services."
+    )
+    parser.add_argument(
+        "--min-gpus",
+        type=int,
+        default=1,
+        help="Minimum number of visible GPUs required (default: 1).",
+    )
+    parser.add_argument(
+        "--min-memory-mib",
+        type=int,
+        default=0,
+        help="Minimum total memory required per GPU (default: disabled).",
+    )
+    parser.add_argument(
+        "--json",
+        action="store_true",
+        help="Print the complete machine-readable report.",
+    )
+    return parser.parse_args()
+
+
+def main() -> int:
+    """Run the GPU preflight and return a shell-friendly status."""
+    args = parse_args()
+    if args.min_gpus < 1:
+        raise SystemExit("--min-gpus must be at least 1")
+    if args.min_memory_mib < 0:
+        raise SystemExit("--min-memory-mib cannot be negative")
+
+    report = build_report(args.min_gpus, args.min_memory_mib)
+    if args.json:
+        print(json.dumps(report, indent=2))
+    else:
+        status = "PASS" if report["available"] else "FAIL"
+        print(f"NVIDIA GPU preflight: {status}")
+        for gpu in report["gpus"]:
+            print(
+                f"  GPU {gpu['index']}: {gpu['name']} "
+                f"({gpu['memory_total_mib']} MiB, driver {gpu['driver_version']})"
+            )
+        if report.get("error"):
+            print(f"  {report['error']}", file=sys.stderr)
+    return 0 if report["available"] else 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

docker/docker-compose-local-embeddings-gpu.yaml

@@ -0,0 +1,159 @@
+# Optional local Qwen3 GPU embedding services.
+#
+# This file is included separately from the CPU service file so the GPU
+# runtime can be reviewed, configured, and enabled without changing CPU
+# service definitions. Every service is still profile-gated.
+
+x-qwen3-gpu-healthcheck: &qwen3_gpu_healthcheck
+  test:
+    - CMD-SHELL
+    - >-
+      curl --fail --silent http://localhost:80/health >/dev/null ||
+      exit 1
+  interval: 15s
+  timeout: 10s
+  retries: 40
+  start_period: 10m
+
+x-qwen3-gpu-common: &qwen3_gpu_common
+  restart: unless-stopped
+  shm_size: 1gb
+  expose:
+    - "80"
+  environment:
+    HF_HOME: /data
+    HF_HUB_DISABLE_TELEMETRY: "1"
+    DO_NOT_TRACK: "1"
+    HF_TOKEN: ${HF_TOKEN:-}
+    TOKENIZERS_PARALLELISM: "false"
+  labels:
+    traefik.enable: "false"
+  stop_grace_period: 30s
+  healthcheck:
+    <<: *qwen3_gpu_healthcheck
+
+x-qwen3-gpu-image: &qwen3_gpu_image
+  ${QWEN3_TEI_GPU_IMAGE:-ghcr.io/huggingface/text-embeddings-inference:cuda-1.9}
+
+services:
+  qwen3-embedding-06b-gpu:
+    <<: *qwen3_gpu_common
+    image: *qwen3_gpu_image
+    container_name: unstract-qwen3-embedding-06b-gpu
+    profiles:
+      - embeddings-gpu
+      - embeddings-both
+    ports:
+      - "127.0.0.1:${QWEN3_06B_GPU_PORT:-8201}:80"
+    volumes:
+      - qwen3_embedding_06b_gpu_cache:/data
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities:
+                - gpu
+    command:
+      - --model-id
+      - Qwen/Qwen3-Embedding-0.6B
+      - --revision
+      - ${QWEN3_06B_REVISION:-97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3}
+      - --served-model-name
+      - qwen3-embedding-06b
+      - --dtype
+      - float16
+      - --pooling
+      - last-token
+      - --max-batch-tokens
+      - ${QWEN3_GPU_MAX_BATCH_TOKENS:-8192}
+      - --max-client-batch-size
+      - ${QWEN3_GPU_MAX_CLIENT_BATCH_SIZE:-32}
+      - --max-concurrent-requests
+      - ${QWEN3_GPU_MAX_CONCURRENT_REQUESTS:-4}
+      - --tokenization-workers
+      - ${QWEN3_GPU_TOKENIZATION_WORKERS:-4}
+
+  qwen3-embedding-4b-gpu:
+    <<: *qwen3_gpu_common
+    image: *qwen3_gpu_image
+    container_name: unstract-qwen3-embedding-4b-gpu
+    profiles:
+      - embeddings-gpu
+      - embeddings-both
+    ports:
+      - "127.0.0.1:${QWEN3_4B_GPU_PORT:-8202}:80"
+    volumes:
+      - qwen3_embedding_4b_gpu_cache:/data
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities:
+                - gpu
+    command:
+      - --model-id
+      - Qwen/Qwen3-Embedding-4B
+      - --revision
+      - ${QWEN3_4B_REVISION:-5cf2132abc99cad020ac570b19d031efec650f2b}
+      - --served-model-name
+      - qwen3-embedding-4b
+      - --dtype
+      - float16
+      - --pooling
+      - last-token
+      - --max-batch-tokens
+      - ${QWEN3_GPU_MAX_BATCH_TOKENS:-8192}
+      - --max-client-batch-size
+      - ${QWEN3_GPU_MAX_CLIENT_BATCH_SIZE:-32}
+      - --max-concurrent-requests
+      - ${QWEN3_GPU_MAX_CONCURRENT_REQUESTS:-4}
+      - --tokenization-workers
+      - ${QWEN3_GPU_TOKENIZATION_WORKERS:-4}
+
+  qwen3-embedding-8b-gpu:
+    <<: *qwen3_gpu_common
+    image: *qwen3_gpu_image
+    container_name: unstract-qwen3-embedding-8b-gpu
+    profiles:
+      - embeddings-gpu
+      - embeddings-both
+    ports:
+      - "127.0.0.1:${QWEN3_8B_GPU_PORT:-8203}:80"
+    volumes:
+      - qwen3_embedding_8b_gpu_cache:/data
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities:
+                - gpu
+    command:
+      - --model-id
+      - Qwen/Qwen3-Embedding-8B
+      - --revision
+      - ${QWEN3_8B_REVISION:-1d8ad4ca9b3dd8059ad90a75d4983776a23d44af}
+      - --served-model-name
+      - qwen3-embedding-8b
+      - --dtype
+      - float16
+      - --pooling
+      - last-token
+      - --max-batch-tokens
+      - ${QWEN3_GPU_MAX_BATCH_TOKENS:-8192}
+      - --max-client-batch-size
+      - ${QWEN3_GPU_MAX_CLIENT_BATCH_SIZE:-32}
+      - --max-concurrent-requests
+      - ${QWEN3_GPU_MAX_CONCURRENT_REQUESTS:-4}
+      - --tokenization-workers
+      - ${QWEN3_GPU_TOKENIZATION_WORKERS:-4}
+
+volumes:
+  qwen3_embedding_06b_gpu_cache:
+  qwen3_embedding_4b_gpu_cache:
+  qwen3_embedding_8b_gpu_cache: