mlflow

Record a script run on an MLflow tracking server.

Lets an example script upload its invocation, configuration, log and outputs so the run can be reproduced from its MLflow entry alone. mlflow is an optional dependency, imported only once tracking is actually enabled.

Classes

MlflowRunLogger

Record one script invocation as an MLflow run.

Functions

current_user

Return the current username, or "unknown" if the uid has no passwd entry.

default_experiment_name

Build an experiment name of the form <user>/<tool>/<model>-<variant>.

validate_tracking_uri

Validate an MLflow tracking URI and return it without a trailing slash.
class MlflowRunLogger

Bases: object

Record one script invocation as an MLflow run.

start() opens the run before the expensive work begins, so a bad URI, a missing token or an unreachable server fails there rather than after hours; it also uploads the invocation and any configuration passed to it, which keeps a crashed run useful. finish() uploads the captured log plus any outputs and closes the run. Everything is a no-op when enabled is false, so callers need no branching.

While the run is open, stdout/stderr are teed to a file that is uploaded as logs/<script>.log. Logging handlers that libraries bound to sys.stderr at import time are re-pointed at the tee for the duration and handed back afterwards.

Failures after the run is open are reported as warnings and never raised: losing a tracking server must not turn a successful job into a failed one.

Note

command.txt masks --*token*-style option values and credentials embedded in a URI, but the captured log is whatever the script printed, so a secret echoed to stdout still reaches the server. Prefer passing credentials via the environment.

tracking_uri must already be validated (see validate_tracking_uri()), experiment_name is created if absent, run_name defaults to the UTC start time YYYYmmdd-HHMMSS, and enabled=False makes every method a no-op – which is how callers skip non-main ranks or an absent flag. required=False additionally downgrades a failure to open the run into a warning: use it when tracking was inferred from the environment rather than asked for, so an uninstalled client or an unreachable server cannot take the job down with it.

Example

>>> logger = MlflowRunLogger(uri, "alice/hf_ptq/Qwen3-0.6B-nvfp4")
>>> logger.start(params={"model": ckpt}, texts={"config.yaml": config_yaml})
>>> status = "FAILED"
>>> try:
...     quantize_and_export()
...     status = "FINISHED"
... finally:
...     logger.finish(status, files={"summary/report.txt": report_path})
__init__(tracking_uri, experiment_name, run_name=None, enabled=True, required=True)

Configure the run without contacting the server; see the class docstring.

Parameters:

  • tracking_uri (str)

  • experiment_name (str)

  • run_name (str | None)

  • enabled (bool)

  • required (bool)

finish(status, texts=None, files=None, metrics=None)

Upload the run’s outputs and close it with status, e.g. "FINISHED".

texts and files both map artifact path to content, from memory and from disk respectively. A files entry is skipped when its file is absent, or was last modified before the run started – so callers can list optional outputs, and a run that produced none of them does not upload a previous run’s leftovers. metrics merges over the default total_time_s.

Parameters:

  • status (str)

  • texts (dict[str, str] | None)

  • files (Mapping[str, Path | str] | None)

  • metrics (dict[str, float] | None)

Return type:

None

property run_url: str

Link to this run in the MLflow UI, or "" before the run is open.

start(params=None, tags=None, texts=None, files=None)

Open the run: capture output, verify the server, upload the inputs.

params are searchable; tags merge over the defaults (user, hostname, ModelOpt version and commit); texts maps artifact path to content, uploaded here rather than at the end so it survives a crash. files names the outputs the run is expected to produce, so finish() can tell them from files that were already there – pass the same mapping to both.

Opening the run is the readiness check: it is MLflow’s own first request, so it honours the client’s TLS and retry configuration rather than second-guessing it. Set MLFLOW_HTTP_REQUEST_MAX_RETRIES to shorten the wait on a dead host.

Raises:

  • ImportError – If mlflow is not installed and required.

  • Exception – Whatever MLflow raises for an unusable server, if required.

Parameters:

  • params (dict[str, Any] | None)

  • tags (dict[str, Any] | None)

  • texts (dict[str, str] | None)

  • files (Mapping[str, Path | str] | None)

Return type:

None

track(params=None, tags=None, texts=None, files=None, metrics=None)

Open the run for the duration of the block, closing it with the right status.

Mirrors mlflow.start_run(). files and metrics are uploaded when the block exits; naming the paths upfront is fine because only files this run actually wrote are uploaded (see finish()).

Example

>>> with logger.track(params={"model": ckpt}, files={"summary.txt": report}):
...     quantize_and_export()
Parameters:

  • params (dict[str, Any] | None)

  • tags (dict[str, Any] | None)

  • texts (dict[str, str] | None)

  • files (Mapping[str, Path | str] | None)

  • metrics (dict[str, float] | None)

Return type:

Iterator[MlflowRunLogger]

current_user()

Return the current username, or "unknown" if the uid has no passwd entry.

Return type:

str

default_experiment_name(tool, model, variant, user=None)

Build an experiment name of the form <user>/<tool>/<model>-<variant>.

Only the basename of model is used, so a local checkpoint directory and an org/name Hugging Face id collapse to the same readable name; variant is whatever distinguishes this run of tool on model, such as a recipe name or a quantization format. Each component is reduced to [A-Za-z0-9._-] so the / separators stay meaningful, and user defaults to the current user.

Example

>>> default_experiment_name("hf_ptq", "/models/Qwen3-0.6B/", "nvfp4", user="alice")
'alice/hf_ptq/Qwen3-0.6B-nvfp4'
Parameters:

  • tool (str)

  • model (str)

  • variant (str)

  • user (str | None)

Return type:

str

validate_tracking_uri(uri)

Validate an MLflow tracking URI and return it without a trailing slash.

Only http(s) servers are accepted; MLflow’s local file: / sqlite: backends are not a useful destination for a shared record of a run.

Raises:

ValueError – If uri is empty, has no host, or is not an http(s) URL.

Parameters:

uri (str)

Return type:

str