2. 工作区与集群工程化配置

在 Databricks 上跑 AI 负载,工作区与集群不是“建好能用”就结束,而是决定训练吞吐、GPU 利用率和单位算力成本的核心工程对象。本节围绕工作区创建、集群策略、自动伸缩、池化、实例选型、Spot 与成本优化,以及面向 AI 负载的调优实践展开,目标是让读者能独立设计一套可复现、可治理、成本可控的集群方案。

一、工作区创建:先定边界,再谈算力

工作区(Workspace)是 Databricks 的组织单元,承载 Notebook、作业、集群、Unity Catalog 元数据与权限。创建时最关键的不是点“新建”,而是先确定三件事:云账号与网络边界、数据平面位置、身份与权限模型。

以 AWS 为例,典型流程是:在 Databricks Account Console 中创建 workspace,选择区域(如 us-east-1),配置 VPC(推荐使用客户自管 VPC,指定私有子网与安全组),绑定 IAM 角色(跨账号角色用于数据访问),再启用 Unity Catalog 作为统一权限层。若数据已在 S3 且要求不出网,应配置 VPC Endpoint 访问 S3 与 STS,避免走公网 NAT 产生额外流量费。

工作区创建后,建议立即做三件治理动作:一是把集群创建权限收敛到管理员或通过策略下发;二是开启集群日志投递到 S3/CloudWatch,便于审计与排障;三是为不同团队建立独立的工作区或至少独立的目录与计算资源标签,避免“一个工作区跑所有负载”导致的资源争抢与成本归属不清。

二、集群策略:把“正确配置”变成默认值

集群策略(Cluster Policy)是一组 JSON 规则,用于限制或强制集群属性。它的价值在于:让普通用户无法创建“8 卡 A100 常驻不释放”的集群,同时让标准负载一键获得合规配置。

策略的核心字段包括:spark_version、node_type_id、autoscale.min_workers/max_workers、autotermination_minutes、custom_tags、aws_attributes.availability(SPOT/ON_DEMAND)、runtime_engine(PHOTON)等。下面是一个面向 AI 训练的策略片段示例:

{
  "spark_version": { "type": "fixed", "value": "15.4.x-gpu-ml-scala2.12" },
  "node_type_id": { "type": "allowlist",
    "values": ["g5.12xlarge", "p4d.24xlarge"] },
  "autoscale": { "min_workers": 2, "max_workers": 8 },
  "autotermination_minutes": { "type": "fixed", "value": 30 },
  "aws_attributes.availability": { "type": "fixed", "value": "SPOT_WITH_FALLBACK" },
  "custom_tags.team": { "type": "fixed", "value": "ai-platform" }
}

策略应分层:基础策略限制运行时版本与标签;GPU 策略限定实例族与最大节点数;作业策略强制自动终止与 Spot 回退。这样既保证灵活性,又避免成本失控。

三、自动伸缩、池化与实例选型

自动伸缩(Autoscaling)根据待处理任务动态增减 worker。对 AI 负载,伸缩的瓶颈往往不在 Spark 任务数,而在 GPU 资源分配与数据加载。建议设置 min_workers 覆盖基线并发,max_workers 不超过配额与数据分区数,避免“扩到 100 节点但每个节点 GPU 空转”。

池化(Instance Pool)预置一组空闲实例,集群创建时直接从池中取,冷启动从数分钟降到数十秒。对交互式 AI 开发、频繁启停的调参作业,池化收益显著。池应设置 idle_instance_autotermination_minutes(如 15 分钟)和 max_capacity,并区分 GPU 池与 CPU 池。

实例选型要匹配负载形态:

  • 数据预处理/ETL:CPU 密集,选计算优化型(如 c5/c6i),内存与 vCPU 比适中。
  • 模型训练:GPU 密集,单机多卡选 p4d/p5,分布式训练关注网络带宽(EFA/InfiniBand)。
  • 推理/微调:显存优先,A10G/L4 适合中等规模,A100/H100 适合大模型。
  • 内存型任务:特征工程、图计算选 r 系列,避免频繁 spill。

选型时还要看 Spark 配置:GPU 集群需设置 spark.task.resource.gpu.amount、spark.executor.resource.gpu.amount,并确保 spark.databricks.delta.preview.enabled 等与运行时匹配。

四、Spot 实例与成本优化

Spot 实例价格通常为按需的 30%–70%,但可能被回收。Databricks 支持 SPOT、SPOT_WITH_FALLBACK、ON_DEMAND 三种模式。对可中断的预处理、特征回填、超参搜索,优先用 Spot;对 driver、关键 checkpoint 写入、长时训练的主节点,用按需或 SPOT_WITH_FALLBACK。

成本优化组合拳:

  1. 用池化降低冷启动,提高 Spot 获取成功率。
  2. 设置 autotermination_minutes,杜绝闲置集群。
  3. 用集群策略限制最大节点数与实例族,防止“选最贵”。
  4. 开启 Photon 加速 SQL/DataFrame 负载,减少同等数据量下的节点数。
  5. 用标签(team/project/job)做成本分摊,按周复盘。
经验值:一个 8 节点 g5.12xlarge 的 GPU 集群,若 24 小时常驻,月成本可达数万美元;改为池化 + Spot + 30 分钟自动终止后,同等吞吐下成本常可下降 50%–70%。具体数值随区域与折扣变化,需以实际账单为准。

五、面向 AI 负载的集群调优实践

AI 负载的调优重点与普通 Spark 不同,核心是“让 GPU 不饿、让数据不堵、让 shuffle 不炸”。

1. 数据加载与格式。 训练数据优先用 Delta Lake 或 Parquet,开启 spark.databricks.io.cache.enabled=true 利用本地 SSD 缓存。小文件先 OPTIMIZE 合并,避免每个 task 读大量小文件。

2. 并行度与分区。 让分区数约为总 vCPU 的 2–4 倍;GPU 任务按 GPU 数设置 spark.task.resource.gpu.amount=1,确保每个 task 独占一张卡。

3. 内存与 shuffle。 提高 spark.sql.shuffle.partitions 到合理值(如 2000),开启 AQE(spark.sql.adaptive.enabled=true)自动合并小分区。对宽表 join,考虑 broadcast 或分桶。

4. 运行时与库。 使用 GPU ML Runtime,预装 CUDA、cuDNN、PyTorch/TensorFlow。自定义库通过 init script 或集群库安装,避免每次启动重装。

5. 监控与判断。 在 Spark UI 看 task 时间分布、shuffle spill、GPU 利用率(通过 Ganglia 或云监控)。若 GPU 利用率长期低于 30%,说明数据加载或 CPU 预处理是瓶颈,应增加 CPU worker 或优化读取;若 shuffle spill 大,应加内存或调整分区。

AI 负载集群工程化配置全景 工作区 Workspace VPC / 子网 IAM 角色 Unity Catalog 日志投递 S3 标签治理 集群策略 Cluster Policy 运行时版本固定 实例族 allowlist 最大节点数限制 自动终止 + Spot 回退 计算资源 Compute 实例池 Pool GPU 池 / CPU 池 空闲 15 分钟回收 冷启动 ↓ 至秒级 自动伸缩 min 2 / max 8 按任务队列扩缩 避免 GPU 空转 实例选型 c 系列:ETL p 系列:训练 r 系列:内存型 Spot 策略 SPOT_WITH_FALLBACK driver 按需 成本 ↓ 50%–70% AI 负载调优要点 Delta + 本地缓存 小文件 OPTIMIZE 分区 = vCPU × 2–4 GPU task 独占 AQE + shuffle 调参 减少 spill GPU ML Runtime init script 装库 Spark UI 监控 GPU 利用率判断 成本反馈

最后强调一个判断原则:先看瓶颈在哪,再调参数。GPU 利用率低,优先查数据读取与 CPU 预处理;shuffle 慢,优先查分区与 AQE;成本高,优先查闲置时间与 Spot 比例。把工作区边界、策略约束、池化与 Spot 组合起来,AI 集群才能既跑得快,又花得值。