Logger utils

WandbConfig

Bases: BaseModel

Note: name controls the exp name is handled by the NeMoLogger so it is ommitted here. directory is also omitted since it is set by the NeMoLogger.

Parameters:

Name	Type	Description	Default
entity		The team posting this run (default: your username or your default team)	required
project		The name of the project to which this run will belong.	required
tags		Tags associated with this run.	required
group		A unique string shared by all runs in a given group	required
offline		Run offline (data can be streamed later to wandb servers).	required
id		Sets the version, mainly used to resume a previous run.	required
anonymous		Enables or explicitly disables anonymous logging.	required

Source code in bionemo/llm/utils/logger_utils.py

class WandbConfig(BaseModel):
    """Note: `name` controls the exp name is handled by the NeMoLogger so it is ommitted here.
    `directory` is also omitted since it is set by the NeMoLogger.

    Args:
        entity: The team posting this run (default: your username or your default team)
        project: The name of the project to which this run will belong.
        tags: Tags associated with this run.
        group: A unique string shared by all runs in a given group
        offline: Run offline (data can be streamed later to wandb servers).
        id: Sets the version, mainly used to resume a previous run.
        anonymous: Enables or explicitly disables anonymous logging.
    """  # noqa: D205

    entity: str | None  # The team posting this run (default: your username or your default team)
    project: str  # The name of the project to which this run will belong.
    # name: #Display name for the run. "This is handled by NeMoLogger"
    # save_dir: #Path where data is saved. "This is handled by NeMoLogger"
    tags: List[str] | None  # Tags associated with this run.
    group: str | None  # A unique string shared by all runs in a given group
    offline: bool  # Run offline (data can be streamed later to wandb servers).
    id: str | None  # Sets the version, mainly used to resume a previous run.
    anonymous: bool  # Enables or explicitly disables anonymous logging.
    log_model: bool  # Save checkpoints in wandb dir to upload on W&B servers.

setup_nemo_lightning_logger(name='default-name', root_dir='./results', initialize_tensorboard_logger=False, wandb_config=None, ckpt_callback=None, **kwargs)

Setup the logger for the experiment.

Parameters:

Name	Type	Description	Default
name	str	The name of the experiment. Results go into root_dir/name	'default-name'
root_dir	str | Path	The root directory to create the name directory in for saving run results.	'./results'
initialize_tensorboard_logger	bool	Whether to initialize the tensorboard logger.	False
wandb_config	Optional[WandbConfig]	The remaining configuration options for the wandb logger.	None
ckpt_callback	Optional[ModelCheckpoint]	The checkpoint callback to use, must be a child of the pytorch lightning ModelCheckpoint callback. NOTE the type annotation in the underlying NeMoCheckpoint constructor is incorrect.	None
**kwargs	Dict[str, Any]	The kwargs for the NeMoLogger.	{}

Returns:

Name	Type	Description
NeMoLogger	NeMoLogger	NeMo logger instance.

Source code in bionemo/llm/utils/logger_utils.py

def setup_nemo_lightning_logger(
    name: str = "default-name",
    root_dir: str | pathlib.Path = "./results",
    initialize_tensorboard_logger: bool = False,
    wandb_config: Optional[WandbConfig] = None,
    ckpt_callback: Optional[nemo_callbacks.ModelCheckpoint] = None,
    **kwargs: Dict[str, Any],
) -> NeMoLogger:
    """Setup the logger for the experiment.

    Arguments:
        name: The name of the experiment. Results go into `root_dir`/`name`
        root_dir: The root directory to create the `name` directory in for saving run results.
        initialize_tensorboard_logger: Whether to initialize the tensorboard logger.
        wandb_config: The remaining configuration options for the wandb logger.
        ckpt_callback: The checkpoint callback to use, must be a child of the pytorch lightning ModelCheckpoint callback.
            NOTE the type annotation in the underlying NeMoCheckpoint constructor is incorrect.
        **kwargs: The kwargs for the NeMoLogger.

    Returns:
        NeMoLogger: NeMo logger instance.
    """
    # The directory that the logger will save to
    save_dir = pathlib.Path(root_dir) / name
    if wandb_config is not None:
        wandb_logger = WandbLogger(save_dir=save_dir, name=name, **wandb_config.model_dump())
    else:
        wandb_logger = None
        logging.warning("WandB is currently turned off.")
    if initialize_tensorboard_logger:
        tb_logger = TensorBoardLogger(save_dir=save_dir, name=name)
    else:
        tb_logger = None
        logging.warning("User-set tensorboard is currently turned off. Internally one may still be set by NeMo2.")
    logger: NeMoLogger = NeMoLogger(
        name=name,
        log_dir=str(root_dir),
        tensorboard=tb_logger,
        wandb=wandb_logger,
        ckpt=ckpt_callback,
        use_datetime_version=False,
        version="dev",
        **kwargs,
    )
    # Needed so that the trainer can find an output directory for the profiler
    logger.save_dir = save_dir
    return logger