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.1 KiB
Python
32 lines
1.1 KiB
Python
"""ML device-selection env parsing (#872). Pure logic — no models/GPU/DB."""
|
|
|
|
from backend.app.services.ml import device
|
|
|
|
|
|
def test_gpu_requested_default_is_auto(monkeypatch):
|
|
monkeypatch.delenv("FC_ML_DEVICE", raising=False)
|
|
assert device.gpu_requested() is True
|
|
|
|
|
|
def test_gpu_requested_modes(monkeypatch):
|
|
for v in ("auto", "cuda", "gpu", "CUDA", " Auto "):
|
|
monkeypatch.setenv("FC_ML_DEVICE", v)
|
|
assert device.gpu_requested() is True
|
|
for v in ("cpu", "CPU", "none", "0"):
|
|
monkeypatch.setenv("FC_ML_DEVICE", v)
|
|
assert device.gpu_requested() is False
|
|
|
|
|
|
def test_onnx_gpu_mem_bytes(monkeypatch):
|
|
monkeypatch.delenv("FC_ML_ONNX_GPU_MEM_GB", raising=False)
|
|
assert device.onnx_gpu_mem_bytes() == 3 * 1024 ** 3
|
|
monkeypatch.setenv("FC_ML_ONNX_GPU_MEM_GB", "2")
|
|
assert device.onnx_gpu_mem_bytes() == 2 * 1024 ** 3
|
|
|
|
|
|
def test_torch_mem_fraction(monkeypatch):
|
|
monkeypatch.delenv("FC_ML_TORCH_MEM_FRACTION", raising=False)
|
|
assert device.torch_mem_fraction() == 0.6
|
|
monkeypatch.setenv("FC_ML_TORCH_MEM_FRACTION", "0.5")
|
|
assert device.torch_mem_fraction() == 0.5
|