def test_model_log(self):
old_uri = mlflow.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) as tmp:
try:
mlflow.set_tracking_uri("test")
if should_start_run:
mlflow.start_run()
artifact_path = "linear"
conda_env = os.path.join(tmp.path(), "conda_env.yaml")
_mlflow_conda_env(conda_env, additional_pip_deps=["sklearn"])
sklearn.log_model(sk_model=self._linear_lr,
artifact_path=artifact_path,
conda_env=conda_env)
x = sklearn.load_model(artifact_path, run_id=mlflow.active_run().info.run_uuid)
model_path = _get_model_log_dir(
artifact_path, mlflow.active_run().info.run_uuid)
model_config = Model.load(os.path.join(model_path, "MLmodel"))
assert pyfunc.FLAVOR_NAME in model_config.flavors
assert pyfunc.ENV in model_config.flavors[pyfunc.FLAVOR_NAME]
env_path = model_config.flavors[pyfunc.FLAVOR_NAME][pyfunc.ENV]
assert os.path.exists(os.path.join(model_path, env_path))
xpred = x.predict(self._X)
np.testing.assert_array_equal(self._linear_lr_predict, xpred)
finally:
mlflow.end_run()
mlflow.set_tracking_uri(old_uri)
def load_and_predict(self, run_id, parameters):
model = model_type.load_model(self.get_model_path(run_id))
model_results = model.predict(
pandas.DataFrame(data=parameters, index=[0]))
return model_results
def test_model_log(self):
with TempDir(chdr=True, remove_on_exit=True):
tracking.start_run()
try:
sklearn.log_model(sk_model=self._linear_lr,
artifact_path="linear")
x = sklearn.load_model(
"linear", run_id=tracking.active_run().info.run_uuid)
xpred = x.predict(self._X)
np.testing.assert_array_equal(self._linear_lr_predict, xpred)
finally:
tracking.end_run()
def test_model_save_load(self):
with TempDir(chdr=True, remove_on_exit=True) as tmp:
model_path = tmp.path("knn.pkl")
with open(model_path, "wb") as f:
pickle.dump(self._knn, f)
path = tmp.path("knn")
sklearn.save_model(self._knn, path=path)
x = sklearn.load_model(path)
xpred = x.predict(self._X)
np.testing.assert_array_equal(self._knn_predict, xpred)
# sklearn should also be stored as a valid pyfunc model
# test pyfunc compatibility
y = pyfunc.load_pyfunc(path)
ypred = y.predict(self._X)
np.testing.assert_array_equal(self._knn_predict, ypred)
def test_model_log(self):
old_uri = tracking.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) as tmp:
try:
tracking.set_tracking_uri("test")
if should_start_run:
tracking.start_run()
sklearn.log_model(sk_model=self._linear_lr, artifact_path="linear")
x = sklearn.load_model("linear", run_id=tracking.active_run().info.run_uuid)
xpred = x.predict(self._X)
np.testing.assert_array_equal(self._linear_lr_predict, xpred)
finally:
tracking.end_run()
tracking.set_tracking_uri(old_uri)
