From 69858e1e21338ae134e4d30200243f64cc677648 Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg <5861166+CompleteDotTech@users.noreply.github.com> Date: Thu, 3 Sep 2026 16:09:40 -0400 Subject: [PATCH 1/2] Add GPU embedding preflight and tuning --- docker/README.md | 13 ++ docker/benchmarks/gpu_preflight.py | 143 ++++++++++++++++++++ docker/docker-compose-local-embeddings.yaml | 24 ++-- 3 files changed, 168 insertions(+), 12 deletions(-) create mode 100644 docker/benchmarks/gpu_preflight.py diff --git a/docker/README.md b/docker/README.md index 0e06aecbe9..a4b28b0db8 100644 --- a/docker/README.md +++ b/docker/README.md @@ -108,6 +108,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 +124,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: diff --git a/docker/benchmarks/gpu_preflight.py b/docker/benchmarks/gpu_preflight.py new file mode 100644 index 0000000000..d11684aae4 --- /dev/null +++ b/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()) diff --git a/docker/docker-compose-local-embeddings.yaml b/docker/docker-compose-local-embeddings.yaml index f870d0e10d..c6ebb883d8 100644 --- a/docker/docker-compose-local-embeddings.yaml +++ b/docker/docker-compose-local-embeddings.yaml @@ -158,13 +158,13 @@ services: - --pooling - last-token - --max-batch-tokens - - ${QWEN3_MAX_BATCH_TOKENS:-8192} + - ${QWEN3_GPU_MAX_BATCH_TOKENS:-8192} - --max-client-batch-size - - ${QWEN3_MAX_CLIENT_BATCH_SIZE:-32} + - ${QWEN3_GPU_MAX_CLIENT_BATCH_SIZE:-32} - --max-concurrent-requests - - ${QWEN3_MAX_CONCURRENT_REQUESTS:-4} + - ${QWEN3_GPU_MAX_CONCURRENT_REQUESTS:-4} - --tokenization-workers - - ${QWEN3_TOKENIZATION_WORKERS:-4} + - ${QWEN3_GPU_TOKENIZATION_WORKERS:-4} qwen3-embedding-4b-gpu: <<: *qwen3_common @@ -197,13 +197,13 @@ services: - --pooling - last-token - --max-batch-tokens - - ${QWEN3_MAX_BATCH_TOKENS:-8192} + - ${QWEN3_GPU_MAX_BATCH_TOKENS:-8192} - --max-client-batch-size - - ${QWEN3_MAX_CLIENT_BATCH_SIZE:-32} + - ${QWEN3_GPU_MAX_CLIENT_BATCH_SIZE:-32} - --max-concurrent-requests - - ${QWEN3_MAX_CONCURRENT_REQUESTS:-4} + - ${QWEN3_GPU_MAX_CONCURRENT_REQUESTS:-4} - --tokenization-workers - - ${QWEN3_TOKENIZATION_WORKERS:-4} + - ${QWEN3_GPU_TOKENIZATION_WORKERS:-4} qwen3-embedding-8b-gpu: <<: *qwen3_common @@ -236,13 +236,13 @@ services: - --pooling - last-token - --max-batch-tokens - - ${QWEN3_MAX_BATCH_TOKENS:-8192} + - ${QWEN3_GPU_MAX_BATCH_TOKENS:-8192} - --max-client-batch-size - - ${QWEN3_MAX_CLIENT_BATCH_SIZE:-32} + - ${QWEN3_GPU_MAX_CLIENT_BATCH_SIZE:-32} - --max-concurrent-requests - - ${QWEN3_MAX_CONCURRENT_REQUESTS:-4} + - ${QWEN3_GPU_MAX_CONCURRENT_REQUESTS:-4} - --tokenization-workers - - ${QWEN3_TOKENIZATION_WORKERS:-4} + - ${QWEN3_GPU_TOKENIZATION_WORKERS:-4} volumes: qwen3_embedding_06b_cpu_cache: From 0f0aecfbcc1d24bf1849ffec5511041bda5fb2be Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg <5861166+CompleteDotTech@users.noreply.github.com> Date: Thu, 3 Sep 2026 16:54:41 -0400 Subject: [PATCH 2/2] Split GPU embeddings into dedicated Compose include --- docker/README.md | 6 + .../docker-compose-local-embeddings-gpu.yaml | 159 ++++++++++++++++++ docker/docker-compose-local-embeddings.yaml | 129 +------------- docker/docker-compose.yaml | 1 + 4 files changed, 169 insertions(+), 126 deletions(-) create mode 100644 docker/docker-compose-local-embeddings-gpu.yaml diff --git a/docker/README.md b/docker/README.md index a4b28b0db8..b4d170e63e 100644 --- a/docker/README.md +++ b/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 | diff --git a/docker/docker-compose-local-embeddings-gpu.yaml b/docker/docker-compose-local-embeddings-gpu.yaml new file mode 100644 index 0000000000..80f85a3560 --- /dev/null +++ b/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: diff --git a/docker/docker-compose-local-embeddings.yaml b/docker/docker-compose-local-embeddings.yaml index c6ebb883d8..aebefbdc31 100644 --- a/docker/docker-compose-local-embeddings.yaml +++ b/docker/docker-compose-local-embeddings.yaml @@ -1,9 +1,9 @@ # Optional local Qwen3 embedding services. # # The services are profile-gated so the normal development stack does not -# download multi-gigabyte model weights. CPU and GPU services deliberately have -# separate names, ports, caches, and containers so their benchmark results can -# be compared without changing an endpoint in place. +# download multi-gigabyte model weights. CPU services live here; GPU services +# are included from docker-compose-local-embeddings-gpu.yaml so either runtime +# variation can be reviewed and operated independently. x-qwen3-healthcheck: &qwen3_healthcheck test: @@ -36,9 +36,6 @@ x-qwen3-common: &qwen3_common x-qwen3-cpu-image: &qwen3_cpu_image ${QWEN3_TEI_CPU_IMAGE:-ghcr.io/huggingface/text-embeddings-inference:cpu-1.9} -x-qwen3-gpu-image: &qwen3_gpu_image - ${QWEN3_TEI_GPU_IMAGE:-ghcr.io/huggingface/text-embeddings-inference:cuda-1.9} - services: qwen3-embedding-06b-cpu: <<: *qwen3_common @@ -127,127 +124,7 @@ services: - --tokenization-workers - ${QWEN3_TOKENIZATION_WORKERS:-4} - qwen3-embedding-06b-gpu: - <<: *qwen3_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_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_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_cpu_cache: qwen3_embedding_4b_cpu_cache: qwen3_embedding_8b_cpu_cache: - qwen3_embedding_06b_gpu_cache: - qwen3_embedding_4b_gpu_cache: - qwen3_embedding_8b_gpu_cache: diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index 6bea2ea1d0..5c444c8e51 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -2,6 +2,7 @@ name: ${COMPOSE_PROJECT_NAME:-docker} include: - docker-compose-dev-essentials.yaml - docker-compose-local-embeddings.yaml + - docker-compose-local-embeddings-gpu.yaml # Reusable host-gateway mapping so containers can reach services on the host # (e.g. host-installed Ollama at http://host.docker.internal:11434).