如何在分布式HPC集群中查看Dask数组persist后的块物理节点位置？

Dask提供了直接查询数组块所在节点的方法，主要通过集群客户端对象结合数组的任务信息来实现，以下是具体方案：

• 通过client.who_has()查询已持久化块的位置
  先获取数组对应的任务键，再用客户端的who_has()方法映射每个块的存储节点，示例代码：

  import dask.array as da
  
  # 假设已初始化client并创建数组，且已执行persist
  arr = da.ones((1000, 1000), chunks=(250, 250)).persist()
  
  # 提取数组所有块对应的任务键
  task_keys = arr.__dask_graph__().keys()
  
  # 查询每个任务键所在的工作节点
  chunk_locations = client.who_has(task_keys)
  
  # 遍历输出结果
  for key, nodes in chunk_locations.items():
      print(f"块 {key} 存储节点: {nodes}")
  

• 通过任务流获取块的执行/存储节点
  如果需要更详细的历史执行信息，可使用client.get_task_stream()：

  # 获取集群任务流数据
  task_stream = client.get_task_stream()
  # 筛选目标数组相关的任务节点信息
  for task in task_stream:
      if any(str(k) in task['key'] for k in task_keys):
          print(f"任务 {task['key']} 在节点 {task['worker']} 执行/存储")
  

• 关键注意事项

  • 只有当数组块被persist()或完成计算后，才会被存储在工作节点上，此时who_has()才能返回有效位置
  • 未计算的块不会在节点上存储，who_has()会返回空列表
  • 动态调度场景下（如节点故障重调度），块的位置可能发生变化，查询结果为当前实时状态

内容的提问来源于stack exchange，提问作者MariusSiuram