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)