如何通过Multiprocessing实现TSNE运行中间输出实时推送到前端

核心问题说明 #

• 原有多进程代码未启用Pipe通信能力：子进程未向连接端写入任何数据，父进程的recv()方法会永久阻塞；且循环内连续调用两次recv()会直接丢弃一半接收数据
• 标准输出默认开启缓冲：Python stdout默认采用全缓冲模式，重定向到StringIO时输出会攒到进程结束才统一返回，无法拿到实时运行日志

完整实现代码 #

from multiprocessing import Process, Pipe
from sklearn.manifold import TSNE
import random
import io
import sys

# 自定义管道输出流：每收到输出就实时发送给父进程
class PipeStdout(io.TextIOBase):
    def __init__(self, conn):
        self.conn = conn
    def write(self, content):
        # 过滤空行输出
        if content.strip():
            self.conn.send(("log", content.strip()))
    def flush(self):
        pass

def TSNE_transform(conn):
    # 重定向标准输出到管道流
    old_stdout = sys.stdout
    sys.stdout = PipeStdout(conn)
    try:
        X = [[random.randint(1,9) for i in range(9)] for j in range(10)]
        model = TSNE(n_components=3, verbose=3)
        result = model.fit_transform(X)
        # 发送最终计算结果
        conn.send(("result", result))
    finally:
        # 恢复标准输出，发送结束标记
        sys.stdout = old_stdout
        conn.send(None)

if __name__ == '__main__':
    parent_conn, child_conn = Pipe()
    # 把子进程的管道端传入处理函数
    p = Process(target=TSNE_transform, args=(child_conn,))
    # 强制关闭Python输出缓冲，避免sklearn内部输出攒批
    p.start(env={"PYTHONUNBUFFERED": "1"})
    
    final_result = None
    while True:
        msg = parent_conn.recv()
        if msg is None:
            break
        msg_type, content = msg
        if msg_type == "log":
            # 此处可直接将日志推送到前端，比如WebSocket/流式HTTP响应
            print(f"[实时日志] {content}")
        elif msg_type == "result":
            final_result = content
    
    p.join()
    print(f"[最终结果] 计算完成，结果长度：{len(final_result)}")

前端对接方案 #

• 短连接场景：采用SSE（服务器推送事件）实现流式响应，父进程每拿到一条日志就按SSE格式data: 日志内容\n\n返回给前端，前端通过EventSource对象接收即可实时渲染
• 长连接场景：采用WebSocket通信，父进程收到日志后直接通过WebSocket连接推送给前端，延迟更低，也支持双向交互

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