控制Dask LocalCluster工作器输出：两种策略详解

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

# 定义一个Dask函数，其中包含打印语句
def dask_function(i):
    print(f'工作器正在处理任务 {i}!') # 这条打印语句我们希望被抑制
    return i**2

# 定义一个Worker插件来抑制输出
class SuppressOutputPlugin(WorkerPlugin):
    def setup(self, worker):
        # 保存原始的stdout，以便在teardown中恢复（如果需要）
        self.original_stdout = sys.stdout
        # 将stdout重定向到os.devnull
        sys.stdout = open(os.devnull, 'w')
        print(f"Worker {worker.name} output suppressed.") # 这条打印会写入devnull

    def teardown(self, worker):
        # 恢复原始的stdout
        sys.stdout.close() # 关闭重定向的文件句柄
        sys.stdout = self.original_stdout
        print(f"Worker {worker.name} output restored.") # 这条打印会显示在控制台

if __name__ == '__main__':
    # 1. 初始化LocalCluster
    # 注意：processes=True 是默认值，但明确指出可以帮助理解
    cluster = LocalCluster(n_workers=2, processes=True, threads_per_worker=1)
    client = Client(cluster)

    print(f"Dask Dashboard link: {client.dashboard_link}")

    # 2. 注册输出抑制插件
    plugin = SuppressOutputPlugin()
    client.register_worker_plugin(plugin)
    print("SuppressOutputPlugin registered successfully.")

    # 3. 提交Dask任务
    dask_delays = []
    for i in range(10):
        dask_delays.append(dask.delayed(dask_function)(i))

    print("\n开始计算，预期工作器内部的'工作器正在处理任务...'打印不会显示...")
    dask_outs = client.compute(dask_delays).result() # .result() 阻塞直到计算完成

    print("\n计算完成。结果：", dask_outs)

    # 4. 关闭客户端和集群
    client.close()
    cluster.close()
    print("\n客户端和集群已关闭。")

# scheduler.py
from distributed import Scheduler
import asyncio

async def start_scheduler():
    scheduler = Scheduler(port=8786, dashboard_address=':8787')
    await scheduler.start()
    print(f"Scheduler started at {scheduler.address}")
    print(f"Dashboard at {scheduler.dashboard_link}")
    await scheduler.finished() # 保持调度器运行直到被外部停止

if __name__ == "__main__":
    asyncio.run(start_scheduler())

# run_dask_with_redirected_workers.py
import subprocess
import time
from distributed import Client
import dask

# 定义一个Dask函数
def dask_function_external(i):
    print(f'外部工作器正在处理任务 {i}!') # 这条打印语句将被重定向
    return i**2

if __name__ == "__main__":
    scheduler_address = "tcp://127.0.0.1:8786"
    num_workers = 2

    # 启动工作器并重定向其stdout/stderr到/dev/null
    worker_processes = []
    for i in range(num_workers):
        # 注意：这里使用了shell=True，这在某些情况下可能存在安全风险，
        # 更好的做法是传递一个列表而不是字符串，但为了演示重定向，这里简化
        # 对于生产环境，请谨慎使用或避免shell=True
        cmd = f"dask worker {scheduler_address} --nthreads 1 > /dev/null 2>&1"
        print(f"启动工作器 {i+1}，命令：{cmd}")
        process = subprocess.Popen(cmd, shell=True)
        worker_processes.append(process)
        time.sleep(1) # 给工作器一点时间启动并连接

    print(f"\n已启动 {num_workers} 个工作器，并重定向其输出。")

    # 连接到调度器
    client = Client(scheduler_address)
    print(f"客户端已连接到调度器：{client.dashboard_link}")

    # 提交Dask任务
    dask_delays = []
    for i in range(10):
        dask_delays.append(dask.delayed(dask_function_external)(i))

    print("\n开始计算，预期工作器内部的'外部工作器正在处理任务...'打印不会显示...")
    dask_outs = client.compute(dask_delays).result()

    print("\n计算完成。结果：", dask_outs)

    # 关闭客户端和工作器进程
    client.close()
    for p in worker_processes:
        p.terminate() # 尝试终止进程
        p.wait()      # 等待进程结束
    print("\n客户端和外部工作器已关闭。")

# 终端 1: 启动调度器
dask scheduler --port 8786 --dashboard-address :8787

# 终端 2: 启动第一个工作器并重定向输出
dask worker tcp://127.0.0.1:8786 --nthreads 1 > /dev/null 2>&1

# 终端 3: 启动第二个工作器并重定向输出
dask worker tcp://127.0.0.1:8786 --nthreads 1 > /dev/null 2>&1

# client_only.py
from distributed import Client
import dask

def dask_function_shell(i):
    print(f'Shell工作器正在处理任务 {i}!')
    return i**2

if __name__ == "__main__":
    client = Client("tcp://127.0.0.1:8786")
    print(f"客户端已连接到调度器：{client.dashboard_link}")

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

    print("\n开始计算，预期工作器内部的'Shell工作器正在处理任务...'打印不会显示...")
    dask_outs = client.compute(dask_delays).result()

    print("\n计算完成。结果：", dask_outs)
    client.close()