如何仅将当前Python脚本文件记录到W&B Code面板？

问题描述 #

我想在运行脚本时，只把当前脚本（比如train.py）记录到W&B的Code面板，而非整个目录。目前使用的实现代码：

wandb.run.log_code(f"./{os.path.basename(__file__)}")

需要确认该方法在不同环境下的可靠性，以及是否有针对此场景的更佳实践。

核心代码片段：

def _main(**kwargs):
    from datetime import datetime
    today = datetime.now().strftime('%Y_m%m_d%d_t%Hh_%Mm_%Ss') # eg '2024_m01_d22_t13h_00m_30s'
    run_name = f'{today}' 
    kwargs = kwargs | {'today': today}
    run = wandb.init(mode=kwargs.get('mode', 'online'), project="putnam-axiom", name=run_name, save_code=True, config=kwargs)
    wandb.run.log_code(f"./{os.path.basename(__file__)}") # maybe logscode immediately
    os.environ['CUDA_VISIBLE_DEVICES'] = str(6)
    output_dir = main(**kwargs)
    run_eval_logic_contamination(output_dir)
    run.alert(title="Run Completed", text=f"Run finished, run url: {run.get_url()}")
    print(f'{run.get_url()=}')
    wandb.finish()

解答 #

一、当前方法的可靠性 #

你的方法在绝大多数环境下是可靠的，但需注意以下细节：

• os.path.basename(__file__)能准确获取当前主脚本的文件名，直接运行脚本时不会有问题；
• 相对路径./依赖当前工作目录，若脚本从其他目录启动（比如用绝对路径执行脚本），可能出现文件找不到的情况，建议改用绝对路径规避：

  wandb.run.log_code(os.path.abspath(__file__))
  

• 脚本文件名含特殊字符时，该方法依然有效，因为W&B的log_code支持单个文件路径作为参数。

二、更佳实践 #

1. 关闭自动代码保存

当前wandb.init中设置了save_code=True，这会自动记录整个目录的代码，和手动log_code单个文件的操作冲突，建议改为save_code=False，避免冗余记录：

run = wandb.init(mode=kwargs.get('mode', 'online'), project="putnam-axiom", name=run_name, save_code=False, config=kwargs)

2. 优先使用绝对路径

用os.path.abspath(__file__)替代相对路径，彻底避免工作目录变化带来的文件定位问题：

import os
wandb.run.log_code(os.path.abspath(__file__))

3. 按需记录依赖脚本

如果主脚本依赖其他本地模块（比如示例中的train目录下的文件），可通过log_code的include参数指定需要记录的文件列表，而非整个目录：

wandb.run.log_code(
    root=".",
    include=[
        os.path.basename(__file__),
        "train/callbacks.py",
        "train/data.py"
    ]
)

4. 尽早调用log_code

你当前在wandb.init后立刻调用log_code是合理的，能确保训练开始前就记录下当前脚本的状态，避免训练过程中脚本被修改导致记录不一致。

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