Models prototype

mlflow/__init__.py

@@ -126,6 +126,7 @@
     active_run,
     autolog,
     create_experiment,
+    create_logged_model,
     delete_experiment,
     delete_run,
     delete_tag,
@@ -136,6 +137,7 @@
     get_artifact_uri,
     get_experiment,
     get_experiment_by_name,
+    get_logged_model,
     get_parent_run,
     get_run,
     last_active_run,
@@ -153,6 +155,7 @@
     log_table,
     log_text,
     search_experiments,
+    search_logged_models,
     search_runs,
     set_experiment,
     set_experiment_tag,
@@ -173,6 +176,7 @@
     "active_run",
     "autolog",
     "create_experiment",
+    "create_logged_model",
     "delete_experiment",
     "delete_run",
     "delete_tag",
@@ -188,6 +192,7 @@
     "get_experiment",
     "get_experiment_by_name",
     "get_last_active_trace",
+    "get_logged_model",
     "get_parent_run",
     "get_registry_uri",
     "get_run",
@@ -212,6 +217,7 @@
     "register_model",
     "run",
     "search_experiments",
+    "search_logged_models",
     "search_model_versions",
     "search_registered_models",
     "search_runs",

mlflow/entities/__init__.py

@@ -11,12 +11,19 @@
 from mlflow.entities.file_info import FileInfo
 from mlflow.entities.input_tag import InputTag
 from mlflow.entities.lifecycle_stage import LifecycleStage
+from mlflow.entities.logged_model import LoggedModel
 from mlflow.entities.metric import Metric
+from mlflow.entities.model_input import ModelInput
+from mlflow.entities.model_output import ModelOutput
+from mlflow.entities.model_param import ModelParam
+from mlflow.entities.model_status import ModelStatus
+from mlflow.entities.model_tag import ModelTag
 from mlflow.entities.param import Param
 from mlflow.entities.run import Run
 from mlflow.entities.run_data import RunData
 from mlflow.entities.run_info import RunInfo
 from mlflow.entities.run_inputs import RunInputs
+from mlflow.entities.run_outputs import RunOutputs
 from mlflow.entities.run_status import RunStatus
 from mlflow.entities.run_tag import RunTag
 from mlflow.entities.source_type import SourceType
@@ -46,6 +53,7 @@
     "InputTag",
     "DatasetInput",
     "RunInputs",
+    "RunOutputs",
     "Span",
     "LiveSpan",
     "NoOpSpan",
@@ -57,4 +65,10 @@
     "TraceInfo",
     "SpanStatusCode",
     "_DatasetSummary",
+    "LoggedModel",
+    "ModelInput",
+    "ModelOutput",
+    "ModelStatus",
+    "ModelTag",
+    "ModelParam",
 ]

mlflow/entities/logged_model.py

@@ -0,0 +1,162 @@
+from typing import Any, Dict, List, Optional, Union
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.entities.metric import Metric
+from mlflow.entities.model_param import ModelParam
+from mlflow.entities.model_status import ModelStatus
+from mlflow.entities.model_tag import ModelTag
+
+
+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,
+        run_id: Optional[str] = None,
+        status: ModelStatus = ModelStatus.READY,
+        status_message: Optional[str] = None,
+        tags: Optional[Union[List[ModelTag], Dict[str, str]]] = None,
+        params: Optional[Union[List[ModelParam], 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._run_id: Optional[str] = run_id
+        self._status: ModelStatus = 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
+
+    @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 run_id(self) -> Optional[str]:
+        """String. MLflow run ID that generated this model."""
+        return self._run_id
+
+    @property
+    def status(self) -> ModelStatus:
+        """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
+
+    @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"] = str(self.status)
+        return model_dict

mlflow/entities/metric.py

@@ -1,3 +1,5 @@
+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
@@ -10,11 +12,31 @@ class Metric(_MlflowObject):
     Metric object.
     """
 
-    def __init__(self, key, value, timestamp, step):
+    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):
@@ -36,16 +58,38 @@ 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
+        # TODO: Add model_id, dataset_name, dataset_digest, and run_id to the proto
         return metric
 
     @classmethod
     def from_proto(cls, proto):
+        # TODO: Add model_id, dataset_name, dataset_digest, and run_id to the proto
         return cls(proto.key, proto.value, proto.timestamp, proto.step)
 
     def __eq__(self, __o):
@@ -69,6 +113,10 @@ def to_dictionary(self):
             "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

mlflow/entities/model_input.py

@@ -0,0 +1,18 @@
+from mlflow.entities._mlflow_object import _MlflowObject
+
+
+class ModelInput(_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

mlflow/entities/model_output.py

@@ -0,0 +1,24 @@
+from mlflow.entities._mlflow_object import _MlflowObject
+
+
+class ModelOutput(_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