Log model uuid (#5149)

* init

Signed-off-by: Weichen Xu <weichen.xu@databricks.com>

* fix

Signed-off-by: Weichen Xu <weichen.xu@databricks.com>

* update

Signed-off-by: Weichen Xu <weichen.xu@databricks.com>

* check uuid valid

Signed-off-by: Weichen Xu <weichen.xu@databricks.com>

* update

Signed-off-by: Weichen Xu <weichen.xu@databricks.com>

* fix tests

Signed-off-by: Weichen Xu <weichen.xu@databricks.com>

* fix test

Signed-off-by: Weichen Xu <weichen.xu@databricks.com>

* update test

Signed-off-by: Weichen Xu <weichen.xu@databricks.com>

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,7 @@ def __init__(
         self.flavors = flavors if flavors is not None else {}
         self.signature = signature
         self.saved_input_example_info = saved_input_example_info
+        self.model_uuid = uuid.uuid4().hex if model_uuid is None else model_uuid
         self.__dict__.update(kwargs)
 
     def __eq__(self, other):

tests/models/test_model.py

@@ -44,6 +44,7 @@ def test_model_save_load():
         saved_input_example_info={"x": 1, "y": 2},
     )
     n.utc_time_created = m.utc_time_created
+    n.model_uuid = m.model_uuid
     assert m == n
     n.signature = None
     assert m != n

tests/pyfunc/test_model_export_with_loader_module_and_data_path.py

@@ -556,6 +556,26 @@ def test_column_schema_enforcement_no_col_names():
     assert pyfunc_model.predict(d).equals(pd.DataFrame(d))
 
 
+def _is_valid_uuid(val):
+    import uuid
+
+    try:
+        uuid.UUID(str(val))
+        return True
+    except ValueError:
+        return False
+
+
+def test_model_uuid():
+    m = Model()
+    assert m.model_uuid is not None
+    assert _is_valid_uuid(m.model_uuid)
+    m_dict = m.to_dict()
+    assert m_dict["model_uuid"] == m.model_uuid
+    m2 = Model.from_dict(m_dict)
+    assert m2.model_uuid == m.model_uuid
+
+
 def test_tensor_schema_enforcement_no_col_names():
     m = Model()
     input_schema = Schema([TensorSpec(np.dtype(np.float32), (-1, 3))])

tests/tracking/test_rest_tracking.py

@@ -386,7 +386,13 @@ def test_log_model(mlflow_client, backend_store_uri):
                 tag = run.data.tags["mlflow.log-model.history"]
                 models = json.loads(tag)
                 model.utc_time_created = models[i]["utc_time_created"]
-                assert models[i] == model.to_dict()
+
+                history_model_meta = models[i].copy()
+                original_model_uuid = history_model_meta.pop("model_uuid")
+                model_meta = model.to_dict().copy()
+                new_model_uuid = model_meta.pop(("model_uuid"))
+                assert history_model_meta == model_meta
+                assert original_model_uuid != new_model_uuid
                 assert len(models) == i + 1
                 for j in range(0, i + 1):
                     assert models[j]["artifact_path"] == model_paths[j]