缩放 Azure Databricks 上的 Ray 群集

项目
09/09/2024

了解如何优化 Ray 群集的大小以获得最佳性能，包括自动缩放、头节点配置、异类群集和资源分配。

在自动缩放模式下创建 Ray 群集

在 Ray 2.8.0 及更高版本中，在 Azure Databricks 上启动的 Ray 群集支持与 Azure Databricks 自动缩放功能的集成。此自动缩放集成在 Azure Databricks 环境内部触发 Azure Databricks 群集自动缩放。

若要启用自动缩放，请运行以下命令：

对于低于 2.10 的 Ray 版本：

from ray.util.spark import setup_ray_cluster

setup_ray_cluster(
  num_worker_nodes=8,
  autoscale=True,
)

对于 Ray 版本 2.10 及更高版本：

from ray.util.spark import setup_ray_cluster, shutdown_ray_cluster

setup_ray_cluster(
  min_worker_nodes=2,
  max_worker_nodes=4,
  num_cpus_per_node=4,
  collect_log_to_path="/dbfs/path/to/ray_collected_logs"
)

# Pass any custom Ray configuration with ray.init
ray.init(ignore_reinit_error=True)

ray.util.spark.setup_ray_cluster API 在 Apache Spark 上创建 Ray 群集。在内部，它会创建一个后台 Apache Spark 作业。该作业中的每个 Apache Spark 任务创建一个 Ray 工作器节点，Ray 头节点是在驱动程序上创建的。参数 min_worker_nodes 和 max_worker_nodes 表示用于创建和使用 Ray 工作负载的 Ray 工作器节点的范围。如果未定义参数 min_worker_nodes，则启动固定大小的 Ray 群集，该群集提供 max_worker_nodes 个工作器。若要指定分配给每个 Ray 工作器节点的 CPU 或 GPU 核心数，请设置参数 num_cpus_worker_node（默认值：1）或 num_gpus_worker_node（默认值：0）。

对于低于 2.10 的 Ray 版本，如果已启用自动缩放，num_worker_nodes 指示 Ray 工作器节点的最大数目。 Ray 工作器节点的默认最小数量为零。此默认设置意味着，当 Ray 群集处于空闲状态时，它会纵向缩减到零个 Ray 工作器节点。这可能并不是在所有方案中快速响应的理想方案，但它可以在启用时显著降低成本。

在自动缩放模式下，无法将 num_worker_nodes 设置为 ray.util.spark.MAX_NUM_WORKER_NODES。

以下参数用于配置纵向扩展和缩减速度：

autoscale_upscaling_speed 表示允许挂起的节点数，为当前节点数的倍数。该值越大，纵向扩展就越激进。例如，如果此值设置为 1.0，则群集大小可以随时以最大 100% 的速度增长。
autoscale_idle_timeout_minutes 表示在自动缩放器删除空闲工作器节点之前需要等待的分钟数。该值越小，纵向缩减就越激进。

使用 Ray 2.9.0 及更高版本，还可以设置 autoscale_min_worker_nodes 以防止 Ray 群集在 Ray 群集空闲时纵向缩减为零个工作器，这将导致群集终止。

配置 Ray 头节点使用的资源

默认情况下，对于 Spark 上的 Ray 配置，Azure Databricks 会将分配给 Ray 头节点的资源限制为：

0 个 CPU 核心
0 个 GPU
128 MB 堆内存
128 MB 对象存储内存

这是因为 Ray 头节点通常仅用于全局协调，而不用于运行 Ray 任务。 Apache Spark 驱动程序节点资源是与多个用户共享的，因此默认设置将资源保存在 Apache Spark 驱动程序端。使用 Ray 2.8.0 及更高版本，可以配置 Ray 头节点使用的资源。在 setup_ray_cluster API 中使用下列参数：

num_cpus_head_node：设置 Ray 头节点使用的 CPU 核心数
num_gpus_head_node：设置 Ray 头节点使用的 GPU
object_store_memory_head_node：按 Ray 头节点设置对象存储内存大小

支持异类群集

你可以为更高效和更具成本效益的训练运行创建一个 Ray on Spark 群集，并在 Ray 头节点和 Ray 工作器节点之间设置不同的配置。但是，所有 Ray 工作器节点都必须具有相同的配置。 Azure Databricks 群集并不完全支持异类群集，但可以通过设置群集策略创建具有不同驱动程序和工作器实例类型的 Azure Databricks 群集。例如：

{
  "node_type_id": {
    "type": "fixed",
    "value": "i3.xlarge"
  },
  "driver_node_type_id": {
    "type": "fixed",
    "value": "g4dn.xlarge"
  },
  "spark_version": {
    "type": "fixed",
    "value": "13.x-snapshot-gpu-ml-scala2.12"
  }
}

优化 Ray 群集配置

每个 Ray 工作器节点的建议配置如下所示：每个 Ray 工作器节点至少 4 个 CPU 内核。每个 Ray 工作器节点至少有 10GB 堆内存。

因此，调用 ray.util.spark.setup_ray_cluster 时，Azure Databricks 建议将 num_cpus_per_node 设置为大于或等于 4 的值。

有关优化每个 Ray 工作器节点的堆内存的详细信息，请参阅下一节。

Ray 工作器节点的内存分配

每个 Ray 工作器节点使用两种类型的内存：堆内存和对象存储内存。

按如下所述确定每种类型的分配内存大小。

分配给每个 Ray 工作器节点的内存总量为：RAY_WORKER_NODE_TOTAL_MEMORY = (SPARK_WORKER_NODE_PHYSICAL_MEMORY / MAX_NUMBER_OF_LOCAL_RAY_WORKER_NODES * 0.8)

MAX_NUMBER_OF_LOCAL_RAY_WORKER_NODES 是可以在 Apache Spark 工作器节点上启动的 Ray 工作器节点的最大数量。此数量由参数 num_cpus_per_node 或 num_gpus_per_node 确定。

如果未设置参数 object_store_memory_per_node，则分配给每个 Ray 工作器节点的堆内存大小和对象存储内存大小为：RAY_WORKER_NODE_HEAP_MEMORY = RAY_WORKER_NODE_TOTAL_MEMORY * 0.7OBJECT_STORE_MEMORY_PER_NODE = RAY_WORKER_NODE_TOTAL_MEMORY * 0.3

如果设置了参数 object_store_memory_per_node：RAY_WORKER_NODE_HEAP_MEMORY = RAY_WORKER_NODE_TOTAL_MEMORY - argument_object_store_memory_per_node

此外，每个 Ray 工作器节点的对象存储内存大小受操作系统共享内存的限制。最大值为：OBJECT_STORE_MEMORY_PER_NODE_CAP = (SPARK_WORKER_NODE_OS_SHARED_MEMORY / MAX_NUMBER_OF_LOCAL_RAY_WORKER_NODES * 0.8)

SPARK_WORKER_NODE_OS_SHARED_MEMORY 是为 Apache Spark 工作器节点配置的 /dev/shm 磁盘大小。

缩放最佳做法

为每个 Ray 工作器节点设置 CPU 和 GPU 编号

建议将参数 num_cpus_worker_node 设置为每个 Apache Spark 工作器节点的 CPU 内核数。同样，将 num_gpus_worker_node 设置为每个 Apache Spark 工作器节点的 GPU 数是最佳的。使用此配置，每个 Apache Spark 工作器节点都会启动一个 Ray 工作器节点，该节点将充分利用每个 Apache Spark 工作器节点的资源。

在启动 Apache Spark 群集时，在 Azure Databricks 群集配置中将环境变量 RAY_memory_monitor_refresh_ms 设置为 0。

Apache Spark 和 Ray 混合工作负载的内存资源配置

如果在 Azure Databricks 群集中运行混合 Spark 和 Ray 工作负载，Azure Databricks 建议将 Spark 执行器内存缩减到较小的值。例如，在 Azure Databricks 群集配置中设置 spark.executor.memory 4g。

Apache Spark 执行程序是一个 Java 进程，它会延迟触发 GC，而 Apache Spark 数据集缓存则使用大量 Apache Spark 执行程序内存。这减少了 Ray 可以使用的可用内存。为避免潜在的内存不足错误，请减小 spark.executor.memory 配置。

Apache Spark 和 Ray 混合工作负载的计算资源配置

如果在 Azure Databricks 群集中运行混合 Spark 和 Ray 工作负载，建议将群集节点或 Ray 工作器节点设置为自动缩放。例如：

如果具有固定数量的工作器节点可用于启动 Azure Databricks 群集，建议启用 Ray-on-Spark 自动缩放。当未运行 Ray 工作负载时，Ray 群集将缩减，从而允许释放资源供 Apache Spark 任务使用。 Apache Spark 任务完成后，再次使用 Ray 时，Ray-on-Spark 群集将再次纵向扩展以满足需求。

此外，还可以将 Azure Databricks 和 Ray-on-spark 群集设为自动缩放。例如，如果将 Azure Databricks 群集的自动缩放节点配置为最多 10 个节点，将 Ray-on-Spark 工作器节点配置为最多四个节点，并将每个 Ray 工作器节点配置为充分利用每个 Apache Spark 工作器的资源，则 Ray 工作负载最多可以在此类群集配置上使用 4 个节点的资源。相比之下，Apache Spark 作业最多可以分配六个节点的资源。

通过