matrice-inference 0.1.2__py3-none-any.whl

matrice_inference/server/proxy_interface.py

@@ -0,0 +1,371 @@
+"""Module providing proxy_interface functionality."""
+
+import logging
+import time
+import threading
+from datetime import datetime, timezone
+from typing import Optional, Set
+import httpx
+import uvicorn
+import asyncio
+
+from fastapi import (
+    FastAPI,
+    HTTPException,
+    Request,
+)
+from fastapi.encoders import jsonable_encoder
+from fastapi.responses import JSONResponse
+from matrice_inference.server.inference_interface import InferenceInterface
+
+
+class MatriceProxyInterface:
+    """Interface for proxying requests to model servers."""
+
+    def __init__(
+        self,
+        session,
+        deployment_id: str,
+        deployment_instance_id: str,
+        external_port: int,
+        inference_interface: InferenceInterface,
+        auth_refresh_interval_minutes: int = 1,
+    ):
+        """Initialize proxy server.
+
+        Args:
+            session: Session object for authentication and RPC
+            deployment_id: ID of the deployment
+            deployment_instance_id: ID of the deployment instance
+            external_port: Port to expose externally
+            inference_interface: Interface for model inference
+            auth_refresh_interval_minutes: Minimum minutes between auth key refreshes
+        """
+        self.session = session
+        self.rpc = self.session.rpc
+        self.deployment_id = deployment_id
+        self.deployment_instance_id = deployment_instance_id
+        self.external_port = external_port
+        self.app = FastAPI()
+        self.logger = logging.getLogger(__name__)
+        self.inference_interface = inference_interface
+        self._shutdown_event = threading.Event()
+        self._server = None
+        self._server_thread = None
+        
+        # Auth key management
+        self.auth_keys: Set[str] = set()
+        self.auth_keys_info = []
+        self.auth_refresh_interval_minutes = auth_refresh_interval_minutes
+        self.last_auth_refresh_time = 0.0
+        self._auth_lock = threading.Lock()
+        
+        # Initialize auth keys on startup
+        self._initialize_auth_keys()
+        self._register_routes()
+
+    def _initialize_auth_keys(self):
+        """Initialize auth keys on startup."""
+        try:
+            self.update_auth_keys()
+            self.logger.info("Auth keys initialized successfully")
+        except Exception as exc:
+            self.logger.error("Failed to initialize auth keys: %s", str(exc))
+            # Continue without auth keys - they will be retried on first request
+
+    def _should_refresh_auth_keys(self) -> bool:
+        """Check if auth keys should be refreshed based on time interval."""
+        current_time = time.time()
+        time_since_last_refresh = current_time - self.last_auth_refresh_time
+        return time_since_last_refresh >= (self.auth_refresh_interval_minutes * 60)
+
+    def _refresh_auth_keys_if_needed(self):
+        """Refresh auth keys if the refresh interval has passed."""
+        if self._should_refresh_auth_keys():
+            with self._auth_lock:
+                # Double-check after acquiring lock
+                if self._should_refresh_auth_keys():
+                    try:
+                        self.update_auth_keys()
+                        self.logger.debug("Auth keys refreshed successfully")
+                    except Exception as exc:
+                        self.logger.error("Failed to refresh auth keys: %s", str(exc))
+
+    def validate_auth_key(self, auth_key):
+        """Validate auth key.
+
+        Args:
+            auth_key: Authentication key to validate
+
+        Returns:
+            bool: True if valid, False otherwise
+        """
+        if not auth_key:
+            return False
+        
+        # Refresh auth keys if needed before validation
+        self._refresh_auth_keys_if_needed()
+        
+        with self._auth_lock:
+            return auth_key in self.auth_keys
+
+    def _parse_expiry_time(self, expiry_time_str: str) -> float:
+        """Parse expiry time string to timestamp."""
+        # Handle different ISO format variations
+        try:
+            # Replace Z with timezone if needed
+            if expiry_time_str.endswith("Z"):
+                expiry_time_str = expiry_time_str.replace("Z", "+00:00")
+            
+            # Normalize ISO format for proper parsing
+            if '.' in expiry_time_str:
+                main, rest = expiry_time_str.split('.', 1)
+                if '+' in rest:
+                    frac, tz = rest.split('+', 1)
+                    frac = (frac + '000000')[:6]  # pad/truncate to 6 digits
+                    expiry_time_str = f"{main}.{frac}+{tz}"
+                elif '-' in rest:
+                    frac, tz = rest.split('-', 1)
+                    frac = (frac + '000000')[:6]  # pad/truncate to 6 digits
+                    expiry_time_str = f"{main}.{frac}-{tz}"
+        except Exception as err:
+            self.logger.error("Error parsing expiry time: %s", str(err))
+            expiry_time_str = expiry_time_str.replace("Z", "+00:00")
+        return datetime.fromisoformat(expiry_time_str).timestamp()
+
+    def update_auth_keys(self) -> None:
+        """Fetch and validate auth keys for the deployment."""      
+        try:
+            response = self.rpc.get(f"/v1/inference/{self.deployment_id}", raise_exception=False)
+            if not response.get("success"):
+                self.logger.error("Failed to fetch auth keys")
+                return
+            
+            if response["data"]["authKeys"]:
+                self.auth_keys_info = response["data"]["authKeys"]
+            else:
+                self.auth_keys_info = []
+            
+            if not self.auth_keys_info:
+                self.logger.debug("No auth keys found for deployment")
+                return
+            
+            current_time = time.time()
+            self.auth_keys.clear()
+            
+            for auth_key_info in self.auth_keys_info:
+                try:
+                    expiry_time = self._parse_expiry_time(auth_key_info["expiryTime"])
+                    if expiry_time > current_time:
+                        self.auth_keys.add(auth_key_info["key"])
+                    else:
+                        self.logger.debug("Skipping expired auth key")
+                except (ValueError, KeyError) as err:
+                    self.logger.error("Invalid auth key data: %s", err)
+                    continue
+            
+            # Update last refresh time
+            self.last_auth_refresh_time = current_time
+            
+            self.logger.debug(
+                "Successfully loaded %d valid auth keys",
+                len(self.auth_keys),
+            )
+        except Exception as err:
+            self.logger.error("Error fetching auth keys: %s", str(err))
+            raise
+
+    def _register_routes(self):
+        """Register proxy routes."""
+
+        @self.app.post("/inference")
+        async def proxy_request(request: Request):
+            # Check if server is shutting down
+            if self._shutdown_event.is_set():
+                raise HTTPException(
+                    status_code=503,
+                    detail="Server is shutting down",
+                )
+            
+            # Parse form data manually
+            try:
+                form_data = await request.form()
+            except Exception as e:
+                raise HTTPException(
+                    status_code=400,
+                    detail=f"Failed to parse form data: {str(e)}",
+                ) from e
+            
+            # Extract parameters from form data
+            auth_key = form_data.get("auth_key") or form_data.get("authKey")
+            input_file = form_data.get("input")
+            input_url_value = form_data.get("input_url") or form_data.get("inputUrl")
+            extra_params = form_data.get("extra_params")
+            apply_post_processing = form_data.get("apply_post_processing", "false")
+            
+            # if not self.validate_auth_key(auth_key): # TODO: enable once fixed to send the external auth key for FR server
+            #     raise HTTPException(
+            #         status_code=401,
+            #         detail="Invalid auth key",
+            #     )
+            
+            # Handle file input
+            input_data = None
+            if input_file and hasattr(input_file, 'read'):
+                input_data = await input_file.read()
+            elif isinstance(input_file, bytes):
+                input_data = input_file
+            
+            if input_url_value:
+                try:
+                    # Use timeout and error handling for URL downloads
+                    timeout = httpx.Timeout(60.0)  # 60 second timeout
+                    async with httpx.AsyncClient(timeout=timeout) as client:
+                        response = await client.get(input_url_value)
+                        response.raise_for_status()  # Raise exception for HTTP errors
+                        input_data = response.content
+                except asyncio.CancelledError:
+                    # Handle shutdown during request
+                    raise HTTPException(
+                        status_code=503,
+                        detail="Request cancelled due to server shutdown",
+                    )
+                except httpx.TimeoutException:
+                    raise HTTPException(
+                        status_code=408,
+                        detail="Timeout fetching input URL",
+                    )
+                except httpx.HTTPStatusError as e:
+                    raise HTTPException(
+                        status_code=400,
+                        detail=f"HTTP error fetching input URL: {e.response.status_code}",
+                    )
+                except Exception as e:
+                    raise HTTPException(
+                        status_code=400,
+                        detail=f"Error fetching input URL: {str(e)}",
+                    )
+            
+            if not input_data:
+                raise HTTPException(
+                    status_code=400,
+                    detail="No input provided",
+                )
+            
+            # Parse apply_post_processing parameter
+            apply_post_processing_flag = False
+            if apply_post_processing:
+                apply_post_processing_flag = apply_post_processing.lower() in ("true", "1", "yes")
+            
+            try:
+                # Check shutdown again before inference
+                if self._shutdown_event.is_set():
+                    raise HTTPException(
+                        status_code=503,
+                        detail="Server is shutting down",
+                    )
+                
+                result, post_processing_result = await self.inference_interface.inference(
+                    input=input_data,
+                    extra_params=extra_params,
+                    apply_post_processing=apply_post_processing_flag
+                )
+                
+
+                response_data = {
+                    "status": 1,
+                    "message": "Request success",
+                    "result": result,
+                }
+                
+                # Include post-processing results if available
+                if post_processing_result is not None:
+                    response_data["post_processing_result"] = post_processing_result
+                    response_data["post_processing_applied"] = True
+                else:
+                    response_data["post_processing_applied"] = False
+
+                return JSONResponse(
+                    content=jsonable_encoder(response_data)
+                )
+            except asyncio.CancelledError:
+                # Handle shutdown during inference
+                self.logger.info("Request cancelled during inference due to shutdown")
+                raise HTTPException(
+                    status_code=503,
+                    detail="Request cancelled due to server shutdown",
+                )
+            except Exception as exc:
+                self.logger.error("Proxy error: %s", str(exc))
+                raise HTTPException(
+                    status_code=500,
+                    detail=str(exc),
+                ) from exc
+
+
+    def start(self):
+        """Start the proxy server in a background thread."""
+        def run_server():
+            """Run the uvicorn server."""
+            try:
+                self.logger.info(
+                    "Starting proxy server on port %d",
+                    self.external_port,
+                )
+                config = uvicorn.Config(
+                    app=self.app,
+                    host="0.0.0.0",
+                    port=self.external_port,
+                    log_level="info",
+                )
+                self._server = uvicorn.Server(config)
+                self._server.run()
+
+            except Exception as exc:
+                if not self._shutdown_event.is_set():
+                    self.logger.error(
+                        "Failed to start proxy server: %s",
+                        str(exc),
+                    )
+                else:
+                    self.logger.info("Proxy server stopped during shutdown")
+        
+        # Start the server in a background thread
+        self._server_thread = threading.Thread(target=run_server, daemon=False, name="ProxyServer")
+        self._server_thread.start()
+        
+        # Wait a moment for the server to start
+        time.sleep(0.5)
+        self.logger.info("Proxy server thread started successfully")
+
+    def stop(self):
+        """Stop the proxy server gracefully."""
+        try:
+            self.logger.info("Stopping proxy server...")
+            
+            # Signal shutdown to prevent new requests
+            self._shutdown_event.set()
+            
+            # Stop the uvicorn server if it exists
+            if self._server:
+                try:
+                    # Force shutdown the server
+                    if hasattr(self._server, 'should_exit'):
+                        self._server.should_exit = True
+                    if hasattr(self._server, 'force_exit'):
+                        self._server.force_exit = True
+                except Exception as exc:
+                    self.logger.warning("Error stopping uvicorn server: %s", str(exc))
+            
+            # Wait for the server thread to finish
+            if self._server_thread and self._server_thread.is_alive():
+                self.logger.info("Waiting for proxy server thread to stop...")
+                self._server_thread.join(timeout=5.0)
+                if self._server_thread.is_alive():
+                    self.logger.warning("Proxy server thread did not stop within timeout")
+                else:
+                    self.logger.info("Proxy server thread stopped successfully")
+            
+            self.logger.info("Proxy server stopped")
+        except Exception as exc:
+            self.logger.error("Error stopping proxy server: %s", str(exc))