terminate context

Lightning-AI · Jan 9, 2025 · 105fa5a · 105fa5a
1 parent a4fc585
commit 105fa5a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git 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
@@ -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()