genesis-flow 1.0.0__py3-none-any.whl

mlflow/entities/logged_model.py

@@ -0,0 +1,228 @@
+from typing import Any, Optional, Union
+
+import mlflow.protos.service_pb2 as pb2
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.entities.logged_model_parameter import LoggedModelParameter
+from mlflow.entities.logged_model_status import LoggedModelStatus
+from mlflow.entities.logged_model_tag import LoggedModelTag
+from mlflow.entities.metric import Metric
+
+
+class LoggedModel(_MlflowObject):
+    """
+    MLflow entity representing a Model logged to an MLflow Experiment.
+    """
+
+    def __init__(
+        self,
+        experiment_id: str,
+        model_id: str,
+        name: str,
+        artifact_location: str,
+        creation_timestamp: int,
+        last_updated_timestamp: int,
+        model_type: Optional[str] = None,
+        source_run_id: Optional[str] = None,
+        status: Union[LoggedModelStatus, int] = LoggedModelStatus.READY,
+        status_message: Optional[str] = None,
+        tags: Optional[Union[list[LoggedModelTag], dict[str, str]]] = None,
+        params: Optional[Union[list[LoggedModelParameter], dict[str, str]]] = None,
+        metrics: Optional[list[Metric]] = None,
+    ):
+        super().__init__()
+        self._experiment_id: str = experiment_id
+        self._model_id: str = model_id
+        self._name: str = name
+        self._artifact_location: str = artifact_location
+        self._creation_time: int = creation_timestamp
+        self._last_updated_timestamp: int = last_updated_timestamp
+        self._model_type: Optional[str] = model_type
+        self._source_run_id: Optional[str] = source_run_id
+        self._status: LoggedModelStatus = (
+            status if isinstance(status, LoggedModelStatus) else LoggedModelStatus.from_int(status)
+        )
+        self._status_message: Optional[str] = status_message
+        self._tags: dict[str, str] = (
+            {tag.key: tag.value for tag in (tags or [])} if isinstance(tags, list) else (tags or {})
+        )
+        self._params: dict[str, str] = (
+            {param.key: param.value for param in (params or [])}
+            if isinstance(params, list)
+            else (params or {})
+        )
+        self._metrics: Optional[list[Metric]] = metrics
+        self._model_uri = f"models:/{self.model_id}"
+
+    def __repr__(self) -> str:
+        return "LoggedModel({})".format(
+            ", ".join(
+                f"{k}={v!r}"
+                for k, v in sorted(self, key=lambda x: x[0])
+                if (
+                    k
+                    not in [
+                        # These fields can be large and take up space on the notebook or terminal
+                        "tags",
+                        "params",
+                        "metrics",
+                    ]
+                )
+            )
+        )
+
+    @property
+    def experiment_id(self) -> str:
+        """String. Experiment ID associated with this Model."""
+        return self._experiment_id
+
+    @experiment_id.setter
+    def experiment_id(self, new_experiment_id: str):
+        self._experiment_id = new_experiment_id
+
+    @property
+    def model_id(self) -> str:
+        """String. Unique ID for this Model."""
+        return self._model_id
+
+    @model_id.setter
+    def model_id(self, new_model_id: str):
+        self._model_id = new_model_id
+
+    @property
+    def name(self) -> str:
+        """String. Name for this Model."""
+        return self._name
+
+    @name.setter
+    def name(self, new_name: str):
+        self._name = new_name
+
+    @property
+    def artifact_location(self) -> str:
+        """String. Location of the model artifacts."""
+        return self._artifact_location
+
+    @artifact_location.setter
+    def artifact_location(self, new_artifact_location: str):
+        self._artifact_location = new_artifact_location
+
+    @property
+    def creation_timestamp(self) -> int:
+        """Integer. Model creation timestamp (milliseconds since the Unix epoch)."""
+        return self._creation_time
+
+    @property
+    def last_updated_timestamp(self) -> int:
+        """Integer. Timestamp of last update for this Model (milliseconds since the Unix
+        epoch).
+        """
+        return self._last_updated_timestamp
+
+    @last_updated_timestamp.setter
+    def last_updated_timestamp(self, updated_timestamp: int):
+        self._last_updated_timestamp = updated_timestamp
+
+    @property
+    def model_type(self) -> Optional[str]:
+        """String. Type of the model."""
+        return self._model_type
+
+    @model_type.setter
+    def model_type(self, new_model_type: Optional[str]):
+        self._model_type = new_model_type
+
+    @property
+    def source_run_id(self) -> Optional[str]:
+        """String. MLflow run ID that generated this model."""
+        return self._source_run_id
+
+    @property
+    def status(self) -> LoggedModelStatus:
+        """String. Current status of this Model."""
+        return self._status
+
+    @status.setter
+    def status(self, updated_status: str):
+        self._status = updated_status
+
+    @property
+    def status_message(self) -> Optional[str]:
+        """String. Descriptive message for error status conditions."""
+        return self._status_message
+
+    @property
+    def tags(self) -> dict[str, str]:
+        """Dictionary of tag key (string) -> tag value for this Model."""
+        return self._tags
+
+    @property
+    def params(self) -> dict[str, str]:
+        """Model parameters."""
+        return self._params
+
+    @property
+    def metrics(self) -> Optional[list[Metric]]:
+        """List of metrics associated with this Model."""
+        return self._metrics
+
+    @property
+    def model_uri(self) -> str:
+        """URI of the model."""
+        return self._model_uri
+
+    @metrics.setter
+    def metrics(self, new_metrics: Optional[list[Metric]]):
+        self._metrics = new_metrics
+
+    @classmethod
+    def _properties(cls) -> list[str]:
+        # aggregate with base class properties since cls.__dict__ does not do it automatically
+        return sorted(cls._get_properties_helper())
+
+    def _add_tag(self, tag):
+        self._tags[tag.key] = tag.value
+
+    def to_dictionary(self) -> dict[str, Any]:
+        model_dict = dict(self)
+        model_dict["status"] = self.status.to_int()
+        # Remove the model_uri field from the dictionary since it is a derived field
+        del model_dict["model_uri"]
+        return model_dict
+
+    def to_proto(self):
+        return pb2.LoggedModel(
+            info=pb2.LoggedModelInfo(
+                experiment_id=self.experiment_id,
+                model_id=self.model_id,
+                name=self.name,
+                artifact_uri=self.artifact_location,
+                creation_timestamp_ms=self.creation_timestamp,
+                last_updated_timestamp_ms=self.last_updated_timestamp,
+                model_type=self.model_type,
+                source_run_id=self.source_run_id,
+                status=self.status.to_proto(),
+                tags=[pb2.LoggedModelTag(key=k, value=v) for k, v in self.tags.items()],
+            ),
+            data=pb2.LoggedModelData(
+                params=[pb2.LoggedModelParameter(key=k, value=v) for (k, v) in self.params.items()],
+                metrics=[m.to_proto() for m in self.metrics] if self.metrics else [],
+            ),
+        )
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(
+            experiment_id=proto.info.experiment_id,
+            model_id=proto.info.model_id,
+            name=proto.info.name,
+            artifact_location=proto.info.artifact_uri,
+            creation_timestamp=proto.info.creation_timestamp_ms,
+            last_updated_timestamp=proto.info.last_updated_timestamp_ms,
+            model_type=proto.info.model_type,
+            source_run_id=proto.info.source_run_id,
+            status=LoggedModelStatus.from_proto(proto.info.status),
+            status_message=proto.info.status_message,
+            tags=[LoggedModelTag.from_proto(tag) for tag in proto.info.tags],
+            params=[LoggedModelParameter.from_proto(param) for param in proto.data.params],
+            metrics=[Metric.from_proto(metric) for metric in proto.data.metrics],
+        )

mlflow/entities/logged_model_input.py

@@ -0,0 +1,26 @@
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.protos.service_pb2 import ModelInput as ProtoModelInput
+
+
+class LoggedModelInput(_MlflowObject):
+    """ModelInput object associated with a Run."""
+
+    def __init__(self, model_id: str):
+        self._model_id = model_id
+
+    def __eq__(self, other: _MlflowObject) -> bool:
+        if type(other) is type(self):
+            return self.__dict__ == other.__dict__
+        return False
+
+    @property
+    def model_id(self) -> str:
+        """Model ID."""
+        return self._model_id
+
+    def to_proto(self):
+        return ProtoModelInput(model_id=self._model_id)
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(proto.model_id)

mlflow/entities/logged_model_output.py

@@ -0,0 +1,32 @@
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.protos.service_pb2 import ModelOutput
+
+
+class LoggedModelOutput(_MlflowObject):
+    """ModelOutput object associated with a Run."""
+
+    def __init__(self, model_id: str, step: int) -> None:
+        self._model_id = model_id
+        self._step = step
+
+    def __eq__(self, other: _MlflowObject) -> bool:
+        if type(other) is type(self):
+            return self.__dict__ == other.__dict__
+        return False
+
+    @property
+    def model_id(self) -> str:
+        """Model ID"""
+        return self._model_id
+
+    @property
+    def step(self) -> str:
+        """Step at which the model was logged"""
+        return self._step
+
+    def to_proto(self):
+        return ModelOutput(model_id=self.model_id, step=self.step)
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(proto.model_id, proto.step)

mlflow/entities/logged_model_parameter.py

@@ -0,0 +1,46 @@
+import sys
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.protos import service_pb2 as pb2
+
+
+class LoggedModelParameter(_MlflowObject):
+    """
+    MLflow entity representing a parameter of a Model.
+    """
+
+    def __init__(self, key, value):
+        if "pyspark.ml" in sys.modules:
+            import pyspark.ml.param
+
+            if isinstance(key, pyspark.ml.param.Param):
+                key = key.name
+                value = str(value)
+        self._key = key
+        self._value = value
+
+    @property
+    def key(self):
+        """String key corresponding to the parameter name."""
+        return self._key
+
+    @property
+    def value(self):
+        """String value of the parameter."""
+        return self._value
+
+    def __eq__(self, __o):
+        if isinstance(__o, self.__class__):
+            return self._key == __o._key
+
+        return False
+
+    def __hash__(self):
+        return hash(self._key)
+
+    def to_proto(self):
+        return pb2.LoggedModelParameter(key=self._key, value=self._value)
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(key=proto.key, value=proto.value)

mlflow/entities/logged_model_status.py

@@ -0,0 +1,74 @@
+from enum import Enum
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos import service_pb2 as pb2
+
+
+class LoggedModelStatus(str, Enum):
+    """Enum for status of an :py:class:`mlflow.entities.LoggedModel`."""
+
+    UNSPECIFIED = "UNSPECIFIED"
+    PENDING = "PENDING"
+    READY = "READY"
+    FAILED = "FAILED"
+
+    def __str__(self):
+        return self.value
+
+    @staticmethod
+    def is_finalized(status) -> bool:
+        """
+        Determines whether or not a LoggedModelStatus is a finalized status.
+        A finalized status indicates that no further status updates will occur.
+        """
+        return status in [LoggedModelStatus.READY, LoggedModelStatus.FAILED]
+
+    def to_proto(self):
+        if self == LoggedModelStatus.UNSPECIFIED:
+            return pb2.LoggedModelStatus.LOGGED_MODEL_STATUS_UNSPECIFIED
+        elif self == LoggedModelStatus.PENDING:
+            return pb2.LoggedModelStatus.LOGGED_MODEL_PENDING
+        elif self == LoggedModelStatus.READY:
+            return pb2.LoggedModelStatus.LOGGED_MODEL_READY
+        elif self == LoggedModelStatus.FAILED:
+            return pb2.LoggedModelStatus.LOGGED_MODEL_UPLOAD_FAILED
+
+        raise MlflowException.invalid_parameter_value(f"Unknown model status: {self}")
+
+    @classmethod
+    def from_proto(cls, proto):
+        if proto == pb2.LoggedModelStatus.LOGGED_MODEL_STATUS_UNSPECIFIED:
+            return LoggedModelStatus.UNSPECIFIED
+        elif proto == pb2.LoggedModelStatus.LOGGED_MODEL_PENDING:
+            return LoggedModelStatus.PENDING
+        elif proto == pb2.LoggedModelStatus.LOGGED_MODEL_READY:
+            return LoggedModelStatus.READY
+        elif proto == pb2.LoggedModelStatus.LOGGED_MODEL_UPLOAD_FAILED:
+            return LoggedModelStatus.FAILED
+
+        raise MlflowException.invalid_parameter_value(f"Unknown model status: {proto}")
+
+    @classmethod
+    def from_int(cls, status_int: int) -> "LoggedModelStatus":
+        if status_int == 0:
+            return cls.UNSPECIFIED
+        elif status_int == 1:
+            return cls.PENDING
+        elif status_int == 2:
+            return cls.READY
+        elif status_int == 3:
+            return cls.FAILED
+
+        raise MlflowException.invalid_parameter_value(f"Unknown model status: {status_int}")
+
+    def to_int(self) -> int:
+        if self == LoggedModelStatus.UNSPECIFIED:
+            return 0
+        elif self == LoggedModelStatus.PENDING:
+            return 1
+        elif self == LoggedModelStatus.READY:
+            return 2
+        elif self == LoggedModelStatus.FAILED:
+            return 3
+
+        raise MlflowException.invalid_parameter_value(f"Unknown model status: {self}")

mlflow/entities/logged_model_tag.py

@@ -0,0 +1,33 @@
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.protos import service_pb2 as pb2
+
+
+class LoggedModelTag(_MlflowObject):
+    """Tag object associated with a Model."""
+
+    def __init__(self, key, value):
+        self._key = key
+        self._value = value
+
+    def __eq__(self, other):
+        if type(other) is type(self):
+            # TODO deep equality here?
+            return self.__dict__ == other.__dict__
+        return False
+
+    @property
+    def key(self):
+        """String name of the tag."""
+        return self._key
+
+    @property
+    def value(self):
+        """String value of the tag."""
+        return self._value
+
+    def to_proto(self):
+        return pb2.LoggedModelTag(key=self._key, value=self._value)
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(key=proto.key, value=proto.value)

mlflow/entities/metric.py

@@ -0,0 +1,200 @@
+from typing import Optional
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.protos.service_pb2 import Metric as ProtoMetric
+from mlflow.protos.service_pb2 import MetricWithRunId as ProtoMetricWithRunId
+
+
+class Metric(_MlflowObject):
+    """
+    Metric object.
+    """
+
+    def __init__(
+        self,
+        key,
+        value,
+        timestamp,
+        step,
+        model_id: Optional[str] = None,
+        dataset_name: Optional[str] = None,
+        dataset_digest: Optional[str] = None,
+        run_id: Optional[str] = None,
+    ):
+        if (dataset_name, dataset_digest).count(None) == 1:
+            raise MlflowException(
+                "Both dataset_name and dataset_digest must be provided if one is provided",
+                INVALID_PARAMETER_VALUE,
+            )
+
+        self._key = key
+        self._value = value
+        self._timestamp = timestamp
+        self._step = step
+        self._model_id = model_id
+        self._dataset_name = dataset_name
+        self._dataset_digest = dataset_digest
+        self._run_id = run_id
+
+    @property
+    def key(self):
+        """String key corresponding to the metric name."""
+        return self._key
+
+    @property
+    def value(self):
+        """Float value of the metric."""
+        return self._value
+
+    @property
+    def timestamp(self):
+        """Metric timestamp as an integer (milliseconds since the Unix epoch)."""
+        return self._timestamp
+
+    @property
+    def step(self):
+        """Integer metric step (x-coordinate)."""
+        return self._step
+
+    @property
+    def model_id(self):
+        """ID of the Model associated with the metric."""
+        return self._model_id
+
+    @property
+    def dataset_name(self) -> Optional[str]:
+        """String. Name of the dataset associated with the metric."""
+        return self._dataset_name
+
+    @property
+    def dataset_digest(self) -> Optional[str]:
+        """String. Digest of the dataset associated with the metric."""
+        return self._dataset_digest
+
+    @property
+    def run_id(self) -> Optional[str]:
+        """String. Run ID associated with the metric."""
+        return self._run_id
+
+    def to_proto(self):
+        metric = ProtoMetric()
+        metric.key = self.key
+        metric.value = self.value
+        metric.timestamp = self.timestamp
+        metric.step = self.step
+        if self.model_id:
+            metric.model_id = self.model_id
+        if self.dataset_name:
+            metric.dataset_name = self.dataset_name
+        if self.dataset_digest:
+            metric.dataset_digest = self.dataset_digest
+        if self.run_id:
+            metric.run_id = self.run_id
+        return metric
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(
+            proto.key,
+            proto.value,
+            proto.timestamp,
+            proto.step,
+            model_id=proto.model_id or None,
+            dataset_name=proto.dataset_name or None,
+            dataset_digest=proto.dataset_digest or None,
+            run_id=proto.run_id or None,
+        )
+
+    def __eq__(self, __o):
+        if isinstance(__o, self.__class__):
+            return self.__dict__ == __o.__dict__
+
+        return False
+
+    def __hash__(self):
+        return hash(
+            (
+                self._key,
+                self._value,
+                self._timestamp,
+                self._step,
+                self._model_id,
+                self._dataset_name,
+                self._dataset_digest,
+                self._run_id,
+            )
+        )
+
+    def to_dictionary(self):
+        """
+        Convert the Metric object to a dictionary.
+
+        Returns:
+            dict: The Metric object represented as a dictionary.
+        """
+        return {
+            "key": self.key,
+            "value": self.value,
+            "timestamp": self.timestamp,
+            "step": self.step,
+            "model_id": self.model_id,
+            "dataset_name": self.dataset_name,
+            "dataset_digest": self.dataset_digest,
+            "run_id": self._run_id,
+        }
+
+    @classmethod
+    def from_dictionary(cls, metric_dict):
+        """
+        Create a Metric object from a dictionary.
+
+        Args:
+            metric_dict (dict): Dictionary containing metric information.
+
+        Returns:
+            Metric: The Metric object created from the dictionary.
+        """
+        required_keys = ["key", "value", "timestamp", "step"]
+        missing_keys = [key for key in required_keys if key not in metric_dict]
+        if missing_keys:
+            raise MlflowException(
+                f"Missing required keys {missing_keys} in metric dictionary",
+                INVALID_PARAMETER_VALUE,
+            )
+
+        return cls(**metric_dict)
+
+
+class MetricWithRunId(Metric):
+    def __init__(self, metric: Metric, run_id):
+        super().__init__(
+            key=metric.key,
+            value=metric.value,
+            timestamp=metric.timestamp,
+            step=metric.step,
+        )
+        self._run_id = run_id
+
+    @property
+    def run_id(self):
+        return self._run_id
+
+    def to_dict(self):
+        return {
+            "key": self.key,
+            "value": self.value,
+            "timestamp": self.timestamp,
+            "step": self.step,
+            "run_id": self.run_id,
+        }
+
+    def to_proto(self):
+        metric = ProtoMetricWithRunId()
+        metric.key = self.key
+        metric.value = self.value
+        metric.timestamp = self.timestamp
+        metric.step = self.step
+        metric.run_id = self.run_id
+        return metric

mlflow/entities/model_registry/__init__.py

@@ -0,0 +1,29 @@
+from mlflow.entities.model_registry.model_version import ModelVersion
+from mlflow.entities.model_registry.model_version_deployment_job_state import (
+    ModelVersionDeploymentJobState,
+)
+from mlflow.entities.model_registry.model_version_search import ModelVersionSearch
+from mlflow.entities.model_registry.model_version_tag import ModelVersionTag
+from mlflow.entities.model_registry.prompt import Prompt
+from mlflow.entities.model_registry.prompt_version import PromptVersion
+from mlflow.entities.model_registry.registered_model import RegisteredModel
+from mlflow.entities.model_registry.registered_model_alias import RegisteredModelAlias
+from mlflow.entities.model_registry.registered_model_deployment_job_state import (
+    RegisteredModelDeploymentJobState,
+)
+from mlflow.entities.model_registry.registered_model_search import RegisteredModelSearch
+from mlflow.entities.model_registry.registered_model_tag import RegisteredModelTag
+
+__all__ = [
+    "Prompt",
+    "PromptVersion",
+    "RegisteredModel",
+    "ModelVersion",
+    "RegisteredModelAlias",
+    "RegisteredModelTag",
+    "ModelVersionTag",
+    "RegisteredModelSearch",
+    "ModelVersionSearch",
+    "ModelVersionDeploymentJobState",
+    "RegisteredModelDeploymentJobState",
+]

mlflow/entities/model_registry/_model_registry_entity.py

@@ -0,0 +1,13 @@
+from abc import abstractmethod
+
+from mlflow.entities._mlflow_object import _MlflowObject
+
+
+class _ModelRegistryEntity(_MlflowObject):
+    @classmethod
+    @abstractmethod
+    def from_proto(cls, proto):
+        pass
+
+    def __eq__(self, other):
+        return dict(self) == dict(other)