#!/usr/bin/env python3 """ Multi-threaded gRPC bridge server for Snakepit. This server variant uses a ThreadPoolExecutor to handle multiple concurrent requests within a single Python process. Designed for Python 3.13+ free-threading mode and CPU-intensive workloads. Key differences from grpc_server.py: - ThreadPoolExecutor with configurable thread count - Thread-safe adapter requirements - Concurrent request handling - Thread safety validation (optional) - Request tracking per thread Requirements: - Python 3.13+ (recommended for free-threading) - Thread-safe adapter implementation - Thread-safe ML libraries (NumPy, PyTorch, etc.) """ import argparse import asyncio import grpc import logging import signal import sys import time import threading import traceback from concurrent import futures from datetime import datetime from typing import Optional, Dict from collections import defaultdict # Add the package to Python path sys.path.insert(0, '.') import snakepit_bridge_pb2 as pb2 import snakepit_bridge_pb2_grpc as pb2_grpc from snakepit_bridge.session_context import SessionContext from snakepit_bridge.serialization import TypeSerializer from google.protobuf.timestamp_pb2 import Timestamp import json import pickle logging.basicConfig( format='%(asctime)s - [%(threadName)s] - %(name)s - %(levelname)s - %(message)s', level=logging.INFO ) logger = logging.getLogger(__name__) class ThreadSafetyMonitor: """ Monitor for tracking thread safety issues during execution. Detects: - Concurrent method access without locking - Thread-unsafe library usage - Race conditions """ def __init__(self, enabled: bool = False): self.enabled = enabled self.access_tracker: Dict[str, set] = defaultdict(set) self.lock = threading.Lock() self.warnings_issued: set = set() def record_access(self, method_name: str): """Record that current thread accessed a method""" if not self.enabled: return thread_id = threading.get_ident() with self.lock: self.access_tracker[method_name].add(thread_id) # Warn if multiple threads accessing if len(self.access_tracker[method_name]) > 1: warning_key = f"{method_name}_concurrent" if warning_key not in self.warnings_issued: logger.warning( f"⚠️ Concurrent access detected: {method_name} accessed by " f"{len(self.access_tracker[method_name])} different threads" ) self.warnings_issued.add(warning_key) def get_stats(self) -> dict: """Get monitoring statistics""" with self.lock: return { "tracked_methods": len(self.access_tracker), "warnings_issued": len(self.warnings_issued), "concurrent_accesses": sum( 1 for threads in self.access_tracker.values() if len(threads) > 1 ) } class ThreadedBridgeServiceServicer(pb2_grpc.BridgeServiceServicer): """ Multi-threaded gRPC bridge service. This servicer can handle multiple concurrent requests by dispatching them to a ThreadPoolExecutor. The adapter MUST be thread-safe. Thread Safety Requirements: - Adapter must implement thread-safe request handling - Shared state must use locks or thread-local storage - ML libraries must release GIL during computation """ def __init__(self, adapter_class, elixir_address: str, max_workers: int, enable_safety_checks: bool = False): logger.info(f"Initializing threaded servicer: max_workers={max_workers}, safety_checks={enable_safety_checks}") self.adapter_class = adapter_class self.elixir_address = elixir_address self.max_workers = max_workers self.server: Optional[grpc.aio.Server] = None # Thread safety monitoring self.safety_monitor = ThreadSafetyMonitor(enabled=enable_safety_checks) # Request tracking self.request_count = 0 self.active_requests = 0 self.start_time = time.time() self.stats_lock = threading.Lock() # Validate adapter thread safety self._validate_adapter_thread_safety() # Create async channel for proxying self.elixir_channel = grpc.aio.insecure_channel(elixir_address) self.elixir_stub = pb2_grpc.BridgeServiceStub(self.elixir_channel) # Create sync channel for SessionContext self.sync_elixir_channel = grpc.insecure_channel(elixir_address) self.sync_elixir_stub = pb2_grpc.BridgeServiceStub(self.sync_elixir_channel) logger.info(f"✅ Threaded servicer initialized. Ready for concurrent requests.") def _validate_adapter_thread_safety(self): """Validate that the adapter declares thread safety""" if not hasattr(self.adapter_class, '__thread_safe__'): logger.warning( f"⚠️ Adapter {self.adapter_class.__name__} does not declare thread safety. " f"Set __thread_safe__ = True if thread-safe, or use process mode instead." ) elif not self.adapter_class.__thread_safe__: logger.error( f"❌ Adapter {self.adapter_class.__name__} explicitly declares it is NOT thread-safe! " f"Use --mode process instead of --mode threaded." ) raise ValueError(f"Adapter {self.adapter_class.__name__} is not thread-safe") def _record_request_start(self): """Thread-safe request start tracking""" with self.stats_lock: self.request_count += 1 self.active_requests += 1 return self.request_count def _record_request_end(self): """Thread-safe request end tracking""" with self.stats_lock: self.active_requests -= 1 def get_stats(self) -> dict: """Get servicer statistics""" with self.stats_lock: uptime = time.time() - self.start_time return { "total_requests": self.request_count, "active_requests": self.active_requests, "max_workers": self.max_workers, "uptime_seconds": uptime, "requests_per_second": self.request_count / uptime if uptime > 0 else 0, "safety_monitor": self.safety_monitor.get_stats() } async def close(self): """Clean up resources""" if self.elixir_channel: await self.elixir_channel.close() if self.sync_elixir_channel: self.sync_elixir_channel.close() # Health & Session Management async def Ping(self, request, context): """Health check endpoint""" self.safety_monitor.record_access("Ping") response = pb2.PingResponse() response.message = f"Pong from threaded Python [{threading.current_thread().name}]: {request.message}" timestamp = Timestamp() timestamp.GetCurrentTime() response.server_time.CopyFrom(timestamp) return response async def InitializeSession(self, request, context): """Initialize session - proxy to Elixir""" self.safety_monitor.record_access("InitializeSession") logger.debug(f"[{threading.current_thread().name}] Proxying InitializeSession: {request.session_id}") return await self.elixir_stub.InitializeSession(request) async def CleanupSession(self, request, context): """Cleanup session - proxy to Elixir""" self.safety_monitor.record_access("CleanupSession") logger.debug(f"[{threading.current_thread().name}] Proxying CleanupSession: {request.session_id}") return await self.elixir_stub.CleanupSession(request) async def GetSession(self, request, context): """Get session - proxy to Elixir""" logger.debug(f"[{threading.current_thread().name}] Proxying GetSession: {request.session_id}") return await self.elixir_stub.GetSession(request) async def Heartbeat(self, request, context): """Heartbeat - proxy to Elixir""" return await self.elixir_stub.Heartbeat(request) # Variable Operations - All proxied to Elixir async def RegisterVariable(self, request, context): return await self.elixir_stub.RegisterVariable(request) async def GetVariable(self, request, context): return await self.elixir_stub.GetVariable(request) async def SetVariable(self, request, context): return await self.elixir_stub.SetVariable(request) async def GetVariables(self, request, context): return await self.elixir_stub.GetVariables(request) async def SetVariables(self, request, context): return await self.elixir_stub.SetVariables(request) async def ListVariables(self, request, context): return await self.elixir_stub.ListVariables(request) async def DeleteVariable(self, request, context): return await self.elixir_stub.DeleteVariable(request) # Tool Execution - Thread-safe concurrent execution async def ExecuteTool(self, request, context): """Execute tool with thread-safe handling""" request_id = self._record_request_start() thread_name = threading.current_thread().name logger.info(f"[{thread_name}] ExecuteTool #{request_id}: {request.tool_name} (session: {request.session_id})") self.safety_monitor.record_access("ExecuteTool") start_time = time.time() try: # Ensure session exists init_request = pb2.InitializeSessionRequest(session_id=request.session_id) try: self.sync_elixir_stub.InitializeSession(init_request) except grpc.RpcError as e: if e.code() != grpc.StatusCode.ALREADY_EXISTS: logger.debug(f"InitializeSession: {e}") # Create ephemeral context and adapter session_context = SessionContext(self.sync_elixir_stub, request.session_id) adapter = self.adapter_class() adapter.set_session_context(session_context) # Register tools if needed if hasattr(adapter, 'register_with_session'): adapter.register_with_session(request.session_id, self.sync_elixir_stub) # Initialize adapter if hasattr(adapter, 'initialize'): if asyncio.iscoroutinefunction(adapter.initialize): await adapter.initialize() else: adapter.initialize() # Decode parameters arguments = {key: TypeSerializer.decode_any(any_msg) for key, any_msg in request.parameters.items()} for key, binary_val in request.binary_parameters.items(): arguments[key] = pickle.loads(binary_val) # Execute tool if not hasattr(adapter, 'execute_tool'): raise NotImplementedError("Adapter does not support execute_tool") import inspect if inspect.iscoroutinefunction(adapter.execute_tool): result_data = await adapter.execute_tool( tool_name=request.tool_name, arguments=arguments, context=session_context ) else: result_data = adapter.execute_tool( tool_name=request.tool_name, arguments=arguments, context=session_context ) # Encode response response = pb2.ExecuteToolResponse(success=True) # Type inference if isinstance(result_data, dict): result_type = "map" elif isinstance(result_data, str): result_type = "string" elif isinstance(result_data, (int, float)): result_type = "float" elif isinstance(result_data, bool): result_type = "boolean" elif isinstance(result_data, list): result_type = "list" else: result_type = "string" any_msg, binary_data = TypeSerializer.encode_any(result_data, result_type) response.result.CopyFrom(any_msg) if binary_data: response.binary_result = binary_data response.execution_time_ms = int((time.time() - start_time) * 1000) logger.info( f"[{thread_name}] ExecuteTool #{request_id} completed in " f"{response.execution_time_ms}ms" ) return response except Exception as e: logger.error(f"[{thread_name}] ExecuteTool #{request_id} failed: {e}", exc_info=True) response = pb2.ExecuteToolResponse() response.success = False response.error_message = str(e) return response finally: self._record_request_end() async def ExecuteStreamingTool(self, request, context): """Execute streaming tool with thread-safe handling""" request_id = self._record_request_start() thread_name = threading.current_thread().name logger.info(f"[{thread_name}] ExecuteStreamingTool #{request_id}: {request.tool_name}") self.safety_monitor.record_access("ExecuteStreamingTool") try: # Similar to ExecuteTool but yields chunks init_request = pb2.InitializeSessionRequest(session_id=request.session_id) try: self.sync_elixir_stub.InitializeSession(init_request) except grpc.RpcError as e: if e.code() != grpc.StatusCode.ALREADY_EXISTS: logger.debug(f"InitializeSession: {e}") session_context = SessionContext(self.sync_elixir_stub, request.session_id) adapter = self.adapter_class() adapter.set_session_context(session_context) if hasattr(adapter, 'register_with_session'): adapter.register_with_session(request.session_id, self.sync_elixir_stub) if hasattr(adapter, 'initialize'): if asyncio.iscoroutinefunction(adapter.initialize): await adapter.initialize() else: adapter.initialize() arguments = {key: TypeSerializer.decode_any(any_msg) for key, any_msg in request.parameters.items()} for key, binary_val in request.binary_parameters.items(): arguments[key] = pickle.loads(binary_val) if not hasattr(adapter, 'execute_tool'): await context.abort(grpc.StatusCode.UNIMPLEMENTED, "Adapter does not support tool execution") return import inspect if inspect.iscoroutinefunction(adapter.execute_tool): stream_iterator = await adapter.execute_tool( tool_name=request.tool_name, arguments=arguments, context=session_context ) else: stream_iterator = adapter.execute_tool( tool_name=request.tool_name, arguments=arguments, context=session_context ) chunk_id_counter = 0 from snakepit_bridge.adapters.showcase.tool import StreamChunk if hasattr(stream_iterator, '__aiter__'): async for chunk_data in stream_iterator: if isinstance(chunk_data, StreamChunk): data_payload = chunk_data.data else: data_payload = chunk_data data_bytes = json.dumps(data_payload).encode('utf-8') chunk_id_counter += 1 chunk = pb2.ToolChunk( chunk_id=f"{request.tool_name}-{chunk_id_counter}", data=data_bytes, is_final=False ) yield chunk elif hasattr(stream_iterator, '__iter__'): for chunk_data in stream_iterator: if isinstance(chunk_data, StreamChunk): data_payload = chunk_data.data else: data_payload = chunk_data data_bytes = json.dumps(data_payload).encode('utf-8') chunk_id_counter += 1 chunk = pb2.ToolChunk( chunk_id=f"{request.tool_name}-{chunk_id_counter}", data=data_bytes, is_final=False ) yield chunk else: data_bytes = json.dumps(stream_iterator).encode('utf-8') yield pb2.ToolChunk( chunk_id=f"{request.tool_name}-1", data=data_bytes, is_final=True ) yield pb2.ToolChunk(is_final=True) logger.info(f"[{thread_name}] ExecuteStreamingTool #{request_id} completed") except Exception as e: logger.error(f"[{thread_name}] ExecuteStreamingTool #{request_id} failed: {e}", exc_info=True) context.abort(grpc.StatusCode.INTERNAL, str(e)) finally: self._record_request_end() # Placeholder methods async def WatchVariables(self, request, context): context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('WatchVariables not implemented') return yield async def AddDependency(self, request, context): return await self.elixir_stub.AddDependency(request) async def StartOptimization(self, request, context): return await self.elixir_stub.StartOptimization(request) async def StopOptimization(self, request, context): return await self.elixir_stub.StopOptimization(request) async def GetVariableHistory(self, request, context): return await self.elixir_stub.GetVariableHistory(request) async def RollbackVariable(self, request, context): return await self.elixir_stub.RollbackVariable(request) def set_server(self, server): """Set server reference for graceful shutdown""" self.server = server async def wait_for_elixir_server(elixir_address: str, max_retries: int = 30, initial_delay: float = 0.05): """Wait for Elixir gRPC server to become available""" delay = initial_delay for attempt in range(1, max_retries + 1): try: channel = grpc.aio.insecure_channel(elixir_address) stub = pb2_grpc.BridgeServiceStub(channel) request = pb2.PingRequest(message="connection_test") await asyncio.wait_for(stub.Ping(request), timeout=1.0) await channel.close() logger.info(f"✅ Connected to Elixir server at {elixir_address} after {attempt} attempt(s)") return True except (grpc.aio.AioRpcError, asyncio.TimeoutError, Exception) as e: if attempt < max_retries: logger.debug(f"Elixir server not ready (attempt {attempt}/{max_retries}), retrying in {delay:.2f}s...") await asyncio.sleep(delay) delay = min(delay * 2, 2.0) else: logger.error(f"❌ Failed to connect to Elixir server after {max_retries} attempts") return False return False async def serve_threaded( port: int, adapter_module: str, elixir_address: str, max_workers: int, enable_safety_checks: bool, shutdown_event: asyncio.Event ): """Start multi-threaded gRPC server""" logger.info(f"🚀 Starting threaded gRPC server: port={port}, max_workers={max_workers}") # Wait for Elixir server if not await wait_for_elixir_server(elixir_address): logger.error("Cannot start: Elixir server unavailable") sys.exit(1) # Import adapter module_parts = adapter_module.split('.') module_name = '.'.join(module_parts[:-1]) class_name = module_parts[-1] try: module = __import__(module_name, fromlist=[class_name]) adapter_class = getattr(module, class_name) except (ImportError, AttributeError) as e: logger.error(f"Failed to import adapter {adapter_module}: {e}") sys.exit(1) # Create server with ThreadPoolExecutor server = grpc.aio.server( futures.ThreadPoolExecutor(max_workers=max_workers), options=[ ('grpc.max_send_message_length', 100 * 1024 * 1024), ('grpc.max_receive_message_length', 100 * 1024 * 1024), ('grpc.max_concurrent_streams', max_workers), ('grpc.so_reuseport', 0), ] ) servicer = ThreadedBridgeServiceServicer( adapter_class, elixir_address, max_workers, enable_safety_checks ) servicer.set_server(server) pb2_grpc.add_BridgeServiceServicer_to_server(servicer, server) # Bind to port try: actual_port = server.add_insecure_port(f'[::]:{port}') if actual_port == 0 and port != 0: logger.error(f"Failed to bind to port {port}") sys.exit(1) except Exception as e: logger.error(f"Exception binding to port {port}: {e}") sys.exit(1) await server.start() # Signal ready logger.info(f"GRPC_READY:{actual_port}") logger.info(f"✅ Threaded gRPC server ready on port {actual_port} with {max_workers} worker threads") # Wait for shutdown server_task = asyncio.create_task(server.wait_for_termination()) try: await shutdown_event.wait() logger.info("Shutdown signal received") finally: server_task.cancel() await servicer.close() await server.stop(0.5) logger.info("Server stopped gracefully") def main(): parser = argparse.ArgumentParser(description='Snakepit Multi-Threaded gRPC Server') parser.add_argument('--port', type=int, default=0, help='Port to listen on') parser.add_argument('--adapter', type=str, required=True, help='Adapter class path') parser.add_argument('--elixir-address', type=str, required=True, help='Elixir server address') parser.add_argument('--max-workers', type=int, default=10, help='Thread pool size') parser.add_argument('--thread-safety-check', action='store_true', help='Enable thread safety checks') parser.add_argument('--snakepit-run-id', type=str, default='', help='Snakepit run ID') args = parser.parse_args() # Signal handlers shutdown_event = None def handle_signal(signum, frame): if shutdown_event and not shutdown_event.is_set(): asyncio.get_running_loop().call_soon_threadsafe(shutdown_event.set) signal.signal(signal.SIGTERM, handle_signal) signal.signal(signal.SIGINT, handle_signal) # Run server loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) shutdown_event = asyncio.Event() try: loop.run_until_complete(serve_threaded( args.port, args.adapter, args.elixir_address, args.max_workers, args.thread_safety_check, shutdown_event )) except BaseException as e: logger.error(f"Unhandled exception: {type(e).__name__}: {e}") traceback.print_exc() sys.exit(1) finally: loop.close() if __name__ == '__main__': main()