如何通过asyncio.Event终止线程中的model.generate并释放GPU资源？

直接硬杀Python线程是不安全的，容易导致GPU显存泄漏、模型状态异常。正确的做法是通过自定义Streamer结合cancel_event来主动中断generate流程，让线程自然退出，同时释放GPU资源。

具体实现步骤 #

1. 自定义带取消检查的Streamer
   继承TextIteratorStreamer，重写put方法，每次接收token后检查cancel_event，如果已触发就抛出异常终止生成流程：

from transformers import TextIteratorStreamer
import asyncio

class CancelableTextStreamer(TextIteratorStreamer):
    def __init__(self, cancel_event: asyncio.Event, **kwargs):
        super().__init__(**kwargs)
        self.cancel_event = cancel_event

    def put(self, value):
        # 每次生成token后检查取消信号
        if self.cancel_event.is_set():
            raise RuntimeError("Generation canceled by user")
        super().put(value)

2. 初始化取消事件与自定义Streamer
   创建cancel_event，将其传入自定义Streamer，再把Streamer加入generation_kwargs：

cancel_event = asyncio.Event()
# 假设tokenizer、model已完成初始化
streamer = CancelableTextStreamer(cancel_event, tokenizer=tokenizer, skip_prompt=True)
generation_kwargs = {
    "inputs": inputs,  # 你的输入张量
    "streamer": streamer,
    # 其他generate参数：max_new_tokens、temperature等
}

3. 启动线程并处理取消逻辑
   启动线程后，当需要取消生成时，调用cancel_event.set()即可触发中断：

import threading

thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
thread.start()

# 示例：模拟异步场景下触发取消
async def trigger_cancel():
    # 模拟用户主动取消操作（根据实际业务场景调整触发时机）
    await asyncio.sleep(5)
    cancel_event.set()
    # 等待线程自然退出
    thread.join()

# 执行取消逻辑（根据你的业务框架调用）
asyncio.run(trigger_cancel())

关键说明 #

• 为什么不杀线程？：Python的threading.Thread没有安全的终止方法，强制终止会导致模型无法正确释放GPU显存、内部状态混乱，甚至需要重启进程才能恢复。
• 异常的作用：model.generate在生成每个token时会调用Streamer的put方法，抛出异常后会终止生成循环，自动释放GPU资源，线程会因未捕获异常自然终止。
• 兼容性：绝大多数HuggingFace Transformers模型的generate方法都能响应这种中途抛出的异常，确保资源正确释放。

内容的提问来源于stack exchange，提问作者A.A