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.
`TeeStream`	Mirror a text stream to sink while passing writes through to stream.

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() verifies the server and opens the run before the expensive work begins, so a bad URI or a missing token fails in seconds 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 records sys.argv verbatim and the log captures everything the script prints, so anything secret on the command line or in the output is uploaded with it. Pass credentials through the environment instead.

Parameters:

tracking_uri – Validated MLflow server URI (see validate_tracking_uri()).
experiment_name – Experiment to log under; created if it does not exist.
run_name – Name for the run. Defaults to the UTC start time, YYYYmmdd-HHMMSS.
enabled – When false, every method returns immediately. Use this to skip non-main ranks or a missing --mlflow flag.

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)

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

Parameters:

tracking_uri (str)
experiment_name (str)
run_name (str | None)
enabled (bool)

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

Upload the run’s outputs and close it with status.

Parameters:

status (str) – MLflow run status, typically "FINISHED" or "FAILED".
texts (dict[str, str] | None) – Artifacts to upload as text, keyed by artifact path.
files (Mapping[str, Path | str] | None) – Artifacts to upload from disk, keyed by artifact path. Entries whose file does not exist are skipped, so callers can list optional outputs.
metrics (dict[str, float] | None) – Extra metrics, merged over the default total_time_s.

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)

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

Parameters:

params (dict[str, Any] | None) – Searchable parameters describing the run’s configuration.
tags (dict[str, Any] | None) – Extra tags, merged over the defaults (user, hostname, ModelOpt version and commit).
texts (dict[str, str] | None) – Artifacts to upload as text, keyed by artifact path. Uploaded here rather than at the end so they survive a crash.

Raises:

ImportError – If mlflow is not installed.
ConnectionError – If the tracking server is unreachable.

Return type:

None

class TeeStream

Bases: object

Mirror a text stream to sink while passing writes through to stream.

Scripts that report progress with bare print() have no log file; wrapping sys.stdout/sys.stderr in this is what produces one. Attribute access falls through to the wrapped stream so isatty() keeps progress bars behaving. Native (C-level) writes go straight to the real file descriptor and are not captured.

__init__(stream, sink): Wrap stream, mirroring everything written to it into the open file sink.

flush()

Flush both the original stream and the sink.

Return type:: None

write(data)

Write to both the original stream and the sink.

Parameters:: data (str)
Return type:: int

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. Each component is reduced to [A-Za-z0-9._-] so the / separators stay meaningful.

Parameters:

tool (str) – Name of the script producing the run, e.g. "hf_ptq".
model (str) – Checkpoint path or Hugging Face model id.
variant (str) – What distinguishes this run of tool on model, e.g. a recipe name or a quantization format.
user (str | None) – Owner of the run. Defaults to the current user.

Returns:

The experiment name.

Return type:

str

Example

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

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.

Parameters:: uri (str) – The tracking URI to validate, e.g. https://mlflow.example.com/.
Returns:: The URI with any trailing slash removed.
Raises:: ValueError – If uri is empty, has no host, or is not an http(s) URL.
Return type:: str