Databricks上Ray+PyTorch Lightning训练遇PyArrow写入错误求助

问题背景 #

在Databricks Runtime 15.2/14.3 ML集群上，使用Ray 2.32.0结合PyTorch Lightning 2.3.3训练神经网络，完全参照官方教程代码实现，但训练仅完成1个epoch就因RayTaskError(OSError: [Errno 22] Invalid argument)终止。错误栈指向PyArrow在复制检查点文件到DBFS路径时失败，已验证DBFS路径读写正常，但问题仍存在。

集群配置代码：

from ray.util.spark import setup_ray_cluster, shutdown_ray_cluster
eps = 1
ray_worker_memory = 410
remote_conn_str = setup_ray_cluster(
  max_worker_nodes=8,
  min_worker_nodes=1,
  memory_worker_node=ray_worker_memory*0.7*10**9-eps,
  object_store_memory_worker_node=ray_worker_memory*0.3*10**9-eps,
  collect_log_to_path="/dbfs/test/ray_collected_logs"
)

核心错误片段：

OSError: [Errno 22] Error writing bytes to file. Detail: [errno 22] Invalid argument

解决方案 #

1. 规避DBFS直接作为Ray存储路径 #

Ray的PyArrow文件复制逻辑与DBFS的FUSE挂载存在兼容性问题，可先将检查点存到本地临时目录，训练后手动同步到DBFS：

修改TorchTrainer的RunConfig：

run_config=ray.train.RunConfig(
    storage_path="/tmp/ray_checkpoints",  # 使用本地临时路径
    name="lightning", 
),

训练完成后同步到DBFS：

dbutils.fs.cp("file:/tmp/ray_checkpoints/lightning", "/dbfs/test/ray_collected_logs/lightning", recurse=True)

2. 禁用自动检查点上报 #

错误触发于RayTrainReportCallback在epoch结束时自动上报检查点，可禁用该功能，仅上报训练指标：

修改训练代码中的回调配置：

callbacks=[ray.train.lightning.RayTrainReportCallback(report_checkpoint=False)],

若需保存检查点，可改用PyTorch Lightning原生的检查点机制，或在自定义回调中手动调用Ray的检查点API。

3. 分离Ray日志与训练存储路径 #

取消setup_ray_cluster中的collect_log_to_path配置，避免Ray日志存储和训练检查点路径冲突：

remote_conn_str = setup_ray_cluster(
  max_worker_nodes=8,
  min_worker_nodes=1,
  memory_worker_node=ray_worker_memory*0.7*10**9-eps,
  object_store_memory_worker_node=ray_worker_memory*0.3*10**9-eps
  # 移除collect_log_to_path配置，使用默认本地日志路径
)

4. 调整依赖版本组合 #

当前Ray与PyTorch Lightning版本可能存在兼容性问题，尝试降级到经过验证的版本：

%pip install ray==2.31.0 pytorch-lightning==2.2.5

5. 统一PyArrow版本 #

Databricks自带的PyArrow版本可能与Ray依赖版本冲突，强制安装兼容版本：

%pip install pyarrow==15.0.0

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