code: def test_mlflow_save_load(URI: Path, tmp_path: Path):
      tracking_db = tmp_path / "mlruns.db"
      mlflow.set_tracking_uri(f"sqlite:///{tracking_db}")
      client = mlflow.tracking.MlflowClient()
      mv = mlflow.register_model(str(URI), "IrisClf")
      client.transition_model_version_stage(
          name="IrisClf",
          version=mv.version,
          stage="Staging",
      )
      bento_model = bentoml.mlflow.import_model(
          MODEL_NAME, str(URI.resolve())
      )
      # make sure the model can be imported with models:/
      model_uri = bentoml.mlflow.import_model(
          MODEL_NAME, "models:/IrisClf/Staging"
      )
  
      pyfunc = bentoml.mlflow.load_model(bento_model.tag)
      np.testing.assert_array_equal(pyfunc.predict(X), res)
      np.testing.assert_array_equal(
          bentoml.mlflow.load_model(model_uri.tag).predict(X), res
      )
  
imports: from __future__ import annotations
  import os
  import typing
  from pathlib import Path
  from typing import TYPE_CHECKING
  import mlflow
  import mlflow.models
  import mlflow.sklearn
  import mlflow.tracking
  import numpy
  import psutil
  import pytest
  from sklearn.datasets import load_iris
  from sklearn.neighbors import KNeighborsClassifier
  import bentoml
  from bentoml._internal.models.model import ModelContext
  from bentoml.exceptions import BentoMLException
  from bentoml.exceptions import NotFound
label_desc: Use image node_id 2 for calling a global function with example usage: test_mlflow_save_load(URI, tmp_path) without return types