diff --git a/mlflow/models/model.py b/mlflow/models/model.py index bb55aa4c1cbfa..5f772ad304fc0 100644 --- a/mlflow/models/model.py +++ b/mlflow/models/model.py @@ -4,6 +4,7 @@ import yaml import os +import uuid from typing import Any, Dict, Optional @@ -41,6 +42,7 @@ def __init__( flavors=None, signature=None, # ModelSignature saved_input_example_info: Dict[str, Any] = None, + model_uuid=None, **kwargs, ): # store model id instead of run_id and path to avoid confusion when model gets exported @@ -52,6 +54,8 @@ def __init__( self.flavors = flavors if flavors is not None else {} self.signature = signature self.saved_input_example_info = saved_input_example_info + if self.model_uuid is None: + self.model_uuid = uuid.uuid4() self.__dict__.update(kwargs) def __eq__(self, other): @@ -98,6 +102,8 @@ def to_dict(self): res["signature"] = self.signature.to_dict() if self.saved_input_example_info is not None: res["saved_input_example_info"] = self.saved_input_example_info + + res["model_uuid"] = self.model_uuid return res def to_yaml(self, stream=None):