如何在Hydra+Databricks工作流中传入非Hydra运行时命令行参数？

问题场景 #

在Databricks工作流中使用Hydra框架时，尝试通过argparse接收非Hydra配置类的运行时命令行参数，传入["--task_values","{'T1':0,'T2':1}"]后触发错误：

usage: set_task_values.py [--help] [--hydra-help] [--version]
                          [--cfg {job,hydra,all}] [--resolve]
                          [--package PACKAGE] [--run] [--multirun]
                          [--shell-completion] [--config-path CONFIG_PATH]
                          [--config-name CONFIG_NAME]
                          [--config-dir CONFIG_DIR]
                          [--experimental-rerun EXPERIMENTAL_RERUN]
                          [--info [{all,config,defaults,defaults-tree,plugins,searchpath}]]
                          [overrides [overrides ...]]
set_task_values.py: error: unrecognized arguments: --task_values

现有脚本代码 #

import hydra
from omegaconf import DictConfig, OmegaConf
from hydra.core.hydra_config import HydraConfig
import argparse

@hydra.main(version_base=None, config_path="conf_tasks", config_name="config_dummy")
def my_app(cfg : DictConfig) -> None:
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument("--task_values", type=str, required=True)
    args, unknown = parser.parse_known_args()
    print("Task Values", "->", args.task_values)
    dbutils.jobs.taskValues.set(key   = "task_values", \
                                value = args.task_values)

if __name__ == "__main__":
    my_app()

现有配置文件config_dummy.yaml #

operand: add
batch_size: 64
learning_rate: 0.01
creation_did: 47
custom:
  email: someone@something.com
  exp_title: Exp_Tasks

hydra:
    run:
      dir: /dbfs/some_path/${custom.email}/${now:%Y-%m-%d}/${custom.exp_title}

问题原因 #

Hydra的@hydra.main装饰器会在脚本启动时优先解析所有命令行参数，任何未被Hydra识别的参数都会直接触发"unrecognized arguments"错误，无法传递给后续的argparse处理。

解决方案 #

方案1：使用Hydra参数分隔符--保留自定义参数 #

Hydra支持用--作为分隔符，分隔符之后的参数会被Hydra忽略，直接留给argparse处理：

• 修改Databricks工作流的命令行参数为：

  ["--", "--task_values", "{'T1':0,'T2':1}"]
  

• 现有脚本无需修改，argparse的parse_known_args会正常获取到--task_values参数。

方案2：将自定义参数整合到Hydra配置（推荐） #

更贴合Hydra的配置优先设计理念，避免参数解析冲突：

1. 在config_dummy.yaml中新增自定义字段：

   # 新增字段，初始值设为空或默认值
   task_values: ""
   operand: add
   # 其余原有配置保持不变
   

2. 修改脚本，直接从Hydra配置中读取参数，移除argparse代码：

   import hydra
   from omegaconf import DictConfig, OmegaConf
   from hydra.core.hydra_config import HydraConfig
   
   @hydra.main(version_base=None, config_path="conf_tasks", config_name="config_dummy")
   def my_app(cfg : DictConfig) -> None:
       print("Task Values", "->", cfg.task_values)
       dbutils.jobs.taskValues.set(key="task_values", value=cfg.task_values)
   
   if __name__ == "__main__":
       my_app()
   

3. 修改Databricks工作流的参数为Hydra的配置覆盖格式：

   ["task_values={'T1':0,'T2':1}"]
   

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