nemo_gym.exporters.mlflow

Module Contents

Classes

Name	Description
`MLflowExporter`	MLflow backend.

Functions

Name	Description
`_elide_key`	Cut an over-long name in the middle, keeping both ends.
`_flatten_config`	Yield (dotted_key, leaf) pairs. MLflow params are flat, unlike W&B’s nested config.
`_normalize_key_path`	Neutralize only the path segments MLflow rejects, leaving the rest of the name intact.
`_sanitize_key`	Map a metric/param/tag name onto what MLflow accepts: character set, path shape, length.

Data

API

class nemo_gym.exporters.mlflow.MLflowExporter(
    global_config_dict: omegaconf.DictConfig
)

Bases: BaseExporter

MLflow backend.

Configured by mlflow_tracking_uri, mlflow_experiment_name and mlflow_run_name in the global config, plus mlflow_tracking_token for servers that require authentication.

Uses MlflowClient with an explicit run id rather than the fluent mlflow.* API, which tracks the active run in thread-local state that Gym’s async call sites don’t share.

CONFIG_ARTIFACT_FILE

str = 'global_config.json'

ROLLOUTS_ARTIFACT_FILE

str = 'rollouts.json'

client

Optional[MlflowClient] = None

config

= MLFlowConfig.model_validate(global_config_dict)

name

str = 'mlflow'

run_id

Optional[str] = None

nemo_gym.exporters.mlflow.MLflowExporter._active() -> tuple[mlflow.MlflowClient, str]

nemo_gym.exporters.mlflow.MLflowExporter._experiment_id(
    name: str
) -> str

nemo_gym.exporters.mlflow.MLflowExporter._log_config(
    config_dict: omegaconf.DictConfig
) -> None

nemo_gym.exporters.mlflow.MLflowExporter._log_metrics(
    metrics: dict[str, typing.Any],
    step: typing.Optional[int] = None
) -> None

nemo_gym.exporters.mlflow.MLflowExporter._log_rollouts(
    rollouts: list[dict[str, typing.Any]]
) -> None

nemo_gym.exporters.mlflow.MLflowExporter.setup() -> None

nemo_gym.exporters.mlflow.MLflowExporter.teardown() -> None

nemo_gym.exporters.mlflow._elide_key(
    key: str,
    limit: int = MAX_ENTITY_KEY_LENGTH
) -> str

Cut an over-long name in the middle, keeping both ends.

Gym names carry their identity at the edges — benchmark and agent in the prefix, the statistic in the suffix (gpqa.../simple_agent/mean/reward) — so the middle is what can be spared.

nemo_gym.exporters.mlflow._flatten_config(
    value: typing.Any,
    prefix: str = ''
) -> typing.Iterator[tuple[str, typing.Any]]

Yield (dotted_key, leaf) pairs. MLflow params are flat, unlike W&B’s nested config.

nemo_gym.exporters.mlflow._normalize_key_path(
    key: str
) -> str

Neutralize only the path segments MLflow rejects, leaving the rest of the name intact.

MLflow requires an already-normalized relative path, so it rejects empty segments (from leading, doubled or trailing slashes) and the traversal forms . and ...

nemo_gym.exporters.mlflow._sanitize_key(
    key: str
) -> typing.Optional[str]

Map a metric/param/tag name onto what MLflow accepts: character set, path shape, length.

Gym emits names MLflow rejects, e.g. pass@1 from pass@k metrics. Returns None for a name made only of separators, leaving nothing to sanitize into. Dropping that one entry is deliberate: MLflow rejects an entire log_batch when any single name is invalid.

nemo_gym.exporters.mlflow._ELISION_MARKER = '___'

nemo_gym.exporters.mlflow._INVALID_KEY_CHARS = re.compile('[^/\\w.\\- :]')

nemo_gym.exporters.mlflow._LEADING_DOTS = re.compile('^\\.+')

nemo_gym.exporters.mlflow.logger = logging.getLogger(__name__)