From 105fa5af6976c6e88f3592150d27dce9afcddfdc Mon Sep 17 00:00:00 2001
From: Aniket Maurya <theaniketmaurya@gmail.com>
Date: Thu, 9 Jan 2025 17:59:43 +0000
Subject: [PATCH] terminate context

---
 src/litserve/loops/base.py  | 2 ++
 src/litserve/loops/loops.py | 3 ++-
 2 files changed, 4 insertions(+), 1 deletion(-)

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()