update the codebase poc ver1
This commit is contained in:
DatTT127
@@ -0,0 +1,184 @@
|
||||
"""Agent tool BFF services — Exa search and Supabase knowledge queries."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
EXA_SEARCH_URL = "https://api.exa.ai/search"
|
||||
EXA_API_KEY = os.getenv("EXA_API_KEY", "").strip()
|
||||
|
||||
SUPABASE_URL = os.getenv("SUPABASE_URL", "").rstrip("/")
|
||||
SUPABASE_SERVICE_ROLE_KEY = os.getenv("SUPABASE_SERVICE_ROLE_KEY", "").strip()
|
||||
|
||||
ALLOWED_EXA_TYPES = frozenset(
|
||||
{"auto", "fast", "instant", "deep-lite", "deep", "deep-reasoning"}
|
||||
)
|
||||
ALLOWED_SUPABASE_RPC = frozenset({"match_semantic_chunks", "get_corpus_citation"})
|
||||
|
||||
|
||||
def _audit(event: str, session_id: str, payload: dict[str, Any]) -> None:
|
||||
query_hash = payload.get("query_hash")
|
||||
logger.info(
|
||||
"[AUDIT] event=%s session=%s query_hash=%s payload_keys=%s",
|
||||
event,
|
||||
session_id,
|
||||
query_hash,
|
||||
list(payload.keys()),
|
||||
)
|
||||
|
||||
|
||||
def _query_hash(text: str) -> str:
|
||||
return hashlib.sha256(text.encode("utf-8")).hexdigest()[:16]
|
||||
|
||||
|
||||
async def exa_search(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
session_id = str(payload.get("session_id", ""))
|
||||
query = str(payload.get("query", "")).strip()
|
||||
if not query:
|
||||
raise ValueError("query is required")
|
||||
if len(query) > 512:
|
||||
raise ValueError("query exceeds 512 characters")
|
||||
|
||||
search_type = str(payload.get("type", "auto"))
|
||||
if search_type not in ALLOWED_EXA_TYPES:
|
||||
raise ValueError(f"unsupported Exa type: {search_type}")
|
||||
|
||||
num_results = int(payload.get("numResults", 10))
|
||||
num_results = max(1, min(num_results, 10))
|
||||
|
||||
if not EXA_API_KEY:
|
||||
raise RuntimeError("EXA_API_KEY is not configured on the backend")
|
||||
|
||||
body: dict[str, Any] = {
|
||||
"query": query,
|
||||
"type": search_type,
|
||||
"numResults": num_results,
|
||||
"contents": {"highlights": True},
|
||||
}
|
||||
|
||||
include_domains = payload.get("includeDomains")
|
||||
exclude_domains = payload.get("excludeDomains")
|
||||
if include_domains:
|
||||
body["includeDomains"] = include_domains
|
||||
if exclude_domains:
|
||||
body["excludeDomains"] = exclude_domains
|
||||
|
||||
_audit(
|
||||
"exa_search",
|
||||
session_id,
|
||||
{"query_hash": _query_hash(query), "type": search_type, "numResults": num_results},
|
||||
)
|
||||
|
||||
headers = {"x-api-key": EXA_API_KEY, "Content-Type": "application/json"}
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.post(EXA_SEARCH_URL, headers=headers, json=body)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
hits = []
|
||||
for index, row in enumerate(data.get("results", [])):
|
||||
hits.append(
|
||||
{
|
||||
"id": row.get("id") or f"exa-{index}",
|
||||
"url": row.get("url") or "",
|
||||
"title": row.get("title") or row.get("url") or "Untitled",
|
||||
"highlights": row.get("highlights") or [],
|
||||
"publishedDate": row.get("publishedDate"),
|
||||
"score": row.get("score"),
|
||||
}
|
||||
)
|
||||
|
||||
return {"hits": hits, "requestId": data.get("requestId")}
|
||||
|
||||
|
||||
async def supabase_query(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
session_id = str(payload.get("session_id", ""))
|
||||
rpc = str(payload.get("rpc", ""))
|
||||
if rpc not in ALLOWED_SUPABASE_RPC:
|
||||
raise ValueError(f"rpc not allowlisted: {rpc}")
|
||||
|
||||
if not SUPABASE_URL or not SUPABASE_SERVICE_ROLE_KEY:
|
||||
raise RuntimeError("SUPABASE_URL and SUPABASE_SERVICE_ROLE_KEY must be configured")
|
||||
|
||||
args = payload.get("args") or {}
|
||||
if rpc == "match_semantic_chunks":
|
||||
query_text = str(args.get("query_text", "")).strip()
|
||||
if not query_text:
|
||||
raise ValueError("args.query_text is required for match_semantic_chunks")
|
||||
|
||||
_audit(
|
||||
"supabase_query",
|
||||
session_id,
|
||||
{"query_hash": _query_hash(query_text), "rpc": rpc},
|
||||
)
|
||||
|
||||
embedding = await _embed_query_text(query_text)
|
||||
rpc_body = {
|
||||
"query_embedding": embedding,
|
||||
"match_count": int(args.get("match_count", 5)),
|
||||
"filter_book_ids": args.get("filter_book_ids"),
|
||||
"filter_edition_ids": args.get("filter_edition_ids"),
|
||||
}
|
||||
else:
|
||||
_audit("supabase_query", session_id, {"rpc": rpc})
|
||||
rpc_body = args
|
||||
|
||||
url = f"{SUPABASE_URL}/rest/v1/rpc/{rpc}"
|
||||
headers = {
|
||||
"apikey": SUPABASE_SERVICE_ROLE_KEY,
|
||||
"Authorization": f"Bearer {SUPABASE_SERVICE_ROLE_KEY}",
|
||||
"Content-Type": "application/json",
|
||||
"Accept-Profile": "knowledge",
|
||||
"Content-Profile": "knowledge",
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.post(url, headers=headers, json=rpc_body)
|
||||
if response.status_code >= 400:
|
||||
detail = response.text
|
||||
try:
|
||||
detail = json.dumps(response.json())
|
||||
except Exception:
|
||||
pass
|
||||
raise RuntimeError(f"Supabase RPC failed ({response.status_code}): {detail}")
|
||||
rows = response.json()
|
||||
|
||||
normalized_rows = []
|
||||
for row in rows if isinstance(rows, list) else []:
|
||||
normalized_rows.append(
|
||||
{
|
||||
"chunk_id": str(row.get("chunk_id", "")),
|
||||
"content": row.get("content") or "",
|
||||
"book_id": row.get("book_id") or "",
|
||||
"parent_title": row.get("parent_title"),
|
||||
"page_start": row.get("page_start"),
|
||||
"page_end": row.get("page_end"),
|
||||
"similarity": row.get("similarity"),
|
||||
}
|
||||
)
|
||||
|
||||
return {"rpc": rpc, "rows": normalized_rows}
|
||||
|
||||
|
||||
async def _embed_query_text(query_text: str) -> list[float]:
|
||||
"""Compute 768-d EmbeddingGemma vector for Supabase RPC.
|
||||
|
||||
PoC: returns NotImplemented until Triton/on-prem embedder is wired.
|
||||
Set EMBED_QUERY_MOCK=1 to return a zero vector for integration testing only.
|
||||
"""
|
||||
if os.getenv("EMBED_QUERY_MOCK") == "1":
|
||||
logger.warning("Using EMBED_QUERY_MOCK zero vector — not for production search quality")
|
||||
return [0.0] * 768
|
||||
|
||||
raise NotImplementedError(
|
||||
"Server-side query embedding is not wired yet. "
|
||||
"Configure Triton EmbeddingGemma or set EMBED_QUERY_MOCK=1 for integration tests."
|
||||
)
|
||||
@@ -0,0 +1,203 @@
|
||||
import logging
|
||||
import httpx
|
||||
import json
|
||||
from typing import AsyncGenerator
|
||||
from datetime import datetime
|
||||
|
||||
from backend.implementation.adapters.redis_adapter import get_redis_client
|
||||
from backend.implementation.adapters.llm_adapter import get_llm_adapter, AuditCallbackHandler
|
||||
from backend.implementation.config import (
|
||||
MODAL_MEDGEMMA_ENDPOINT,
|
||||
VERTEX_AI_GEMINI_ENDPOINT,
|
||||
GCP_ACCESS_TOKEN,
|
||||
PROJECT_ID,
|
||||
LOCATION,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
redis_client = get_redis_client()
|
||||
llm_adapter = get_llm_adapter()
|
||||
|
||||
|
||||
def _set_consult_mode(session_id: str, mode: str):
|
||||
redis_client.set(f"consult_mode:{session_id}", mode, ex=7200)
|
||||
|
||||
|
||||
async def _verify_consent(session_id: str) -> bool:
|
||||
consent_key = f"consent:{session_id}"
|
||||
return bool(await asyncio.to_thread(redis_client.exists, consent_key))
|
||||
|
||||
|
||||
async def _verify_session_ownership(session_id: str, user_id: str) -> bool:
|
||||
owner_key = f"session_owner:{session_id}"
|
||||
owner_id = await asyncio.to_thread(redis_client.get, owner_key)
|
||||
if not owner_id:
|
||||
return False
|
||||
return owner_id == user_id
|
||||
|
||||
|
||||
def _audit_event(session_id: str, event_type: str, payload: dict):
|
||||
key = f"audit:{session_id}:{event_type}"
|
||||
redis_client.set(key, json.dumps(payload), ex=86400)
|
||||
logger.info(
|
||||
"[AUDIT] event=%s session=%s payload=%s",
|
||||
event_type,
|
||||
session_id,
|
||||
payload,
|
||||
)
|
||||
|
||||
|
||||
async def route_gemini_request(payload: dict, user_id: str) -> dict:
|
||||
session_id = payload.get("session_id", "")
|
||||
task_type = payload.get("task_type", "orchestration")
|
||||
prompt = payload.get("prompt", "")
|
||||
redaction_hash = payload.get("redaction_hash")
|
||||
|
||||
if not await _verify_consent(session_id):
|
||||
raise PermissionError("User consent for cloud LLM egress is required.")
|
||||
|
||||
if not await _verify_session_ownership(session_id, user_id):
|
||||
raise PermissionError("You do not own this session.")
|
||||
|
||||
_audit_event(session_id, "egress_consent_gemini", {
|
||||
"user_id": user_id,
|
||||
"task_type": task_type,
|
||||
"redaction_hash": redaction_hash,
|
||||
"ts": datetime.utcnow().isoformat(),
|
||||
})
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {GCP_ACCESS_TOKEN}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
vertex_payload = {
|
||||
"contents": [{"parts": [{"text": prompt}]}],
|
||||
"generationConfig": {
|
||||
"temperature": 0.2,
|
||||
"topP": 0.8,
|
||||
"topK": 40,
|
||||
"maxOutputTokens": 1024,
|
||||
},
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient(timeout=22.0) as client:
|
||||
response = await client.post(
|
||||
VERTEX_AI_GEMINI_ENDPOINT,
|
||||
headers=headers,
|
||||
json=vertex_payload,
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
_audit_event(session_id, "egress_response_gemini", {
|
||||
"task_type": task_type,
|
||||
"status": "success",
|
||||
"ts": datetime.utcnow().isoformat(),
|
||||
})
|
||||
_set_consult_mode(session_id, "tier_2")
|
||||
|
||||
return {
|
||||
"text": result.get("candidates", [{}])[0].get("content", {}).get("parts", [{}])[0].get("text", ""),
|
||||
"tier": "gemini",
|
||||
"task_type": task_type,
|
||||
}
|
||||
|
||||
|
||||
async def route_medgemma_request(payload: dict, user_id: str) -> dict | AsyncGenerator[str, None]:
|
||||
session_id = payload.get("session_id", "")
|
||||
task_type = payload.get("task_type", "clinical_deep_reasoning")
|
||||
prompt = payload.get("prompt", "")
|
||||
stream = payload.get("stream", False)
|
||||
redaction_hash = payload.get("redaction_hash")
|
||||
|
||||
if not await _verify_consent(session_id):
|
||||
raise PermissionError("User consent for cloud LLM egress is required.")
|
||||
|
||||
if not await _verify_session_ownership(session_id, user_id):
|
||||
raise PermissionError("You do not own this session.")
|
||||
|
||||
_audit_event(session_id, "egress_consent_medgemma", {
|
||||
"user_id": user_id,
|
||||
"task_type": task_type,
|
||||
"redaction_hash": redaction_hash,
|
||||
"ts": datetime.utcnow().isoformat(),
|
||||
})
|
||||
|
||||
modal_payload = {
|
||||
"model": payload.get("model", "medgemma:4b"),
|
||||
"prompt": prompt,
|
||||
"stream": stream,
|
||||
"options": {
|
||||
"temperature": 0.1,
|
||||
"top_p": 0.8,
|
||||
"top_k": 40,
|
||||
"num_predict": 2048,
|
||||
},
|
||||
}
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
if stream:
|
||||
return _stream_medgemma(session_id, modal_payload, headers, task_type)
|
||||
|
||||
async with httpx.AsyncClient(timeout=22.0) as client:
|
||||
response = await client.post(
|
||||
f"{MODAL_MEDGEMMA_ENDPOINT}/api/generate",
|
||||
headers=headers,
|
||||
json=modal_payload,
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
_audit_event(session_id, "egress_response_medgemma", {
|
||||
"task_type": task_type,
|
||||
"status": "success",
|
||||
"ts": datetime.utcnow().isoformat(),
|
||||
})
|
||||
_set_consult_mode(session_id, "tier_3")
|
||||
|
||||
return {
|
||||
"text": result.get("response", ""),
|
||||
"tier": "medgemma",
|
||||
"task_type": task_type,
|
||||
}
|
||||
|
||||
|
||||
async def _stream_medgemma(
|
||||
session_id: str,
|
||||
modal_payload: dict,
|
||||
headers: dict,
|
||||
task_type: str,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
async with httpx.AsyncClient(timeout=22.0) as client:
|
||||
async with client.stream(
|
||||
"POST",
|
||||
f"{MODAL_MEDGEMMA_ENDPOINT}/api/generate",
|
||||
headers=headers,
|
||||
json=modal_payload,
|
||||
) as response:
|
||||
response.raise_for_status()
|
||||
async for line in response.aiter_lines():
|
||||
if not line.startswith("data:"):
|
||||
continue
|
||||
data_str = line[len("data:"):].strip()
|
||||
if not data_str:
|
||||
continue
|
||||
try:
|
||||
data = json.loads(data_str)
|
||||
chunk = data.get("response", "")
|
||||
if chunk:
|
||||
yield chunk
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
_audit_event(session_id, "egress_response_medgemma_stream", {
|
||||
"task_type": task_type,
|
||||
"status": "success",
|
||||
"ts": datetime.utcnow().isoformat(),
|
||||
})
|
||||
_set_consult_mode(session_id, "tier_3")
|
||||
@@ -0,0 +1,336 @@
|
||||
"""Spec-compliant CV inference orchestration — Sprint 1–2 architecture §7."""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import io
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from backend.implementation.config import get_angle_type, get_model_name, get_segmentation_model
|
||||
from backend.implementation.pipeline.cv_spec_pipeline import (
|
||||
BRANCH_ANGLE_CLASSES,
|
||||
build_segmentation_skipped,
|
||||
build_severity_zero,
|
||||
)
|
||||
from backend.implementation.postprocessing.calibration import (
|
||||
CalibrationConfig,
|
||||
calibration_config_from_params,
|
||||
interpret_angle_logits,
|
||||
interpret_inflammation_logits,
|
||||
)
|
||||
from backend.implementation.postprocessing.measurement import calculate_thickness
|
||||
from backend.implementation.postprocessing.overlay import COLOR_MAP_POST, COLOR_MAP_SUP, create_overlay
|
||||
from backend.implementation.postprocessing.severity import calculate_severity
|
||||
from backend.implementation.preprocessing.clahe import apply_clahe
|
||||
from backend.services import cv_result_cache
|
||||
from backend.services import triton_runtime_service as triton_runtime
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
SEGMENT_CLASSES_SUP = {
|
||||
0: "background",
|
||||
1: "effusion",
|
||||
2: "fat",
|
||||
3: "fat-pat",
|
||||
4: "femur",
|
||||
5: "synovium",
|
||||
6: "tendon",
|
||||
}
|
||||
SEGMENT_CLASSES_POST = {
|
||||
0: "background",
|
||||
1: "fat",
|
||||
2: "tendon",
|
||||
3: "muscle",
|
||||
4: "femur",
|
||||
5: "artery",
|
||||
6: "baker's cyst",
|
||||
}
|
||||
|
||||
_triton_pipeline_lock = asyncio.Lock()
|
||||
|
||||
|
||||
@dataclass
|
||||
class CvInferenceOptions:
|
||||
calibration: CalibrationConfig | None = None
|
||||
model_versions: dict[str, str] | None = None
|
||||
use_cache: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class CvBatchResult:
|
||||
results: list[dict[str, Any]]
|
||||
triton_infer_calls: int
|
||||
triton_infer_modes: list[str]
|
||||
|
||||
|
||||
def _encode_image_to_data_url(image_pil: Image.Image) -> str:
|
||||
buffered = io.BytesIO()
|
||||
image_pil.save(buffered, format="PNG")
|
||||
encoded = base64.b64encode(buffered.getvalue()).decode()
|
||||
return f"data:image/png;base64,{encoded}"
|
||||
|
||||
|
||||
def _build_inflammation_payload(logits_row: np.ndarray, config: CalibrationConfig) -> dict:
|
||||
interpreted = interpret_inflammation_logits(logits_row, config)
|
||||
return {
|
||||
"detected": interpreted["detected"],
|
||||
"confidence": interpreted["confidence"],
|
||||
"calibration": interpreted["calibration"],
|
||||
}
|
||||
|
||||
|
||||
def _logits_to_masks(
|
||||
logits: np.ndarray,
|
||||
angle_class: str,
|
||||
image_size: tuple[int, int],
|
||||
batch_idx: int = 0,
|
||||
) -> dict[str, np.ndarray]:
|
||||
logits_arr = np.asarray(logits, dtype=np.float32)
|
||||
while logits_arr.ndim < 4:
|
||||
logits_arr = np.expand_dims(logits_arr, 0)
|
||||
if logits_arr.ndim != 4:
|
||||
raise ValueError(f"Unexpected segmentation logits shape: {logits_arr.shape}")
|
||||
if batch_idx >= logits_arr.shape[0]:
|
||||
raise IndexError(f"batch_idx {batch_idx} out of range for shape {logits_arr.shape}")
|
||||
|
||||
preds_lowres = logits_arr.argmax(axis=1)[batch_idx]
|
||||
width, height = image_size
|
||||
preds = cv2.resize(
|
||||
preds_lowres.astype(np.uint8),
|
||||
(width, height),
|
||||
interpolation=cv2.INTER_NEAREST,
|
||||
)
|
||||
|
||||
angle_type = get_angle_type(angle_class)
|
||||
class_map = SEGMENT_CLASSES_SUP if angle_type == "sup" else SEGMENT_CLASSES_POST
|
||||
|
||||
masks: dict[str, np.ndarray] = {}
|
||||
for class_id, class_name in class_map.items():
|
||||
masks[class_name] = (preds == class_id).astype(np.uint8)
|
||||
return masks
|
||||
|
||||
|
||||
def _build_color_legend(classes_detected: list[str], angle_type: str) -> dict[str, list[int]]:
|
||||
color_map = COLOR_MAP_SUP if angle_type == "sup" else COLOR_MAP_POST
|
||||
legend: dict[str, list[int]] = {}
|
||||
for class_name in classes_detected:
|
||||
if class_name in color_map:
|
||||
legend[class_name] = color_map[class_name]
|
||||
return legend
|
||||
|
||||
|
||||
def _build_segmentation_result(
|
||||
image_pil: Image.Image,
|
||||
logits: np.ndarray,
|
||||
angle_class: str,
|
||||
seg_model: str,
|
||||
*,
|
||||
frame_id: str | None,
|
||||
inflammation: dict,
|
||||
angle_payload: dict,
|
||||
enhanced_data_url: str,
|
||||
inflam_model: str,
|
||||
angle_model: str,
|
||||
) -> dict:
|
||||
angle_type = get_angle_type(angle_class)
|
||||
masks = _logits_to_masks(logits, angle_class, image_pil.size)
|
||||
measurement = calculate_thickness(masks, image_pil.size)
|
||||
severity = calculate_severity(masks, image_pil.size)
|
||||
overlay = create_overlay(image_pil, masks, measurement, angle_type)
|
||||
|
||||
classes_detected = [name for name, mask in masks.items() if np.sum(mask) > 0 and name != "background"]
|
||||
color_legend = _build_color_legend(classes_detected, angle_type)
|
||||
|
||||
result: dict[str, Any] = {
|
||||
"success": True,
|
||||
"angle": angle_payload,
|
||||
"inflammation": inflammation,
|
||||
"measurement": measurement,
|
||||
"severity": severity,
|
||||
"segmentation": {
|
||||
"performed": True,
|
||||
"angle_type": angle_type,
|
||||
"classes_detected": classes_detected,
|
||||
"color_legend": color_legend,
|
||||
},
|
||||
"images": {
|
||||
"enhanced": enhanced_data_url,
|
||||
"segmented": _encode_image_to_data_url(overlay),
|
||||
},
|
||||
"models_used": {
|
||||
"angle": angle_model,
|
||||
"inflammation": inflam_model,
|
||||
"segmentation": seg_model,
|
||||
},
|
||||
}
|
||||
if frame_id is not None:
|
||||
result["frame_id"] = frame_id
|
||||
return result
|
||||
|
||||
|
||||
async def _run_spec_cv_pipeline_single(
|
||||
image_pil: Image.Image,
|
||||
*,
|
||||
frame_id: str | None,
|
||||
options: CvInferenceOptions,
|
||||
) -> tuple[dict[str, Any], str, int]:
|
||||
"""
|
||||
Per-image spec path:
|
||||
CLAHE → angle → (post-trans|sup-up-long only) inflammation → conditional segmentation.
|
||||
"""
|
||||
config = options.calibration or CalibrationConfig()
|
||||
model_versions = options.model_versions
|
||||
angle_model = get_model_name("angle", model_versions)
|
||||
inflam_model = get_model_name("inflammation", model_versions)
|
||||
|
||||
triton_calls = 0
|
||||
modes: list[str] = []
|
||||
|
||||
enhanced_pil = apply_clahe(image_pil)
|
||||
enhanced_data_url = _encode_image_to_data_url(enhanced_pil)
|
||||
|
||||
angle_logits, angle_mode, angle_calls = await triton_runtime.infer_angle_logits_single(
|
||||
image_pil, angle_model
|
||||
)
|
||||
modes.append(angle_mode)
|
||||
triton_calls += angle_calls
|
||||
|
||||
angle_interpreted = interpret_angle_logits(angle_logits, config)
|
||||
angle_payload = {
|
||||
"class": angle_interpreted["class"],
|
||||
"confidence": angle_interpreted["confidence"],
|
||||
"calibration": angle_interpreted["calibration"],
|
||||
}
|
||||
|
||||
result: dict[str, Any] = {
|
||||
"success": True,
|
||||
"angle": angle_payload,
|
||||
"models_used": {"angle": angle_model},
|
||||
"images": {"enhanced": enhanced_data_url},
|
||||
}
|
||||
if frame_id is not None:
|
||||
result["frame_id"] = frame_id
|
||||
|
||||
angle_class = angle_interpreted["class"]
|
||||
if angle_class not in BRANCH_ANGLE_CLASSES:
|
||||
result["segmentation"] = build_segmentation_skipped("angle_only")
|
||||
result["severity"] = build_severity_zero("angle_only")
|
||||
return result, "+".join(modes), triton_calls
|
||||
|
||||
inflam_logits, inflam_mode, inflam_calls = await triton_runtime.infer_inflammation_logits_single(
|
||||
image_pil, inflam_model
|
||||
)
|
||||
modes.append(inflam_mode)
|
||||
triton_calls += inflam_calls
|
||||
|
||||
inflammation = _build_inflammation_payload(inflam_logits, config)
|
||||
result["inflammation"] = inflammation
|
||||
result["models_used"]["inflammation"] = inflam_model
|
||||
|
||||
if not inflammation.get("detected"):
|
||||
result["segmentation"] = build_segmentation_skipped("no_inflammation")
|
||||
result["severity"] = build_severity_zero("no_inflammation")
|
||||
return result, "+".join(modes), triton_calls
|
||||
|
||||
seg_model = get_segmentation_model(angle_class, model_versions)
|
||||
seg_logits, seg_mode, seg_calls = await triton_runtime.infer_segmentation_logits_single(
|
||||
image_pil, seg_model
|
||||
)
|
||||
modes.append(seg_mode)
|
||||
triton_calls += seg_calls
|
||||
|
||||
seg_result = _build_segmentation_result(
|
||||
image_pil,
|
||||
seg_logits,
|
||||
angle_class,
|
||||
seg_model,
|
||||
frame_id=frame_id,
|
||||
inflammation=inflammation,
|
||||
angle_payload=angle_payload,
|
||||
enhanced_data_url=enhanced_data_url,
|
||||
inflam_model=inflam_model,
|
||||
angle_model=angle_model,
|
||||
)
|
||||
return seg_result, "+".join(modes), triton_calls
|
||||
|
||||
|
||||
async def run_single(
|
||||
image: Image.Image,
|
||||
*,
|
||||
frame_id: str | None = None,
|
||||
options: CvInferenceOptions | None = None,
|
||||
) -> dict[str, Any]:
|
||||
opts = options or CvInferenceOptions()
|
||||
result, _, _ = await _run_spec_cv_pipeline_single(image, frame_id=frame_id, options=opts)
|
||||
return result
|
||||
|
||||
|
||||
async def _run_batch_uncached(
|
||||
images: list[Image.Image],
|
||||
frame_ids: list[str],
|
||||
options: CvInferenceOptions,
|
||||
) -> CvBatchResult:
|
||||
if not images:
|
||||
return CvBatchResult(results=[], triton_infer_calls=0, triton_infer_modes=[])
|
||||
if len(frame_ids) != len(images):
|
||||
raise ValueError("frame_ids length must match images length")
|
||||
|
||||
async with _triton_pipeline_lock:
|
||||
results: list[dict[str, Any]] = []
|
||||
infer_modes: list[str] = []
|
||||
triton_call_count = 0
|
||||
for image_pil, fid in zip(images, frame_ids, strict=True):
|
||||
item, mode, calls = await _run_spec_cv_pipeline_single(
|
||||
image_pil,
|
||||
frame_id=fid,
|
||||
options=options,
|
||||
)
|
||||
results.append(item)
|
||||
infer_modes.append(mode)
|
||||
triton_call_count += calls
|
||||
return CvBatchResult(
|
||||
results=results,
|
||||
triton_infer_calls=triton_call_count,
|
||||
triton_infer_modes=infer_modes,
|
||||
)
|
||||
|
||||
|
||||
async def run_batch(
|
||||
images: list[Image.Image],
|
||||
frame_ids: list[str],
|
||||
options: CvInferenceOptions | None = None,
|
||||
) -> CvBatchResult:
|
||||
opts = options or CvInferenceOptions()
|
||||
if not opts.use_cache or not images:
|
||||
return await _run_batch_uncached(images, frame_ids, opts)
|
||||
|
||||
image_hashes = []
|
||||
for image in images:
|
||||
buf = io.BytesIO()
|
||||
image.save(buf, format="PNG")
|
||||
image_hashes.append(cv_result_cache.hash_image_bytes(buf.getvalue()))
|
||||
|
||||
cache_key = cv_result_cache.analyze_batch_cache_key(frame_ids, image_hashes)
|
||||
|
||||
async def compute():
|
||||
return await _run_batch_uncached(images, frame_ids, opts)
|
||||
|
||||
return await cv_result_cache.with_result_cache(cache_key, compute, enabled=opts.use_cache)
|
||||
|
||||
|
||||
def options_from_params(params: dict[str, Any] | None) -> CvInferenceOptions:
|
||||
params = params or {}
|
||||
calibration = calibration_config_from_params(params)
|
||||
model_versions = params.get("model_versions")
|
||||
use_cache = params.get("use_cache", True)
|
||||
return CvInferenceOptions(
|
||||
calibration=calibration,
|
||||
model_versions=model_versions,
|
||||
use_cache=use_cache,
|
||||
)
|
||||
@@ -0,0 +1,72 @@
|
||||
"""In-memory CV inference result cache with in-flight request coalescing."""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any, Awaitable, Callable, TypeVar
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
CV_PIPELINE_VERSION = os.getenv("CV_PIPELINE_VERSION", "poc-v2-spec-cv-seg-norm")
|
||||
CV_RESULT_CACHE_TTL_S = float(os.getenv("CV_RESULT_CACHE_TTL_S", "3600"))
|
||||
CV_CACHE_ENABLED = os.getenv("CV_CACHE_ENABLED", "true").lower() in {"1", "true", "yes"}
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
_result_cache: dict[str, tuple[float, Any]] = {}
|
||||
_inflight: dict[str, asyncio.Future] = {}
|
||||
|
||||
|
||||
def hash_image_bytes(raw: bytes) -> str:
|
||||
return hashlib.sha256(raw).hexdigest()
|
||||
|
||||
|
||||
def analyze_batch_cache_key(frame_ids: list[str], image_hashes: list[str]) -> str:
|
||||
pairs = sorted(zip(frame_ids, image_hashes, strict=True), key=lambda item: item[0])
|
||||
payload = "|".join(f"{frame_id}:{digest}" for frame_id, digest in pairs)
|
||||
return f"analyze|{CV_PIPELINE_VERSION}|{payload}"
|
||||
|
||||
|
||||
async def with_result_cache(cache_key: str, compute: Callable[[], Awaitable[T]], *, enabled: bool = True) -> T:
|
||||
if not enabled or not CV_CACHE_ENABLED:
|
||||
return await compute()
|
||||
|
||||
now = time.monotonic()
|
||||
cached = _result_cache.get(cache_key)
|
||||
if cached and cached[0] > now:
|
||||
logger.info("CV cache HIT: %s", cache_key[:96])
|
||||
return cached[1]
|
||||
|
||||
inflight = _inflight.get(cache_key)
|
||||
if inflight is not None:
|
||||
logger.info("CV in-flight coalesce: %s", cache_key[:96])
|
||||
return await inflight
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
fut: asyncio.Future = loop.create_future()
|
||||
_inflight[cache_key] = fut
|
||||
try:
|
||||
result = await compute()
|
||||
_result_cache[cache_key] = (time.monotonic() + CV_RESULT_CACHE_TTL_S, result)
|
||||
fut.set_result(result)
|
||||
logger.info("CV cache STORE: %s", cache_key[:96])
|
||||
return result
|
||||
except Exception as exc:
|
||||
if not fut.done():
|
||||
fut.set_exception(exc)
|
||||
raise
|
||||
finally:
|
||||
_inflight.pop(cache_key, None)
|
||||
|
||||
|
||||
def cache_stats() -> dict[str, int | bool | float | str]:
|
||||
return {
|
||||
"cache_enabled": CV_CACHE_ENABLED,
|
||||
"pipeline_version": CV_PIPELINE_VERSION,
|
||||
"cache_ttl_s": CV_RESULT_CACHE_TTL_S,
|
||||
"cache_entries": len(_result_cache),
|
||||
"inflight_batches": len(_inflight),
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
"""EmbeddingGemma-compatible embed endpoint for episodic memory and RAG queries."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
from typing import Literal
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
EMBEDDING_DIMENSIONS = 768
|
||||
|
||||
|
||||
def _format_embed_input(text: str, task: str, title: str | None = None) -> str:
|
||||
if task == "retrieval-document":
|
||||
title_value = title.strip() if title and title.strip() else "none"
|
||||
return f"title: {title_value} | text: {text}"
|
||||
return f"task: search result | query: {text}"
|
||||
|
||||
|
||||
def deterministic_embed(text: str, dimensions: int = EMBEDDING_DIMENSIONS) -> list[float]:
|
||||
"""PoC fallback — matches gemma4_e2b deterministicEmbed (SHA-256 + char histogram)."""
|
||||
vec = [0.0] * dimensions
|
||||
normalized = text.lower().strip()
|
||||
if not normalized:
|
||||
return vec
|
||||
|
||||
for index, char in enumerate(normalized):
|
||||
code = ord(char)
|
||||
bucket = (code * (index + 17)) % dimensions
|
||||
vec[bucket] += 1.0
|
||||
|
||||
digest = hashlib.sha256(normalized.encode("utf-8")).digest()
|
||||
for i in range(dimensions):
|
||||
vec[i] += digest[i % len(digest)] / 255.0
|
||||
|
||||
norm = math.sqrt(sum(value * value for value in vec))
|
||||
if norm == 0:
|
||||
return vec
|
||||
return [value / norm for value in vec]
|
||||
|
||||
|
||||
def _try_gemma_embedder(formatted: str) -> list[float] | None:
|
||||
"""Optional real EmbeddingGemma via knowledge ingestion pipeline."""
|
||||
if os.getenv("EMBED_QUERY_MOCK") == "1":
|
||||
return None
|
||||
try:
|
||||
from knowledge.implementation.ingestion.embedding import GemmaEmbedder, EmbedTask
|
||||
|
||||
embedder = GemmaEmbedder()
|
||||
task = (
|
||||
EmbedTask.RETRIEVAL_DOCUMENT
|
||||
if formatted.startswith("title:")
|
||||
else EmbedTask.RETRIEVAL_QUERY
|
||||
)
|
||||
raw = formatted
|
||||
title = None
|
||||
if task == EmbedTask.RETRIEVAL_DOCUMENT and "| text: " in formatted:
|
||||
prefix, body = formatted.split("| text: ", 1)
|
||||
title = prefix.replace("title:", "").strip()
|
||||
raw = body
|
||||
vector = embedder.embed(raw, task, title=title if title and title != "none" else None)
|
||||
return vector.tolist()
|
||||
except Exception as exc:
|
||||
logger.debug("GemmaEmbedder unavailable: %s", exc)
|
||||
return None
|
||||
|
||||
|
||||
async def embed_text(
|
||||
text: str,
|
||||
task: Literal["retrieval-query", "retrieval-document"] = "retrieval-query",
|
||||
title: str | None = None,
|
||||
) -> dict[str, object]:
|
||||
formatted = _format_embed_input(text, task, title)
|
||||
vector = _try_gemma_embedder(formatted)
|
||||
if vector is not None:
|
||||
return {"vector": vector, "model": "embeddinggemma-300m", "source": "gemma"}
|
||||
|
||||
if os.getenv("EMBED_QUERY_MOCK") == "1":
|
||||
logger.warning("Using deterministic embed fallback (EMBED_QUERY_MOCK or no embedder)")
|
||||
|
||||
return {
|
||||
"vector": deterministic_embed(formatted),
|
||||
"model": "embeddinggemma-300m-deterministic",
|
||||
"source": "deterministic",
|
||||
}
|
||||
@@ -0,0 +1,384 @@
|
||||
"""Triton inference runtime — lock, retry, batching with batched→sequential fallback."""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Literal
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
from backend.implementation.adapters.triton_adapter import TritonAdapter
|
||||
from backend.implementation.preprocessing.tensor_prep import (
|
||||
prepare_angle_tensor,
|
||||
prepare_inflammation_tensor,
|
||||
prepare_segmentation_tensor,
|
||||
)
|
||||
from backend.implementation.triton_batch import TRITON_MAX_BATCH_SIZE, chunk_sequence
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
INPUT_NAME = "input_image"
|
||||
OUTPUT_NAME = "logits"
|
||||
|
||||
TRITON_INFER_TIMEOUT = float(os.getenv("TRITON_INFER_TIMEOUT", "120"))
|
||||
TRITON_INFER_RETRIES = int(os.getenv("TRITON_INFER_RETRIES", "5"))
|
||||
TRITON_RETRY_BASE_S = float(os.getenv("TRITON_RETRY_BASE_S", "4"))
|
||||
TRITON_USE_BATCH_INFER = os.getenv("TRITON_USE_BATCH_INFER", "true").lower()
|
||||
RETRYABLE_STATUS = {429, 502, 503, 504}
|
||||
|
||||
_triton_infer_lock = asyncio.Lock()
|
||||
_adapter: TritonAdapter | None = None
|
||||
_adapter_endpoint: str | None = None
|
||||
|
||||
|
||||
def get_triton_endpoint() -> str:
|
||||
return os.getenv("TRITON_ENDPOINT", "http://localhost:8000").rstrip("/")
|
||||
|
||||
|
||||
def _get_adapter() -> TritonAdapter:
|
||||
global _adapter, _adapter_endpoint
|
||||
endpoint = get_triton_endpoint()
|
||||
if _adapter is None or _adapter_endpoint != endpoint:
|
||||
_adapter = TritonAdapter(endpoint_url=endpoint, timeout=TRITON_INFER_TIMEOUT)
|
||||
_adapter_endpoint = endpoint
|
||||
return _adapter
|
||||
|
||||
|
||||
def _retry_backoff_seconds(attempt: int) -> float:
|
||||
return TRITON_RETRY_BASE_S * (2 ** (attempt - 1))
|
||||
|
||||
|
||||
def _should_try_batched_infer(image_count: int) -> bool:
|
||||
if image_count <= 1:
|
||||
return True
|
||||
if TRITON_USE_BATCH_INFER in {"1", "true", "yes"}:
|
||||
return True
|
||||
if TRITON_USE_BATCH_INFER in {"0", "false", "no"}:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _logits_array_from_adapter(result: dict) -> np.ndarray:
|
||||
logits = result.get(OUTPUT_NAME, [])
|
||||
if not logits:
|
||||
raise ValueError(f"Empty {OUTPUT_NAME} in Triton response")
|
||||
return np.asarray(logits, dtype=np.float32)
|
||||
|
||||
|
||||
def _infer_sync_with_retry(
|
||||
model_name: str,
|
||||
batch_tensor: np.ndarray,
|
||||
*,
|
||||
operation: str,
|
||||
max_retries: int | None = None,
|
||||
) -> np.ndarray:
|
||||
adapter = _get_adapter()
|
||||
attempts = max_retries if max_retries is not None else TRITON_INFER_RETRIES
|
||||
last_error: Exception | None = None
|
||||
|
||||
inputs = {
|
||||
INPUT_NAME: {
|
||||
"data": batch_tensor,
|
||||
"shape": list(batch_tensor.shape),
|
||||
"datatype": "FP32",
|
||||
}
|
||||
}
|
||||
|
||||
for attempt in range(1, attempts + 1):
|
||||
try:
|
||||
result = adapter._infer_sync(model_name, inputs, outputs=[OUTPUT_NAME])
|
||||
if attempt > 1:
|
||||
logger.info("%s succeeded on attempt %s/%s", operation, attempt, attempts)
|
||||
return _logits_array_from_adapter(result)
|
||||
except requests.HTTPError as exc:
|
||||
status = exc.response.status_code if exc.response is not None else None
|
||||
if status not in RETRYABLE_STATUS:
|
||||
raise
|
||||
last_error = exc
|
||||
except (requests.ConnectionError, requests.Timeout) as exc:
|
||||
last_error = exc
|
||||
|
||||
if attempt >= attempts:
|
||||
logger.error("%s failed on final attempt %s/%s: %s", operation, attempt, attempts, last_error)
|
||||
break
|
||||
wait_s = _retry_backoff_seconds(attempt)
|
||||
logger.warning(
|
||||
"%s attempt %s/%s failed (%s); exponential retry in %.1fs",
|
||||
operation,
|
||||
attempt,
|
||||
attempts,
|
||||
last_error,
|
||||
wait_s,
|
||||
)
|
||||
time.sleep(wait_s)
|
||||
|
||||
assert last_error is not None
|
||||
raise last_error
|
||||
|
||||
|
||||
def _stack_angle_tensors(images: list[Image.Image]) -> np.ndarray:
|
||||
if not images:
|
||||
raise ValueError("images must not be empty")
|
||||
if len(images) > TRITON_MAX_BATCH_SIZE:
|
||||
raise ValueError(f"At most {TRITON_MAX_BATCH_SIZE} images per Triton batch, got {len(images)}")
|
||||
tensors = [prepare_angle_tensor(img) for img in images]
|
||||
return np.concatenate(tensors, axis=0).astype(np.float32)
|
||||
|
||||
|
||||
def _stack_inflammation_tensors(images: list[Image.Image]) -> np.ndarray:
|
||||
if not images:
|
||||
raise ValueError("images must not be empty")
|
||||
if len(images) > TRITON_MAX_BATCH_SIZE:
|
||||
raise ValueError(f"At most {TRITON_MAX_BATCH_SIZE} images per Triton batch, got {len(images)}")
|
||||
tensors = [prepare_inflammation_tensor(img) for img in images]
|
||||
return np.concatenate(tensors, axis=0).astype(np.float32)
|
||||
|
||||
|
||||
def _stack_segmentation_tensors(images: list[Image.Image]) -> np.ndarray:
|
||||
if not images:
|
||||
raise ValueError("images must not be empty")
|
||||
if len(images) > TRITON_MAX_BATCH_SIZE:
|
||||
raise ValueError(f"At most {TRITON_MAX_BATCH_SIZE} images per Triton batch, got {len(images)}")
|
||||
tensors = [prepare_segmentation_tensor(img) for img in images]
|
||||
return np.concatenate(tensors, axis=0).astype(np.float32)
|
||||
|
||||
|
||||
def _normalize_batched_angle_logits(logits: np.ndarray, expected_count: int) -> np.ndarray:
|
||||
if logits.ndim == 1:
|
||||
logits = np.expand_dims(logits, axis=0)
|
||||
if logits.ndim != 2:
|
||||
raise ValueError(f"Unexpected batched angle logits shape: {logits.shape}")
|
||||
if logits.shape[0] != expected_count:
|
||||
raise ValueError(
|
||||
f"Triton returned batch {logits.shape[0]} but expected {expected_count} angle rows",
|
||||
)
|
||||
return logits
|
||||
|
||||
|
||||
def _normalize_batched_segmentation_logits(logits: np.ndarray, expected_count: int) -> np.ndarray:
|
||||
if logits.ndim == 3:
|
||||
logits = np.expand_dims(logits, axis=0)
|
||||
if logits.ndim != 4:
|
||||
raise ValueError(f"Unexpected batched segmentation logits shape: {logits.shape}")
|
||||
if logits.shape[0] != expected_count:
|
||||
raise ValueError(
|
||||
f"Triton returned batch {logits.shape[0]} but expected {expected_count} images",
|
||||
)
|
||||
return logits
|
||||
|
||||
|
||||
def _infer_angle_logits_batch(
|
||||
images: list[Image.Image],
|
||||
model_name: str,
|
||||
*,
|
||||
max_retries: int | None = None,
|
||||
) -> np.ndarray:
|
||||
batch_tensor = _stack_angle_tensors(images)
|
||||
logits = _infer_sync_with_retry(
|
||||
model_name,
|
||||
batch_tensor,
|
||||
operation=f"Triton angle batch×{len(images)} ({model_name})",
|
||||
max_retries=max_retries,
|
||||
)
|
||||
return _normalize_batched_angle_logits(logits, len(images))
|
||||
|
||||
|
||||
def _infer_angle_logits_sequential(images: list[Image.Image], model_name: str) -> np.ndarray:
|
||||
rows: list[np.ndarray] = []
|
||||
for index, image in enumerate(images, start=1):
|
||||
logits = _infer_angle_logits_batch([image], model_name)
|
||||
row = logits[0] if logits.ndim == 2 else logits
|
||||
rows.append(row)
|
||||
logger.info("Triton sequential angle infer %s/%s complete for %s", index, len(images), model_name)
|
||||
return np.stack(rows, axis=0)
|
||||
|
||||
|
||||
def _infer_angle_logits_chunk(images: list[Image.Image], model_name: str) -> tuple[np.ndarray, Literal["batched", "sequential"]]:
|
||||
if not _should_try_batched_infer(len(images)):
|
||||
return _infer_angle_logits_sequential(images, model_name), "sequential"
|
||||
try:
|
||||
return _infer_angle_logits_batch(images, model_name, max_retries=1), "batched"
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Batched angle infer×%s failed (%s); falling back to sequential single-image calls",
|
||||
len(images),
|
||||
exc,
|
||||
)
|
||||
return _infer_angle_logits_sequential(images, model_name), "sequential"
|
||||
|
||||
|
||||
def _infer_inflammation_logits_batch(
|
||||
images: list[Image.Image],
|
||||
model_name: str,
|
||||
*,
|
||||
max_retries: int | None = None,
|
||||
) -> np.ndarray:
|
||||
batch_tensor = _stack_inflammation_tensors(images)
|
||||
logits = _infer_sync_with_retry(
|
||||
model_name,
|
||||
batch_tensor,
|
||||
operation=f"Triton inflammation batch×{len(images)} ({model_name})",
|
||||
max_retries=max_retries,
|
||||
)
|
||||
return _normalize_batched_angle_logits(logits, len(images))
|
||||
|
||||
|
||||
def _infer_inflammation_logits_sequential(images: list[Image.Image], model_name: str) -> np.ndarray:
|
||||
rows: list[np.ndarray] = []
|
||||
for index, image in enumerate(images, start=1):
|
||||
logits = _infer_inflammation_logits_batch([image], model_name)
|
||||
row = logits[0] if logits.ndim == 2 else logits
|
||||
rows.append(row)
|
||||
logger.info(
|
||||
"Triton sequential inflammation infer %s/%s complete for %s",
|
||||
index,
|
||||
len(images),
|
||||
model_name,
|
||||
)
|
||||
return np.stack(rows, axis=0)
|
||||
|
||||
|
||||
def _infer_inflammation_logits_chunk(
|
||||
images: list[Image.Image],
|
||||
model_name: str,
|
||||
) -> tuple[np.ndarray, Literal["batched", "sequential"]]:
|
||||
if not _should_try_batched_infer(len(images)):
|
||||
return _infer_inflammation_logits_sequential(images, model_name), "sequential"
|
||||
try:
|
||||
return _infer_inflammation_logits_batch(images, model_name, max_retries=1), "batched"
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Batched inflammation infer×%s failed (%s); falling back to sequential",
|
||||
len(images),
|
||||
exc,
|
||||
)
|
||||
return _infer_inflammation_logits_sequential(images, model_name), "sequential"
|
||||
|
||||
|
||||
def _infer_segmentation_logits_batch(
|
||||
images: list[Image.Image],
|
||||
model_name: str,
|
||||
*,
|
||||
max_retries: int | None = None,
|
||||
) -> np.ndarray:
|
||||
batch_tensor = _stack_segmentation_tensors(images)
|
||||
logits = _infer_sync_with_retry(
|
||||
model_name,
|
||||
batch_tensor,
|
||||
operation=f"Triton segmentation batch×{len(images)} ({model_name})",
|
||||
max_retries=max_retries,
|
||||
)
|
||||
return _normalize_batched_segmentation_logits(logits, len(images))
|
||||
|
||||
|
||||
def _infer_segmentation_logits_sequential(images: list[Image.Image], model_name: str) -> np.ndarray:
|
||||
rows: list[np.ndarray] = []
|
||||
for index, image in enumerate(images, start=1):
|
||||
logits = _infer_segmentation_logits_batch([image], model_name)
|
||||
row = logits[0] if logits.ndim == 4 else logits
|
||||
rows.append(row)
|
||||
logger.info(
|
||||
"Triton sequential segmentation infer %s/%s complete for %s",
|
||||
index,
|
||||
len(images),
|
||||
model_name,
|
||||
)
|
||||
return np.stack(rows, axis=0)
|
||||
|
||||
|
||||
def _infer_segmentation_logits_chunk(
|
||||
images: list[Image.Image],
|
||||
model_name: str,
|
||||
) -> tuple[np.ndarray, Literal["batched", "sequential"]]:
|
||||
if not _should_try_batched_infer(len(images)):
|
||||
return _infer_segmentation_logits_sequential(images, model_name), "sequential"
|
||||
try:
|
||||
return _infer_segmentation_logits_batch(images, model_name, max_retries=1), "batched"
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Batched segmentation infer×%s failed (%s); falling back to sequential",
|
||||
len(images),
|
||||
exc,
|
||||
)
|
||||
return _infer_segmentation_logits_sequential(images, model_name), "sequential"
|
||||
|
||||
|
||||
async def infer_angle_logits(
|
||||
images: list[Image.Image],
|
||||
model_name: str,
|
||||
) -> tuple[np.ndarray, Literal["batched", "sequential"], int]:
|
||||
"""Run angle classification under the global Triton lock. Returns (logits, mode, call_count)."""
|
||||
async with _triton_infer_lock:
|
||||
all_logits: list[np.ndarray] = []
|
||||
modes: list[str] = []
|
||||
call_count = 0
|
||||
for chunk in chunk_sequence(images):
|
||||
logits, mode = await asyncio.to_thread(_infer_angle_logits_chunk, chunk, model_name)
|
||||
all_logits.append(logits)
|
||||
modes.append(mode)
|
||||
call_count += 1 if mode == "batched" else len(chunk)
|
||||
combined = np.concatenate(all_logits, axis=0) if len(all_logits) > 1 else all_logits[0]
|
||||
infer_mode: Literal["batched", "sequential"] = (
|
||||
"sequential" if any(m == "sequential" for m in modes) else "batched"
|
||||
)
|
||||
return combined, infer_mode, call_count
|
||||
|
||||
|
||||
async def infer_inflammation_logits(
|
||||
images: list[Image.Image],
|
||||
model_name: str,
|
||||
) -> tuple[np.ndarray, Literal["batched", "sequential"], int]:
|
||||
async with _triton_infer_lock:
|
||||
all_logits: list[np.ndarray] = []
|
||||
modes: list[str] = []
|
||||
call_count = 0
|
||||
for chunk in chunk_sequence(images):
|
||||
logits, mode = await asyncio.to_thread(_infer_inflammation_logits_chunk, chunk, model_name)
|
||||
all_logits.append(logits)
|
||||
modes.append(mode)
|
||||
call_count += 1 if mode == "batched" else len(chunk)
|
||||
combined = np.concatenate(all_logits, axis=0) if len(all_logits) > 1 else all_logits[0]
|
||||
infer_mode: Literal["batched", "sequential"] = (
|
||||
"sequential" if any(m == "sequential" for m in modes) else "batched"
|
||||
)
|
||||
return combined, infer_mode, call_count
|
||||
|
||||
|
||||
async def infer_segmentation_logits(
|
||||
images: list[Image.Image],
|
||||
model_name: str,
|
||||
) -> tuple[np.ndarray, Literal["batched", "sequential"], int]:
|
||||
async with _triton_infer_lock:
|
||||
all_logits: list[np.ndarray] = []
|
||||
modes: list[str] = []
|
||||
call_count = 0
|
||||
for chunk in chunk_sequence(images):
|
||||
logits, mode = await asyncio.to_thread(_infer_segmentation_logits_chunk, chunk, model_name)
|
||||
all_logits.append(logits)
|
||||
modes.append(mode)
|
||||
call_count += 1 if mode == "batched" else len(chunk)
|
||||
combined = np.concatenate(all_logits, axis=0) if len(all_logits) > 1 else all_logits[0]
|
||||
infer_mode: Literal["batched", "sequential"] = (
|
||||
"sequential" if any(m == "sequential" for m in modes) else "batched"
|
||||
)
|
||||
return combined, infer_mode, call_count
|
||||
|
||||
|
||||
async def infer_angle_logits_single(image: Image.Image, model_name: str) -> tuple[np.ndarray, str, int]:
|
||||
logits, mode, calls = await infer_angle_logits([image], model_name)
|
||||
return logits[0], f"angle:{mode}", calls
|
||||
|
||||
|
||||
async def infer_inflammation_logits_single(image: Image.Image, model_name: str) -> tuple[np.ndarray, str, int]:
|
||||
logits, mode, calls = await infer_inflammation_logits([image], model_name)
|
||||
return logits[0], f"inflam:{mode}", calls
|
||||
|
||||
|
||||
async def infer_segmentation_logits_single(image: Image.Image, model_name: str) -> tuple[np.ndarray, str, int]:
|
||||
logits, mode, calls = await infer_segmentation_logits([image], model_name)
|
||||
row = logits[0] if logits.ndim == 4 else logits
|
||||
return row, f"seg:{mode}", calls
|
||||
@@ -0,0 +1,48 @@
|
||||
"""Warm up Modal Triton on CV service startup — reduces cold-start 502s and first-frame latency."""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from backend.implementation.config import get_model_name, get_segmentation_model
|
||||
from backend.services import triton_runtime_service as triton_runtime
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _warmup_model_versions() -> dict[str, str]:
|
||||
versions: dict[str, str] = {}
|
||||
if angle := os.getenv("ANGLE_MODEL"):
|
||||
versions["angle"] = angle
|
||||
else:
|
||||
versions["angle"] = "angle_classify_resnet50"
|
||||
if inflam := os.getenv("INFLAMMATION_MODEL"):
|
||||
versions["inflammation"] = inflam
|
||||
return versions
|
||||
|
||||
|
||||
async def warmup_triton_models() -> None:
|
||||
if os.getenv("CV_SKIP_TRITON_WARMUP", "").lower() in {"1", "true", "yes"}:
|
||||
logger.info("Triton warmup skipped (CV_SKIP_TRITON_WARMUP)")
|
||||
return
|
||||
|
||||
model_versions = _warmup_model_versions()
|
||||
angle_model = get_model_name("angle", model_versions)
|
||||
inflam_model = get_model_name("inflammation", model_versions)
|
||||
seg_model = get_segmentation_model("sup-up-long", model_versions)
|
||||
|
||||
img224 = Image.new("RGB", (224, 224), color=(128, 128, 128))
|
||||
img512 = Image.new("RGB", (512, 512), color=(128, 128, 128))
|
||||
|
||||
logger.info(
|
||||
"Warming up Triton (angle=%s, inflammation=%s, segmentation=%s)…",
|
||||
angle_model,
|
||||
inflam_model,
|
||||
seg_model,
|
||||
)
|
||||
await triton_runtime.infer_angle_logits_single(img224, angle_model)
|
||||
await triton_runtime.infer_inflammation_logits_single(img224, inflam_model)
|
||||
await triton_runtime.infer_segmentation_logits_single(img512, seg_model)
|
||||
logger.info("Triton warmup complete")
|
||||
Reference in New Issue
Block a user