Merge pull request #18 from saichandrapandraju/global-state

feat: Refactor Detector lifecycle to use FastAPI app.state

Author: ruivieira

detectors/common/app.py

@@ -94,6 +94,17 @@ async def http_exception_handler(self, request, exc):
             content={"code": exc.status_code, "message": exc.detail},
         )
 
+    def set_detector(self, detector) -> None:
+        """Store detector in app.state"""
+        self.state.detector = detector
+        
+    def get_detector(self):
+        """Retrieve detector from app.state"""
+        return getattr(self.state, 'detector', None)
+    
+    def cleanup_detector(self) -> None:
+        """Clean up detector resources"""
+        self.state.detector = None
 
 async def health():
     return "ok"

detectors/huggingface/app.py

@@ -15,15 +15,16 @@
     Error,
 )
 
-detector_objects = {}
-
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    detector_objects["detector"] = Detector()
+    app.set_detector(Detector())
     yield
     # Clean up the ML models and release the resources
-    detector_objects.clear()
+    detector: Detector = app.get_detector()
+    if detector and hasattr(detector, 'close'):
+        detector.close()
+    app.cleanup_detector()
 
 
 app = FastAPI(lifespan=lifespan, dependencies=[])
@@ -44,4 +45,7 @@ async def detector_unary_handler(
     request: ContentAnalysisHttpRequest,
     detector_id: Annotated[str, Header(example="en_syntax_slate.38m.hap")],
 ):
-    return ContentsAnalysisResponse(root=detector_objects["detector"].run(request))
+    detector: Detector = app.get_detector()
+    if not detector:
+        raise RuntimeError("Detector is not initialized")
+    return ContentsAnalysisResponse(root=detector.run(request))

detectors/huggingface/detector.py

@@ -16,7 +16,7 @@
     ContentAnalysisResponse,
     ContentsAnalysisResponse,
 )
-
+import gc
 
 class Detector:
     risk_names = [
@@ -280,3 +280,20 @@ def run(self, input: ContentAnalysisHttpRequest) -> ContentsAnalysisResponse:
                 raise ValueError("Unsupported model type for analysis.")
             contents_analyses.append(analyses)
         return contents_analyses
+
+
+    def close(self) -> None:
+        """Clean up model and tokenizer resources."""
+        
+        if self.model:
+            if hasattr(self.model, 'to') and hasattr(self.model, 'device') and self.model.device.type != "cpu":
+                self.model = self.model.to(torch.device("cpu"))
+            self.model = None
+
+        if self.tokenizer:
+            self.tokenizer = None
+
+        gc.collect()
+
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()

detectors/llm_judge/app.py

@@ -13,20 +13,18 @@
     Error,
 )
 
-detector_objects: Dict[str, LLMJudgeDetector] = {}
-
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     """Application lifespan management."""
-    try:
-        detector_objects["detector"] = LLMJudgeDetector()
-        yield
-    finally:
-        # Clean up resources
-        if "detector" in detector_objects:
-            await detector_objects["detector"].close()
-        detector_objects.clear()
+
+    app.set_detector(LLMJudgeDetector())
+    yield
+    # Clean up resources
+    detector: LLMJudgeDetector = app.get_detector()
+    if detector and hasattr(detector, 'close'):
+        await detector.close()
+    app.cleanup_detector()
 
 
 app = FastAPI(lifespan=lifespan, dependencies=[])
@@ -49,7 +47,8 @@ async def detector_unary_handler(
     detector_id: Annotated[str, Header(example="llm_judge_safety")],
 ):
     """Analyze content using LLM-as-Judge evaluation."""
-    return ContentsAnalysisResponse(root=await detector_objects["detector"].run(request))
+    detector: LLMJudgeDetector = app.get_detector()
+    return ContentsAnalysisResponse(root=await detector.run(request))
 
 
 @app.get(
@@ -62,7 +61,7 @@ async def detector_unary_handler(
 )
 async def list_metrics():
     """List all available evaluation metrics."""
-    detector = detector_objects.get("detector")
+    detector: LLMJudgeDetector = app.get_detector()
     if not detector:
         return {"metrics": [], "total": 0}
 

detectors/llm_judge/detector.py

@@ -16,7 +16,7 @@ class LLMJudgeDetector:
 
     def __init__(self) -> None:
         """Initialize the LLM Judge Detector."""
-        self.judge = None
+        self.judge: Judge = None
         self.available_metrics = set(BUILTIN_METRICS.keys())
 
         # Get configuration from environment