STAC¶

def create_base_models_collection() -> pystac.Collection:
    """base-models: model blueprints contributed via PR."""
    return pystac.Collection(
        id=BASE_MODELS_COLLECTION,
        description="Model blueprints contributed via PR. Each item is a complete model card.",
        extent=pystac.Extent(
            spatial=pystac.SpatialExtent(bboxes=[[-180, -90, 180, 90]]),
            temporal=pystac.TemporalExtent(intervals=[[datetime(2026, 1, 1, tzinfo=UTC), None]]),
        ),
        license="various",
        stac_extensions=BASE_MODEL_EXTENSIONS,
    )

def create_local_models_collection() -> pystac.Collection:
    """local-models: finetuned models, only promoted versions."""
    return pystac.Collection(
        id=LOCAL_MODELS_COLLECTION,
        description="Finetuned models produced by ZenML pipelines. Only promoted versions appear here.",
        extent=pystac.Extent(
            spatial=pystac.SpatialExtent(bboxes=[[-180, -90, 180, 90]]),
            temporal=pystac.TemporalExtent(intervals=[[datetime(2026, 1, 1, tzinfo=UTC), None]]),
        ),
        license="various",
        stac_extensions=LOCAL_MODEL_EXTENSIONS,
    )

def create_datasets_collection() -> pystac.Collection:
    """datasets: training data registered via fAIr UI/backend."""
    return pystac.Collection(
        id=DATASETS_COLLECTION,
        description="Training data registered via fAIr UI/backend.",
        extent=pystac.Extent(
            spatial=pystac.SpatialExtent(bboxes=[[-180, -90, 180, 90]]),
            temporal=pystac.TemporalExtent(intervals=[[datetime(2026, 1, 1, tzinfo=UTC), None]]),
        ),
        license="various",
        stac_extensions=DATASET_EXTENSIONS,
    )

def initialize_catalog(catalog_path: str) -> pystac.Catalog:
    """Create catalog.json + 3 empty collections. Saves to disk.

    returns existing catalog if already present.
    """
    if os.path.exists(catalog_path):
        return pystac.Catalog.from_file(catalog_path)

    catalog = pystac.Catalog(
        id="fair-models",
        description="fAIr model registry and dataset catalog",
    )

    catalog.add_child(create_base_models_collection())
    catalog.add_child(create_local_models_collection())
    catalog.add_child(create_datasets_collection())

    catalog_dir = os.path.dirname(catalog_path) or "."
    catalog.normalize_hrefs(catalog_dir)
    catalog.save(catalog_type=CatalogType.SELF_CONTAINED)

    return catalog

def prediction_output(properties: dict) -> dict | None:
    """The `mlm:output` entry describing post-processed prediction features through MLM `variables`."""
    return next((output for output in properties.get("mlm:output", []) if output.get("variables")), None)

def is_trainable(assets: Mapping[str, Any]) -> bool:
    """A base model without a training runtime serves predictions only and cannot be finetuned."""
    return TRAINING_ASSET in assets

def batch_runtime(assets: Mapping[str, Any]) -> Any | None:
    """The full toolchain image that runs batch predictions; the distroless mlm:inference image cannot."""
    return assets.get(TRAINING_ASSET) or assets.get(BATCH_INFERENCE_ASSET)

def validate_inference_params(params: dict[str, Any], model_item: pystac.Item) -> None:
    """Raise ValueError unless every key is a declared, spec'd inference parameter with a valid value."""
    props = model_item.properties
    declared = inference_params(props.get("mlm:hyperparameters", {}))
    spec_keys = {entry["key"] for entry in props.get("fair:hyperparameters_spec", [])}
    unknown = sorted(key for key in params if key not in declared or key not in spec_keys)
    if unknown:
        raise ValueError(f"Unknown inference parameter(s) for '{model_item.id}': {', '.join(unknown)}")
    if errors := validate_hyperparameters(params, model_item):
        raise ValueError("; ".join(errors))

def normalize_version_link_hrefs(
    item: pystac.Item,
    item_href_factory,
    default_collection_id: str,
) -> None:
    """Rewrite version-link relative hrefs (`../../coll/id/id.json`) to the backend's `item_href` form.

    pystac's local-catalog write produces these relative hrefs. Use raw `link.target` so we
    don't trigger pystac's root resolution (which makes HTTP calls).
    """
    for link in item.links:
        if link.rel not in _VERSION_RELS:
            continue
        target = link.target if isinstance(link.target, str) else ""
        if target.startswith(("http://", "https://")):
            continue
        parts = target.replace("\\", "/").split("/")
        json_parts = [p for p in parts if p.endswith(".json")]
        if not json_parts:
            continue
        target_item_id = json_parts[-1].removesuffix(".json")
        target_coll = parts[-3] if len(parts) >= 3 else default_collection_id
        link.target = item_href_factory(target_coll, target_item_id)

def relativize_catalog_links(item: pystac.Item, catalog_url: str) -> None:
    """Store links into this catalog relative to its root; the STAC API resolves them against each request's host."""
    prefix = f"{catalog_url.rstrip('/')}/"
    for link in item.links:
        # The catalog root itself has no relative form; the API regenerates root links on read.
        if isinstance(link.target, str) and link.target.startswith(prefix) and link.target != prefix:
            link.target = link.target.removeprefix(prefix)

def prepare_api_links(item: pystac.Item, catalog_url: str) -> None:
    """Shape links for a STAC API: catalog links relative to its root, no self, no latest-version on a current item.

    The API infers self per request, and pystac turns a relative self href into a local file path.
    """
    relativize_catalog_links(item, catalog_url)
    item.links = [lnk for lnk in item.links if lnk.rel != "self" and not _is_current_latest_link(item, lnk)]

def delete_item_with_archives(backend: StacBackend, collection_id: str, item_id: str) -> list[str]:
    """Delete the item and its deprecated `{id}-v{n}` archives; return the deleted ids."""
    version = int(backend.get_item(collection_id, item_id).properties.get("version", "1"))
    # A dataset archive carries the item's own version, since a rebuild under the same title gets a new id.
    candidates = [f"{item_id}-v{n}" for n in range(1, version + 1)]
    deleted = [
        archived_id
        for archived_id in candidates
        if backend.item_exists(collection_id, archived_id)
        and backend.get_item(collection_id, archived_id).properties.get("deprecated")
    ]
    for archived_id in deleted:
        backend.delete_item(collection_id, archived_id)
    backend.delete_item(collection_id, item_id)
    log.info("Deleted %s/%s and %d archived versions", collection_id, item_id, len(deleted))
    return [*deleted, item_id]

def archive_previous_version(
    backend: StacBackend,
    collection_id: str,
    old_item: pystac.Item,
    new_item_href: str,
) -> pystac.Item:
    """Copy the old item to a version-suffixed ID and deprecate it.

    Preserves the old version as a discoverable, deprecated STAC item
    per the STAC Version Extension v1.2.0 spec.
    """
    old_version = old_item.properties.get("version", "1")
    archived_id = f"{old_item.id}-v{old_version}"

    archived = old_item.clone()
    archived.id = archived_id
    archived.properties["deprecated"] = True
    archived.links = [lnk for lnk in archived.links if lnk.rel not in {"latest-version", "self"}]
    archived.add_link(pystac.Link(rel="successor-version", target=new_item_href))

    archived_href = backend.item_href(collection_id, archived_id)
    archived.add_link(pystac.Link(rel="self", target=archived_href, media_type="application/geo+json"))

    published = backend.publish_item(collection_id, archived)
    log.info("Archived %s/%s as deprecated", collection_id, archived_id)
    return published