Update examples to use model_info (mlflow#14636)

Signed-off-by: serena-ruan <[email protected]>

Author: serena-ruan

examples/catboost/train.py

@@ -28,11 +28,10 @@
 with mlflow.start_run() as run:
     signature = infer_signature(eval_data, model.predict(eval_data))
     mlflow.log_params(params)
-    mlflow.catboost.log_model(model, artifact_path="model", signature=signature)
-    model_uri = mlflow.get_artifact_uri("model")
+    model_info = mlflow.catboost.log_model(model, artifact_path="model", signature=signature)
 
 # Load model
-loaded_model = mlflow.catboost.load_model(model_uri)
+loaded_model = mlflow.catboost.load_model(model_info.model_uri)
 
 # Get predictions
 preds = loaded_model.predict(eval_data)

examples/diviner/train.py

@@ -47,7 +47,7 @@ def generate_data(location_data, start_dt) -> pd.DataFrame:
     return pd.concat(generated_listing).reset_index().drop("index", axis=1)
 
 
-def grouped_prophet_example(locations, start_dt, artifact_path):
+def grouped_prophet_example(locations, start_dt):
     print("Generating data...\n")
     data = generate_data(location_data=locations, start_dt=start_dt)
     grouping_keys = ["country", "city"]
@@ -73,7 +73,7 @@ def grouped_prophet_example(locations, start_dt, artifact_path):
     )
     print(f"Cross Validation Metrics: \n{metrics.to_string()}")
 
-    mlflow.diviner.log_model(diviner_model=model, artifact_path=artifact_path)
+    model_info = mlflow.diviner.log_model(diviner_model=model, artifact_path="diviner_model")
 
     # As an Alternative to saving metrics and params directly with a `log_dict()` function call,
     # Serializing the DataFrames to local as a .csv can be done as well, without requiring
@@ -99,7 +99,7 @@ def grouped_prophet_example(locations, start_dt, artifact_path):
 
     mlflow.log_dict(metrics.to_dict(), "metrics.json")
 
-    return mlflow.get_artifact_uri(artifact_path=artifact_path)
+    return model_info.model_uri
 
 
 if __name__ == "__main__":
@@ -112,10 +112,9 @@ def grouped_prophet_example(locations, start_dt, artifact_path):
         ("MX", "MexicoCity"),
     ]
     start_dt = "2022-02-01 04:11:35"
-    artifact_path = "diviner_model"
 
     with mlflow.start_run():
-        uri = grouped_prophet_example(locations, start_dt, artifact_path)
+        uri = grouped_prophet_example(locations, start_dt)
 
     loaded_model = mlflow.diviner.load_model(model_uri=uri)
 

examples/evaluation/evaluate_on_binary_classifier.py

@@ -24,12 +24,11 @@
 
 with mlflow.start_run() as run:
     # Log the XGBoost binary classifier model to MLflow
-    mlflow.sklearn.log_model(model, "model", signature=signature)
-    model_uri = mlflow.get_artifact_uri("model")
+    model_info = mlflow.sklearn.log_model(model, "model", signature=signature)
 
     # Evaluate the logged model
     result = mlflow.evaluate(
-        model_uri,
+        model_info.model_uri,
         eval_data,
         targets="label",
         model_type="classifier",

examples/evaluation/evaluate_on_multiclass_classifier.py

@@ -4,17 +4,15 @@
 
 import mlflow
 
-mlflow.sklearn.autolog()
-
 X, y = make_classification(n_samples=10000, n_classes=10, n_informative=5, random_state=1)
 
 X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)
 
 with mlflow.start_run() as run:
     model = LogisticRegression(solver="liblinear").fit(X_train, y_train)
-    model_uri = mlflow.get_artifact_uri("model")
+    model_info = mlflow.sklearn.log_model(model, "model")
     result = mlflow.evaluate(
-        model_uri,
+        model_info.model_uri,
         X_test,
         targets=y_test,
         model_type="classifier",

examples/evaluation/evaluate_on_regressor.py

@@ -4,8 +4,6 @@
 
 import mlflow
 
-mlflow.sklearn.autolog()
-
 california_housing_data = fetch_california_housing()
 
 X_train, X_test, y_train, y_test = train_test_split(
@@ -14,10 +12,10 @@
 
 with mlflow.start_run() as run:
     model = LinearRegression().fit(X_train, y_train)
-    model_uri = mlflow.get_artifact_uri("model")
+    model_info = mlflow.sklearn.log_model(model, "model")
 
     result = mlflow.evaluate(
-        model_uri,
+        model_info.model_uri,
         X_test,
         targets=y_test,
         model_type="regressor",

examples/evaluation/evaluate_with_custom_metrics.py

@@ -61,10 +61,9 @@ def prediction_target_scatter(eval_df, _builtin_metrics, artifacts_dir):
 
 
 with mlflow.start_run() as run:
-    mlflow.sklearn.log_model(lin_reg, "model", signature=signature)
-    model_uri = mlflow.get_artifact_uri("model")
+    model_info = mlflow.sklearn.log_model(lin_reg, "model", signature=signature)
     result = mlflow.evaluate(
-        model=model_uri,
+        model=model_info.model_uri,
         data=eval_data,
         targets="target",
         model_type="regressor",

examples/evaluation/evaluate_with_custom_metrics_comprehensive.py

@@ -62,10 +62,9 @@ def custom_artifact(eval_df, builtin_metrics, _artifacts_dir):
 
 
 with mlflow.start_run() as run:
-    mlflow.sklearn.log_model(lin_reg, "model", signature=signature)
-    model_uri = mlflow.get_artifact_uri("model")
+    model_info = mlflow.sklearn.log_model(lin_reg, "model", signature=signature)
     result = mlflow.evaluate(
-        model=model_uri,
+        model=model_info.model_uri,
         data=eval_data,
         targets="target",
         model_type="regressor",

examples/pip_requirements/pip_requirements.py

@@ -19,14 +19,12 @@ def read_lines(path):
         return f.read().splitlines()
 
 
-def get_pip_requirements(run_id, artifact_path, return_constraints=False):
-    req_path = download_artifacts(run_id=run_id, artifact_path=f"{artifact_path}/requirements.txt")
+def get_pip_requirements(artifact_uri, return_constraints=False):
+    req_path = download_artifacts(artifact_uri=f"{artifact_uri}/requirements.txt")
     reqs = read_lines(req_path)
 
     if return_constraints:
-        con_path = download_artifacts(
-            run_id=run_id, artifact_path=f"{artifact_path}/constraints.txt"
-        )
+        con_path = download_artifacts(artifact_uri=f"{artifact_uri}/constraints.txt")
         cons = read_lines(con_path)
         return set(reqs), set(cons)
 
@@ -43,30 +41,28 @@ def main():
     xgb_req = f"xgboost=={xgb.__version__}"
     sklearn_req = f"scikit-learn=={sklearn.__version__}"
 
-    with mlflow.start_run() as run:
-        run_id = run.info.run_id
-
+    with mlflow.start_run():
         # Default (both `pip_requirements` and `extra_pip_requirements` are unspecified)
         artifact_path = "default"
-        mlflow.xgboost.log_model(model, artifact_path, signature=signature)
-        pip_reqs = get_pip_requirements(run_id, artifact_path)
+        model_info = mlflow.xgboost.log_model(model, artifact_path, signature=signature)
+        pip_reqs = get_pip_requirements(model_info.model_uri)
         assert xgb_req in pip_reqs, pip_reqs
 
         # Overwrite the default set of pip requirements using `pip_requirements`
         artifact_path = "pip_requirements"
-        mlflow.xgboost.log_model(
+        model_info = mlflow.xgboost.log_model(
             model, artifact_path, pip_requirements=[sklearn_req], signature=signature
         )
-        pip_reqs = get_pip_requirements(run_id, artifact_path)
+        pip_reqs = get_pip_requirements(model_info.model_uri)
         assert sklearn_req in pip_reqs, pip_reqs
 
         # Add extra pip requirements on top of the default set of pip requirements
         # using `extra_pip_requirements`
         artifact_path = "extra_pip_requirements"
-        mlflow.xgboost.log_model(
+        model_info = mlflow.xgboost.log_model(
             model, artifact_path, extra_pip_requirements=[sklearn_req], signature=signature
         )
-        pip_reqs = get_pip_requirements(run_id, artifact_path)
+        pip_reqs = get_pip_requirements(model_info.model_uri)
         assert pip_reqs.issuperset({xgb_req, sklearn_req}), pip_reqs
 
         # Specify pip requirements using a requirements file
@@ -76,21 +72,21 @@ def main():
 
             # Path to a pip requirements file
             artifact_path = "requirements_file_path"
-            mlflow.xgboost.log_model(
+            model_info = mlflow.xgboost.log_model(
                 model, artifact_path, pip_requirements=f.name, signature=signature
             )
-            pip_reqs = get_pip_requirements(run_id, artifact_path)
+            pip_reqs = get_pip_requirements(model_info.model_uri)
             assert sklearn_req in pip_reqs, pip_reqs
 
             # List of pip requirement strings
             artifact_path = "requirements_file_list"
-            mlflow.xgboost.log_model(
+            model_info = mlflow.xgboost.log_model(
                 model,
                 artifact_path,
                 pip_requirements=[xgb_req, f"-r {f.name}"],
                 signature=signature,
             )
-            pip_reqs = get_pip_requirements(run_id, artifact_path)
+            pip_reqs = get_pip_requirements(model_info.model_uri)
             assert pip_reqs.issuperset({xgb_req, sklearn_req}), pip_reqs
 
         # Using a constraints file
@@ -99,15 +95,13 @@ def main():
             f.flush()
 
             artifact_path = "constraints_file"
-            mlflow.xgboost.log_model(
+            model_info = mlflow.xgboost.log_model(
                 model,
                 artifact_path,
                 pip_requirements=[xgb_req, f"-c {f.name}"],
                 signature=signature,
             )
-            pip_reqs, pip_cons = get_pip_requirements(
-                run_id, artifact_path, return_constraints=True
-            )
+            pip_reqs, pip_cons = get_pip_requirements(model_info.model_uri, return_constraints=True)
             assert pip_reqs.issuperset({xgb_req, "-c constraints.txt"}), pip_reqs
             assert pip_cons == {sklearn_req}, pip_cons
 

examples/pmdarima/train.py

@@ -48,16 +48,15 @@ def calculate_cv_metrics(model, endog, metric, cv):
     predictions = arima.predict(n_periods=30, return_conf_int=False)
     signature = infer_signature(train, predictions)
 
-    mlflow.pmdarima.log_model(
+    model_info = mlflow.pmdarima.log_model(
         pmdarima_model=arima, artifact_path=ARTIFACT_PATH, signature=signature
     )
     mlflow.log_params(parameters)
     mlflow.log_metrics(metrics)
-    model_uri = mlflow.get_artifact_uri(ARTIFACT_PATH)
 
-    print(f"Model artifact logged to: {model_uri}")
+    print(f"Model artifact logged to: {model_info.model_uri}")
 
-loaded_model = mlflow.pmdarima.load_model(model_uri)
+loaded_model = mlflow.pmdarima.load_model(model_info.model_uri)
 
 forecast = loaded_model.predict(30)