Ray Tune 自定义训练函数中的检查点机制：tune.report 与 Checkpoint 的完整实践
【免费下载链接】rayRay is an AI compute engine. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.项目地址: https://gitcode.com/gh_mirrors/ra/ray
本文以 Ray 仓库中的示例脚本 custom_func_checkpointing.py（由文档片段 custom_func_checkpointing.rst 通过literalinclude直接引入）为主线，讲解如何在不使用 PyTorch/TensorFlow 等官方集成、仅使用纯 Python 自定义训练函数的情况下，通过ray.tune.report(metrics, checkpoint=...)API 实现断点保存与恢复（resume），并结合 trainable_fn_utils.py 与 Checkpoint 类实现 说明其底层机制。读完后你将掌握：自定义训练函数中检查点的写入与读取方式、Checkpoint对象的目录语义，以及Tuner层面超参搜索配置的完整可运行代码。
核心思路：用tune.report(checkpoint=...)报告检查点
当训练逻辑是自定义函数（而非 Ray 的 PyTorch/TF 等集成）时，Tune 无法替你感知训练状态，因此需要你主动做两件事：
- 写入：在训练循环中周期性地把状态（如当前 step、模型权重）序列化到一个目录中，然后通过
tune.report(..., checkpoint=Checkpoint.from_directory(...))交给 Tune 持久化； - 恢复：在训练函数入口处调用
tune.get_checkpoint()获取最近一次已报告的检查点，反序列化后从断点继续。
示例完整代码如下（来自 custom_func_checkpointing.py）：
import argparse import json import os import tempfile import time from ray import tune from ray.tune import Checkpoint def evaluation_fn(step, width, height): time.sleep(0.1) return (0.1 + width * step / 100) ** (-1) + height * 0.1 def train_func(config): step = 0 width, height = config["width"], config["height"] checkpoint = tune.get_checkpoint() if checkpoint: with checkpoint.as_directory() as checkpoint_dir: with open(os.path.join(checkpoint_dir, "checkpoint.json")) as f: state = json.load(f) step = state["step"] + 1 for current_step in range(step, 100): intermediate_score = evaluation_fn(current_step, width, height) with tempfile.TemporaryDirectory() as temp_checkpoint_dir: with open(os.path.join(temp_checkpoint_dir, "checkpoint.json"), "w") as f: json.dump({"step": current_step}, f) tune.report( {"iterations": current_step, "mean_loss": intermediate_score}, checkpoint=Checkpoint.from_directory(temp_checkpoint_dir), ) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( "--smoke-test", action="store_true", help="Finish quickly for testing" ) args, _ = parser.parse_known_args() tuner = tune.Tuner( train_func, run_config=tune.RunConfig( name="hyperband_test", stop={"training_iteration": 1 if args.smoke_test else 10}, ), tune_config=tune.TuneConfig( metric="mean_loss", mode="min", num_samples=5, ), param_space={ "steps": 10, "width": tune.randint(10, 100), "height": tune.loguniform(10, 100), }, ) results = tuner.fit() best_result = results.get_best_result() print("Best hyperparameters: ", best_result.config) best_checkpoint = best_result.checkpoint print("Best checkpoint: ", best_checkpoint)脚本支持--smoke-test命令行参数：传入后stop中的training_iteration从默认的 10 降为 1，便于快速验证流水线。
train_func的恢复逻辑：tune.get_checkpoint()的语义
训练函数开头的关键代码是：
checkpoint = tune.get_checkpoint() if checkpoint: with checkpoint.as_directory() as checkpoint_dir: with open(os.path.join(checkpoint_dir, "checkpoint.json")) as f: state = json.load(f) step = state["step"] + 1从源码结构看，trainable_fn_utils.py 中get_checkpoint()的实现是：
@PublicAPI(stability="stable") @_warn_session_misuse() def get_checkpoint() -> Optional[Checkpoint]: """Access the latest reported checkpoint to resume from if one exists.""" return get_session().loaded_checkpoint即它返回的是当前 Train 会话（session）中最近加载的检查点；首次运行时没有任何历史检查点，返回None，训练函数从step = 0开始。当实验因故障恢复或断点续跑时，Tune 会把上一次报告的检查点装载进会话——这一点在 function_trainable.py 中可以看到闭环：
session.loaded_checkpoint = checkpoint_result.checkpoint也就是说，每一次tune.report(..., checkpoint=...)的结果都会写回session.loaded_checkpoint，形成"报告 → 持久化 → 恢复时加载"的完整链路。
as_directory()：把检查点当作只读本地目录
Checkpoint.as_directory()是上下文管理器。从 Checkpoint 实现 的文档说明可以确认其两种行为：
- 本地目录检查点：直接返回原目录路径，不做拷贝，退出上下文后不做清理；
- 远端存储检查点（如 S3 URI）：下载到本地临时目录后返回路径，退出上下文时清理临时目录；
- 若同一节点上多个进程并发访问同一检查点，只有一个进程真正执行下载，其余进程等待后共享同一份数据（通过
TempFileLock文件锁实现）。
官方文档同时强调：返回的目录应视为只读，因为临时数据可能在退出上下文后被删除。示例中读取checkpoint.json恢复 step 的做法正是这种"目录即状态文件集合"的典型用法——检查点的物理形态就是一个目录，用户完全自由决定其中放什么（JSON、pickle、模型权重文件等）。
写入检查点：tune.report+Checkpoint.from_directory
训练循环中每步都执行一次检查点报告：
with tempfile.TemporaryDirectory() as temp_checkpoint_dir: with open(os.path.join(temp_checkpoint_dir, "checkpoint.json"), "w") as f: json.dump({"step": current_step}, f) tune.report( {"iterations": current_step, "mean_loss": intermediate_score}, checkpoint=Checkpoint.from_directory(temp_checkpoint_dir), )这里有两个值得注意的细节：
- 先写临时目录，再包装成
Checkpoint对象：Checkpoint.from_directory(path)（源码）会以pyarrow.fs.LocalFileSystem()为后端构造Checkpoint，把"本地目录"变成"可被存储层持久化/迁移的检查点引用"。tune.report的 docstring 明确说明：提供checkpoint时，它会被持久化到配置的存储位置（persistent storage）。 - metrics 与 checkpoint 在同一次调用中报告：
tune.report的 源码 表明每次调用都会自动递增底层的training_iteration计数——示例main块中stop={"training_iteration": 10}正是以"报告次数"为停止条件，而非 epoch 数。docstring 也提醒：这个"iteration"的物理含义由用户按调用report的频率自行定义，不一定对应一个 epoch。
另外注意 trainable_fn_utils.py 中的实现细节：
@_copy_doc(TrainCheckpoint) class Checkpoint(TrainCheckpoint): # NOTE: This is just a pass-through wrapper around `ray.train.Checkpoint` # in order to detect whether the import module was correct `ray.tune.Checkpoint`. passray.tune.Checkpoint只是ray.train.Checkpoint的透传包装类，专门用于检测导入模块是否正确；若你在report中传入了错误来源的Checkpoint实例，会触发 v2 迁移弃用警告（见 report 函数中的类型检查），因此务必从ray.tune导入Checkpoint，与示例文件开头的from ray.tune import Checkpoint保持一致。
Checkpoint类的更多能力（示例未用到但实用）
从 Checkpoint 完整实现 可以看到，该类还支持：
- 远端构造：
Checkpoint("s3://bucket/path/to/checkpoint")会根据 URI scheme 推断文件系统，无需显式传入 filesystem； to_directory(path=None)：把检查点内容写到指定本地目录（或自动生成的临时目录），适用于需要持久保留下载内容的场景；- 元数据 API：
set_metadata/get_metadata/update_metadata会把键值对以.metadata.json形式随检查点持久化，适合存放超参摘要、预处理器配置等； - 误用防护：
__fspath__被刻意实现为抛TypeError，强制你使用to_directory()/as_directory()而非把Checkpoint当普通路径拼接使用。
Tuner配置逐项解析
示例main块展示了完整的tune.Tuner调用，各部分含义如下：
| 配置项 | 取值 | 作用 |
|---|---|---|
train_func（位置参数） | 上面的训练函数 | Tune 自动包装为 Function Trainable |
RunConfig.name | "hyperband_test" | 实验名，用于标识运行目录 |
RunConfig.stop | {"training_iteration": 1}（smoke）/10（默认） | 每个 trial 报告满指定 iteration 次数即停止 |
TuneConfig.metric | "mean_loss" | 用于排序与早停的目标指标，每次report中提供 |
TuneConfig.mode | "min" | 目标越小越好（loss 场景） |
TuneConfig.num_samples | 5 | 每个 trial 重复采样 5 次，以抑制单次随机性 |
param_space["steps"] | 10（常量） | 非随机超参，仅随配置下发，示例中实际未使用 |
param_space["width"] | tune.randint(10, 100) | 在 [10, 100] 均匀取整 |
param_space["height"] | tune.loguniform(10, 100) | 在 [10, 100] 对数均匀采样，适合量级跨度大的参数 |
训练结束后通过results.get_best_result()取最优 trial，并打印其config与checkpoint——best_result.checkpoint即最后一次报告的Checkpoint对象，可直接用as_directory()加载其中的checkpoint.json（或真实场景中的模型权重）用于部署/推理。
运行方式与适用前提
python python/ray/tune/examples/custom_func_checkpointing.py # 完整跑 10 次 iteration python python/ray/tune/examples/custom_func_checkpointing.py --smoke-test # 1 次 iteration 快速验证适用前提与限制：
- 训练函数运行在 Ray 集群的 worker 中，
tune.report/tune.get_checkpoint只能在 Train 会话内调用（源码中由_warn_session_misuse()装饰器检测误用）； - 检查点内容本身只是目录数据，Tune 不负责解释其结构，恢复逻辑（如
step + 1）完全由用户编写； - 该模式与官方集成（如 cifar10_pytorch.py、pbt_convnet_function_example.py 中对
tune.get_checkpoint()的同款用法）遵循同一套 API，因此从自定义函数迁移到框架集成时检查点代码基本可以平移。
小结
Ray Tune 的自定义函数检查点机制可归纳为三步闭环：用Checkpoint.from_directory把任意本地目录包装为检查点 → 随tune.report一起上报以触发持久化 → 恢复时经tune.get_checkpoint().as_directory()以只读目录形式读回。整套机制不依赖任何深度学习框架，检查点可以是 JSON、pickle 或模型权重文件的任意组合；而training_iteration随每次report自动递增的特性，使得RunConfig.stop可以按"报告次数"精确控制试验时长。
【免费下载链接】rayRay is an AI compute engine. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.项目地址: https://gitcode.com/gh_mirrors/ra/ray
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
` 方法详解：取匹配节点集合的最后一个元素)
