fix model_loaded race condition with lock

Author: eliteprox

pytrickle/client.py

@@ -76,6 +76,7 @@ def __init__(
 
         # Guard to ensure model is loaded once per client lifecycle
         self._model_loaded = False
+        self._model_load_lock = asyncio.Lock()
 
     async def start(self, request_id: str = "default"):
         """Start the trickle client."""
@@ -183,21 +184,26 @@ def set_target_fps(self, target_fps: Optional[float]):
 
     async def _ensure_model_loaded(self):
         """Load the model once on the same event loop before processing begins."""
-        if not self._model_loaded:
-            # Transition to LOADING while model warms up, if state is available
-            try:
-                if getattr(self.frame_processor, "state", None) is not None:
-                    self.frame_processor.state.set_state(PipelineState.LOADING)
-                await self.frame_processor.load_model()
-                # Mark startup complete; this moves LOADING → IDLE per state machine
-                if getattr(self.frame_processor, "state", None) is not None:
-                    self.frame_processor.state.set_startup_complete()
-                self._model_loaded = True
-            except Exception as e:
-                # Reflect error in state if available, then propagate
-                if getattr(self.frame_processor, "state", None) is not None:
-                    self.frame_processor.state.set_error(str(e))
-                raise
+        # Transition to LOADING while model warms up, if state is available
+        try:
+            if getattr(self.frame_processor, "state", None) is not None:
+                self.frame_processor.state.set_state(PipelineState.LOADING)
+            
+            # Use the thread-safe wrapper
+            await self.frame_processor.ensure_model_loaded()
+            
+            # Mark startup complete; this moves LOADING → IDLE per state machine
+            if getattr(self.frame_processor, "state", None) is not None:
+                self.frame_processor.state.set_startup_complete()
+            
+            # Update our local flag for consistency
+            self._model_loaded = True
+            
+        except Exception as e:
+            # Reflect error in state if available, then propagate
+            if getattr(self.frame_processor, "state", None) is not None:
+                self.frame_processor.state.set_error(str(e))
+            raise
 
     async def _on_protocol_error(self, error_type: str, exception: Optional[Exception] = None):
         """Handle protocol errors and shutdown events."""

pytrickle/frame_processor.py

@@ -4,6 +4,7 @@
 making it easy to integrate AI models and async pipelines with PyTrickle.
 """
 
+import asyncio
 import logging
 from abc import ABC, abstractmethod
 from typing import Optional, Any, Dict, List
@@ -55,11 +56,25 @@ def __init__(
         """
         self.error_callback = error_callback
         self.state: Optional[StreamState] = None
+        
+        # Model loading protection
+        self._model_load_lock = asyncio.Lock()
+        self._model_loaded = False
 
     def attach_state(self, state: StreamState) -> None:
         """Attach a pipeline state manager and set IDLE if model already loaded."""
         self.state = state
 
+    async def ensure_model_loaded(self, **kwargs):
+        """Thread-safe wrapper that ensures model is loaded exactly once."""
+        async with self._model_load_lock:
+            if not self._model_loaded:
+                await self.load_model(**kwargs)
+                self._model_loaded = True
+                logger.debug(f"Model loaded for {self.__class__.__name__}")
+            else:
+                logger.debug(f"Model already loaded for {self.__class__.__name__}")
+
     @abstractmethod
     async def load_model(self, **kwargs):
         """

pytrickle/stream_processor.py

@@ -63,6 +63,10 @@ def __init__(
         self.frame_skip_config = frame_skip_config
         self.server_kwargs = server_kwargs
 
+        # Add lock for model loading protection
+        self._model_load_lock = asyncio.Lock()
+        self._model_loaded = False
+        
         # Create internal frame processor
         self._frame_processor = _InternalFrameProcessor(
             video_processor=video_processor,
@@ -86,7 +90,7 @@ def __init__(
             self._frame_processor.attach_state(self.server.state)
         except Exception:
             # If attach fails for any reason, log and continue (non-fatal)
-            logger.debug("Failed to attach server state to frame processor")
+            logger.warning("Failed to attach server state to frame processor")
 
         # Register server startup hook to preload model on same event loop
         async def _on_startup(_app):
@@ -95,10 +99,17 @@ async def _background_preload():
                 try:
                     if getattr(self._frame_processor, "state", None) is not None:
                         self._frame_processor.state.set_state(PipelineState.LOADING)
-                    await self._frame_processor.load_model()
+                    
+                    # Use the thread-safe wrapper
+                    await self._frame_processor.ensure_model_loaded()
+                    
                     if getattr(self._frame_processor, "state", None) is not None:
                         self._frame_processor.state.set_startup_complete()
+                    
+                    # Update our local flag for consistency
+                    self._model_loaded = True
                     logger.info(f"StreamProcessor '{self.name}' model preloaded on server startup")
+                    
                 except Exception as e:
                     if getattr(self._frame_processor, "state", None) is not None:
                         self._frame_processor.state.set_error(str(e))