使用Dask Worker插件管理LocalCluster的标准输出

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

# 配置Dask的日志级别，以避免Dask自身的冗余日志
logging.basicConfig(level=logging.WARNING)

class NullWriter:
    """
    一个空写入器，用于抑制输出。
    它模拟文件对象的write和flush方法，但实际上不做任何事情。
    """
    def write(self, s):
        pass # 忽略所有写入操作
    def flush(self):
        pass # 忽略所有刷新操作

class OutputRedirectorPlugin(WorkerPlugin):
    """
    Dask Worker插件，用于重定向工作器的标准输出。
    默认情况下，它会将sys.stdout重定向到一个NullWriter，从而抑制所有输出。
    """
    def __init__(self, target_stdout=None):
        self.original_stdout = None
        # 如果未指定目标，则使用NullWriter抑制输出
        self.target_stdout = target_stdout if target_stdout is not None else NullWriter()
        self.worker_id = None # 用于识别是哪个worker实例

    async def setup(self, worker):
        """
        在工作器启动时调用。
        在此方法中，我们保存原始的sys.stdout并将其替换为目标输出流。
        """
        self.worker_id = worker.name
        print(f"[{self.worker_id}] Plugin setup: Redirecting stdout...")
        self.original_stdout = sys.stdout
        sys.stdout = self.target_stdout
        # 也可以在此处重定向sys.stderr
        # self.original_stderr = sys.stderr
        # sys.stderr = NullWriter() # 或其他目标

    async def teardown(self, worker):
        """
        在工作器关闭时调用。
        在此方法中，我们将sys.stdout恢复到其原始状态。
        """
        if self.original_stdout:
            sys.stdout = self.original_stdout
            print(f"[{self.worker_id}] Plugin teardown: Restored stdout.")
        # if self.original_stderr:
        #     sys.stderr = self.original_stderr

# 示例函数，包含打印语句
def dask_function(i):
    # 这条打印语句应该被抑制
    print(f'Worker {os.getpid()} processing item {i}. This print should be suppressed!')
    return i**2

# 主执行流程
if __name__ == "__main__":
    print("--------------------------------------------------")
    print("正在启动Dask LocalCluster...")
    # n_workers: 工作器数量
    # processes=True: 每个工作器作为独立进程运行，这是我们遇到stdout问题的场景
    cluster = LocalCluster(n_workers=2, processes=True, threads_per_worker=1)
    client = Client(cluster)
    print(f"Dask Dashboard链接: {client.dashboard_link}")

    # 实例化插件，并将其注册到Dask Client。
    # 注册后，所有由该Client管理的工作器在启动时都会应用此插件。
    # 这里我们使用默认的NullWriter来抑制输出。
    plugin = OutputRedirectorPlugin()
    client.register_worker_plugin(plugin)
    print("OutputRedirectorPlugin已注册。")
    print("--------------------------------------------------")

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

    print("\n正在计算Dask任务 (工作器打印应该被抑制):")
    # 注意：这里的client.compute()会阻塞直到所有任务完成
    # .result()用于获取计算结果
    dask_outs = client.compute(dask_delays).result()
    print("计算完成。结果:", dask_outs)
    print("--------------------------------------------------")

    # 如果需要，可以在任务完成后取消注册插件
    # client.unregister_worker_plugin(plugin)
    # print("\n插件已取消注册。后续任务将再次打印:")
    # dask_delays_unreg = [dask.delayed(dask_function)(i) for i in range(5, 7)]
    # dask_outs_unreg = client.compute(dask_delays_unreg).result()
    # print("未注册插件的计算完成。结果:", dask_outs_unreg)
    # print("--------------------------------------------------")

    # 关闭Dask集群
    client.close()
    cluster.close()
    print("\nDask LocalCluster已关闭。")
    print("--------------------------------------------------")

# ... (NullWriter和dask_function定义不变)

# 修改插件实例化方式
if __name__ == "__main__":
    # ... (集群和客户端启动代码)

    # 重定向到文件
    log_file_path = "dask_worker_output.log"
    with open(log_file_path, "w", buffering=1) as f: # buffering=1 实现行缓冲，实时写入
        plugin_to_file = OutputRedirectorPlugin(target_stdout=f)
        client.register_worker_plugin(plugin_to_file)
        print(f"OutputRedirectorPlugin已注册，输出将写入到 '{log_file_path}'。")

        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()
        print("计算完成。结果:", dask_outs)

        client.close()
        cluster.close()

    print(f"\n请检查文件 '{log_file_path}' 获取工作器输出。")
    # ... (其他代码)

class PerWorkerFileRedirectorPlugin(WorkerPlugin):
    def __init__(self, log_base_path="dask_worker_logs"):
        self.log_base_path = log_base_path
        self.original_stdout = None
        self.log_file = None
        self.worker_id = None

    async def setup(self, worker):
        self.worker_id = worker.name
        log_dir = os.path.join(os.getcwd(), self.log_base_path)
        os.makedirs(log_dir, exist_ok=True)
        log_filename = os.path.join(log_dir, f"{self.worker_id}.log")

        self.log_file = open(log_filename, "w", buffering=1) # 在worker进程中打开文件
        self.original_stdout = sys.stdout
        sys.stdout = self.log_file
        print(f"[{self.worker_id}] Plugin setup: Redirecting stdout to {log_filename}")

    async def teardown(self, worker):
        if self.original_stdout:
            sys.stdout = self.original_stdout
        if self.log_file:
            self.log_file.close()
            print(f"[{self.worker_id}] Plugin teardown: Log file closed.")

# 使用方法：
# plugin_per_worker_file = PerWorkerFileRedirectorPlugin()
# client.register_worker_plugin(plugin_per_worker_file)

class CombinedOutputRedirectorPlugin(WorkerPlugin):
    def __init__(self, target_stdout=None, target_stderr=None):
        self.original_stdout = None
        self.original_stderr = None
        self.target_stdout = target_stdout if target_stdout is not None else NullWriter()
        self.target_stderr = target_stderr if target_stderr is not None else NullWriter()

    async def setup(self, worker):
        self.original_stdout = sys.stdout
        sys.stdout = self.target_stdout
        self.original_stderr = sys.stderr
        sys.stderr = self.target_stderr

    async def teardown(self, worker):
        if self.original_stdout:
            sys.stdout = self.original_stdout
        if self.original_stderr:
            sys.stderr = self.original_stderr