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