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
job_type		Type of run, which is useful when you're grouping runs together into larger experiments.	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
        job_type: Type of run, which is useful when you're grouping runs together into larger experiments.
        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.
    job_type: str | None = (
        None  # Type of run, which is useful when you're grouping runs together into larger experiments.
    )
    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