From 0b661bae853650ec8782b4bd627cf356c30e41b8 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 8 Aug 2026 13:12:40 +0800 Subject: [PATCH] fix: remove blocking cleanup from streaming generator - stream finally froze main thread on cache.task_free - scheduler handles cleanup in next loop iteration instead --- astrai/inference/engine.py | 26 +++++++++++--------------- 1 file changed, 11 insertions(+), 15 deletions(-) diff --git a/astrai/inference/engine.py b/astrai/inference/engine.py index 6de91e8..0c739df 100644 --- a/astrai/inference/engine.py +++ b/astrai/inference/engine.py @@ -204,21 +204,17 @@ class InferenceEngine: def gen(): nonlocal remaining - try: - while remaining > 0: - items = result.pop_all() - for idx, token in items: - if token is STOP: - if not finished[idx]: - finished[idx] = True - remaining -= 1 - else: - yield (idx, token) if is_batch else token - if remaining > 0: - result.wait(timeout=0.05) - finally: - for tid in task_ids: - self.scheduler.remove_task(tid) + while remaining > 0: + items = result.pop_all() + for idx, token in items: + if token is STOP: + if not finished[idx]: + finished[idx] = True + remaining -= 1 + else: + yield (idx, token) if is_batch else token + if remaining > 0: + result.wait(timeout=0.05) return gen()