Dask LocalCluster 工作器输出重定向指南

import dask
from dask.distributed import Client, LocalCluster

def dask_function(i):
    print(f'Worker processing {i}: Ignore me!')
    return i**2

if __name__ == "__main__":
    cluster = LocalCluster(n_workers=2, processes=True, threads_per_worker=1)
    client = Client(cluster)

    dask_delays = []
    for i in range(5):
        dask_delays.append(dask.delayed(dask_function)(i))

    print("开始计算Dask任务...")
    dask_outs = client.compute(dask_delays).result()
    print("Dask任务计算完成。")

    client.close()
    cluster.close()

import sys
import os
from dask.distributed import WorkerPlugin

class SuppressPrintPlugin(WorkerPlugin):
    """
    一个Dask工作器插件，用于在工作器运行时抑制其标准输出。
    """
    def setup(self, worker):
        """
        在工作器启动时调用。
        将sys.stdout重定向到操作系统的空设备（/dev/null 或 NUL）。
        """
        self.original_stdout = sys.stdout
        # 打开一个指向空设备的写模式文件句柄
        # os.devnull 在不同操作系统上会自动解析为 /dev/null (Unix) 或 NUL (Windows)
        self.devnull_fd = open(os.devnull, 'w')
        sys.stdout = self.devnull_fd
        # print(f"Worker {worker.name} stdout redirected.")
        # 注意：这行代码将不会打印，因为它在重定向之后执行

    def teardown(self, worker):
        """
        在工作器关闭时调用。
        恢复sys.stdout到其原始状态，并关闭空设备的文件句柄。
        """
        if hasattr(self, 'original_stdout'):
            sys.stdout = self.original_stdout
        if hasattr(self, 'devnull_fd') and not self.devnull_fd.closed:
            self.devnull_fd.close()
        # print(f"Worker {worker.name} stdout restored.")
        # 这行代码将正常打印，因为它在恢复之后执行

import dask
from dask.distributed import Client, LocalCluster
import sys
import os

# 1. 定义Dask任务函数
def dask_function(i):
    print(f'Worker processing {i}: Ignore me!') # 包含打印输出
    return i**2

# 2. 定义SuppressPrintPlugin（如上所示）
class SuppressPrintPlugin(WorkerPlugin):
    def setup(self, worker):
        self.original_stdout = sys.stdout
        self.devnull_fd = open(os.devnull, 'w')
        sys.stdout = self.devnull_fd

    def teardown(self, worker):
        if hasattr(self, 'original_stdout'):
            sys.stdout = self.original_stdout
        if hasattr(self, 'devnull_fd') and not self.devnull_fd.closed:
            self.devnull_fd.close()

if __name__ == "__main__":
    print("启动Dask LocalCluster...")
    # 为了演示效果，使用较少的工作器
    cluster = LocalCluster(n_workers=2, processes=True, threads_per_worker=1)
    client = Client(cluster)
    print(f"Dask Dashboard链接: {client.dashboard_link}")

    # 注册插件
    client.register_worker_plugin(SuppressPrintPlugin())
    print("SuppressPrintPlugin已成功注册。")

    dask_delays = []
    for i in range(5):
        dask_delays.append(dask.delayed(dask_function)(i))

    print("\n开始计算Dask任务（工作器输出将被抑制）...")
    dask_outs = client.compute(dask_delays).result() # .result() 阻塞直到计算完成
    print("Dask任务计算完成。")

    print("\n计算结果:", dask_outs)

    # 清理Dask资源
    client.close()
    cluster.close()
    print("Dask LocalCluster已关闭。")

    # 验证插件关闭后，主进程的打印功能恢复正常
    print("\n此打印语句应在Dask关闭后正常显示。")