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.
`Tool`	What distinguishes one script's tracking from another's.

Functions

`add_mlflow_args`	Add `--mlflow`, `--mlflow_experiment` and `--mlflow_run_name` to parser.
`command_text`	The invocation, as a copy-pasteable line, with credentials masked.
`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>`.
`default_run_name`	The UTC start time as `YYYYmmdd-HHMMSS`, which is what the flags document.
`describe_run`	The keyword arguments `MlflowRunLogger.track()` takes, for this run.
`drop_experiment_json`	Remove a provenance pointer an untracked export would otherwise inherit.
`log_active_run_experiment_json`	Record MLflow's currently active run as the producer of a checkpoint.
`mask_tracking_uri`	Mask any `user:token@` a tracking URI carries, for printing.
`masked_args`	A copy of args whose tracking URI cannot leak credentials into a printed namespace.
`resolve_mlflow_args`	Settle where tracking is configured from, and name the experiment, in place.
`resolve_tracking_uri`	Settle the tracking URI from `--mlflow` and the environment.
`resolved_recipe_texts`	`{artifact path: content}` for recipe, or `{}` when the run used none.
`run_tags`	This run's join keys, shared with whatever is later done with what it produced.
`split_tracking_credentials`	Move any `user:token@` out of uri into MLflow's own credential variables.
`tracked_run`	Track one invocation of tool for the duration of the block.
`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

log_experiment_json(checkpoint_dir=None)#

Record which MLflow run produced a checkpoint, on the server and in the checkpoint.

Tags point from a run to the checkpoint it wrote; this is the reverse, so a checkpoint found on disk can be traced back to the run that produced it without searching the server. The artifact goes up for any run that opened, so a failure is traceable from the server side too.

checkpoint_dir also writes the JSON there as EXPERIMENT_JSON. Pass it only once the checkpoint is really on disk, since the file claims authorship of the weights sitting next to it: an output directory existing proves nothing, as it may hold a checkpoint from an earlier attempt whose weights this run never touched.

After the checkpoint is written, the pointer beside it is this run’s or absent – never a previous run’s. So a run that never opened removes the pointer rather than leaving one: tracking can disable itself mid-flight (an unreachable server or an uninstalled client, which a URI inherited from the environment tolerates by design), and the caller’s untracked cleanup was skipped because tracking looked configured.

Parameters:: checkpoint_dir (Path | str | None)
Return type:: None

log_text(artifact_path, text)#

Upload text as an artifact while the run is open, best-effort.

For a value that is only settled midway through the run and is worth having even if the run later crashes – the quantization config a calibration is about to apply, say. start() and finish() cover everything known at the two ends.

Parameters:

artifact_path (str)
text (str)

Return type:

None

property run_info: dict[str, str]#

Identity of this run on the server, or {} before it is open.

Enough for a consumer holding only this run’s outputs to find it again: run_id is MLflow’s own identifier for the run, a uuid4 hex, unique across experiments. Every field is read back off the run the server returned rather than off what was requested, so a run MLflow resolved differently is reported as it really is.

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]

class Tool#

Bases: object

What distinguishes one script’s tracking from another’s.

A script declares one of these and the functions below do the rest, so the flags, the experiment-name convention, the join tags and the provenance pointer are settled in one place rather than per script.

model and checkpoint read the arguments naming this run’s input and output; source is what it consumed when that differs from the model, which is what the next run in a chain joins on; variant names what this run did, for the default experiment name. texts and outputs are artifacts to upload, from memory and from disk, and metrics is read once the block ends. non_params names arguments to keep out of the params – the tracking settings and a script’s own bookkeeping. settles_pointer is false when something other than tracked_run() writes the pointer – a script that writes several checkpoints in one run and points each of them at the run itself, or one whose checkpoint is saved by a training loop.

__init__(name, tracks, variant_help, variant, model, checkpoint=<function Tool.<lambda>>, source=None, texts=<function Tool.<lambda>>, outputs=<function Tool.<lambda>>, metrics=<function Tool.<lambda>>, non_params=frozenset({}), settles_pointer=True)#

Parameters:

name (str)
tracks (str)
variant_help (str)
variant (Callable[[Namespace], str])
model (Callable[[Namespace], str])
checkpoint (Callable[[Namespace], str | None])
source (Callable[[Namespace], str] | None)
texts (Callable[[Namespace], dict[str, str]])
outputs (Callable[[Namespace], dict[str, Path]])
metrics (Callable[[Namespace], dict[str, float]])
non_params (frozenset[str])
settles_pointer (bool)

Return type:

None

checkpoint()#

metrics()#

model: Callable[[Namespace], str]#

name: str#

non_params: frozenset[str] = frozenset({})#

outputs()#

settles_pointer: bool = True#

source: Callable[[Namespace], str] | None = None#

texts()#

tracks: str#

variant: Callable[[Namespace], str]#

variant_help: str#

add_mlflow_args(parser, tool)#

Add --mlflow, --mlflow_experiment and --mlflow_run_name to parser.

The help text comes from tool: its tracks describes what this script uploads and its variant_help says what the experiment name’s variant is derived from. Pair with resolve_mlflow_args().

The multi-word flags are registered under both the underscored and the dashed spelling: vLLM’s FlexibleArgumentParser rewrites every --foo_bar on the command line to --foo-bar before matching, so the dashed spelling has to exist for the flag to be reachable there at all, and a user moving between the example scripts should not have to remember which spelling each one took.

Parameters:

parser (ArgumentParser)
tool (Tool)

Return type:

None

command_text(argv=None)#

The invocation, as a copy-pasteable line, with credentials masked.

argv defaults to this process’s own sys.argv. Pass another process’s argv when the run is opened somewhere the user never typed a command – a worker subprocess, whose own sys.argv is an implementation detail rather than a reproducible invocation.

Parameters:: argv (list[str] | None)
Return type:: str

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

default_run_name()#

The UTC start time as YYYYmmdd-HHMMSS, which is what the flags document.

Used by MlflowRunLogger and by a caller handing the name to something else that opens the run, so both honour the documented default rather than MLflow’s random one.

Return type:: str

describe_run(args, tool, world_size=1)#

The keyword arguments MlflowRunLogger.track() takes, for this run.

Every command-line argument becomes a searchable param, so a flag added later is tracked without touching this. world_size is recorded separately because the parallelism flags say how a run was laid out but not how many processes it took.

Parameters:

args (Namespace)
tool (Tool)
world_size (int)

Return type:

dict

drop_experiment_json(checkpoint_dir)#

Remove a provenance pointer an untracked export would otherwise inherit.

A fresh checkpoint written into a reused output directory would keep the previous run’s pointer, and one produced from a tracked source checkpoint could be handed that source’s pointer. Either way the file would name a run that did not produce these weights. Call it only for a completed export; a failed run leaves whatever checkpoint was already there, pointer included.

Parameters:: checkpoint_dir (Path | str)
Return type:: None

log_active_run_experiment_json(checkpoint_dir)#

Record MLflow’s currently active run as the producer of a checkpoint.

For a caller whose run is opened and owned by something else – Megatron-Bridge’s LoggerConfig opens the run for a training job, on its last rank – so the checkpoint still names the run that produced it, in the format MlflowRunLogger.log_experiment_json() writes. Call it from the rank that owns the run, once the checkpoint is on disk.

When no run can be found – mlflow absent, tracking off, the run never opened, or the server unreachable when asked – any pointer already beside the checkpoint is removed rather than left: the weights are new, so a previous run’s pointer would misname their author. Failures are warnings: losing the pointer must not fail a finished job, and the file beside the weights is written before the upload is attempted so a server that goes away cannot cost it.

Returns whether a pointer was written, so a caller that asked for tracking can tell that apart from an untracked job, which reaches this too and is quiet by design.

Parameters:: checkpoint_dir (Path | str)
Return type:: bool

mask_tracking_uri(uri)#

Mask any user:token@ a tracking URI carries, for printing.

Credentials in the URI are a supported form, so everything this module prints or uploads masks them – command.txt, the logged params, MlflowRunLogger.run_url. A caller that prints the URI itself (a script echoing its parsed arguments, say) has to do the same, or the secret reaches a console log that is routinely archived.

Parameters:: uri (str | None)
Return type:: str | None

masked_args(args, attr='mlflow')#

A copy of args whose tracking URI cannot leak credentials into a printed namespace.

For a script that echoes its parsed arguments: a user:token@ in the URI is a supported form that this module masks wherever it prints or uploads one, and a job log is routinely archived. Uploaded artifacts are unaffected – command_text() redacts, and the URI is not worth logging as a param.

Parameters:

args (Namespace)
attr (str)

Return type:

Namespace

resolve_mlflow_args(args, parser, tool)#

Settle where tracking is configured from, and name the experiment, in place.

Sets args.mlflow to the validated URI or None, args.mlflow_required to whether the flag asked for it, and defaults args.mlflow_experiment from tool. Pair with add_mlflow_args().

Parameters:

args (Namespace)
parser (ArgumentParser)
tool (Tool)

Return type:

None

resolve_tracking_uri(uri, parser)#

Settle the tracking URI from --mlflow and the environment.

Returns (uri or None, required), where required records that the flag was passed. Only the flag is a deliberate request, so only the flag is fatal when the URI is unusable: the environment variable is commonly exported for unrelated tooling and must not fail a job that would otherwise have worked.

Parameters:

uri (str | None)
parser (ArgumentParser)

Return type:

tuple[str | None, bool]

resolved_recipe_texts(recipe)#

{artifact path: content} for recipe, or {} when the run used none.

The resolved recipe, not the source file: a recipe may be a directory or use $imports, and only the resolved form stands alone.

Parameters:: recipe (str | None)
Return type:: dict[str, str]

run_tags(args, tool)#

This run’s join keys, shared with whatever is later done with what it produced.

checkpoint_path is the checkpoint the run writes, because that is what an export or an evaluation is later pointed at (NEL takes deployment.checkpoint_path), and source_checkpoint_path is what it consumed, so a chain of runs joins on the pair: a distillation’s source is the checkpoint it continues from, not the model that was quantized. Both are resolved, since a relative path is useless as a join key – except a source that names no directory, such as a Hub org/name id.

Parameters:

args (Namespace)
tool (Tool)

Return type:

dict[str, str]

split_tracking_credentials(uri)#

Move any user:token@ out of uri into MLflow’s own credential variables.

For a caller that hands the URI to something which records it. Megatron-Bridge logs its resolved config as MLflow params and serialises it into the checkpoint it writes, so a credential left in the URI becomes durable in two places; masking is not an option there because the value is also what authenticates. MLflow reads MLFLOW_TRACKING_USERNAME/MLFLOW_TRACKING_PASSWORD itself, so moving them across keeps the request working. Variables the caller already exported win, since those were set deliberately.

Parameters:: uri (str)
Return type:: str

tracked_run(args, tool, is_main, exported, world_size=1)#

Track one invocation of tool for the duration of the block.

Inert unless --mlflow settled a URI and this is the rank that records it, so the caller needs no branching. is_main is that rank, and also gates the writes every rank would otherwise race on; exported is read on the way out, once the run knows whether it wrote the checkpoint its pointer would claim.

Example

>>> with tracked_run(args, HF_PTQ, is_main, lambda: args.exported, world_size):
...     quantize_and_export(args)

Parameters:

args (Namespace)
tool (Tool)
is_main (bool)
exported (Callable[[], bool])
world_size (int)

Return type:

Iterator[MlflowRunLogger]

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