kempnerforge.metrics.mlflow.backend

The MLflow logging backend, for Databricks or any MLflow tracking server.

MetricsTracker builds MLflowBackend on rank 0 when metrics.enable_mlflow is on. Its server calls run on the threads of kempnerforge.metrics.mlflow.worker, each under a deadline.

Classes

MLflowBackend

MLflow logging backend (Databricks-hosted or any MLflow tracking server).

class kempnerforge.metrics.mlflow.backend.MLflowBackend

Bases: _LoggingBackend

MLflow logging backend (Databricks-hosted or any MLflow tracking server).

Lazy init on first log; run ID written back to config for checkpoint resume. Every server call has a deadline: the run is set up and later ended on one owner thread (the active run is thread-local), and metric, param, tag and artifact calls go through a best-effort worker, so a slow server never stalls training. A new run (not a resumed one) also gets its description and tags, the git commit, and the resolved config (plus the launch TOML) as artifacts. With mlflow_register_checkpoints, the same worker registers each committed checkpoint as an external model, so a save never waits on the server for it.

The tracking URI resolves as config, then $MLFLOW_TRACKING_URI, then “databricks”. Databricks credentials are not pre-checked: init attempts the server calls, and a failure disables the backend with a warning that names the auth routes.

__init__(config, job_config=None)

Parameters:

• config (MetricsConfig)

• job_config (JobConfig | None)

Return type:

None

log(metrics, step)

Parameters:

• metrics (dict[str, float])

• step (int)

Return type:

None

tracking_entry()

The live run as its checkpoint reader needs it: all four fields non-blank, or none.

Return type:

dict[str, str] | None

checkpoint_fields(step)

The model name of the checkpoint saved at step, plus its id once registered.

Only with mlflow_register_checkpoints on and a live run; answered from memory.

Parameters:

step (int)

Return type:

dict[str, str] | None

register_checkpoint(step, directory, record)

Find or make the external model of a committed checkpoint, then record its id.

Only with mlflow_register_checkpoints on and a live run, and it never waits: a step registered before is recorded at once from memory, and any other is queued on the worker, which records the id when the server answers. A failure warns (rate limited) and costs only the id; a later commit of the same step tries again.

Parameters:

• step (int)

• directory (Path)

• record (Callable[[dict[str, str]], None])

Return type:

None

close()

Return type:

None