hssm.track

hssm.track records an inference workflow — HSSM(...), sample() or vi(), and save_model() — as one MLflow run: the model and its priors, the sampler settings, convergence diagnostics, the traces, and any parameters, metrics or figures you add yourself. Where the likelihood is a published network, the run also carries that network's provenance. Tracking is opt-in and requires the optional tracking extra (uv add "hssm[tracking]"). See Track fits with MLflow for usage.

hssm.tracking.track ¶

track(
    experiment: str | None = None,
    run_name: str | None = None,
    *,
    tracking_uri: str | None = None,
    log_artifacts: bool | Literal["all"] = True,
    lineage_id: str | None = None,
    tags: dict[str, str] | None = None,
    params: dict[str, Any] | None = None,
    dataset_name: str | None = None,
) -> Iterator[Tracker]

Record the enclosed HSSM workflow as one MLflow run (phase=infer).

One block records one fit. Fitting twice inside a single block is not supported: only the first sample() or vi() is recorded, and a later one is skipped with a warning. Open one block per fit.

Parameters:

experiment (str | None, default: None ) –

MLflow experiment name. Defaults to MLFLOW_EXPERIMENT_NAME; the ecosystem convention is infer/<model>.
run_name (str | None, default: None ) –

Optional run name.
tracking_uri (str | None, default: None ) –

Tracking server. Defaults to MLFLOW_TRACKING_URI (the lab's shared server), then MLflow's own default. Switching to a different server without naming an experiment records the run in that server's Default experiment.
log_artifacts (bool | Literal['all'], default: True ) –

True attaches traces.nc, summary.csv, model_spec.json and data.parquet — the data the model was fit to, so note that it is copied into every run, including runs sent to a shared server; "all" additionally attaches model.pkl when save_model is called inside the block; False logs params/metrics only.
lineage_id (str | None, default: None ) –

Override the lineage id. By default it is read from the HuggingFace manifest.json entry of the network in use, else minted.
tags (dict[str, str] | None, default: None ) –

Extra tags. schema_version, phase and lineage_id are reserved and rejected.
params (dict[str, Any] | None, default: None ) –

Your own parameters for this run — the dimensions you are varying, such as a dataset version or a prior set — so runs can be compared on them in the MLflow UI. Keys HSSM records itself are rejected; see :data:RESERVED_PARAMS. More can be added later with :meth:Tracker.log_param.
dataset_name (str | None, default: None ) –

Name to register the observed data under, which fills the run's "Datasets" panel with the column schema and row count and lets the UI filter by dataset. Omitted, no dataset is recorded: a DataFrame carries no name of its own, and an invented one would look the same for every study. The data is identified either way by the data_sha256 tag.

Raises:

ImportError –

If MLflow is not installed (uv add "hssm[tracking]").
RuntimeError –

If the block is nested inside another.
ValueError –

If tags or params collide with what HSSM records itself.
Exception –

Whatever MLflow raises when the tracking server or the experiment cannot be set up, before the block runs. Once the run is open, tracking failures are logged and never raised.

hssm.tracking.load_run ¶

load_run(run_id: str, tracking_uri: str | None = None) -> HSSM

Rebuild a fitted model from a run recorded by :func:track.

The model is constructed again from the run's model_spec.json and data.parquet, and its traces.nc and vi_traces.nc are attached when the run has them.

Parameters:

run_id (str) –

The MLflow run to rebuild.
tracking_uri (str | None, default: None ) –

Tracking server to read from. Defaults to the current MLflow tracking URI, which is left unchanged either way.

Returns:

HSSM –

The rebuilt model, with its posterior restored.

Raises:

ImportError –

If MLflow is not installed (uv add "hssm[tracking]").
NotImplementedError –

If the run fitted a model class other than HSSM.
ValueError –

If the run kept no data or spec (it was logged with log_artifacts=False), or if its spec holds arguments that were recorded as text, such as bmb.Prior objects or custom functions.

hssm.tracking.list_runs ¶

list_runs(tracking_uri: str | None = None) -> pd.DataFrame

List the runs recorded by :func:track, newest first.

Runs logged by other tools on the same server are left out.

Parameters:

tracking_uri (str | None, default: None ) –

Tracking server to read from. Defaults to the current MLflow tracking URI, which is left unchanged either way.

Returns:

DataFrame –

One row per run, with columns run_id, run_name, experiment, start_time, status, user, model, loglik_kind, dataset_name, n_trials and restorable, which says whether :func:load_run can rebuild the run.

Raises:

ImportError –

If MLflow is not installed (uv add "hssm[tracking]").

hssm.tracking.Tracker ¶

Tracker(
    mlflow: Any,
    run_id: str,
    *,
    log_artifacts: bool | Literal["all"],
    lineage_id: str | None,
    dataset_name: str | None = None,
)

State for one tracked inference run. Obtained from :func:track.

Methods:

log_param –

Attach one of your own parameters to the run.
log_params –

Attach several of your own parameters at once. See :meth:log_param.
log_metric –

Attach an extra metric (e.g. loo_elpd) to the run.
log_figure –

Attach a matplotlib figure, e.g. log_figure(fig, "plots/trace.png").

hssm.tracking.Tracker.log_param ¶

log_param(key: str, value: Any) -> None

Attach one of your own parameters to the run.

Use this for the dimensions you are varying — a dataset version, a prior set, which subjects were included — so runs can be compared on them in the MLflow UI. Keys this module logs itself are refused; see :data:RESERVED_PARAMS.

hssm.tracking.Tracker.log_params ¶

log_params(params: dict[str, Any]) -> None

Attach several of your own parameters at once. See :meth:log_param.

hssm.tracking.Tracker.log_metric ¶

log_metric(key: str, value: float) -> None

Attach an extra metric (e.g. loo_elpd) to the run.

hssm.tracking.Tracker.log_figure ¶

log_figure(figure: Any, artifact_file: str) -> None

Attach a matplotlib figure, e.g. log_figure(fig, "plots/trace.png").