# COMMAND ----------
import mlflow
from mlflow.tracking.client import MlflowClient
client = MlflowClient()
model_name = "linear-regression-model"
artifact_path = "best_model"
model_uri = "runs:/{run_id}/{artifact_path}".format(run_id=run_id, artifact_path=artifact_path)
registered_model = mlflow.register_model(model_uri=model_uri, name=model_name, await_registration_for=120)
#Add model and model version descriptions to Model Registry
client.update_model_version(
name=registered_model.name,
version=registered_model.version,
description="This predicts the age of a customer using transaction history."
)
#Transition a model version to Staging/Prod/Archived
client.transition_model_version_stage(
name=registered_model.name,
version=registered_model.version,
stage='Staging',
)
#Get model version details
model_version = client.get_model_version(
name=registered_model.name,
version=registered_model.version,
)
# COMMAND ----------
Utils.wait_until_ready(model_name, model_version)
# COMMAND ----------
# MAGIC %md ## Add a description to the new model version
# COMMAND ----------
from mlflow.tracking.client import MlflowClient
client = MlflowClient()
client.update_model_version(
name=model_name,
version=model_version,
description=
"This is the best model version of a random forest that was trained in scikit-learn."
)
# COMMAND ----------
# MAGIC %md ## Test the new model version in `Staging`
# MAGIC
# MAGIC Before deploying a model to a production application, it is often best practice to test it in a staging environment. The following cells transition the new model version to `Staging` and evaluate its performance.
# COMMAND ----------
client.transition_model_version_stage(
name=model_name,
version=model_version,
stage="Staging",
model_uri="runs:/{}/model".format(latest_model.run_id)
latest_sk_model = mlflow.sklearn.load_model(model_uri)
# COMMAND ----------
from mlflow.tracking.client import MlflowClient
client = MlflowClient()
client.update_registered_model(
name=model_name,
description="This model forecasts the wine quality based on the characteristics."
)
client.update_model_version(
name=model_name,
version=model_version,
description="This model version was built using sklearn."
)
# COMMAND ----------
client.transition_model_version_stage(
name=model_name,
version=model_version,
stage=stage,
archive_existing_versions=True
)
model_version_details = client.get_model_version(
name=model_name,
version=model_version,
client = MlflowClient()
client.update_registered_model(
name=model_details.name,
description="This model forecasts the power output of a wind farm based on weather data. The weather data consists of three features: wind speed, wind direction, and air temperature."
)
# COMMAND ----------
# MAGIC %md Add a model version description with information about the model architecture and machine learning framework.
# COMMAND ----------
client.update_model_version(
name=model_details.name,
version=model_details.version,
description="This model version was built using Keras. It is a feed-forward neural network with one hidden layer."
)
# COMMAND ----------
# MAGIC %md ### Perform a model stage transition
# MAGIC
# MAGIC The MLflow Model Registry defines several model stages: **None**, **Staging**, **Production**, and **Archived**. Each stage has a unique meaning. For example, **Staging** is meant for model testing, while **Production** is for models that have completed the testing or review processes and have been deployed to applications.
# MAGIC
# MAGIC Users with appropriate permissions can transition models between stages. In private preview, any user can transition a model to any stage. In the near future, administrators in your organization will be able to control these permissions on a per-user and per-model basis.
# MAGIC
# MAGIC If you have permission to transition a model to a particular stage, you can make the transition directly by using the `MlflowClient.update_model_version()` function. If you do not have permission, you can request a stage transition using the REST API; for example:
# MAGIC
# MAGIC ```
# MAGIC %sh curl -i -X POST -H "X-Databricks-Org-Id: <YOUR_ORG_ID>" -H "Authorization: Bearer <YOUR_ACCESS_TOKEN>" https://<YOUR_DATABRICKS_WORKSPACE_URL>/api/2.0/preview/mlflow/transition-requests/create -d '{"comment": "Please move this model into production!", "model_version": {"version": 1, "registered_model": {"name": "power-forecasting-model"}}, "stage": "Production"}'
