From 11cade0107129bb3d01a85112ed219e390b8191d Mon Sep 17 00:00:00 2001 From: WeichenXu Date: Mon, 13 Dec 2021 17:31:54 +0800 Subject: [PATCH] Log model uuid (#5149) * init Signed-off-by: Weichen Xu * fix Signed-off-by: Weichen Xu * update Signed-off-by: Weichen Xu * check uuid valid Signed-off-by: Weichen Xu * update Signed-off-by: Weichen Xu * fix tests Signed-off-by: Weichen Xu * fix test Signed-off-by: Weichen Xu * update test Signed-off-by: Weichen Xu --- mlflow/models/model.py | 3 +++ tests/models/test_model.py | 1 + ...export_with_loader_module_and_data_path.py | 20 +++++++++++++++++++ tests/tracking/test_rest_tracking.py | 8 +++++++- 4 files changed, 31 insertions(+), 1 deletion(-) diff --git a/mlflow/models/model.py b/mlflow/models/model.py index bb55aa4c1cbfa..1ada7241bd406 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,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): diff --git a/tests/models/test_model.py b/tests/models/test_model.py index df77f020be64f..3f058c765b87e 100644 --- a/tests/models/test_model.py +++ b/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 diff --git a/tests/pyfunc/test_model_export_with_loader_module_and_data_path.py b/tests/pyfunc/test_model_export_with_loader_module_and_data_path.py index 8b4f395096bc8..6e9baca1e20df 100644 --- a/tests/pyfunc/test_model_export_with_loader_module_and_data_path.py +++ b/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))]) diff --git a/tests/tracking/test_rest_tracking.py b/tests/tracking/test_rest_tracking.py index 263e13229c1b9..b3b8adafb21d3 100644 --- a/tests/tracking/test_rest_tracking.py +++ b/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]