nemo_automodel.components.loggers

Submodules

Package 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).

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

Comet API key. None reads from COMET_API_KEY env var.

auto_metric_logging

bool = False

Enable Comet’s automatic metric logging.

experiment_name

str | None = None

Display name for this experiment run.

project_name

str = 'automodel'

Comet project name.

tags

list[str] = field(default_factory=list)

List of string tags attached to the experiment.

workspace

str | None = None

Comet workspace. None uses the default.

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

Root artifact store URI for the experiment.

description

str | None = None

Free-text description shown in the MLflow UI (sets the mlflow.note.content tag).

experiment_name

str = 'automodel-experiment'

MLflow experiment name.

flatten_depth

int | None = 1

Nesting depth for mlflow.log_params. 1 (default) splits one level; None is fully recursive.

resume

bool = True

When True (default), look for a mlflow_run_id sidecar in the checkpoint dir and resume that run.

run_name

str = ''

Display name for the run.

tags

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

Dictionary of string tags attached to the run.

tracking_uri

str | None = None

MLflow tracking server URI. None delegates to MLFLOW_TRACKING_URI and then MLflow’s local default (sqlite:///mlflow.db in current releases).

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,
    swanlab: nemo_automodel.components.loggers.loggers.SwanLabConfig | 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

WandB team / entity. None uses the default from the wandb config or WANDB_ENTITY env var.

extra

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

Any additional wandb.init() kwargs (mode, dir, …) carried through unfiltered.

group

str | None = None

Group name for related runs.

name

str = ''

Display name for the run. When empty, build derives one from the model name.

notes

str | None = None

Free-text notes shown in the WandB UI.

project

str = 'automodel'

WandB project name.

swanlab

SwanLabConfig | None = None

Optional SwanLabConfig that mirrors this run to SwanLab (the wandb.swanlab: sub-block).