bvandeusen
db7e1f2b59
Step 1 of GPU enablement (code only — CPU-safe, CI-green; the CUDA image is a separate step pending the host driver version). - New services/ml/device.py: FC_ML_DEVICE (auto|cuda|cpu) intent + VRAM knobs (FC_ML_ONNX_GPU_MEM_GB, FC_ML_TORCH_MEM_FRACTION). Per-worker-host bootstrap → env, not a DB setting (the GPU host runs CUDA, others CPU). - tagger: use CUDAExecutionProvider (with gpu_mem_limit) when requested AND the provider is actually present (onnxruntime-gpu), else CPUExecutionProvider. Logs the active providers. - embedder: move model + inputs to cuda when requested AND torch.cuda is available; cap torch's VRAM share; .detach().cpu() before numpy. fp32 kept so GPU embeddings stay in the same space as existing CPU ones. Both AND the env intent with the framework's real availability, so on CPU (CI / CPU onnxruntime / no GPU) they fall back cleanly — behavior unchanged. The 8GB P4 is shared by both frameworks, hence the conservative default caps. Tests: device env parsing. (tagger/embedder GPU paths are operator-verified on the GPU host — models aren't in CI.) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
32 lines
1.2 KiB
Python
32 lines
1.2 KiB
Python
"""ML device selection (#872 — GPU enablement for the ml-worker).
|
|
|
|
The ml-worker is GPU-capable but must run unchanged on CPU (CI, non-GPU hosts).
|
|
Selection is a per-worker-HOST bootstrap concern (the GPU host runs CUDA, others
|
|
CPU), so it's an env var, not a DB setting — different workers need different
|
|
values. Each framework still ANDs this intent with its OWN runtime availability
|
|
(onnxruntime providers / torch.cuda), so "want GPU but none present" falls back
|
|
to CPU cleanly.
|
|
|
|
Env:
|
|
FC_ML_DEVICE auto (default) | cuda | gpu -> try GPU; cpu -> force CPU
|
|
FC_ML_ONNX_GPU_MEM_GB ONNX CUDA arena cap, GB (default 3) — the P4 is 8GB
|
|
total and torch shares it, so keep headroom.
|
|
FC_ML_TORCH_MEM_FRACTION fraction of total VRAM torch may use (default 0.6).
|
|
"""
|
|
|
|
import os
|
|
|
|
|
|
def gpu_requested() -> bool:
|
|
return os.environ.get("FC_ML_DEVICE", "auto").strip().lower() in (
|
|
"auto", "cuda", "gpu",
|
|
)
|
|
|
|
|
|
def onnx_gpu_mem_bytes() -> int:
|
|
return int(float(os.environ.get("FC_ML_ONNX_GPU_MEM_GB", "3")) * 1024 ** 3)
|
|
|
|
|
|
def torch_mem_fraction() -> float:
|
|
return float(os.environ.get("FC_ML_TORCH_MEM_FRACTION", "0.6"))
|