Dask LocalCluster工作器输出重定向与抑制策略

# Linux/macOS: 将输出重定向到文件
dask scheduler --port 8786 > scheduler.log 2>&1 &

# Linux/macOS: 完全抑制输出
dask scheduler --port 8786 > /dev/null 2>&1 &

# Windows: 将输出重定向到文件
start /B dask scheduler --port 8786 > scheduler.log 2>&1

# Windows: 完全抑制输出
start /B dask scheduler --port 8786 > NUL 2>&1

# Linux/macOS: 将工作器输出重定向到文件
dask worker tcp://localhost:8786 --nprocs 4 > worker.log 2>&1 &

# Linux/macOS: 完全抑制工作器输出
dask worker tcp://localhost:8786 --nprocs 4 > /dev/null 2>&1 &

# Windows: 将工作器输出重定向到文件
start /B dask worker tcp://localhost:8786 --nprocs 4 > worker.log 2>&1

# Windows: 完全抑制工作器输出
start /B dask worker tcp://localhost:8786 --nprocs 4 > NUL 2>&1

from dask.distributed import Client
import dask

def dask_function(i):
    print(f'Worker processing {i}. This will be redirected!') # 此处的打印会被重定向
    return i**2

# 连接到手动启动的调度器
client = Client("tcp://localhost:8786")

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

dask_outs = client.compute(dask_delays).result()
print("Computation finished. Results:", dask_outs)

client.close()
# 注意：手动启动的调度器和工作器需要手动关闭
# 例如，通过 `kill` 命令或任务管理器

import sys
import os
from distributed.diagnostics.plugin import WorkerPlugin
from distributed import LocalCluster, Client
import dask

class SuppressPrintsPlugin(WorkerPlugin):
    """
    一个Dask工作器插件，用于在工作器启动时重定向其标准输出和标准错误。
    """
    def __init__(self, redirect_path=os.devnull):
        self.redirect_path = redirect_path
        self._original_stdout = None
        self._original_stderr = None
        self._redirected_stdout_file = None
        self._redirected_stderr_file = None

    def setup(self, worker):
        """
        在工作器启动时调用，用于重定向sys.stdout和sys.stderr。
        """
        # 存储原始的sys.stdout和sys.stderr，以便在teardown时恢复
        self._original_stdout = sys.stdout
        self._original_stderr = sys.stderr

        # 打开目标文件（如/dev/null或自定义日志文件）
        # 'w' 模式会覆盖文件，如果希望追加，请使用 'a'
        self._redirected_stdout_file = open(self.redirect_path, 'w')
        self._redirected_stderr_file = open(self.redirect_path, 'w')

        # 重定向sys.stdout和sys.stderr
        sys.stdout = self._redirected_stdout_file
        sys.stderr = self._redirected_stderr_file

        # 可以在此处添加工作器启动时的日志信息
        # self._original_stdout.write(f"Worker {worker.name} stdout/stderr redirected to {self.redirect_path}\n")

    def teardown(self, worker):
        """
        在工作器关闭时调用，用于恢复原始的sys.stdout和sys.stderr。
        """
        # 恢复原始的sys.stdout和sys.stderr
        if self._original_stdout:
            sys.stdout = self._original_stdout
        if self._original_stderr:
            sys.stderr = self._original_stderr

        # 关闭重定向文件
        if self._redirected_stdout_file:
            self._redirected_stdout_file.close()
        if self._redirected_stderr_file:
            self._redirected_stderr_file.close()

        # self._original_stdout.write(f"Worker {worker.name} stdout/stderr restored.\n")

# 示例函数，包含打印语句
def dask_function(i):
    print(f'Worker processing {i}. This message should be suppressed!')
    return i**2

# 初始化LocalCluster和Client
# n_workers 参数可以控制启动的工作器数量
cluster = LocalCluster(n_workers=4, processes=True, dashboard_address=None)
client = Client(cluster)

# 实例化插件，可以将输出重定向到 /dev/null 进行完全抑制
# 或者重定向到特定文件，例如：SuppressPrintsPlugin(f'/tmp/dask_worker_output_{client.scheduler_info["id"]}.log')
plugin = SuppressPrintsPlugin(redirect_path=os.devnull) 

# 注册插件。插件会在每个工作器启动时被激活
client.register_worker_plugin(plugin)

# 创建Dask延迟任务
dask_delays = []
for i in range(10):
    dask_delays.append(dask.delayed(dask_function)(i))

# 执行计算
dask_outs = client.compute(dask_delays).result()
print("Computation finished. Results:", dask_outs) # 这条打印会正常显示在控制台

# 关闭客户端和集群
client.close()
cluster.close()