ZenML Integration¶

def generate_inference_config(
    model_item: pystac.Item,
    input_images: Mapping[str, str],
    overrides: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """`input_images` maps each of the model's input names to its chips path."""
    # Works for both base-model and local-model items (same MLM structure)
    props = model_item.properties

    model_asset = model_item.assets.get("model")
    if model_asset is None:
        msg = f"Model item '{model_item.id}' missing 'model' (ONNX) asset"
        raise KeyError(msg)

    params = inference_params(props.get("mlm:hyperparameters", {}))
    if overrides:
        validate_inference_params(overrides, model_item)
        params.update(overrides)

    names = require_imagery_for_inputs(model_item, input_images)
    chips = {name: http_url_to_s3_uri(input_images[name]) for name in names}
    parameters: dict[str, Any] = {
        "model_uri": http_url_to_s3_uri(model_asset.href),
        "input_images": chips,
        "inference_params": params,
    }

    config: dict[str, Any] = {
        "parameters": parameters,
        "tags": [
            f"model:{model_item.id}",
        ],
    }

    runtime = batch_runtime(model_item.assets)
    if runtime is None:
        msg = (
            f"Model item '{model_item.id}' has neither '{TRAINING_ASSET}' nor '{BATCH_INFERENCE_ASSET}'; "
            "batch inference needs a full toolchain image"
        )
        raise KeyError(msg)
    if runtime.media_type == OCI_IMAGE_INDEX_TYPE:
        docker_cfg: dict[str, Any] = {
            "parent_image": _normalize_container_href(runtime.href),
            "skip_build": True,
        }
        config["settings"] = {"docker": docker_cfg}

    k8s = _scheduling_settings(model_item, "inference")
    if k8s:
        config.setdefault("steps", {}).setdefault("run_inference", {}).setdefault("settings", {}).update(k8s)

    return config

@step
def load_model(
    model_uri: str,
    zenml_artifact_version_id: str = "",
) -> Any:
    """Resolve model from ZenML artifact store. Framework-agnostic via materializer."""
    client = Client()
    if zenml_artifact_version_id:
        art = client.get_artifact_version(zenml_artifact_version_id)
    else:
        results = client.list_artifact_versions(uri=model_uri)
        if not results:
            msg = f"No artifact found for URI: {model_uri}"
            raise RuntimeError(msg)
        art = results[0]
    return art.load()

@contextmanager
def mlflow_training_context(
    hyperparameters: dict[str, Any],
    model_name: str | None = None,
    base_model_id: str | None = None,
    dataset_id: str | None = None,
):
    """Context manager that instruments a training step with MLflow.

    Handles autolog, param logging, tag setting, and wall-clock timing.
    Contributors use this instead of manual MLflow calls.
    """
    import mlflow

    mlflow.autolog()  # ty: ignore[possibly-missing-attribute]
    mlflow.log_params(  # ty: ignore[possibly-missing-attribute]
        {k: v for k, v in hyperparameters.items() if not isinstance(v, (dict, list))}
    )

    tags: dict[str, str] = {}
    if model_name:
        tags["fair.model_name"] = model_name
    if base_model_id:
        tags["fair.base_model"] = base_model_id
    if dataset_id:
        tags["fair.dataset"] = dataset_id
    if tags:
        mlflow.set_tags(tags)  # ty: ignore[possibly-missing-attribute]

    run_meta: dict[str, Any] = {f"fair/{k}": v for k, v in hyperparameters.items() if not isinstance(v, (dict, list))}
    if model_name:
        run_meta["fair/model_name"] = model_name
    if base_model_id:
        run_meta["fair/base_model_id"] = base_model_id
    if dataset_id:
        run_meta["fair/dataset_id"] = dataset_id
    if run_meta:
        log_run_metadata(run_meta)

    wall_start = time.perf_counter()
    yield
    wall_seconds = time.perf_counter() - wall_start
    log_training_wall_time(wall_seconds)

def log_evaluation_results(metrics: dict[str, Any]) -> None:
    """Log evaluation metrics to both MLflow and ZenML fair-prefixed metadata."""
    import mlflow

    scalar_metrics = {k: v for k, v in metrics.items() if isinstance(v, (int, float))}
    mlflow.log_metrics(scalar_metrics)  # ty: ignore[possibly-missing-attribute]
    log_fair_metrics(metrics)