适用于 Apache Spark 的深度学习流水线

该仓库仅包含用于本地 CI 的 HorovodRunner 代码和 API 文档。若要使用 HorovodRunner 进行分布式训练，请使用 Databricks Runtime for Machine Learning。 有关详细信息，请参阅 Databricks 文档中的 HorovodRunner：使用 Horovod 进行分布式深度学习。

如需使用包含 Spark 深度学习流水线 API 的旧版本，请访问 Spark Packages 页面。

API 文档

类 sparkdl.HorovodRunner(*, np, driver_log_verbosity='all')

基类：object

HorovodRunner 使用 Horovod 运行分布式深度学习训练作业。

在 Databricks Runtime 5.0 ML 及更高版本上，它会将 Horovod 作业作为分布式 Spark 作业启动。 通过管理集群设置并与 Spark 集成，它使在 Databricks 上运行 Horovod 变得更加容易。 请查阅 Databricks 文档以查看端到端示例和性能调优提示。

开源版本仅在同一 Python 进程内本地运行作业，仅供本地开发使用。

注意：Horovod 是由 Uber 开发的分布式训练框架。

• 参数

  • np - 用于 Horovod 作业的并行进程数。 此参数仅在 Databricks Runtime 5.0 ML 及更高版本上生效。 在开源版本中会被忽略。 在 Databricks 上，每个进程将占用一个可用的任务槽位， 该槽位对应于 GPU 集群上的 GPU 或 CPU 集群上的 CPU 核心。 允许的取值如下：

    • 如果 <0，则会在驱动节点上派生 -np 个子进程，在本地运行 Horovod。 训练的 stdout 和 stderr 消息会输出到笔记本单元格中，并且在单元格输出被截断时， 也会记录在驱动程序日志中。这有助于调试，我们建议您首先在此模式下测试代码。 不过，在共享的 Databricks 集群上，请谨慎使用 Spark 驱动程序资源。 注意，np < -1 仅在 Databricks Runtime 5.5 ML 及更高版本上支持。
    • 如果 >0，则会启动一个包含 np 个任务的 Spark 作业，所有任务同时开始， 并在这些任务节点上运行 Horovod 作业。 它会等待直到有 np 个任务槽位可用才启动作业。 如果 np 大于集群中任务槽位的总数，则作业会失败。自 Databricks Runtime 5.4 ML 起， 训练的 stdout 和 stderr 消息会输出到笔记本单元格中。如果单元格输出被截断， 则完整的日志可以在由 HorovodRunner 启动的第二个 Spark 作业中，任务 0 的 stderr 流中找到， 您可以通过 Spark UI 查看这些日志。
    • 如果 0，则会使用集群中的所有任务槽位来启动作业。 .. 警告:: 将 np 设置为 0 已弃用，并将在下一个主要的 Databricks Runtime 版本中移除。 根据运行时的任务槽位总数来选择 np 值是不可靠的，因为执行器可能会动态注册。 请明确指定所需的并行进程数量。

  • driver_log_verbosity - 驱动程序日志的详细程度：此参数仅在 Databricks Runtime 上可用。

run(main, **kwargs)

运行一个调用 main(**kwargs) 的 Horovod 训练作业。

开源版本仅在同一 Python 进程内调用 main(**kwargs)。 在 Databricks Runtime 5.0 ML 及更高版本上，它会根据 np 的文档行为启动 Horovod 作业。 主函数和关键字参数都会使用 cloudpickle 序列化，并分发给集群中的工作节点。

• 参数

  • main – 包含 Horovod 训练代码的 Python 函数。 预期签名为 def main(**kwargs) 或兼容形式。 由于函数会被序列化并分发给工作节点，因此请在函数内部更改全局状态， 例如设置日志级别，并注意 pickling 的限制。 避免在函数中引用大型对象，否则可能导致序列化数据过大， 从而延长作业启动时间。

  • kwargs – 调用时传递给主函数的关键字参数。

• 返回值

  主函数的返回值。 当 np≥0 时，返回的是 rank 0 进程的值。请注意，返回值必须能够使用 cloudpickle 序列化。

发布版本

请访问 Github 发布页面 查看发布说明。

许可证

• 源代码根据 Apache License 2.0 发布（详见 LICENSE 文件）。

spark-deep-learning 快速上手指南

spark-deep-learning 是 Databricks 提供的用于在 Apache Spark 上运行深度学习流水线的工具库。请注意：当前开源仓库主要包含用于本地开发和 CI 测试的 HorovodRunner 代码。若需在生产环境中进行分布式训练，官方推荐使用 Databricks Runtime for Machine Learning。开源版本仅支持在单个 Python 进程中本地运行，适用于代码调试和本地开发。

环境准备

• 操作系统: Linux, macOS, Windows (仅限本地开发)
• Python: 3.6+
• 核心依赖:
  • Apache Spark (本地模式)
  • PySpark
  • Horovod (用于本地单进程模拟)
  • cloudpickle
• 推荐环境: 建议使用 Conda 或虚拟环境隔离依赖。

注意：开源版本不包含完整的 Spark Deep Learning Pipelines API（如预定义的深度学习转换器），该部分功能需通过 Spark Packages 获取或在 Databricks 平台上使用。

安装步骤

由于该工具主要针对 Databricks 平台优化，开源版本可通过 pip 直接安装其核心组件用于本地测试。

pip install sparkdl

如果需要手动安装特定版本的依赖（例如配合特定版本的 PySpark）：

pip install pyspark horovod cloudpickle

国内加速提示：如果下载速度较慢，可使用国内镜像源安装：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple sparkdl

基本使用

开源版本的 HorovodRunner 主要用于在本地单机环境中模拟和调试 Horovod 训练代码。它不会真正启动分布式集群，而是在当前 Python 进程中执行训练函数。

最简单的使用示例

以下示例展示如何定义一个训练函数并通过 HorovodRunner 在本地运行：

from sparkdl import HorovodRunner

# 1. 定义包含 Horovod 训练逻辑的主函数
def train_func(**kwargs):
    import horovod.keras as hvd
    import tensorflow as tf
    
    # 初始化 Horovod
    hvd.init()
    
    # 设置 GPU 可见性 (如果是多卡本地模拟，需配合 np 参数，但在开源版中通常只跑单进程)
    gpus = tf.config.experimental.list_physical_devices('GPU')
    if gpus:
        tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'GPU')

    # 构建简单的模型和数据 (此处仅为示意)
    model = ... 
    optimizer = ...
    
    # 包装 Optimizer
    optimizer = hvd.DistributedOptimizer(optimizer)
    
    # 训练循环
    # model.fit(...)
    
    print(f"Training completed on rank {hvd.rank()}")
    return "success"

# 2. 实例化 HorovodRunner
# np 参数在开源版本中被忽略，仅在 Databricks 集群生效
# 这里主要用于保持代码与云端部署的一致性
hr = HorovodRunner(np=1)

# 3. 运行训练任务
# 在开源版本中，这将直接在当前进程中调用 train_func
result = hr.run(train_func)

print(f"Job result: {result}")

关键说明

• 运行模式：在开源版本中，hr.run() 仅仅是序列化并本地调用 main 函数，不会启动多个子进程或 Spark 任务。
• 代码迁移：编写此代码的主要目的是为了无缝迁移到 Databricks 平台。在 Databricks ML Runtime 上，只需更改 np 参数（例如 np=4 或 np=-1），相同的代码即可自动利用集群资源进行分布式训练。
• 序列化限制：传递给 run 方法的 main 函数及其闭包变量会被 cloudpickle 序列化。请避免在函数外部引用大型对象，并确保函数内部处理所有全局状态（如日志级别）。