dispcraft.model_registry¶

def latest_version(name):
    """Highest version number currently registered under `name`.

    Uses `get_registered_model(name).latest_versions`, not
    `search_model_versions` -- the latter is a *fourth* GitLab quirk found
    while validating this module: `SearchModelVersions` 404s outright on
    this registry (unlike `search_registered_models`, which works), while
    `GetRegisteredModel`'s `latest_versions` field returns the real,
    per-name version list correctly.
    """
    client = MlflowClient()
    versions = client.get_registered_model(name).latest_versions
    if not versions:
        raise ValueError(f"no registered model versions found for {name!r}")
    return max(int(v.version) for v in versions)

def list_models():
    """Names of every model currently registered (sorted).

    Bypasses `MlflowClient.search_registered_models()` -- see quirk 5 in the
    module docstring -- with a direct, unfiltered REST call instead.
    """
    url = f"{mlflow.get_tracking_uri()}/api/2.0/mlflow/registered-models/search"
    token = os.environ.get("MLFLOW_TRACKING_TOKEN")
    headers = {"Authorization": f"Bearer {token}"} if token else {}
    response = requests.get(url, headers=headers)
    response.raise_for_status()
    return sorted(m["name"] for m in response.json().get("registered_models", []))

def load_model(name, version=None):
    """Load a registered `LitResidualRegressor` by name.

    Parameters
    ----------
    name    : str -- registered model name
    version : int, optional -- defaults to `latest_version(name)`

    Returns
    -------
    (lit_model, extra) : (dispcraft.ml.LitResidualRegressor, dict) -- the
        model, in eval mode, ready for `dispcraft.ml.predict`; `extra` is
        whatever dict was passed to `register_model` (empty if none).
    """
    if version is None:
        version = latest_version(name)
    with tempfile.TemporaryDirectory() as tmp:
        local_dir = mlflow.artifacts.download_artifacts(artifact_uri=f"models:/{name}/{version}", dst_path=tmp)
        ckpt_path = next(Path(local_dir).rglob(_CHECKPOINT_NAME))
        ckpt = torch.load(ckpt_path, weights_only=False)

    hp = ckpt["hparams"]
    lit_model = LitResidualRegressor(
        n_inputs=hp["n_inputs"], n_outputs=hp["n_outputs"], hidden_layer_sizes=hp["hidden_layer_sizes"],
        activation=hp["activation"], alpha=hp["alpha"], lr=hp["lr"], batch_size=hp["batch_size"],
    )
    lit_model.load_state_dict(ckpt["state_dict"])
    lit_model.eval()
    return lit_model, ckpt.get("extra", {})

def model_web_url(name, version=None):
    """(model_url, version_url) -- GitLab UI links for a registered model,
    for documentation/`models/*.toml` where a human wants to click through
    and inspect a model, as opposed to `load_model` (name/version alone is
    enough to fetch weights programmatically).

    Reads the URL straight off `ModelVersion.source` (see quirk 6 in the
    module docstring) -- no extra API call beyond what `latest_version`
    already makes.

    Parameters
    ----------
    name    : str -- registered model name
    version : int, optional -- defaults to `latest_version(name)`; an older
        version's URL is derived by substituting the trailing
        `/versions/<n>` segment (`get_model_version` 404s on this registry,
        quirk 6), so this assumes every version shares the same
        `.../-/ml/models/<model_id>` prefix, true by construction (one
        registered model, many versions).

    Returns
    -------
    (model_url, version_url) : (str, str)
    """
    client = MlflowClient()
    versions = client.get_registered_model(name).latest_versions
    if not versions:
        raise ValueError(f"no registered model versions found for {name!r}")
    latest = versions[0]
    model_url = latest.source.rsplit("/versions/", 1)[0]
    version = version if version is not None else latest.version
    return model_url, f"{model_url}/versions/{version}"

def register_model(lit_model, name, params=None, extra=None, run_name=None, experiment_name=DEFAULT_EXPERIMENT):
    """Register a trained `LitResidualRegressor` as a new version of the
    model registry entry `name` (creating both the experiment and the
    registered model on first use).

    Parameters
    ----------
    lit_model : dispcraft.ml.LitResidualRegressor -- already trained (eval mode)
    name      : str -- registered model name, e.g. "rgs000_0-hybrid-mlp-y"
    params    : dict, optional -- hyperparameters to log alongside the run
        (e.g. the dict passed to `dispcraft.ml.train_residual_mlp`)
    extra     : dict, optional -- arbitrary picklable data needed to use the
        model that isn't part of its own state (e.g. a fitted
        `StandardScaler`'s `mean_`/`scale_` for the RGS/BGS hybrid MLPs,
        which standardize `(y, z, wavelength)` before prediction). Returned
        back by `load_model`. Not needed for the zeroth-order NNs, whose
        standardization is a fixed formula (`standardize_field`), not fitted.
    run_name  : str, optional -- MLflow run name, defaults to `name`
    experiment_name : str, default `DEFAULT_EXPERIMENT`

    Returns
    -------
    str -- the new model version number
    """
    mlflow.set_experiment(experiment_name)
    with mlflow.start_run(run_name=run_name or name) as run:
        if params:
            mlflow.log_params(params)
        with tempfile.TemporaryDirectory() as tmp:
            ckpt_path = Path(tmp) / _CHECKPOINT_NAME
            torch.save({"state_dict": lit_model.state_dict(), "hparams": dict(lit_model.hparams),
                        "extra": extra or {}}, ckpt_path)
            mlflow.log_artifact(str(ckpt_path), artifact_path="model")
        run_id = run.info.run_id
        artifact_uri = mlflow.get_artifact_uri("model")

    client = MlflowClient()
    try:
        client.get_registered_model(name)
    except mlflow.exceptions.MlflowException:
        client.create_registered_model(name)
    mv = client.create_model_version(name=name, source=artifact_uri, run_id=run_id)
    return mv.version