Skip to content

Logging

Functions:

init_wandb(cfg, run_dir, config_for_wandb)

Initialize W&B.

config_for_wandb is the configuration that will be logged to W&B.

Source code in tapeagents/finetune/logging_.py
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def init_wandb(
    cfg: DictConfig,
    run_dir: Path,
    config_for_wandb: DictConfig | dict,
) -> wandb_run.Run:
    """Initialize W&B.

    config_for_wandb is the configuration that will be logged to W&B.

    """
    if config_for_wandb is None:
        config_for_wandb = cfg.dict()

    wandb_id = cfg.finetune.wandb_id

    if cfg.finetune.wandb_resume == "always":
        resume = True
    elif cfg.finetune.wandb_resume == "if_not_interactive":
        resume = not cfg.finetune.force_restart
    else:
        raise ValueError(f"Unknown value for wandb_resume: {cfg.finetune.wandb_resume}")
    wandb_name = run_dir.name if cfg.finetune.wandb_use_basename else str(run_dir)

    if len(wandb_name) > 128:
        logger.warning(f"wandb_name: {wandb_name} is longer than 128 characters. Truncating to 128 characters.")

    logging.info(f"Initializing W&B with name: {wandb_name[:128]}, resume: {resume}")
    run = wandb.init(
        name=wandb_name[:128],  # wandb limits name to 128 characters
        entity=cfg.finetune.wandb_entity_name,
        project=cfg.finetune.wandb_project_name,
        config=config_for_wandb,  # type: ignore
        resume=resume,
        id=wandb_id,
        tags=cfg.finetune.tags,
    )
    if not isinstance(run, wandb_run.Run):
        raise ValueError("W&B init failed")
    return run