nemo_gym.artifacts.exporters.mlflow

Module Contents

Classes

Name	Description
`MLflowExporter`	MLflow backend.

Functions

Name	Description
`_flatten_config`	Yield (dotted_key, leaf) pairs. MLflow params are flat, unlike W&B’s nested config.
`_landed`	Whether the rollout landed in the output: a failed attempt carries its failure class, and a no-persist one is
`_rollout_key`	-
`_sanitize_key`	Map a metric/param/tag name onto what MLflow accepts: character set, path shape, length.

Data

logger

API

class nemo_gym.artifacts.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'

RESUME_COUNT_TAG

str = 'gym_resume_count'

ROLLOUTS_ARTIFACT_FILE

str = 'rollouts.json'

SEGMENT_CONFIG_ARTIFACT_DIR

str = 'config'

_delivered_config_digest

bytes | None = None

_delivered_metrics

dict[str, Any] = {}

_delivered_rollouts

Counter[bytes] = Counter()

_pending_operations

list[Any] = []

_prior_resumes

Optional[str] = None

_prior_status

Optional[str] = None

client

Optional[MlflowClient] = None

config

= MLFlowConfig.model_validate(global_config_dict)

config_model

type[RunExportConfig] = MLflowExportConfig

name

str = 'mlflow'

run_id

Optional[str] = None

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

nemo_gym.artifacts.exporters.mlflow.MLflowExporter._batch_mode() -> dict[str, bool]

log_batch arguments: a post-run export waits for each batch.

With MLflow’s async logging (MLFLOW_ENABLE_ASYNC_LOGGING) a batch call returns before the server has the batch. The post-run receipt says the whole run reached the backend, so the export logs synchronously and a failed batch raises before the receipt is written. A live run keeps MLflow’s default.

nemo_gym.artifacts.exporters.mlflow.MLflowExporter._close_reopened() -> None

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

nemo_gym.artifacts.exporters.mlflow.MLflowExporter._is_active(
    run_id: str
) -> bool

Whether the recorded run still exists and is not deleted; a new run takes the place of one that is not.

nemo_gym.artifacts.exporters.mlflow.MLflowExporter._landed_rollout_keys(
    client: mlflow.MlflowClient,
    run_id: str
) -> set[tuple[typing.Any, typing.Any]]

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

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

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

nemo_gym.artifacts.exporters.mlflow.MLflowExporter._payload_digest(
    payload: typing.Any
) -> bytes

staticmethod

nemo_gym.artifacts.exporters.mlflow.MLflowExporter._retain_operation(
    operation: typing.Any
) -> None

nemo_gym.artifacts.exporters.mlflow.MLflowExporter.finalize_live(
    artifacts: nemo_gym.artifacts.run_artifacts.RunArtifacts,
    version: dict[str, str | None]
) -> None

Certify the same fresh live run only after the caller verifies final collection and health files.

nemo_gym.artifacts.exporters.mlflow.MLflowExporter.link_superseding_export(
    live_run_id: str,
    run_id: str
) -> None

Tag the live run superseded_by=<run_id> and the post-run export’s run supersedes=<live_run_id>.

A reader that finds two runs for one evaluation picks the one the receipt names. Only tags are set, with this live activation’s client: metrics, params and status are untouched, and MLflow accepts tags on a finished run. Each failure is logged and ignored; the receipt is already written.

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

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

nemo_gym.artifacts.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.artifacts.exporters.mlflow._landed(
    rollout: dict[str, typing.Any]
) -> bool

Whether the rollout landed in the output: a failed attempt carries its failure class, and a no-persist one is written nowhere.

nemo_gym.artifacts.exporters.mlflow._rollout_key(
    rollout: dict[str, typing.Any]
) -> tuple[typing.Any, typing.Any]

nemo_gym.artifacts.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.artifacts.exporters.mlflow.logger = logging.getLogger(__name__)