def test_log_model_calls_register_model(sklearn_knn_model,
main_scoped_model_class):
register_model_patch = mock.patch("mlflow.register_model")
with register_model_patch:
sklearn_artifact_path = "sk_model_no_run"
with kiwi.start_run():
kiwi.sklearn.log_model(sk_model=sklearn_knn_model,
artifact_path=sklearn_artifact_path)
sklearn_model_uri = "runs:/{run_id}/{artifact_path}".format(
run_id=kiwi.active_run().info.run_id,
artifact_path=sklearn_artifact_path)
def test_predict(sk_model, model_input):
return sk_model.predict(model_input) * 2
pyfunc_artifact_path = "pyfunc_model"
assert kiwi.active_run() is None
kiwi.pyfunc.log_model(
artifact_path=pyfunc_artifact_path,
artifacts={"sk_model": sklearn_model_uri},
python_model=main_scoped_model_class(test_predict),
registered_model_name="AdsModel1")
model_uri = "runs:/{run_id}/{artifact_path}".format(
run_id=kiwi.active_run().info.run_id,
artifact_path=pyfunc_artifact_path)
kiwi.register_model.assert_called_once_with(model_uri, "AdsModel1")
kiwi.end_run()
def test_pyfunc_model_log_load_no_active_run(sklearn_knn_model,
main_scoped_model_class,
iris_data):
sklearn_artifact_path = "sk_model_no_run"
with kiwi.start_run():
kiwi.sklearn.log_model(sk_model=sklearn_knn_model,
artifact_path=sklearn_artifact_path)
sklearn_model_uri = "runs:/{run_id}/{artifact_path}".format(
run_id=kiwi.active_run().info.run_id,
artifact_path=sklearn_artifact_path)
def test_predict(sk_model, model_input):
return sk_model.predict(model_input) * 2
pyfunc_artifact_path = "pyfunc_model"
assert kiwi.active_run() is None
kiwi.pyfunc.log_model(artifact_path=pyfunc_artifact_path,
artifacts={"sk_model": sklearn_model_uri},
python_model=main_scoped_model_class(test_predict))
pyfunc_model_uri = "runs:/{run_id}/{artifact_path}".format(
run_id=kiwi.active_run().info.run_id,
artifact_path=pyfunc_artifact_path)
loaded_pyfunc_model = kiwi.pyfunc.load_pyfunc(model_uri=pyfunc_model_uri)
np.testing.assert_array_equal(
loaded_pyfunc_model.predict(iris_data[0]),
test_predict(sk_model=sklearn_knn_model, model_input=iris_data[0]))
kiwi.end_run()
def test_model_log(h2o_iris_model):
h2o_model = h2o_iris_model.model
old_uri = kiwi.get_tracking_uri()
# should_start_run tests whether or not calling log_model() automatically starts a run.
for should_start_run in [False, True]:
with TempDir(chdr=True, remove_on_exit=True):
try:
artifact_path = "gbm_model"
kiwi.set_tracking_uri("test")
if should_start_run:
kiwi.start_run()
kiwi.h2o.log_model(h2o_model=h2o_model,
artifact_path=artifact_path)
model_uri = "runs:/{run_id}/{artifact_path}".format(
run_id=kiwi.active_run().info.run_id,
artifact_path=artifact_path)
# Load model
h2o_model_loaded = kiwi.h2o.load_model(model_uri=model_uri)
assert all(
h2o_model_loaded.predict(h2o_iris_model.inference_data).
as_data_frame() == h2o_model.predict(
h2o_iris_model.inference_data).as_data_frame())
finally:
kiwi.end_run()
kiwi.set_tracking_uri(old_uri)
def test_sparkml_estimator_model_log(tmpdir, spark_model_estimator):
# Print the coefficients and intercept for multinomial logistic regression
old_tracking_uri = kiwi.get_tracking_uri()
cnt = 0
# should_start_run tests whether or not calling log_model() automatically starts a run.
for should_start_run in [False, True]:
for dfs_tmp_dir in [None, os.path.join(str(tmpdir), "test")]:
print("should_start_run =", should_start_run, "dfs_tmp_dir =", dfs_tmp_dir)
try:
tracking_dir = os.path.abspath(str(tmpdir.join("mlruns")))
kiwi.set_tracking_uri("file://%s" % tracking_dir)
if should_start_run:
kiwi.start_run()
artifact_path = "model%d" % cnt
cnt += 1
sparkm.log_model(
artifact_path=artifact_path,
spark_model=spark_model_estimator.model,
dfs_tmpdir=dfs_tmp_dir)
model_uri = "runs:/{run_id}/{artifact_path}".format(
run_id=kiwi.active_run().info.run_id,
artifact_path=artifact_path)
# test reloaded model
reloaded_model = sparkm.load_model(model_uri=model_uri, dfs_tmpdir=dfs_tmp_dir)
preds_df = reloaded_model.transform(spark_model_estimator.spark_df)
preds = [x.prediction for x in preds_df.select("prediction").collect()]
assert spark_model_estimator.predictions == preds
finally:
kiwi.end_run()
kiwi.set_tracking_uri(old_tracking_uri)
x = dfs_tmp_dir or sparkm.DFS_TMP
shutil.rmtree(x)
shutil.rmtree(tracking_dir)
def test_autologging_does_not_start_run(spark_session, format_to_file_path):
try:
kiwi.spark.autolog()
data_format = list(format_to_file_path.keys())[0]
file_path = format_to_file_path[data_format]
df = spark_session.read.format(data_format).option("header", "true"). \
option("inferSchema", "true").load(file_path)
df.collect()
time.sleep(1)
active_run = kiwi.active_run()
assert active_run is None
assert len(kiwi.search_runs()) == 0
finally:
kiwi.end_run()
def test_delete_tag():
"""
Confirm that fluent API delete tags actually works
:return:
"""
kiwi.set_tag('a', 'b')
run = MlflowClient().get_run(kiwi.active_run().info.run_id)
print(run.info.run_id)
assert 'a' in run.data.tags
kiwi.delete_tag('a')
run = MlflowClient().get_run(kiwi.active_run().info.run_id)
assert 'a' not in run.data.tags
with pytest.raises(MlflowException):
kiwi.delete_tag('a')
with pytest.raises(MlflowException):
kiwi.delete_tag('b')
kiwi.end_run()
def fit(self, *args, **kwargs):
if not kiwi.active_run():
auto_end_run = True
else:
auto_end_run = False
original = gorilla.get_original_attribute(Estimator, "fit")
if len(args) >= 4:
l = list(args)
l[3] += [__MLflowGluonCallback()]
args = tuple(l)
elif "event_handlers" in kwargs:
kwargs["event_handlers"] += [__MLflowGluonCallback()]
else:
kwargs["event_handlers"] = [__MLflowGluonCallback()]
result = original(self, *args, **kwargs)
if auto_end_run:
kiwi.end_run()
return result
def test_model_log_evaluate_pyfunc_format(onnx_model, data, predicted):
import onnx
import kiwi.onnx
x, y = data
# should_start_run tests whether or not calling log_model() automatically starts a run.
for should_start_run in [False, True]:
try:
if should_start_run:
kiwi.start_run()
artifact_path = "onnx_model"
kiwi.onnx.log_model(onnx_model=onnx_model, artifact_path=artifact_path)
model_uri = "runs:/{run_id}/{artifact_path}".format(
run_id=kiwi.active_run().info.run_id,
artifact_path=artifact_path)
# Loading pyfunc model
pyfunc_loaded = kiwi.pyfunc.load_pyfunc(model_uri=model_uri)
assert np.allclose(pyfunc_loaded.predict(x).values, predicted,
rtol=1e-05, atol=1e-05)
finally:
kiwi.end_run()
def test_model_log(model, data, predicted):
x, _ = data
# should_start_run tests whether or not calling log_model() automatically starts a run.
for should_start_run in [False, True]:
try:
if should_start_run:
kiwi.start_run()
artifact_path = "keras_model"
kiwi.keras.log_model(model, artifact_path=artifact_path)
model_uri = "runs:/{run_id}/{artifact_path}".format(
run_id=kiwi.active_run().info.run_id,
artifact_path=artifact_path)
# Load model
model_loaded = kiwi.keras.load_model(model_uri=model_uri)
assert all(model_loaded.predict(x) == predicted)
# Loading pyfunc model
pyfunc_loaded = kiwi.pyfunc.load_model(model_uri=model_uri)
assert all(pyfunc_loaded.predict(x).values == predicted)
finally:
kiwi.end_run()
def test_model_log(onnx_model, onnx_custom_env):
# pylint: disable=unused-argument
import onnx
import kiwi.onnx
# should_start_run tests whether or not calling log_model() automatically starts a run.
for should_start_run in [False, True]:
try:
if should_start_run:
kiwi.start_run()
artifact_path = "onnx_model"
kiwi.onnx.log_model(onnx_model=onnx_model,
artifact_path=artifact_path,
conda_env=onnx_custom_env)
model_uri = "runs:/{run_id}/{artifact_path}".format(
run_id=kiwi.active_run().info.run_id,
artifact_path=artifact_path)
# Load model
onnx.checker.check_model = mock.Mock()
kiwi.onnx.load_model(model_uri)
assert onnx.checker.check_model.called
finally:
kiwi.end_run()
def test_model_log_load_no_active_run(sklearn_knn_model, iris_data, tmpdir):
sk_model_path = os.path.join(str(tmpdir), "knn.pkl")
with open(sk_model_path, "wb") as f:
pickle.dump(sklearn_knn_model, f)
pyfunc_artifact_path = "pyfunc_model"
assert kiwi.active_run() is None
kiwi.pyfunc.log_model(artifact_path=pyfunc_artifact_path,
data_path=sk_model_path,
loader_module=os.path.basename(__file__)[:-3],
code_path=[__file__])
pyfunc_model_path = _download_artifact_from_uri(
"runs:/{run_id}/{artifact_path}".format(
run_id=kiwi.active_run().info.run_id,
artifact_path=pyfunc_artifact_path))
model_config = Model.load(os.path.join(pyfunc_model_path, "MLmodel"))
assert kiwi.pyfunc.FLAVOR_NAME in model_config.flavors
assert kiwi.pyfunc.PY_VERSION in model_config.flavors[
kiwi.pyfunc.FLAVOR_NAME]
reloaded_model = kiwi.pyfunc.load_pyfunc(pyfunc_model_path)
np.testing.assert_array_equal(sklearn_knn_model.predict(iris_data[0]),
reloaded_model.predict(iris_data[0]))
kiwi.end_run()
def manual_run(request):
if request.param:
kiwi.start_run()
yield
kiwi.end_run()
def start_run():
kiwi.start_run()
yield
kiwi.end_run()
def manual_run(request, tracking_uri_mock):
if request.param:
kiwi.start_run()
yield
kiwi.end_run()