File size: 4,426 Bytes
c794b6b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | """
GPU / provider detection and vision-engine mode selection for CEPHEUS.
CEPHEUS_CLOUD=1 without CEPHEUS_GPU_VISION uses the lightweight stub.
CEPHEUS_GPU_VISION=1 on cloud uses the full vision stack (for Cloud Run + L4).
CEPHEUS_FORCE_FULL_VISION=1 on cloud forces the full vision stack even on CPU.
Local dev (CEPHEUS_CLOUD unset/0) always uses the full vision stack.
"""
from __future__ import annotations
import logging
import os
import threading
from typing import Any
logger = logging.getLogger(__name__)
_warmload_lock = threading.Lock()
_warmload_state: dict[str, Any] = {
"started": False,
"complete": False,
"error": None,
"result": None,
}
def _env_truthy(name: str) -> bool:
return os.getenv(name, "").strip().lower() in ("1", "true", "yes")
def use_full_vision_engine() -> bool:
"""Return True when YOLO + InsightFace should load (local or GPU cloud)."""
if not _env_truthy("CEPHEUS_CLOUD"):
return True
return _env_truthy("CEPHEUS_GPU_VISION") or _env_truthy("CEPHEUS_FORCE_FULL_VISION")
def detect_acceleration() -> dict[str, Any]:
"""Probe PyTorch CUDA and ONNX Runtime providers without crashing."""
info: dict[str, Any] = {
"cuda_available": False,
"torch_cuda": False,
"onnx_providers": [],
"provider": "cpu",
"device_name": None,
"fallback_reason": None,
}
try:
import onnxruntime as ort
info["onnx_providers"] = list(ort.get_available_providers())
except Exception as exc:
logger.debug("ONNX Runtime probe skipped: %s", exc)
try:
import torch
info["torch_cuda"] = bool(torch.cuda.is_available())
if info["torch_cuda"]:
info["cuda_available"] = True
info["provider"] = "cuda"
try:
info["device_name"] = torch.cuda.get_device_name(0)
except Exception:
info["device_name"] = "cuda:0"
except ImportError:
if "CUDAExecutionProvider" in info["onnx_providers"]:
info["cuda_available"] = True
info["provider"] = "cuda"
info["device_name"] = "onnxruntime-cuda"
if not info["cuda_available"]:
if _env_truthy("CEPHEUS_FORCE_CPU"):
info["fallback_reason"] = "CEPHEUS_FORCE_CPU=1"
elif not info["onnx_providers"]:
info["fallback_reason"] = "ONNX Runtime unavailable"
elif "CUDAExecutionProvider" not in info["onnx_providers"] and not info["torch_cuda"]:
info["fallback_reason"] = (
"CUDA not available (install torch+cuda or onnxruntime-gpu for GPU inference)"
)
else:
info["fallback_reason"] = "CUDA not available"
return info
def insightface_ctx_id() -> int:
"""InsightFace ctx_id: 0 = GPU, -1 = CPU."""
if _env_truthy("CEPHEUS_FORCE_CPU"):
return -1
accel = detect_acceleration()
return 0 if accel["cuda_available"] else -1
def get_warmload_state() -> dict[str, Any]:
with _warmload_lock:
return dict(_warmload_state)
def mark_warmload_started() -> bool:
"""Return False if warmload was already started."""
with _warmload_lock:
if _warmload_state["started"]:
return False
_warmload_state["started"] = True
return True
def mark_warmload_complete(result: dict[str, Any] | None = None) -> None:
with _warmload_lock:
_warmload_state["complete"] = True
_warmload_state["result"] = result
def mark_warmload_failed(error: str) -> None:
with _warmload_lock:
_warmload_state["error"] = error
_face_ready_probe: Any = None
def register_face_ready_probe(probe: Any) -> None:
"""Optional callback from main after VisionEngine init — avoids import cycles."""
global _face_ready_probe
_face_ready_probe = probe
def live_status_payload() -> dict[str, str]:
state = get_warmload_state()
if state.get("complete"):
return {"status": "ready", "message": "Instance ready"}
if _face_ready_probe:
try:
if _face_ready_probe():
return {"status": "ready", "message": "Face engine loaded"}
except Exception:
pass
if state.get("started"):
return {"status": "warming", "message": "Model warmload in progress"}
return {"status": "warming", "message": "Instance initialization triggered"}
|