Return Model from log_model

mlflow/keras.py

@@ -376,7 +376,7 @@ def log_model(
         with mlflow.start_run() as run:
             mlflow.keras.log_model(keras_model, "models")
     """
-    Model.log(
+    return Model.log(
         artifact_path=artifact_path,
         flavor=mlflow.keras,
         keras_model=keras_model,

mlflow/lightgbm.py

@@ -264,7 +264,7 @@ def log_model(
     :param extra_pip_requirements: {{ extra_pip_requirements }}
     :param kwargs: kwargs to pass to `lightgbm.Booster.save_model`_ method.
     """
-    Model.log(
+    return Model.log(
         artifact_path=artifact_path,
         flavor=mlflow.lightgbm,
         registered_model_name=registered_model_name,

mlflow/models/model.py

@@ -1,3 +1,4 @@
+from collections import namedtuple
 from datetime import datetime
 import json
 import logging
@@ -22,12 +23,29 @@
     "Logging model metadata to the tracking server has failed, possibly due older "
     "server version. The model artifacts have been logged successfully under %s. "
     "In addition to exporting model artifacts, MLflow clients 1.7.0 and above "
-    "attempt to record model metadata to the  tracking store. If logging to a "
-    "mlflow server via REST, consider  upgrading the server version to MLflow "
+    "attempt to record model metadata to the tracking store. If logging to a "
+    "mlflow server via REST, consider upgrading the server version to MLflow "
     "1.7.0 or above."
 )
 
 
+ModelInfo = namedtuple(
+    "ModelInfo",
+    [
+        "run_id",
+        "artifact_path",
+        "model_uri",
+        "utc_time_created",
+        "flavors",
+        "model_uuid",
+        "saved_input_example_info",
+        "signature",
+        "input_schema",
+        "output_schema",
+    ],
+)
+
+
 class Model(object):
     """
     An MLflow Model that can support multiple model flavors. Provides APIs for implementing
@@ -109,6 +127,21 @@ def saved_input_example_info(self, value: Dict[str, Any]):
         # pylint: disable=attribute-defined-outside-init
         self._saved_input_example_info = value
 
+    def get_model_info(self) -> ModelInfo:
+        """Create a ModelInfo instance that contains the model metadata."""
+        return ModelInfo(
+            run_id=self.run_id,
+            artifact_path=self.artifact_path,
+            model_uri="runs:/{}/{}".format(self.run_id, self.artifact_path),
+            utc_time_created=self.utc_time_created,
+            flavors=self.flavors,
+            model_uuid=self.model_uuid,
+            saved_input_example_info=self.saved_input_example_info,
+            signature=self.signature,
+            input_schema=self.get_input_schema(),
+            output_schema=self.get_output_schema(),
+        )
+
     def to_dict(self):
         """Serialize the model to a dictionary."""
         res = {k: v for k, v in self.__dict__.items() if not k.startswith("_")}
@@ -223,3 +256,4 @@ def log(
                     registered_model_name,
                     await_registration_for=await_registration_for,
                 )
+        return mlflow_model.get_model_info()

mlflow/pytorch/__init__.py

@@ -283,7 +283,7 @@ def gen_data():
         PyTorch logged models
     """
     pickle_module = pickle_module or mlflow_pytorch_pickle_module
-    Model.log(
+    return Model.log(
         artifact_path=artifact_path,
         flavor=mlflow.pytorch,
         pytorch_model=pytorch_model,

mlflow/sklearn/__init__.py

@@ -371,7 +371,7 @@ def log_model(
         # log model
         mlflow.sklearn.log_model(sk_model, "sk_models")
     """
-    Model.log(
+    return Model.log(
         artifact_path=artifact_path,
         flavor=mlflow.sklearn,
         sk_model=sk_model,

mlflow/xgboost/__init__.py

@@ -241,7 +241,7 @@ def log_model(
     :param extra_pip_requirements: {{ extra_pip_requirements }}
     :param kwargs: kwargs to pass to `xgboost.Booster.save_model`_ method.
     """
-    Model.log(
+    return Model.log(
         artifact_path=artifact_path,
         flavor=mlflow.xgboost,
         registered_model_name=registered_model_name,