def handle_model_uri(model_uri, service_name):
"""
Handle the various types of model uris we could receive.
:param model_uri:
:type model_uri: str
:param service_name:
:type service_name: str
:return:
:rtype:
"""
client = MlflowClient()
if model_uri.startswith("models:/"):
model_name = model_uri.split("/")[-2]
model_stage_or_version = model_uri.split("/")[-1]
if model_stage_or_version in client.get_model_version_stages(None, None):
# TODO: Add exception handling for no models found with specified stage
model_version = client.get_latest_versions(model_name, [model_stage_or_version])[0].version
else:
model_version = model_stage_or_version
elif (model_uri.startswith("runs:/") or model_uri.startswith("file://")) \
and get_tracking_uri().startswith("azureml") and get_registry_uri().startswith("azureml"):
# We will register the model for the user
model_name = service_name + "-model"
mlflow_model = mlflow_register_model(model_uri, model_name)
model_version = mlflow_model.version
_logger.info(
"Registered an Azure Model with name: `%s` and version: `%s`",
mlflow_model.name,
mlflow_model.version,
)
else:
raise MlflowException("Unsupported model uri provided, or tracking or registry uris are not set to "
"an AzureML uri.")
return model_name, model_version
import mlflow.sklearn
from mlflow.tracking import MlflowClient
from sklearn.ensemble import RandomForestRegressor
if __name__ == "__main__":
mlflow.set_tracking_uri("sqlite:///mlruns.db")
params = {"n_estimators": 3, "random_state": 42}
name = "RandomForestRegression"
rfr = RandomForestRegressor(**params).fit([[0, 1]], [1])
# Log MLflow entities
with mlflow.start_run() as run:
mlflow.log_params(params)
mlflow.sklearn.log_model(rfr, artifact_path="models/sklearn-model")
# Register model name in the model registry
client = MlflowClient()
client.create_registered_model(name)
# Create a new version of the rfr model under the registered model name
model_uri = "runs:/{}/models/sklearn-model".format(run.info.run_id)
mv = client.create_model_version(name, model_uri, run.info.run_id)
stages = client.get_model_version_stages(name, mv.version)
print("Model list of valid stages: {}".format(stages))
