from __future__ import annotations
import logging
import os
import shutil
import tempfile
import warnings
from typing import Any, ClassVar
from semantic_version import Version
from bailo.core.client import Client
from bailo.core.entry import Entry
from bailo.core.enums import CollaboratorEntry, EntryKind, MinimalSchema, ModelVisibility
from bailo.core.exceptions import BailoException
from bailo.core.utils import NestedDict
from bailo.helper.entry import ReleaseMixin
from bailo.helper.release import Release
try:
import mlflow
ML_FLOW = True
except ImportError:
ML_FLOW = False
logger = logging.getLogger(__name__)
[docs]
class Model(ReleaseMixin, Entry):
"""Represent a model within Bailo.
:param client: A client object used to interact with Bailo
:param model_id: A unique ID for the model
:param name: Name of model
:param description: Description of model
:param organisation: Organisation responsible for the model, defaults to None
:param state: Development readiness of the model, defaults to None
:param tags: Tags to assign to the model, defaults to None
:param collaborators: List of CollaboratorEntry to define who the model's collaborators (a.k.a. model access) are, defaults to None
:param visibility: Visibility of model, using ModelVisibility enum (e.g Public or Private), defaults to None
"""
entry_kind: ClassVar[EntryKind] = EntryKind.MODEL
_id_alias: ClassVar[str] = "model_id"
def __init__(
self,
client: Client,
model_id: str,
name: str,
description: 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.MODEL,
visibility=visibility,
organisation=organisation,
state=state,
tags=tags,
collaborators=collaborators,
)
[docs]
@classmethod
def from_id(cls, client: Client, model_id: str) -> Model:
"""Return an existing model from Bailo.
:param client: A client object used to interact with Bailo
:param model_id: A unique model ID
:return: A model object
"""
return cls._from_id(client, model_id)
[docs]
@classmethod
def from_mlflow(
cls,
client: Client,
mlflow_uri: str,
name: str,
schema_id: str = MinimalSchema.MODEL,
version: str | None = None,
files: bool = True,
visibility: ModelVisibility | None = None,
organisation: str | None = None,
state: str | None = None,
tags: list[str] | None = None,
collaborators: list[CollaboratorEntry] | None = None,
) -> Model:
"""Import an MLFlow Model into Bailo.
:param client: A client object used to interact with Bailo
:param mlflow_uri: MLFlow server URI
:param name: Name of model (on MLFlow). Same name will be used on Bailo
:param schema_id: A unique schema ID, only required when files is True, defaults to minimal-general-v10
:param version: Specific MLFlow model version to import, defaults to None
:param files: Import files?, defaults to True
:param visibility: Visibility of model on Bailo, using ModelVisibility enum (e.g Public or Private), defaults to None
:param organisation: Organisation responsible for the model, defaults to None
:param state: Development readiness of the model, defaults to None
:param tags: Tags to assign to the model, defaults to None
:param collaborators: List of CollaboratorEntry to define who the model's collaborators (a.k.a. model access) are, defaults to None
:return: A model object
"""
if not ML_FLOW:
raise ImportError("Optional MLFlow dependencies (needed for this method) are not installed.")
mlflow_client = mlflow.tracking.MlflowClient(tracking_uri=mlflow_uri) # type: ignore[reportPrivateImportUsage]
mlflow.set_tracking_uri(mlflow_uri)
filter_string = f"name = '{name}'"
res = mlflow_client.search_model_versions(filter_string=filter_string, order_by=["version_number DESC"])
if not res:
raise BailoException("No MLFlow models found. Are you sure the name/alias/version provided is correct?")
sel_model = None
if version:
for model in res:
if model.version == version:
sel_model = model
else:
sel_model = res[0]
if sel_model is None:
raise BailoException("No MLFlow model found. Are you sure the name/alias/version provided is correct?")
name = sel_model.name
description = str(sel_model.description) + " Imported from MLFlow."
model = cls.create(
client=client,
name=name,
description=description,
visibility=visibility,
organisation=organisation,
state=state,
tags=tags,
collaborators=collaborators,
)
logger.info("MLFlow model successfully imported to Bailo with ID %s", model.id)
if files:
model.card_from_schema(schema_id=schema_id)
release = model.create_release(version=Version.coerce(str(sel_model.version)), notes=" ")
run_id = sel_model.run_id
if run_id is None:
raise BailoException(
"MLFlow model does not have an associated run_id, therefore artifacts cannot be transferred."
)
mlflow_run = mlflow_client.get_run(run_id)
artifact_uri: str = str(mlflow_run.info.artifact_uri)
if artifact_uri is None:
raise BailoException("Artifact URI could not be found, therefore artifacts cannot be transferred.")
if mlflow.artifacts.list_artifacts(artifact_uri=artifact_uri) is not None: # type: ignore[reportPrivateImportUsage]
temp_dir = os.path.join(tempfile.gettempdir(), "mlflow_model")
mlflow_dir = os.path.join(temp_dir, f"mlflow_{run_id}")
mlflow.artifacts.download_artifacts(artifact_uri=artifact_uri, dst_path=mlflow_dir) # type: ignore[reportPrivateImportUsage]
release.upload(mlflow_dir)
return model
[docs]
def update_model_card(self, model_card: dict[str, Any] | None = None) -> None:
"""Upload and retrieve any changes to the 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 create_experiment(
self,
) -> Experiment:
"""Create an experiment locally
:return: An experiment object
"""
return Experiment.create(model=self)
[docs]
def create_release(
self,
version: Version | str,
notes: str,
files: list[str] | None = None,
images: list[str] | None = None,
minor: bool = False,
draft: bool = True,
) -> Release:
"""Call the Release.create method to build a release from Bailo and upload it.
:param version: A semantic version for the release
:param notes: Notes on release
:param files: A list of files for release, defaults to None
:param images: A list of images for release, defaults to None
:param minor: Is a minor release?, defaults to False
:param draft: Is a draft release?, defaults to True
:return: Release object
"""
if self.model_card_schema:
return Release.create(
client=self.client,
model_id=self.model_id,
version=version,
notes=notes,
model_card_version=self.model_card_version,
files=files,
images=images,
minor=minor,
draft=draft,
)
raise BailoException("Create a model card before creating a release")
@property
def model_card(self) -> dict[str, Any] | None:
"""Get the data of the model card.
:return: Model card data.
"""
return self.card
@model_card.setter
def model_card(self, value: dict[str, Any] | None) -> None:
"""Set the data of the model card.
:param value: The data to set.
"""
self.card = value
@property
def model_card_version(self) -> int | None:
"""Get the version of the model card.
:return: Model card version.
"""
return self.card_version
@model_card_version.setter
def model_card_version(self, value: int | None) -> None:
"""Set the version of the model card.
:param value: The version to set.
"""
self.card_version = value
@property
def model_card_schema(self) -> str | None:
"""Get the schema of the model card.
:return: Model card schema.
"""
return self.card_schema
@model_card_schema.setter
def model_card_schema(self, value: str | None) -> None:
"""Set the schema of the model card.
:param value: The schema to set.
"""
self.card_schema = value
[docs]
class Experiment:
"""Represent an experiment locally.
:param model: A Bailo model object which the experiment is being run on
:param raw: Raw information about the experiment runs
.. code-block:: python
experiment = model.create_experiment()
for x in range(5):
experiment.start_run()
experiment.log_params({"lr": 0.01})
### INSERT MODEL TRAINING HERE ###
experiment.log_metrics("accuracy": 0.86)
experiment.log_artifacts(["weights.pth"])
experiment.publish(mc_loc="performance.performanceMetrics", run_id=1)
"""
def __init__(
self,
model: Model,
):
self.model = model
self.raw = []
self.run = -1
self.temp_dir = os.path.join(tempfile.gettempdir(), "bailo_runs")
self.published = False
self.run_data = {}
[docs]
@classmethod
def create(
cls,
model: Model,
) -> Experiment:
"""Create an experiment locally.
:param model: A Bailo model object which the experiment is being run on
:return: Experiment object
"""
return cls(model=model)
[docs]
def start_run(self, is_mlflow: bool = False):
"""Starts a new experiment run.
:param is_mlflow: Marks a run as MLFlow
"""
self.run += 1
self.run_data = {
"run": self.run,
"params": [],
"metrics": [],
"artifacts": [],
"dataset": "",
}
self.raw.append(self.run_data)
if not is_mlflow:
logger.info("Bailo tracking run %s.", self.run)
[docs]
def log_params(self, params: dict[str, Any]):
"""Logs parameters to the current run.
:param params: Dictionary of parameters to be logged
"""
self.run_data["params"].append(params)
[docs]
def log_metrics(self, metrics: dict[str, Any]):
"""Logs metrics to the current run.
:param metrics: Dictionary of metrics to be logged
"""
self.run_data["metrics"].append(metrics)
[docs]
def log_artifacts(self, artifacts: list):
"""Logs artifacts to the current run.
:param artifacts: A list of artifact paths to be logged
"""
self.run_data["artifacts"].extend(artifacts)
[docs]
def log_dataset(self, dataset: str):
"""Logs a dataset to the current run.
:param dataset: Arbitrary title of dataset
"""
self.run_data["dataset"] = dataset
[docs]
def from_mlflow(self, tracking_uri: str, experiment_id: str):
"""Imports information from an MLFlow Tracking experiment.
:param tracking_uri: MLFlow Tracking server URI
:param experiment_id: MLFlow Tracking experiment ID
:raises ImportError: Import error if MLFlow not installed
"""
if not ML_FLOW:
raise ImportError("Optional MLFlow dependencies (needed for this method) are not installed.")
client = mlflow.tracking.MlflowClient(tracking_uri=tracking_uri) # type: ignore[reportPrivateImportUsage]
runs = client.search_runs([experiment_id])
if len(runs):
logger.info(
"Successfully retrieved MLFlow experiment %s from tracking server. %d were found.",
experiment_id,
len(runs),
)
else:
warnings.warn(
f"MLFlow experiment {experiment_id} does not have any runs and publishing requires at least one valid run. Are you sure the ID is correct?",
stacklevel=2,
)
for run in runs:
data = run.data
info = run.info
inputs = run.inputs
artifact_uri: str = str(info.artifact_uri)
run_id = info.run_id
status = info.status
datasets = inputs.dataset_inputs
datasets_str = [dataset.dataset.name for dataset in datasets]
artifacts = []
# MLFlow run must be status FINISHED
if status != "FINISHED":
continue
if mlflow.artifacts.list_artifacts(artifact_uri=artifact_uri) is not None: # type: ignore[reportPrivateImportUsage]
mlflow_dir = os.path.join(self.temp_dir, f"mlflow_{run_id}")
mlflow.artifacts.download_artifacts(artifact_uri=artifact_uri, dst_path=mlflow_dir) # type: ignore[reportPrivateImportUsage]
artifacts.append(mlflow_dir)
logger.info(
"Successfully downloaded artifacts for MLFlow experiment %s to %s.",
experiment_id,
mlflow_dir,
)
self.start_run(is_mlflow=True)
self.log_params(data.params)
self.log_metrics(data.metrics)
self.log_artifacts(artifacts)
self.log_dataset("".join(datasets_str))
self.run_data["run"] = info.run_id
logger.info("Successfully imported MLFlow experiment %s.", experiment_id)
[docs]
def publish( # noqa: PLR0912
self,
mc_loc: str,
semver: str = "0.1.0",
notes: str = "",
run_id: str | None = None,
select_by: str | None = None,
):
"""Publishes a given experiments results to the model card.
:param mc_loc: Location of metrics in the model card (e.g. performance.performanceMetrics)
:param semver: Semantic version of release to create (if artifacts present), defaults to 0.1.0 or next
:param notes: Notes for release, defaults to ""
:param run_id: Local experiment run ID to be selected, defaults to None
:param select_by: String describing experiment to be selected (e.g. "accuracy MIN|MAX"), defaults to None
.. note:: mc_loc is dependent on the model card schema being used
.. warning:: User must specify either run_id or select_by, otherwise the code will error
"""
# Check if already published, can only published once
if self.published:
raise BailoException("This experiment has already been published.")
mc = self.model.model_card
if mc is None:
raise BailoException("Model card needs to be populated before publishing an experiment.")
mc = NestedDict(mc)
if len(self.raw) == 0:
raise BailoException("This experiment has no runs to publish.")
if (select_by is None) and (run_id is None):
raise BailoException(
"Either select_by (e.g. 'accuracy MIN|MAX') or run_id is required to publish an experiment run."
)
sel_run: dict[Any, Any]
if (select_by is not None) and (run_id is None):
sel_run = self.__select_run(select_by=select_by)
if run_id is not None:
for run in self.raw:
if run["run"] == run_id:
sel_run = run
break
else:
raise NameError(f"Run {run_id} does not exist.")
values = []
for metric in sel_run["metrics"]:
for k, v in metric.items():
values.append({"name": k, "value": v})
# Updating the model card
parsed_values = [{"dataset": sel_run["dataset"], "datasetMetrics": values}]
mc[tuple(mc_loc.split("."))] = parsed_values
self.model.update_model_card(model_card=mc)
# Creating a release and uploading artifacts (if artifacts present)
artifacts = sel_run["artifacts"]
if len(artifacts):
# Create new release
try:
release_latest_version = self.model.get_latest_release().version
release_new_version = release_latest_version.next_minor()
except Exception:
release_new_version = semver
run_id = sel_run["run"]
notes = f"{notes} (Run ID: {run_id})"
release_new = self.model.create_release(version=release_new_version, minor=True, notes=notes)
logger.info(
"Uploading %d artifacts to version %s of model %s.",
len(artifacts),
str(release_new_version),
self.model.model_id,
)
for artifact in artifacts:
release_new.upload(path=artifact)
self.published = True
if os.path.exists(self.temp_dir) and os.path.isdir(self.temp_dir):
shutil.rmtree(self.temp_dir)
logger.info(
"Successfully published experiment run %s to model %s.",
str(run_id),
self.model.model_id,
)
def __select_run(self, select_by: str) -> dict:
# Parse target and order from select_by string
select_by_split = select_by.split(" ")
if len(select_by_split) != 2: # noqa: PLR2004
raise BailoException("Invalid select_by string. Expected format is 'metric_name MIN|MAX'.")
order_str = select_by_split[1].upper()
order_opt = {"MIN": 0, "MAX": -1}
if order_str not in order_opt:
raise BailoException(f"Runs can only be ordered by MIN or MAX, not {order_str}.")
target_str = select_by_split[0]
# Extract target value for each run
runs = self.raw
for run in runs:
metrics = run["metrics"]
for metric in metrics:
target_value = metric.get(target_str, None)
if target_value is not None:
run["target"] = target_value
else:
raise BailoException(
f"Target '{target_str}' could not be found in at least one experiment run, or is not an integer. Therefore ordering cannot take place."
)
# Sort experiment runs by target value into ascending order, and select first or last depending on order_str
ordered_runs = sorted(runs, key=lambda run: run["target"])
return ordered_runs[order_opt[order_str]]