from __future__ import annotations
import logging
import warnings
from typing import Any
from semantic_version import Version
# isort: split
from bailo.core.client import Client
from bailo.core.enums import CollaboratorEntry, EntryKind, ModelVisibility
from bailo.core.exceptions import BailoException
from bailo.helper.entry import Entry
from bailo.helper.release import Release
logger = logging.getLogger(__name__)
[docs]
class MirroredModel(Entry):
"""Represent a mirrored model within Bailo.
:param client: A client object used to interact with Bailo
:param model_id: A unique ID for the mirrored model
:param name: Name of mirrored model
:param description: Description of mirrored model
:param sourceModelId: Used for linking a mirrored model to its source model
:param organisation: Organisation responsible for the mirrored model, defaults to None
:param state: Development readiness of the mirrored model, defaults to None
:param tags: Tags to assign to the mirrored model, defaults to None
:param collaborators: List of CollaboratorEntry to define who the mirrored model's collaborators (a.k.a. mirrored model access) are, defaults to None
:param visibility: Visibility of the mirrored model, using ModelVisibility enum (e.g Public or Private), defaults to None
"""
def __init__(
self,
client: Client,
model_id: str,
name: str,
description: str,
sourceModelId: str,
organisation: str | None = None,
state: str | None = None,
tags: list[str] | None = None,
collaborators: list[CollaboratorEntry] | None = None,
visibility: ModelVisibility | None = None,
) -> None:
super().__init__(
client=client,
id=model_id,
name=name,
description=description,
kind=EntryKind.MIRRORED_MODEL,
visibility=visibility,
organisation=organisation,
state=state,
tags=tags,
collaborators=collaborators,
)
self.sourceModelId = sourceModelId
self._original_source_model_id = sourceModelId
self.model_id = model_id
self._mirrored_card = None
self._mirrored_card_version = None
[docs]
@classmethod
def create(
cls,
client: Client,
name: str,
description: str,
sourceModelId: str,
organisation: str | None = None,
state: str | None = None,
tags: list[str] | None = None,
collaborators: list[CollaboratorEntry] | None = None,
visibility: ModelVisibility | None = None,
) -> MirroredModel:
"""Build a mirrored model from Bailo and upload it.
:param client: A client object used to interact with Bailo
:param name: Name of mirrored model
:param description: Description of mirrored model
:param sourceModelId: Used for linking a mirrored model to its source model
:param organisation: Organisation responsible for the mirrored model, defaults to None
:param state: Development readiness of the mirrored model, defaults to None
:param tags: Tags to assign to the mirrored model, defaults to None
:param collaborators: List of CollaboratorEntry to define who the mirrored model's collaborators (a.k.a. model access) are, defaults to None
:param visibility: Visibility of the mirrored model, using ModelVisibility enum (e.g Public or Private), defaults to None
:return: MirroredModel object
"""
res = client.post_model(
name=name,
kind=EntryKind.MIRRORED_MODEL,
description=description,
sourceModelId=sourceModelId,
visibility=visibility,
organisation=organisation,
state=state,
tags=tags,
collaborators=collaborators,
)
model_id = res["model"]["id"]
logger.info("Model successfully created on server with ID %s.", model_id)
model = cls(
client=client,
model_id=model_id,
name=name,
description=description,
sourceModelId=sourceModelId,
visibility=visibility,
organisation=organisation,
state=state,
tags=tags,
collaborators=collaborators,
)
model._unpack(res["model"])
return model
[docs]
def update(self) -> None:
"""Upload and retrieve any changes to the mirrored model summary on Bailo.
Merges the sourceModelId into settings when it has changed, then delegates to the base update.
"""
if self.sourceModelId != self._original_source_model_id:
if self.settings is None:
self.settings = {}
self.settings.setdefault("mirror", {})["sourceModelId"] = self.sourceModelId
super().update()
self._original_source_model_id = self.sourceModelId
[docs]
def _unpack(self, res):
"""Update mirrored model attributes from API response.
:param res: Response dictionary containing model information.
"""
super()._unpack(res)
source_id = res.get("settings", {}).get("mirror", {}).get("sourceModelId", self.sourceModelId)
self.sourceModelId = source_id
self._original_source_model_id = source_id
[docs]
@classmethod
def from_id(cls, client: Client, model_id: str) -> MirroredModel:
"""Return an existing mirrored model from Bailo.
:param client: A client object used to interact with Bailo
:param model_id: A unique mirrored model ID
:return: A mirrored model object
"""
res = client.get_model(model_id=model_id)["model"]
if res["kind"] != EntryKind.MIRRORED_MODEL:
raise BailoException(
f"ID {model_id} does not belong to a mirrored model. Did you mean to use MirroredModel.from_id()?"
)
logger.info("Model %s successfully retrieved from server.", model_id)
model = cls(
client=client,
model_id=model_id,
name=res["name"],
description=res["description"],
sourceModelId=res["settings"]["mirror"]["sourceModelId"],
collaborators=res["collaborators"],
organisation=res.get("organisation"),
state=res.get("state"),
tags=res.get("tags"),
)
model._unpack(res)
model.get_card_latest()
return model
[docs]
@classmethod
def search(
cls,
client: Client,
task: str | None = None,
libraries: list[str] | None = None,
filters: list[str] | None = None,
search: str = "",
organisations: list[str] | None = None,
states: list[str] | None = None,
allow_templating: bool | None = None,
schema_id: str | None = None,
admin_access: bool | None = None,
peers: list[str] | None = None,
title_only: bool | None = None,
) -> list[MirroredModel]:
"""Return a list of mirrored model objects from Bailo, based on search parameters.
:param client: A client object used to interact with Bailo
:param task: Mirrored model task (e.g. image classification), defaults to None
:param libraries: Mirrored model library (e.g. TensorFlow), defaults to None
:param filters: List of collaborator role filters. Special value `"mine"` restricts results to
models where the current user is a collaborator. Otherwise, values are treated as collaborator
roles, defaults to None
:param search: Free-text search string. Always performs a partial, case-insensitive match against
the mirrored model name. If `title_only` is False, a full-text search across mirrored model
content is also performed, defaults to ""
:param organisations: List of organisation identifiers to restrict results, defaults to None
:param states: List of mirrored model lifecycle states to restrict results, defaults to None
:param allow_templating: If True, restricts results to models with templating enabled, defaults to None
:param schema_id: Schema ID to restrict results to models using that schema, defaults to None
:param admin_access: If True, returns models requiring admin access. The caller must
have the Admin role or the request will be rejected by the backend, defaults to None
:param peers: List of peer identifiers to include remote search results from, defaults to None
:param title_only: If True, limits searching to mirrored model titles only and disables
full-text search, defaults to None
:return: List of mirrored model objects
"""
res = client.get_models(
task=task,
libraries=libraries,
filters=filters,
search=search,
kind=EntryKind.MIRRORED_MODEL,
organisations=organisations,
states=states,
allow_templating=allow_templating,
schema_id=schema_id,
admin_access=admin_access,
peers=peers,
title_only=title_only,
)
models = []
for model in res["models"]:
res_model = client.get_model(model_id=model["id"])["model"]
model_obj = cls(
client=client,
model_id=model["id"],
name=model["name"],
description=model["description"],
sourceModelId=res_model["settings"]["mirror"]["sourceModelId"],
collaborators=model["collaborators"],
organisation=model.get("organisation"),
state=model.get("state"),
tags=model.get("tags"),
)
model_obj._unpack(res_model)
model_obj.get_card_latest()
models.append(model_obj)
return models
[docs]
def get_releases(self) -> list[Release]:
"""Get all releases for the mirrored model.
:return: List of Release objects
"""
res = self.client.get_all_releases(model_id=self.model_id)
releases = []
for release in res["releases"]:
releases.append(self.get_release(version=release["semver"]))
logger.info("Successfully retrieved all releases for mirrored model %s.", self.model_id)
return releases
[docs]
def get_release(self, version: Version | str) -> Release:
"""Call the Release.from_version method to return an existing release from Bailo.
:param version: A semantic version for the release
:return: Release object
"""
return Release.from_version(self.client, self.model_id, version)
[docs]
def get_latest_release(self):
"""Get the latest release for the mirrored model from Bailo.
:return: Release object
"""
releases = self.get_releases()
if not releases:
raise BailoException("This mirrored model has no releases.")
latest_release = max(releases)
logger.info(
"latest_release (%s) for %s retrieved successfully.",
str(latest_release.version),
self.model_id,
)
return max(releases)
[docs]
def get_images(self):
"""Get all model image references for the mirrored model.
:return: List of images
"""
res = self.client.get_all_images(model_id=self.model_id)
logger.info("Images for %s retrieved successfully.", self.model_id)
return res["images"]
[docs]
def get_image(self):
"""Get a model image reference.
:raises NotImplementedError: Not implemented error.
"""
raise NotImplementedError
[docs]
def update_model_card(self, model_card: dict[str, Any] | None = None) -> None:
"""Upload and retrieve any changes to the editable mirrored model card on Bailo.
:param model_card: Model card dictionary, defaults to None
.. note:: If a model card is not provided, the current model card attribute value is used
"""
self._update_card(card=model_card)
[docs]
def get_card_latest(self) -> None:
"""Get the latest card from Bailo."""
res = self.client.get_model(model_id=self.id)
if "card" in res["model"]:
self._unpack_card(res["model"]["card"])
logger.info("Latest card for ID %s successfully retrieved.", self.id)
else:
warnings.warn(f"ID {self.id} does not have any associated model card.")
if "mirroredCard" in res["model"]:
self._unpack_card(res["model"]["mirroredCard"], True)
else:
warnings.warn(f"ID {self.id} does not have any associated additional information.")
def _unpack_card(self, res, mirrored=False) -> None:
if mirrored:
super()._unpack_card(res)
else:
self._mirrored_card_version = res["version"]
try:
self._mirrored_card = res["metadata"]
except KeyError:
self._mirrored_card = None
@property
def model_card(self):
"""Get the data of the model card.
:return: Model card data.
"""
return {"card": self._card, "additional_information": self._mirrored_card}
@property
def model_card_version(self):
"""Get the version of the mirrored model card.
:return: Model card version.
"""
return {
"card": self._card_version,
"additional_information": self._mirrored_card_version,
}
@property
def model_card_schema(self):
"""Get the schema of the mirrored model card.
:return: Model card schema.
"""
return self._card_schema
def __repr__(self) -> str:
"""Return a developer-oriented string representation of the mirrored model.
:return: String representation with class and ID
"""
return f"{self.__class__.__name__}({str(self)})"
def __str__(self) -> str:
"""Return the human-readable string representation of the mirrored model.
:return: String representation of the mirrored model.
"""
return f"{self.model_id}"