用于多进程python输出的类屏幕分割屏幕

2024-09-24 02:25:16 发布

您现在位置:Python中文网/ 问答频道 /正文

我经常遇到这样的情况:我必须在大量输入上运行一些时间密集型代码,并希望加速并行运行多个代码实例(在不同的CPU内核或Cuda设备上)。你知道吗

到目前为止,我是手工完成的:

  • 打开屏幕/tmux
  • 创建拆分平面
  • 在每个窗格中只使用一小部分输入来运行代码

然而,这很麻烦,容易出错,并且在代码的非并行部分的情况下不起作用。因此我有时多处理池并行运行代码,但如果我想/需要能够读取不同部分的输出,这就行不通了,如果并行运行的(第三方)代码包含类似tqdm的内容,这会变得一团糟。你知道吗

我想创建一个分割屏幕,并将每个不同进程的stdout/stderr写入其中一个分割窗格。然而,我还没有找到任何包装/技巧来正确地做到这一点。 我使用了urwind、ncurses和prompt\u工具包,但它们不是用来服务stdout/stderr的。你知道吗

我目前最好的解决方案是使用libtmux:

import os
import sys
import tempfile
import time
from multiprocessing import Pool, current_process

import libtmux
from tqdm import tqdm

tmpdir = tempfile.mkdtemp()
fifo1_name = os.path.join(tmpdir, "fifo1")
fifo2_name = os.path.join(tmpdir, "fifo2")
os.mkfifo(fifo1_name)
os.mkfifo(fifo2_name)
session_name = "session_name"


def setup(para):
    """Change stdout/stderr to point to the fifo connected to the right tmux plane"""
    pid = int(current_process().name.split("-")[-1])
    if pid % 2 == 0:
        fifo_write = open(fifo1_name, "w", 1)
    else:
        fifo_write = open(fifo2_name, "w", 1)
    sys.stdout = fifo_write
    sys.stderr = fifo_write
    run(para, pid)


def run(para, pid):
    """Dummy function"""
    print("start run nr.", para, "on", pid)
    for _ in tqdm(range(10), desc=str(para)):
        time.sleep(1)
    print("finished run nr.", para)


parameter_space = [1, 2, 3, 4, 5, 6]

# Set up tmux
server = libtmux.Server()
session = server.find_where({"session_name": session_name})
if session is None:
    session = server.new_session(
        session_name=session_name, window_command="cat " + fifo1_name
    )
    window = session.windows[0]
    window.panes[0]
    window.split_window(shell="cat " + fifo2_name)
# session.attach_session() Does not work for unknown reasons the planes stay empty

# Run always 2 runs in parallel
with Pool(2) as p:
    p.map(setup, parameter_space)

if session:
    session.kill_session()

这仍然有一个问题,我必须手动附加到会话,因为attach_session()会导致空窗格。而且,这感觉像是一个黑客和丑陋的解决办法。你知道吗

有更好/更好的方法吗?你知道吗


Tags: 代码nameimportossessionstderrstdoutwindow