shutdown bug

microsoft · Nov 8, 2023 · 3115397 · 3115397
1 parent aaf202b
commit 3115397
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/mii/batching/ragged_batching.py b/mii/batching/ragged_batching.py
@@ -713,8 +713,7 @@ def __call__(self) -> None:
         while True:
             self.generate()
 
-            if (self.stop_thread and self.request_queue.empty()
-                    and all(q.empty() for q in self.result_queues.values())):
+            if self.stop_thread:  # and self.request_queue.empty() and all(q.empty() for q in self.result_queues.values())):
                 break
 
     def _get_uid(self, session_id: Union[str, None]):
@@ -758,8 +757,8 @@ def put_request(self,
                 self.request_queue.put(r)
 
         # Temporary hack to avoid non-rank 0 processes not shutting down. See related TODO above.
-        if self.is_rank_0:
-            self.request_queue.empty()
+        if not self.is_rank_0:
+            self.request_queue.queue.clear()
 
         return uid
 
@@ -814,4 +813,4 @@ def destroy_session(self,
                         post_processing=None,
                         stream=None,
                     ))
-                self.uids.remove(uid)
+            self.uids.remove(uid)