genesis-flow 1.0.0__py3-none-any.whl

mlflow/entities/run.py

@@ -0,0 +1,97 @@
+from typing import Any, Optional
+
+from mlflow.entities._mlflow_object import _MlflowObject
+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.exceptions import MlflowException
+from mlflow.protos.service_pb2 import Run as ProtoRun
+
+
+class Run(_MlflowObject):
+    """
+    Run object.
+    """
+
+    def __init__(
+        self,
+        run_info: RunInfo,
+        run_data: RunData,
+        run_inputs: Optional[RunInputs] = None,
+        run_outputs: Optional[RunOutputs] = None,
+    ) -> None:
+        if run_info is None:
+            raise MlflowException("run_info cannot be None")
+        self._info = run_info
+        self._data = run_data
+        self._inputs = run_inputs
+        self._outputs = run_outputs
+
+    @property
+    def info(self) -> RunInfo:
+        """
+        The run metadata, such as the run id, start time, and status.
+
+        :rtype: :py:class:`mlflow.entities.RunInfo`
+        """
+        return self._info
+
+    @property
+    def data(self) -> RunData:
+        """
+        The run data, including metrics, parameters, and tags.
+
+        :rtype: :py:class:`mlflow.entities.RunData`
+        """
+        return self._data
+
+    @property
+    def inputs(self) -> RunInputs:
+        """
+        The run inputs, including dataset inputs.
+
+        :rtype: :py:class:`mlflow.entities.RunInputs`
+        """
+        return self._inputs
+
+    @property
+    def outputs(self) -> RunOutputs:
+        """
+        The run outputs, including model outputs.
+
+        :rtype: :py:class:`mlflow.entities.RunOutputs`
+        """
+        return self._outputs
+
+    def to_proto(self):
+        run = ProtoRun()
+        run.info.MergeFrom(self.info.to_proto())
+        if self.data:
+            run.data.MergeFrom(self.data.to_proto())
+        if self.inputs:
+            run.inputs.MergeFrom(self.inputs.to_proto())
+        if self.outputs:
+            run.outputs.MergeFrom(self.outputs.to_proto())
+        return run
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(
+            RunInfo.from_proto(proto.info),
+            RunData.from_proto(proto.data),
+            RunInputs.from_proto(proto.inputs) if proto.inputs else None,
+            RunOutputs.from_proto(proto.outputs) if proto.outputs else None,
+        )
+
+    def to_dictionary(self) -> dict[Any, Any]:
+        run_dict = {
+            "info": dict(self.info),
+        }
+        if self.data:
+            run_dict["data"] = self.data.to_dictionary()
+        if self.inputs:
+            run_dict["inputs"] = self.inputs.to_dictionary()
+        if self.outputs:
+            run_dict["outputs"] = self.outputs.to_dictionary()
+        return run_dict

mlflow/entities/run_data.py

@@ -0,0 +1,84 @@
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.entities.metric import Metric
+from mlflow.entities.param import Param
+from mlflow.entities.run_tag import RunTag
+from mlflow.protos.service_pb2 import Param as ProtoParam
+from mlflow.protos.service_pb2 import RunData as ProtoRunData
+from mlflow.protos.service_pb2 import RunTag as ProtoRunTag
+
+
+class RunData(_MlflowObject):
+    """
+    Run data (metrics and parameters).
+    """
+
+    def __init__(self, metrics=None, params=None, tags=None):
+        """Construct a new mlflow.entities.RunData instance.
+
+        Args:
+            metrics: List of mlflow.entities.Metric.
+            params: List of mlflow.entities.Param.
+            tags: List of mlflow.entities.RunTag.
+
+        """
+        # Maintain the original list of metrics so that we can easily convert it back to
+        # protobuf
+        self._metric_objs = metrics or []
+        self._metrics = {metric.key: metric.value for metric in self._metric_objs}
+        self._params = {param.key: param.value for param in (params or [])}
+        self._tags = {tag.key: tag.value for tag in (tags or [])}
+
+    @property
+    def metrics(self):
+        """
+        Dictionary of string key -> metric value for the current run.
+        For each metric key, the metric value with the latest timestamp is returned. In case there
+        are multiple values with the same latest timestamp, the maximum of these values is returned.
+        """
+        return self._metrics
+
+    @property
+    def params(self):
+        """Dictionary of param key (string) -> param value for the current run."""
+        return self._params
+
+    @property
+    def tags(self):
+        """Dictionary of tag key (string) -> tag value for the current run."""
+        return self._tags
+
+    def _add_metric(self, metric):
+        self._metrics[metric.key] = metric.value
+        self._metric_objs.append(metric)
+
+    def _add_param(self, param):
+        self._params[param.key] = param.value
+
+    def _add_tag(self, tag):
+        self._tags[tag.key] = tag.value
+
+    def to_proto(self):
+        run_data = ProtoRunData()
+        run_data.metrics.extend([m.to_proto() for m in self._metric_objs])
+        run_data.params.extend([ProtoParam(key=key, value=val) for key, val in self.params.items()])
+        run_data.tags.extend([ProtoRunTag(key=key, value=val) for key, val in self.tags.items()])
+        return run_data
+
+    def to_dictionary(self):
+        return {
+            "metrics": self.metrics,
+            "params": self.params,
+            "tags": self.tags,
+        }
+
+    @classmethod
+    def from_proto(cls, proto):
+        run_data = cls()
+        # iterate proto and add metrics, params, and tags
+        for proto_metric in proto.metrics:
+            run_data._add_metric(Metric.from_proto(proto_metric))
+        for proto_param in proto.params:
+            run_data._add_param(Param.from_proto(proto_param))
+        for proto_tag in proto.tags:
+            run_data._add_tag(RunTag.from_proto(proto_tag))
+        return run_data

mlflow/entities/run_info.py

@@ -0,0 +1,188 @@
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.entities.lifecycle_stage import LifecycleStage
+from mlflow.entities.run_status import RunStatus
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.protos.service_pb2 import RunInfo as ProtoRunInfo
+
+
+def check_run_is_active(run_info):
+    if run_info.lifecycle_stage != LifecycleStage.ACTIVE:
+        raise MlflowException(
+            f"The run {run_info.run_id} must be in 'active' lifecycle_stage.",
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+
+
+class searchable_attribute(property):
+    # Wrapper class over property to designate some of the properties as searchable
+    # run attributes
+    pass
+
+
+class orderable_attribute(property):
+    # Wrapper class over property to designate some of the properties as orderable
+    # run attributes
+    pass
+
+
+class RunInfo(_MlflowObject):
+    """
+    Metadata about a run.
+    """
+
+    def __init__(
+        self,
+        run_id,
+        experiment_id,
+        user_id,
+        status,
+        start_time,
+        end_time,
+        lifecycle_stage,
+        artifact_uri=None,
+        run_name=None,
+    ):
+        if experiment_id is None:
+            raise Exception("experiment_id cannot be None")
+        if user_id is None:
+            raise Exception("user_id cannot be None")
+        if status is None:
+            raise Exception("status cannot be None")
+        if start_time is None:
+            raise Exception("start_time cannot be None")
+        self._run_id = run_id
+        self._experiment_id = experiment_id
+        self._user_id = user_id
+        self._status = status
+        self._start_time = start_time
+        self._end_time = end_time
+        self._lifecycle_stage = lifecycle_stage
+        self._artifact_uri = artifact_uri
+        self._run_name = run_name
+
+    def __eq__(self, other):
+        if type(other) is type(self):
+            # TODO deep equality here?
+            return self.__dict__ == other.__dict__
+        return False
+
+    def _copy_with_overrides(self, status=None, end_time=None, lifecycle_stage=None, run_name=None):
+        """A copy of the RunInfo with certain attributes modified."""
+        proto = self.to_proto()
+        if status:
+            proto.status = status
+        if end_time:
+            proto.end_time = end_time
+        if lifecycle_stage:
+            proto.lifecycle_stage = lifecycle_stage
+        if run_name:
+            proto.run_name = run_name
+        return RunInfo.from_proto(proto)
+
+    @searchable_attribute
+    def run_id(self):
+        """String containing run id."""
+        return self._run_id
+
+    @property
+    def experiment_id(self):
+        """String ID of the experiment for the current run."""
+        return self._experiment_id
+
+    @searchable_attribute
+    def run_name(self):
+        """String containing run name."""
+        return self._run_name
+
+    def _set_run_name(self, new_name):
+        self._run_name = new_name
+
+    @searchable_attribute
+    def user_id(self):
+        """String ID of the user who initiated this run."""
+        return self._user_id
+
+    @searchable_attribute
+    def status(self):
+        """
+        One of the values in :py:class:`mlflow.entities.RunStatus`
+        describing the status of the run.
+        """
+        return self._status
+
+    @searchable_attribute
+    def start_time(self):
+        """Start time of the run, in number of milliseconds since the UNIX epoch."""
+        return self._start_time
+
+    @searchable_attribute
+    def end_time(self):
+        """End time of the run, in number of milliseconds since the UNIX epoch."""
+        return self._end_time
+
+    @searchable_attribute
+    def artifact_uri(self):
+        """String root artifact URI of the run."""
+        return self._artifact_uri
+
+    @property
+    def lifecycle_stage(self):
+        """
+        One of the values in :py:class:`mlflow.entities.lifecycle_stage.LifecycleStage`
+        describing the lifecycle stage of the run.
+        """
+        return self._lifecycle_stage
+
+    def to_proto(self):
+        proto = ProtoRunInfo()
+        proto.run_uuid = self.run_id
+        proto.run_id = self.run_id
+        if self.run_name is not None:
+            proto.run_name = self.run_name
+        proto.experiment_id = self.experiment_id
+        proto.user_id = self.user_id
+        proto.status = RunStatus.from_string(self.status)
+        proto.start_time = self.start_time
+        if self.end_time:
+            proto.end_time = self.end_time
+        if self.artifact_uri:
+            proto.artifact_uri = self.artifact_uri
+        proto.lifecycle_stage = self.lifecycle_stage
+        return proto
+
+    @classmethod
+    def from_proto(cls, proto):
+        end_time = proto.end_time
+        # The proto2 default scalar value of zero indicates that the run's end time is absent.
+        # An absent end time is represented with a NoneType in the `RunInfo` class
+        if end_time == 0:
+            end_time = None
+        return cls(
+            run_id=proto.run_id,
+            run_name=proto.run_name,
+            experiment_id=proto.experiment_id,
+            user_id=proto.user_id,
+            status=RunStatus.to_string(proto.status),
+            start_time=proto.start_time,
+            end_time=end_time,
+            lifecycle_stage=proto.lifecycle_stage,
+            artifact_uri=proto.artifact_uri,
+        )
+
+    @classmethod
+    def get_searchable_attributes(cls):
+        return sorted(
+            [p for p in cls.__dict__ if isinstance(getattr(cls, p), searchable_attribute)]
+        )
+
+    @classmethod
+    def get_orderable_attributes(cls):
+        # Note that all searchable attributes are also orderable.
+        return sorted(
+            [
+                p
+                for p in cls.__dict__
+                if isinstance(getattr(cls, p), (searchable_attribute, orderable_attribute))
+            ]
+        )

mlflow/entities/run_inputs.py

@@ -0,0 +1,59 @@
+from typing import Any, Optional
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.entities.dataset_input import DatasetInput
+from mlflow.entities.logged_model_input import LoggedModelInput
+from mlflow.protos.service_pb2 import RunInputs as ProtoRunInputs
+
+
+class RunInputs(_MlflowObject):
+    """RunInputs object."""
+
+    def __init__(
+        self,
+        dataset_inputs: list[DatasetInput],
+        model_inputs: Optional[list[LoggedModelInput]] = None,
+    ) -> None:
+        self._dataset_inputs = dataset_inputs
+        self._model_inputs = model_inputs or []
+
+    def __eq__(self, other: _MlflowObject) -> bool:
+        if type(other) is type(self):
+            return self.__dict__ == other.__dict__
+        return False
+
+    @property
+    def dataset_inputs(self) -> list[DatasetInput]:
+        """Array of dataset inputs."""
+        return self._dataset_inputs
+
+    @property
+    def model_inputs(self) -> list[LoggedModelInput]:
+        """Array of model inputs."""
+        return self._model_inputs
+
+    def to_proto(self):
+        run_inputs = ProtoRunInputs()
+        run_inputs.dataset_inputs.extend(
+            [dataset_input.to_proto() for dataset_input in self.dataset_inputs]
+        )
+        run_inputs.model_inputs.extend(
+            [model_input.to_proto() for model_input in self.model_inputs]
+        )
+        return run_inputs
+
+    def to_dictionary(self) -> dict[str, Any]:
+        return {
+            "model_inputs": self.model_inputs,
+            "dataset_inputs": [d.to_dictionary() for d in self.dataset_inputs],
+        }
+
+    @classmethod
+    def from_proto(cls, proto):
+        dataset_inputs = [
+            DatasetInput.from_proto(dataset_input) for dataset_input in proto.dataset_inputs
+        ]
+        model_inputs = [
+            LoggedModelInput.from_proto(model_input) for model_input in proto.model_inputs
+        ]
+        return cls(dataset_inputs, model_inputs)

mlflow/entities/run_outputs.py

@@ -0,0 +1,43 @@
+from typing import Any
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.entities.logged_model_output import LoggedModelOutput
+from mlflow.protos.service_pb2 import RunOutputs as ProtoRunOutputs
+
+
+class RunOutputs(_MlflowObject):
+    """RunOutputs object."""
+
+    def __init__(self, model_outputs: list[LoggedModelOutput]) -> None:
+        self._model_outputs = model_outputs
+
+    def __eq__(self, other: _MlflowObject) -> bool:
+        if type(other) is type(self):
+            return self.__dict__ == other.__dict__
+        return False
+
+    @property
+    def model_outputs(self) -> list[LoggedModelOutput]:
+        """Array of model outputs."""
+        return self._model_outputs
+
+    def to_proto(self):
+        run_outputs = ProtoRunOutputs()
+        run_outputs.model_outputs.extend(
+            [model_output.to_proto() for model_output in self.model_outputs]
+        )
+
+        return run_outputs
+
+    def to_dictionary(self) -> dict[Any, Any]:
+        return {
+            "model_outputs": self.model_outputs,
+        }
+
+    @classmethod
+    def from_proto(cls, proto):
+        model_outputs = [
+            LoggedModelOutput.from_proto(model_output) for model_output in proto.model_outputs
+        ]
+
+        return cls(model_outputs)

mlflow/entities/run_status.py

@@ -0,0 +1,41 @@
+from mlflow.protos.service_pb2 import RunStatus as ProtoRunStatus
+
+
+class RunStatus:
+    """Enum for status of an :py:class:`mlflow.entities.Run`."""
+
+    RUNNING = ProtoRunStatus.Value("RUNNING")
+    SCHEDULED = ProtoRunStatus.Value("SCHEDULED")
+    FINISHED = ProtoRunStatus.Value("FINISHED")
+    FAILED = ProtoRunStatus.Value("FAILED")
+    KILLED = ProtoRunStatus.Value("KILLED")
+
+    _STRING_TO_STATUS = {k: ProtoRunStatus.Value(k) for k in ProtoRunStatus.keys()}
+    _STATUS_TO_STRING = {value: key for key, value in _STRING_TO_STATUS.items()}
+    _TERMINATED_STATUSES = {FINISHED, FAILED, KILLED}
+
+    @staticmethod
+    def from_string(status_str):
+        if status_str not in RunStatus._STRING_TO_STATUS:
+            raise Exception(
+                f"Could not get run status corresponding to string {status_str}. Valid run "
+                f"status strings: {list(RunStatus._STRING_TO_STATUS.keys())}"
+            )
+        return RunStatus._STRING_TO_STATUS[status_str]
+
+    @staticmethod
+    def to_string(status):
+        if status not in RunStatus._STATUS_TO_STRING:
+            raise Exception(
+                f"Could not get string corresponding to run status {status}. Valid run "
+                f"statuses: {list(RunStatus._STATUS_TO_STRING.keys())}"
+            )
+        return RunStatus._STATUS_TO_STRING[status]
+
+    @staticmethod
+    def is_terminated(status):
+        return status in RunStatus._TERMINATED_STATUSES
+
+    @staticmethod
+    def all_status():
+        return list(RunStatus._STATUS_TO_STRING.keys())

mlflow/entities/run_tag.py

@@ -0,0 +1,36 @@
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.protos.service_pb2 import RunTag as ProtoRunTag
+
+
+class RunTag(_MlflowObject):
+    """Tag object associated with a run."""
+
+    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):
+        param = ProtoRunTag()
+        param.key = self.key
+        param.value = self.value
+        return param
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(proto.key, proto.value)

mlflow/entities/source_type.py

@@ -0,0 +1,31 @@
+class SourceType:
+    """Enum for originating source of a :py:class:`mlflow.entities.Run`."""
+
+    NOTEBOOK, JOB, PROJECT, LOCAL, UNKNOWN = range(1, 6)
+
+    _STRING_TO_SOURCETYPE = {
+        "NOTEBOOK": NOTEBOOK,
+        "JOB": JOB,
+        "PROJECT": PROJECT,
+        "LOCAL": LOCAL,
+        "UNKNOWN": UNKNOWN,
+    }
+    SOURCETYPE_TO_STRING = {value: key for key, value in _STRING_TO_SOURCETYPE.items()}
+
+    @staticmethod
+    def from_string(status_str):
+        if status_str not in SourceType._STRING_TO_SOURCETYPE:
+            raise Exception(
+                f"Could not get run status corresponding to string {status_str}. Valid run "
+                f"status strings: {list(SourceType._STRING_TO_SOURCETYPE.keys())}"
+            )
+        return SourceType._STRING_TO_SOURCETYPE[status_str]
+
+    @staticmethod
+    def to_string(status):
+        if status not in SourceType.SOURCETYPE_TO_STRING:
+            raise Exception(
+                f"Could not get string corresponding to run status {status}. Valid run "
+                f"statuses: {list(SourceType.SOURCETYPE_TO_STRING.keys())}"
+            )
+        return SourceType.SOURCETYPE_TO_STRING[status]