ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

Hydra Callbacks 事件机制深度指南:在 Hydra 配置框架中接入自定义生命周期钩子

Hydra Callbacks 事件机制深度指南:在 Hydra 配置框架中接入自定义生命周期钩子 Hydra Callbacks 事件机制深度指南在 Hydra 配置框架中接入自定义生命周期钩子【免费下载链接】hydraHydra is a framework for elegantly configuring complex applications项目地址: https://gitcode.com/GitHub_Trending/hyd/hydraHydra 的 Callback回调机制允许开发者编写自定义代码在 Hydra 运行生命周期的各类事件如任务启动前、任务结束后、multirun 扫描开始/结束时被自动触发是优雅配置复杂应用的核心扩展点之一。本文基于 Hydra 1.3 官方文档与源码完整讲解 Callback 接口的六个事件方法、注册配置方式、多个回调的调用顺序以及仓库内置的实用示例回调帮助你快速实现诸如自动上传产物、远程集群错误日志、任务复现等真实场景。一、什么是 Hydra CallbackCallback 接口Callback interface使得用户自定义代码能够在 Hydra 的各个事件节点被触发执行。它属于 Hydra 的 experimental 模块位于 hydra/experimental/callback.py核心是一个定义了六个事件钩子的抽象类from hydra.types import TaskFunction class Callback: def on_run_start(self, config: DictConfig, **kwargs: Any) - None: Called in RUN mode before job/application code starts. config is composed with overrides. Some hydra.runtime configs are not populated yet. See hydra.core.utils.run_job for more info. ... def on_run_end(self, config: DictConfig, **kwargs: Any) - None: Called in RUN mode after job/application code returns. ... def on_multirun_start(self, config: DictConfig, **kwargs: Any) - None: Called in MULTIRUN mode before any job starts. When using a launcher, this will be executed on local machine before any Sweeper/Launcher is initialized. ... def on_multirun_end(self, config: DictConfig, **kwargs: Any) - None: Called in MULTIRUN mode after all jobs returns. When using a launcher, this will be executed on local machine. ... def on_job_start( self, config: DictConfig, *, task_function: TaskFunction, **kwargs: Any ) - None: Called in both RUN and MULTIRUN modes, once for each Hydra job (before running application code). This is called from within hydra.core.utils.run_job. In the case of remote launching, this will be executed on the remote server along with your application code. The task_function argument is the function decorated with hydra.main. ... def on_job_end( self, config: DictConfig, job_return: JobReturn, **kwargs: Any ) - None: Called in both RUN and MULTIRUN modes, once for each Hydra job (after running application code). This is called from within hydra.core.utils.run_job. In the case of remote launching, this will be executed on the remote server after your application code. job_return contains info that could be useful for logging or post-processing. See hydra.core.utils.JobReturn for more. ...使用回调 API 的第一步是导入Callback类from hydra.experimental.callback import Callback随后用户创建Callback的子类覆写其中的一个或多个方法。子类必须在hydra.callbacks配置中注册注册方式见下文示例Hydra 才会在合适的时机调用这些方法。二、六个事件方法的语义与触发时机Callback 接口定义了 6 个事件方法按运行模式分为两类RUN 模式与 MULTIRUN 模式。RUN 模式单任务运行on_run_start(config, **kwargs)在 RUN 模式下、任务/应用代码开始之前调用。此时传入的config已经完成了组合composed并应用了 overrides但部分hydra.runtime配置项尚未填充。相关实现可参考 hydra/core/utils.py 中的run_job。on_run_end(config, **kwargs)在 RUN 模式下、任务/应用代码返回之后调用。MULTIRUN 模式多任务扫描运行on_multirun_start(config, **kwargs)在 MULTIRUN 模式下、任何 job 开始之前调用。使用 launcher 时该方法在本地机器上、任何 Sweeper/Launcher 初始化之前执行。on_multirun_end(config, **kwargs)在 MULTIRUN 模式下、所有 job 都返回之后调用。使用 launcher 时在本地机器上执行。每个 Job 的事件两种模式共用on_job_start(config, *, task_function, **kwargs)在 RUN 与 MULTIRUN 模式下都会调用每个 Hydra job 各触发一次在运行应用代码之前。该方法从hydra.core.utils.run_job内部调用在远程启动场景下它将与应用代码一起在远程服务器上执行。task_function即被hydra.main装饰的函数。on_job_end(config, job_return, **kwargs)在 RUN 与 MULTIRUN 模式下都会调用每个 Hydra job 各触发一次在运行应用代码之后同样从hydra.core.utils.run_job内部调用远程启动时在远程服务器上执行。job_return是JobReturn对象包含可用于日志记录或后处理的信息详见 hydra/core/utils.py 中的JobReturn定义。触发时序全景综合 hydra/_internal/hydra.py 中run第 159-186 行与multirun第 208-221 行的实现以及 hydra/core/utils.py 中run_job第 174 行与第 194 行的调用点完整的事件流如下RUN 模式on_run_start→on_job_start→应用代码→on_job_end→on_run_endMULTIRUN 模式on_multirun_start→on_job_start每个 job 一次→应用代码→on_job_end每个 job 一次→on_multirun_end仓库测试 tests/test_callbacks.py 中的test_app_with_callbacks用例完整验证了这一时序RUN 模式输出依次为on_run_start → on_job_start → on_job_end → on_run_endMULTIRUN 模式输出依次为on_multirun_start → on_job_start → on_job_end → on_multirun_end。注意其中有一个细节on_job_start的方法签名包含一个仅限关键字的参数task_function测试用例on_job_start_accepts_task_function验证了该参数确实能被回调接收。三、注册与配置 Callback回调子类必须通过配置注册到hydra.callbacks节点Hydra 才会在生命周期中实例化并调用它。注册原理位于 hydra/_internal/callbacks.py 的Callbacks.__init__它通过OmegaConf.select(config, hydra.callbacks)读取配置若存在则遍历config.hydra.callbacks.values()对每个参数集调用hydra.utils.instantiate实例化回调对象。因此每个回调配置本质上是一个 instantiate 节点使用_target_指向可导入的回调类其他字段作为构造函数参数传入。3.1 最小配置示例假设我们有MyCallback希望在每次 job 结束后把某个文件上传到 S3 bucket。为简化演示这里将回调类直接写在应用文件里真实项目中应放在独立文件中。应用代码my_app.pyclass MyCallback(Callback): def __init__(self, bucket: str, file_path: str) - None: self.bucket bucket self.file_path file_path def on_job_end(self, config: DictConfig, **kwargs: Any) - None: print(fJob ended,uploading...) # uploading... hydra.main(version_baseNone, config_pathconf, config_nameconfig) def my_app(cfg: DictConfig) - None: print(OmegaConf.to_yaml(cfg)) if __name__ __main__: my_app()运行结果$ python my_app.py foo: bar Job ended,uploading...配置目录结构conf ├── config.yaml └── hydra └── callbacks └── my_callback.yaml主配置conf/config.yamldefaults: - /hydra/callbacks: - my_callback foo: bar回调配置conf/hydra/callbacks/my_callback.yaml# package _global_ hydra: callbacks: my_callback: _target_: my_app.MyCallback bucket: my_s3_bucket file_path: ./test.pt3.2 配置要点解析通过 defaults list 引入回调hydra.callbacks是一个可选的配置组在 hydra/conf/init.py 的HydraConf中对应callbacks: Dict[str, Any]字段且defaults列表中包含{callbacks: None}表示默认不启用任何回调。在主配置的defaults中声明- /hydra/callbacks: [my_callback]Hydra 会在组合配置时将conf/hydra/callbacks/下的同名配置合并进来。# package _global_头回调配置文件使用_global_包从而把hydra.callbacks键合并到全局配置的hydra节点下这是回调配置生效的关键。_target_指向回调类值为my_app.MyCallback这样可导入的module.Class路径其余字段bucket、file_path会作为MyCallback.__init__的关键字参数传入。回调名称唯一hydra.callbacks下的键如my_callback是该回调实例的标识可用多个键注册多个回调。仓库测试应用 tests/test_apps/app_with_callbacks/custom_callback/my_app.py 展示了完整的多事件回调类实现其配置config_with_two_callbacks.yaml同时注册了callback_1与callback_2两个回调实例是学习配置写法的极佳参考。四、Callback 的调用顺序on_run_start或on_multirun_start最先被调用随后是on_job_start每个 job 调用一次。每个 job 结束后调用on_job_end最后在应用退出前调用一次on_run_end或on_multirun_end。在配置的hydra.callbacks一节中可以使用多个键注册多个回调。这些回调在组合后的最终顺序下对start类事件按正序调用对end类事件按逆序调用。例如假设组合后的配置如下# package hydra.callbacks my_callback1: _target_: my_app.MyCallback1 param1: val1 my_callback2: _target_: my_app.MyCallback2 param2: val2每个 job 开始前先调用MyCallback1.on_job_start再调用MyCallback2.on_job_start每个 job 结束后先调用MyCallback2.on_job_end再调用MyCallback1.on_job_end。这一顺序由 hydra/_internal/callbacks.py 的_notify方法实现on_run_end、on_multirun_end、on_job_end三个“结束类”事件以reverseTrue调用内部对回调列表取反其余事件按注册顺序正序遍历。仓库测试 tests/test_callbacks.py 的two_custom_callbacks用例精确验证了这一点期望输出为[HYDRA] Init callback_1 [HYDRA] Init callback_2 [HYDRA] callback_1 on_run_start [HYDRA] callback_2 on_run_start [JOB] callback_1 on_job_start [JOB] callback_2 on_job_start [JOB] {} [JOB] callback_2 on_job_end [JOB] callback_1 on_job_end [JOB] callback_2 on_run_end [JOB] callback_1 on_run_end可以清楚看到start事件按callback_1 → callback_2正序end事件按callback_2 → callback_1逆序。这一设计保证了“后注册的回调先进行收尾”类似栈式的对称语义便于实现资源清理类逻辑。五、仓库内置的示例回调仓库在 hydra/experimental/callbacks.py 中提供了两个开箱即用的示例回调5.1 LogJobReturnCallback记录任务返回值与错误LogJobReturnCallback在每次 job 结束时记录任务结果尤其适合在远程集群如 Slurm上排查运行错误——当任务在远端失败且日志难以追踪时它会把异常栈打印到本地日志class LogJobReturnCallback(Callback): Log the jobs return value or error upon job end def __init__(self) - None: self.log logging.getLogger(f{__name__}.{self.__class__.__name__}) def on_job_end( self, config: DictConfig, job_return: JobReturn, **kwargs: Any ) - None: if job_return.status JobStatus.COMPLETED: self.log.info(fSucceeded with return value: {job_return.return_value}) elif job_return.status JobStatus.FAILED: self.log.error(, exc_infojob_return._return_value) else: self.log.error(Status unknown. This should never happen.)其判定逻辑基于JobReturn.statusJobStatus.COMPLETED/JobStatus.FAILED成功时以 INFO 级别记录返回值失败时以 ERROR 级别结合_return_value输出异常栈。对应的测试位于 tests/test_callbacks.py 的test_save_job_return_callback验证了 job 失败时错误信息确实被写入 job 日志文件。5.2 PickleJobInfoCallback序列化任务信息用于复现PickleJobInfoCallback将 job 配置与返回值以 pickle 形式保存到输出目录可用于复现一个 Hydra job复现命令见 experimental/rerun 文档on_job_start将 job 配置 pickle 到${output_dir}/config.pickleon_job_end将JobReturnpickle 到${output_dir}/job_return.pickle。class PickleJobInfoCallback(Callback): Pickle the job config/return-value in ${output_dir}/{config,job_return}.pickle output_dir: Path def __init__(self) - None: self.log logging.getLogger(f{__name__}.{self.__class__.__name__}) def on_job_start(self, config: DictConfig, **kwargs: Any) - None: Pickle the jobs config in ${output_dir}/config.pickle. self.output_dir Path(config.hydra.runtime.output_dir) / Path( config.hydra.output_subdir ) filename config.pickle self._save_pickle(objconfig, filenamefilename, output_dirself.output_dir) self.log.info(fSaving job configs in {self.output_dir / filename}) def on_job_end( self, config: DictConfig, job_return: JobReturn, **kwargs: Any ) - None: Pickle the jobs return value in ${output_dir}/job_return.pickle. filename job_return.pickle self._save_pickle(objjob_return, filenamefilename, output_dirself.output_dir) self.log.info(fSaving job_return in {self.output_dir / filename}) def _save_pickle(self, obj: Any, filename: str, output_dir: Path) - None: output_dir.mkdir(parentsTrue, exist_okTrue) assert output_dir is not None with open(str(output_dir / filename), wb) as file: pickle.dump(obj, file, protocol4)注意其输出目录取自config.hydra.runtime.output_dir与config.hydra.output_subdir默认.hydra即最终落在job输出目录/.hydra/下。仓库测试 tests/test_callbacks.py 的test_experimental_save_job_info_callback验证了RUN 模式下文件生成于run.dir/.hydra/MULTIRUN 模式下生成于sweep.dir/0/.hydra/且 pickle 内容与运行中的task_cfg、hydra_cfg完全一致。对应的测试应用为 tests/test_apps/app_with_pickle_job_info_callback/my_app.py。六、综合实战一个可复制的多回调示例将以上知识综合起来下面是一个同时使用内置回调与自定义回调的完整应用骨架可直接对照仓库测试应用app_with_callbacks见 tests/test_apps/app_with_callbacks/custom_callback/my_app.py与app_with_pickle_job_info_callback来验证。应用my_app.py覆写全部 6 个事件方法import logging from typing import Any from omegaconf import DictConfig, OmegaConf import hydra from hydra.core.utils import JobReturn from hydra.experimental.callback import Callback log logging.getLogger(__name__) class CustomCallback(Callback): def __init__(self, callback_name: str) - None: self.name callback_name log.info(fInit {self.name}) def on_job_start(self, config: DictConfig, **kwargs: Any) - None: log.info(f{self.name} on_job_start) def on_job_end( self, config: DictConfig, job_return: JobReturn, **kwargs: Any ) - None: log.info(f{self.name} on_job_end) def on_run_start(self, config: DictConfig, **kwargs: Any) - None: log.info(f{self.name} on_run_start) def on_run_end(self, config: DictConfig, **kwargs: Any) - None: log.info(f{self.name} on_run_end) def on_multirun_start(self, config: DictConfig, **kwargs: Any) - None: log.info(f{self.name} on_multirun_start) def on_multirun_end(self, config: DictConfig, **kwargs: Any) - None: log.info(f{self.name} on_multirun_end) hydra.main(config_path., config_nameconfig) def my_app(cfg: DictConfig) - None: log.info(OmegaConf.to_yaml(cfg)) if __name__ __main__: my_app()配置config_with_two_callbacks.yaml注册两个回调实例hydra: callbacks: callback_1: _target_: my_app.CustomCallback callback_name: callback_1 callback_2: _target_: my_app.CustomCallback callback_name: callback_2以--config-name config_with_two_callbacks运行即可观察两个回调按“start 正序、end 逆序”的完整调用链。若在配置中加入内置示例hydra: callbacks: log_job_return: _target_: hydra.experimental.callbacks.LogJobReturnCallback pickle_job_info: _target_: hydra.experimental.callbacks.PickleJobInfoCallback则每个 job 结束时日志会记录返回值同时输出目录/.hydra/下会生成config.pickle与job_return.pickle用于后续任务复现与审计。七、注意事项与最佳实践注册是生效前提仅覆写Callback方法而不在hydra.callbacks中注册回调不会被实例化见 hydra/_internal/callbacks.py 的遍历实例化逻辑。回调类需可导入_target_指向的类必须位于可导入模块中。文档示例将回调写在应用内真实场景建议独立成文件并配合hydra.utils.target_whitelist显式放行自定义模块仓库测试应用即使用了target_whitelist(my_app.*)见 tests/test_apps/app_with_callbacks/custom_callback/my_app.py 第 48-49 行。结束类事件按逆序触发多个回调同时存在时on_run_end、on_multirun_end、on_job_end会按注册顺序的反向调用编写依赖关系时需注意。on_job_start拥有task_function该参数仅在on_job_start中提供关键字参数可用于记录被调度的目标函数信息on_job_end则额外获得job_return携带任务结果与元数据。远程启动的执行位置on_job_start/on_job_end在远程服务器上与应用代码一同执行而on_multirun_start/on_multirun_end始终在本地机器执行——在设计上传、聚合等逻辑时要区分执行环境。回调异常不中断主流程从源码 hydra/_internal/callbacks.py 的_notify可见单个回调抛出的异常会被捕获并转为warnings.warn不会影响任务本身KeyboardInterrupt场景下仓库测试 tests/test_callbacks.py 的test_callbacks_on_keyboard_interrupt还验证了中断时on_job_end与on_run_end仍会各执行一次并携带完整的JobReturn。依赖关系提示Callback API 属于hydra.experimental模块接口未来可能演进升级 Hydra 版本时建议关注 CHANGELOG 中与 callbacks 相关的说明。【免费下载链接】hydraHydra is a framework for elegantly configuring complex applications项目地址: https://gitcode.com/GitHub_Trending/hyd/hydra创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表