python线程无法手动关闭

2024-06-14  本文已影响0人  KyoDante

业务场景是:大模型每次推理都是新建一个线程进行推理,如果用户要取消回答,或者遇到异常的时候,需要停止线程;主要针对的是第一种情况,流失推理实际上就是用一个队列保存推理之后的结果,然后用另外一个线程不断地从这个队列里面取推理结果返回,达到所谓的“打字机”效果;以下是模拟的场景:


from queue import Queue, Empty
from threading import Thread
from multiprocessing import Process
import time


class Streamer:
    def __init__ (self, _text_queue):
        self.text_queue = _text_queue
        self.stop_signal = "stop"
    
    def put(self, value):
        self.text_queue.put(value, timeout=0.5)
    
    def __iter__(self):
        return self
    
    def __next__(self):
        try:
            value = self.text_queue.get(timeout=0.5)
        except Empty as empty:
            value = self.stop_signal
        if value == self.stop_signal:
            print("stop here!")
            raise StopIteration()
        return value

def test_func():
    def _inference():
        count = 0
        while True:
            if count < 10:
                streamer.put("shit")
                time.sleep(20) # 例如卡在执行.so,比如调用模型的推理
                print("put shit & getting gem")
                time.sleep(0.2 * count)  # 后面推理超时
                count += 1
            else:
                print("breaking!")
                break
    t_queue = Queue()
    streamer = Streamer(t_queue)
    
    # 模式1 Thread daemon
    # inference_thread = Thread(target=_inference, daemon=True)
    # inference_thread.start()
    # for idx, i in enumerate(streamer):
    #     print(f"get {i}{idx}")
    
    # 模式2 Thread+stop
    # https://www.cnblogs.com/conscience-remain/p/16930488.html
    # stop_thread(inference_thread) # 54行会导致停止失败

    # 模式3 ThreadPoolExecutor的cancel
    # import concurrent.futures
    # with concurrent.futures.ThreadPoolExecutor(max_workers=1) as tpe:
    #     future = tpe.submit(_inference)
    #     for idx, i in enumerate(streamer):
    #         print(i, idx)
    #     # getting jammed or finished
    #     if future.running():
    #         print("canceling here!")
    #         future.cancel() # 取消线程?如果正在运行的,并不会生效
    #     try:
    #         future.result(timeout=1)
    #     except concurrent.futures.TimeoutError as ex:
    #         print(f"ex: {ex}")

    # 模式4 trace 主线程都退出了,子线程还gam
    # thread = thread_with_trace(target=_inference, daemon=True)
    # thread.start()
    # for idx, i in enumerate(streamer):
    #     print(i, idx)
    # thread.kill()

def worker():
    worker = Thread(target=test_func)
    worker.start()
    worker.join()

if __name__ == "__main__":
    worker()
    # time.sleep(5) # 主线程如果不退出,daemon也不会退出
    print("out")

以上几种方式,都无法正常地停止掉当前的推理线程。因此无法达到停止当前推理的功能;
另外,涉及到Python中Thread和ThreadPoolExecutor的相关用法,另外可以了解什么是daemon线程。

上一篇 下一篇

猜你喜欢

热点阅读