如何确保Python逻辑仅在Apache Ray工作节点上运行？

我用Apache Ray搭建自定义集群运行业务逻辑，但通过ray.remote提交的任务总是跑在驱动节点，而非初始化时配置的工作节点，怎么确保逻辑只在工作节点执行？

核心思路 #

Ray默认允许驱动节点利用空闲资源执行任务，要强制任务调度到工作节点，需要通过节点标签标记工作节点，再给远程函数添加调度约束，限制任务只能在带标签的工作节点上运行。

具体修改步骤 #

1. 给工作节点添加专属标签

在调用setup_ray_cluster时，通过worker_node_args参数给所有工作节点添加自定义标签，以此区分驱动节点。

2. 给远程函数指定调度约束

在@ray.remote装饰器中，用resources参数指定必须匹配工作节点的标签，Ray会自动将任务调度到符合标签的节点。

修改后的完整代码 #

import os
os.environ["RAY_PROFILING"] = "1"
os.environ["RAY_task_events_report_interval_ms"] = "0"

import ray
from ray.util.spark import setup_ray_cluster, shutdown_ray_cluster
from ray.util import get_node_ip_address

RAY_LOGS_VOLUME_PATH = "***/temp/apache_ray_logs"
RAY_TIMELINE_VOLUME_PATH = "***/temp/apache_ray_logs/timeline.json"
MAX_WORKER_NODES = 2

try:
  shutdown_ray_cluster()
except:
  pass

try:
  ray.shutdown()
except:
  pass

if not os.path.exists(RAY_LOGS_VOLUME_PATH):
  os.makedirs(RAY_LOGS_VOLUME_PATH)

# 给工作节点添加专属标签
_, cluster_address = setup_ray_cluster(
  max_worker_nodes = MAX_WORKER_NODES,
  memory_worker_node = 28 * 2**30,
  num_cpus_per_node = 4,
  num_gpus_worker_node = 0,
  collect_log_to_path = RAY_LOGS_VOLUME_PATH,
  # 给所有工作节点打上标签
  worker_node_args={"node_labels": {"worker_node": "true"}}
  )

ray.init(
  address = cluster_address,
  ignore_reinit_error = True
  )

ray.timeline(filename=RAY_TIMELINE_VOLUME_PATH)

# 指定任务只能在带有worker_node=true标签的节点运行
@ray.remote(resources={"worker_node": 0.1})
def some_function(input):
  print(f"Node used: {get_node_ip_address()}")
  # 你的业务逻辑

sources = ["input1", "input2"]
results = ray.get([some_function.remote(x) for x in sources])

关键说明 #

• worker_node_args={"node_labels": {"worker_node": "true"}}：给所有工作节点打上专属标签，和驱动节点做区分。
• @ray.remote(resources={"worker_node": 0.1})：要求任务必须占用带有worker_node标签节点的0.1单位资源（仅做标记用，不消耗实际资源），Ray会自动过滤驱动节点，只将任务调度到符合标签的工作节点。

验证方法 #

可以在驱动节点打印自身IP，对比任务输出的节点IP，确认任务运行位置：

# 在驱动节点打印自身IP
print(f"Driver node IP: {get_node_ip_address()}")

运行后若任务输出的IP与驱动节点IP不同，说明任务已在工作节点执行。

内容的提问来源于stack exchange，提问作者question.it