nemo_automodel.components.loggers.loggers

Typed configs for remote loggers (WandB, MLflow, Comet).

Each logger config is a plain dataclass exposing its YAML-configurable fields plus a build(...) method that initialises and returns the logger / run object. Loggers are a closed, section-named set (no _target_ dispatch), so there is no free builder function — config.build(...) is the entry point.

Module Contents

Classes

Name	Description
`CometConfig`	User-facing Comet ML configuration.
`MLflowConfig`	User-facing MLflow configuration (maps to the YAML `mlflow:` block).
`WandbConfig`	User-facing WandB configuration (maps to the YAML `wandb:` block).

Data

__all__

API

class nemo_automodel.components.loggers.loggers.CometConfig(
    project_name: str = 'automodel',
    workspace: str | None = None,
    api_key: str | None = None,
    experiment_name: str | None = None,
    tags: list[str] = list(),
    auto_metric_logging: bool = False
)

Dataclass

User-facing Comet ML configuration.

api_key

str | None = None

auto_metric_logging

bool = False

experiment_name

str | None = None

project_name

str = 'automodel'

tags

list[str] = field(default_factory=list)

workspace

str | None = None

nemo_automodel.components.loggers.loggers.CometConfig.build(
    model_name: str | None = None
) -> typing.Any

Initialise Comet ML and return the logger (active on rank 0).

When model_name is provided a model:<name> tag is appended and, if experiment_name is empty, one is derived from the model name.

Parameters:

model_name

str | NoneDefaults to None

Optional model name used to tag the run and derive an experiment name when none is set.

Returns: Any

A CometLogger instance.

class nemo_automodel.components.loggers.loggers.MLflowConfig(
    experiment_name: str = 'automodel-experiment',
    run_name: str = '',
    tracking_uri: str | None = None,
    artifact_location: str | None = None,
    tags: dict[str, str] = dict(),
    resume: bool = True,
    description: str | None = None,
    flatten_depth: int | None = 1
)

Dataclass

User-facing MLflow configuration (maps to the YAML mlflow: block).

artifact_location

str | None = None

description

str | None = None

experiment_name

str = 'automodel-experiment'

flatten_depth

int | None = 1

resume

bool = True

run_name

str = ''

tags

dict[str, str] = field(default_factory=dict)

tracking_uri

str | None = None

nemo_automodel.components.loggers.loggers.MLflowConfig.build(
    checkpoint_dir: str | None = None,
    run_config: collections.abc.Mapping[str, typing.Any] | None = None
) -> typing.Any

Initialise MLflow on rank 0 and start (or resume) a run.

Installs a sys.excepthook so crashed jobs report as FAILED rather than FINISHED. On non-rank-0 processes returns None.

Parameters:

checkpoint_dir

str | NoneDefaults to None

Checkpoint directory used to persist / read the mlflow_run_id sidecar for run resumption.

run_config

Mapping[str, Any] | NoneDefaults to None

Full training config dict logged as MLflow params and as a config.yaml artifact.

Returns: Any

Active mlflow.entities.Run on rank 0, or None.

class nemo_automodel.components.loggers.loggers.WandbConfig(
    project: str = 'automodel',
    entity: str | None = None,
    name: str = '',
    group: str | None = None,
    tags: list[str] = list(),
    notes: str | None = None,
    extra: dict[str, typing.Any] = dict()
)

Dataclass

User-facing WandB configuration (maps to the YAML wandb: block).

The named fields are the common ones; any other key under the YAML wandb: block (e.g. mode, dir, resume) is a valid wandb.init() kwarg and is preserved verbatim in extra and forwarded to wandb.init().

entity

str | None = None

extra

dict[str, Any] = field(default_factory=dict)

group

str | None = None

name

str = ''

notes

str | None = None

project

str = 'automodel'

tags

list[str] = field(default_factory=list)

nemo_automodel.components.loggers.loggers.WandbConfig.build(
    run_config: collections.abc.Mapping[str, typing.Any] | None = None,
    model_name: str | None = None
) -> typing.Any

Initialise WandB and return the run.

Parameters:

run_config

Mapping[str, Any] | NoneDefaults to None

Full training config dict logged to the WandB run.

model_name

str | NoneDefaults to None

Optional model name used to derive the run name when name is empty.

Returns: Any

Initialised wandb.Run.

nemo_automodel.components.loggers.loggers.WandbConfig.from_kwargs(
    kwargs: typing.Any = {}
) -> 'WandbConfig'

classmethod

Build from a flat kwargs dict, routing unknown keys to extra.

Keys matching a named field are assigned directly; everything else (valid wandb.init params such as mode/dir that are not first-class fields here) is preserved in extra.

nemo_automodel.components.loggers.loggers.__all__ = ['CometConfig', 'MLflowConfig', 'WandbConfig']