def register_model(model_uri, name):
"""
Create a new model version in model registry for the model files specified by ``model_uri``.
Note that this method assumes the model registry backend URI is the same as that of the
tracking backend.
:param model_uri: URI referring to the MLmodel directory. Use a ``runs:/`` URI if you want to
record the run ID with the model in model registry. ``models:/`` URIs are
currently not supported.
:param name: Name of the registered model under which to create a new model version. If a
registered model with the given name does not exist, it will be created
automatically.
:return: Single :py:class:`mlflow.entities.model_registry.ModelVersion` object created by
backend.
"""
client = MlflowClient()
try:
client.create_registered_model(name)
except MlflowException as e:
if e.error_code == ErrorCode.Name(RESOURCE_ALREADY_EXISTS):
eprint(
"Registered model %s already exists. Using it to create a new version."
% name)
else:
raise e
if RunsArtifactRepository.is_runs_uri(model_uri):
source = RunsArtifactRepository.get_underlying_uri(model_uri)
(run_id, _) = RunsArtifactRepository.parse_runs_uri(model_uri)
return client.create_model_version(name, source, run_id)
else:
return client.create_model_version(name, source=model_uri, run_id=None)
def register(run_id: str, experiment_id: int, model_name: str, block: bool):
client = MlflowClient(tracking_uri="databricks")
if not any([m.name == model_name for m in client.list_registered_models()]):
client.create_registered_model(model_name)
source = f"dbfs:/databricks/mlflow-tracking/{experiment_id}/{run_id}/artifacts/model"
model_details = client.create_model_version(model_name, source, run_id)
if block:
def wait_until_ready() -> bool:
for _ in range(60):
model_version_details = client.get_model_version(
name=model_name,
version=model_details.version,
)
status = ModelVersionStatus.from_string(model_version_details.status)
print("Model status: %s" % ModelVersionStatus.to_string(status))
if status == ModelVersionStatus.READY:
return True
time.sleep(5)
return False
if not wait_until_ready():
print(f"Timeout waiting on registration of model. Will not stage model.")
return
client.transition_model_version_stage(model_name, model_details.version, "Staging")
def test_client_can_be_serialized_with_pickle(tmpdir):
"""
Verifies that instances of `MlflowClient` can be serialized using pickle, even if the underlying
Tracking and Model Registry stores used by the client are not serializable using pickle
"""
class MockUnpickleableTrackingStore(SqlAlchemyTrackingStore):
pass
class MockUnpickleableModelRegistryStore(SqlAlchemyModelRegistryStore):
pass
backend_store_path = tmpdir.join("test.db").strpath
artifact_store_path = tmpdir.join("artfiacts").strpath
mock_tracking_store = MockUnpickleableTrackingStore(
"sqlite:///" + backend_store_path, artifact_store_path
)
mock_model_registry_store = MockUnpickleableModelRegistryStore(
"sqlite:///" + backend_store_path
)
# Verify that the mock stores cannot be pickled because they are defined within a function
# (i.e. the test function)
with pytest.raises(AttributeError, match="<locals>.MockUnpickleableTrackingStore'"):
pickle.dumps(mock_tracking_store)
with pytest.raises(AttributeError, match="<locals>.MockUnpickleableModelRegistryStore'"):
pickle.dumps(mock_model_registry_store)
_tracking_store_registry.register("pickle", lambda *args, **kwargs: mock_tracking_store)
_model_registry_store_registry.register(
"pickle", lambda *args, **kwargs: mock_model_registry_store
)
# Create an MlflowClient with the store that cannot be pickled, perform
# tracking & model registry operations, and verify that the client can still be pickled
client = MlflowClient("pickle://foo")
client.create_experiment("test_experiment")
client.create_registered_model("test_model")
pickle.dumps(client)
def register_model(model_uri,
name,
await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS):
"""
Create a new model version in model registry for the model files specified by ``model_uri``.
Note that this method assumes the model registry backend URI is the same as that of the
tracking backend.
:param model_uri: URI referring to the MLmodel directory. Use a ``runs:/`` URI if you want to
record the run ID with the model in model registry. ``models:/`` URIs are
currently not supported.
:param name: Name of the registered model under which to create a new model version. If a
registered model with the given name does not exist, it will be created
automatically.
:param await_registration_for: Number of seconds to wait for the model version to finish
being created and is in ``READY`` status. By default, the function
waits for five minutes. Specify 0 or None to skip waiting.
:return: Single :py:class:`mlflow.entities.model_registry.ModelVersion` object created by
backend.
"""
client = MlflowClient()
try:
create_model_response = client.create_registered_model(name)
eprint("Successfully registered model '%s'." %
create_model_response.name)
except MlflowException as e:
if e.error_code == ErrorCode.Name(RESOURCE_ALREADY_EXISTS):
eprint(
"Registered model '%s' already exists. Creating a new version of this model..."
% name)
else:
raise e
if RunsArtifactRepository.is_runs_uri(model_uri):
source = RunsArtifactRepository.get_underlying_uri(model_uri)
(run_id, _) = RunsArtifactRepository.parse_runs_uri(model_uri)
create_version_response = client.create_model_version(
name, source, run_id)
else:
create_version_response = client.create_model_version(
name,
source=model_uri,
run_id=None,
await_creation_for=await_registration_for)
eprint("Created version '{version}' of model '{model_name}'.".format(
version=create_version_response.version,
model_name=create_version_response.name))
return create_version_response
def test_client_registry_operations_raise_exception_with_unsupported_registry_store():
"""
This test case ensures that Model Registry operations invoked on the `MlflowClient`
fail with an informative error message when the registry store URI refers to a
store that does not support Model Registry features (e.g., FileStore).
"""
with TempDir() as tmp:
client = MlflowClient(registry_uri=tmp.path())
expected_failure_functions = [
client._get_registry_client,
lambda: client.create_registered_model("test"),
lambda: client.get_registered_model("test"),
lambda: client.create_model_version("test", "source", "run_id"),
lambda: client.get_model_version("test", 1),
]
for func in expected_failure_functions:
with pytest.raises(MlflowException) as exc:
func()
assert exc.value.error_code == ErrorCode.Name(FEATURE_DISABLED)
def register_model(model_uri,
name,
await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS):
"""
Create a new model version in model registry for the model files specified by ``model_uri``.
Note that this method assumes the model registry backend URI is the same as that of the
tracking backend.
:param model_uri: URI referring to the MLmodel directory. Use a ``runs:/`` URI if you want to
record the run ID with the model in model registry. ``models:/`` URIs are
currently not supported.
:param name: Name of the registered model under which to create a new model version. If a
registered model with the given name does not exist, it will be created
automatically.
:param await_registration_for: Number of seconds to wait for the model version to finish
being created and is in ``READY`` status. By default, the function
waits for five minutes. Specify 0 or None to skip waiting.
:return: Single :py:class:`mlflow.entities.model_registry.ModelVersion` object created by
backend.
.. code-block:: python
:caption: Example
import mlflow.sklearn
from sklearn.ensemble import RandomForestRegressor
mlflow.set_tracking_uri("sqlite:////tmp/mlruns.db")
params = {"n_estimators": 3, "random_state": 42}
# Log MLflow entities
with mlflow.start_run() as run:
rfr = RandomForestRegressor(**params).fit([[0, 1]], [1])
mlflow.log_params(params)
mlflow.sklearn.log_model(rfr, artifact_path="sklearn-model")
model_uri = "runs:/{}/sklearn-model".format(run.info.run_id)
mv = mlflow.register_model(model_uri, "RandomForestRegressionModel")
print("Name: {}".format(mv.name))
print("Version: {}".format(mv.version))
.. code-block:: text
:caption: Output
Name: RandomForestRegressionModel
Version: 1
"""
client = MlflowClient()
try:
create_model_response = client.create_registered_model(name)
eprint("Successfully registered model '%s'." %
create_model_response.name)
except MlflowException as e:
if e.error_code == ErrorCode.Name(RESOURCE_ALREADY_EXISTS):
eprint(
"Registered model '%s' already exists. Creating a new version of this model..."
% name)
else:
raise e
if RunsArtifactRepository.is_runs_uri(model_uri):
source = RunsArtifactRepository.get_underlying_uri(model_uri)
(run_id, _) = RunsArtifactRepository.parse_runs_uri(model_uri)
create_version_response = client.create_model_version(
name, source, run_id, await_creation_for=await_registration_for)
else:
create_version_response = client.create_model_version(
name,
source=model_uri,
run_id=None,
await_creation_for=await_registration_for)
eprint("Created version '{version}' of model '{model_name}'.".format(
version=create_version_response.version,
model_name=create_version_response.name))
return create_version_response
# MAGIC %md Register model
# MAGIC
# MAGIC https://www.mlflow.org/docs/latest/model-registry.html#registering-a-model
# MAGIC https://docs.microsoft.com/fr-fr/azure/databricks/applications/machine-learning/manage-model-lifecycle/
# COMMAND ----------
result = mlflow.register_model("runs:<model-path>", "<model-name>")
# COMMAND ----------
from mlflow.tracking import MlflowClient
client = MlflowClient()
name = "spark-lr-registered-model"
client.create_registered_model(name)
desc = "A new version of the model"
model_uri = "runs:/{}/sklearn-model".format(run.info.run_id)
mv = client.create_model_version(name,
model_uri,
run.info.run_id,
description=desc)
# COMMAND ----------
# MAGIC %md mlflow.pyfunc
# COMMAND ----------
import mlflow.pyfunc
import mlflow
from mlflow.tracking import MlflowClient
if __name__ == "__main__":
def print_model_info(models):
for m in models:
print("--")
print("name: {}".format(m.name))
print("tags: {}".format(m.tags))
print("description: {}".format(m.description))
mlflow.set_tracking_uri("sqlite:///mlruns.db")
client = MlflowClient()
# Register a couple of models with respective names, tags, and descriptions
for name, tags, desc in [("name1", {
"t1": "t1"
}, 'description1'), ("name2", {
"t2": "t2"
}, 'description2')]:
client.create_registered_model(name, tags, desc)
# Fetch all registered models
print_model_info(client.list_registered_models())
# COMMAND ----------
client = MlflowClient()
# Register the model with the run URI and unique name
#model_uri = "runs:/{}/model".format(runId)
#model_registry_id = spark.conf.get("com.databricks.tmp.uniqueid") # ensures we have a uniqueID for each student
#model_registry_name = "bank_singlenode_{}".format(model_registry_id)
#model_details = mlflow.register_model(model_uri=model_uri, name=model_registry_name)
# COMMAND ----------
## Create ModelRegistry (only have to do this once) TODO: Can probably delete this code
client = MlflowClient()
client.create_registered_model(modelRegistryName)
# COMMAND ----------
# MAGIC %md
# MAGIC ### Register the model associated to the runID to our model Registry
# COMMAND ----------
mlflow.register_model("runs:/" + runId + "/model", modelRegistryName)
# COMMAND ----------
# DBTITLE 1,Push old model(s) to archive stage if it exists
model_list = client.search_model_versions("name='%s'" % modelRegistryName)
# MAGIC %md ##### Call register_model() using the remote client, using the new DBFS location as “source”.
# MAGIC Note that the "source" URI, which is part of the metadata for the newly created model version, may contain the original Run ID if you used the default artifact store location when you logged the model from the Run.
# COMMAND ----------
### USER INPUT
model_name = "my_remote_model"
# COMMAND ----------
import posixpath
source = posixpath.join(
artifact_uri, artifact_path
) # we preserved the model artifact dbfs path from the source workspace
try:
remote_client.create_registered_model(model_name)
except Exception as e:
if e.error_code == 'RESOURCE_ALREADY_EXISTS':
print(e)
else:
throw(e)
mv = remote_client.create_model_version(
model_name, source,
run_id) # `source` must point to the DBFS location in the new workspace
print(mv)
# COMMAND ----------
# MAGIC %md At this point, if you log into the registry workspace you should see the new model version.
# COMMAND ----------
mlflow.set_tracking_uri("sqlite:///mlruns.db")
with mlflow.start_run(run_name="YOUR_RUN_NAME") as run:
params = {"n_estimators": 5, "random_state": 42}
sk_learn_rfr = RandomForestRegressor(**params)
# Log parameters and metrics using the MLflow APIs
mlflow.log_params(params)
mlflow.log_param("param_1", randint(0, 100))
mlflow.log_metrics({"metric_1": random(), "metric_2": random()})
# log the sklearn model and register as version 1
mlflow.sklearn.log_model(sk_model=sk_learn_rfr,
artifact_path="sklearn-model")
client = MlflowClient()
client.create_registered_model("sk-learn-random_forest-reg-model")
client.create_model_version(name="sk-learn-random-forest-reg-model",
source="",
run_id="")
# Fetching an mlflow model from the model registry
import mlflow.pyfunc
model_name = "sk-learn-random-forest-reg-model"
model_version = 1
data = ""
model = mlflow.pyfunc.load_model(
import mlflow
from mlflow.tracking import MlflowClient
from sklearn.linear_model import LogisticRegression, LinearRegression
if __name__ == "__main__":
mlflow.set_tracking_uri("sqlite:///mlruns.db")
client = MlflowClient()
client.create_registered_model("RegressionModels")
# Create two models under a same run
m = LinearRegression()
m2 = LogisticRegression()
with mlflow.start_run() as run:
mlflow.log_metric("m", 2.3)
mlflow.sklearn.log_model(m, "sk_learn_m")
mlflow.log_metric("m2", 2.5)
mlflow.sklearn.log_model(m2, "sk_learn_m2")
# Register different versions in the register under the same run
mlflow.register_model("runs:/{}/sk_learn_m".format(run.info.run_id),
"RegressionModels")
mlflow.register_model("runs:/{}/sk_learn_m2".format(run.info.run_id),
"RegressionModels")
#Save model to Model registry
model_name = "lr_trip_duration_model"
run_id = run.info.run_id #"c5218874277e4644a6536affee9b3ba0"
model_uri = f"runs:/{run_id}/model"
model_details = mlflow.register_model(model_uri, model_name)
# COMMAND ----------
#https://www.mlflow.org/docs/latest/model-registry.html
from mlflow.tracking import MlflowClient
client = MlflowClient()
client.create_registered_model("spark-lr-model")
#client.log_artifacts(run.info.run_id, "/FileStore/spark-model", artifact_path=mlflow.get_artifact_uri())
# COMMAND ----------
artifacts = [f.path for f in client.list_artifacts(run.info.run_id, "model")]
print("artifacts: {}".format(artifacts))
# COMMAND ----------
mlflow.get_artifact_uri()
# COMMAND ----------
# MAGIC %fs ls /databricks/mlflow-tracking/