diff --git a/src/litserve/loops/base.py b/src/litserve/loops/base.py index 78ff6640..c79fa76b 100644 --- a/src/litserve/loops/base.py +++ b/src/litserve/loops/base.py @@ -129,6 +129,7 @@ def run( """ + zmq_ctx: Optional[zmq.Context] = None socket: Optional[zmq.Socket] = None def pre_setup(self, lit_api: LitAPI, spec: Optional[LitSpec]): @@ -265,6 +266,7 @@ def put_error_response( def __del__(self): if self.socket: self.socket.close() + self.zmq_ctx.term() class DefaultLoop(LitLoop): diff --git a/src/litserve/loops/loops.py b/src/litserve/loops/loops.py index d9f375dc..23667d3a 100644 --- a/src/litserve/loops/loops.py +++ b/src/litserve/loops/loops.py @@ -84,6 +84,7 @@ def inference_worker( socket = ctx.socket(zmq.PUB) logger.debug(f"Inference worker binding to port {zmq_port}") socket.bind(f"tcp://*:{zmq_port}") + loop.zmq_context = ctx loop( lit_api, @@ -101,4 +102,4 @@ def inference_worker( ) if use_zmq: socket.close() - ctx.destroy() + loop.zmq_context.term()