feat(ml): GPU-capable tagger + embedder with CPU fallback (#872)
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>
This commit is contained in:
@@ -0,0 +1,31 @@
|
||||
"""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"))
|
||||
Reference in New Issue
Block a user