def test_sparkml_model_log_without_specified_conda_env_uses_default_env_with_expected_dependencies( spark_model_iris): artifact_path = "model" with kiwi.start_run(): sparkm.log_model( spark_model=spark_model_iris.model, artifact_path=artifact_path, conda_env=None) model_uri = "runs:/{run_id}/{artifact_path}".format( run_id=kiwi.active_run().info.run_id, artifact_path=artifact_path) model_path = _download_artifact_from_uri(artifact_uri=model_uri) pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME) conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]) with open(conda_env_path, "r") as f: conda_env = yaml.safe_load(f) assert conda_env == sparkm.get_default_conda_env()
def test_model_save_persists_specified_conda_env_in_mlflow_model_directory( spacy_model_with_data, model_path, spacy_custom_env): kiwi.spacy.save_model(spacy_model=spacy_model_with_data.model, path=model_path, conda_env=spacy_custom_env) pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME) saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]) assert os.path.exists(saved_conda_env_path) assert saved_conda_env_path != spacy_custom_env with open(spacy_custom_env, "r") as f: spacy_custom_env_text = f.read() with open(saved_conda_env_path, "r") as f: saved_conda_env_text = f.read() assert saved_conda_env_text == spacy_custom_env_text
def test_model_save_persists_specified_conda_env_in_mlflow_model_directory( sklearn_knn_model, model_path, sklearn_custom_env): kiwi.sklearn.save_model(sk_model=sklearn_knn_model.model, path=model_path, conda_env=sklearn_custom_env) pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME) saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]) assert os.path.exists(saved_conda_env_path) assert saved_conda_env_path != sklearn_custom_env with open(sklearn_custom_env, "r") as f: sklearn_custom_env_parsed = yaml.safe_load(f) with open(saved_conda_env_path, "r") as f: saved_conda_env_parsed = yaml.safe_load(f) assert saved_conda_env_parsed == sklearn_custom_env_parsed
def test_save_model_accepts_conda_env_as_dict(saved_tf_iris_model, model_path): conda_env = dict(kiwi.tensorflow.get_default_conda_env()) conda_env["dependencies"].append("pytest") kiwi.tensorflow.save_model( tf_saved_model_dir=saved_tf_iris_model.path, tf_meta_graph_tags=saved_tf_iris_model.meta_graph_tags, tf_signature_def_key=saved_tf_iris_model.signature_def_key, path=model_path, conda_env=conda_env) pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME) saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]) assert os.path.exists(saved_conda_env_path) with open(saved_conda_env_path, "r") as f: saved_conda_env_parsed = yaml.safe_load(f) assert saved_conda_env_parsed == conda_env
def test_save_model_persists_specified_conda_env_in_mlflow_model_directory( saved_tf_iris_model, model_path, tf_custom_env): kiwi.tensorflow.save_model( tf_saved_model_dir=saved_tf_iris_model.path, tf_meta_graph_tags=saved_tf_iris_model.meta_graph_tags, tf_signature_def_key=saved_tf_iris_model.signature_def_key, path=model_path, conda_env=tf_custom_env) pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME) saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]) assert os.path.exists(saved_conda_env_path) assert saved_conda_env_path != tf_custom_env with open(tf_custom_env, "r") as f: tf_custom_env_text = f.read() with open(saved_conda_env_path, "r") as f: saved_conda_env_text = f.read() assert saved_conda_env_text == tf_custom_env_text
def test_model_log_persists_specified_conda_env_in_mlflow_model_directory( sequential_model, pytorch_custom_env): artifact_path = "model" with kiwi.start_run(): kiwi.pytorch.log_model(pytorch_model=sequential_model, artifact_path=artifact_path, conda_env=pytorch_custom_env) model_path = _download_artifact_from_uri("runs:/{run_id}/{artifact_path}".format( run_id=kiwi.active_run().info.run_id, artifact_path=artifact_path)) pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME) saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]) assert os.path.exists(saved_conda_env_path) assert saved_conda_env_path != pytorch_custom_env with open(pytorch_custom_env, "r") as f: pytorch_custom_env_text = f.read() with open(saved_conda_env_path, "r") as f: saved_conda_env_text = f.read() assert saved_conda_env_text == pytorch_custom_env_text
def _get_and_parse_flavor_configuration(model_path): """ :param path: Local filesystem path to the MLflow Model with the ``tensorflow`` flavor. :return: A triple containing the following elements: - ``tf_saved_model_dir``: The local filesystem path to the underlying TensorFlow SavedModel directory. - ``tf_meta_graph_tags``: A list of tags identifying the TensorFlow model's metagraph within the serialized ``SavedModel`` object. - ``tf_signature_def_key``: A string identifying the input/output signature associated with the model. This is a key within the serialized ``SavedModel``'s signature definition mapping. """ flavor_conf = _get_flavor_configuration(model_path=model_path, flavor_name=FLAVOR_NAME) tf_saved_model_dir = os.path.join(model_path, flavor_conf['saved_model_dir']) tf_meta_graph_tags = flavor_conf['meta_graph_tags'] tf_signature_def_key = flavor_conf['signature_def_key'] return tf_saved_model_dir, tf_meta_graph_tags, tf_signature_def_key
def test_model_log_persists_specified_conda_env_in_mlflow_model_directory(onnx_model, onnx_custom_env): import kiwi.onnx artifact_path = "model" with kiwi.start_run(): kiwi.onnx.log_model( onnx_model=onnx_model, artifact_path=artifact_path, conda_env=onnx_custom_env) model_path = _download_artifact_from_uri("runs:/{run_id}/{artifact_path}".format( run_id=kiwi.active_run().info.run_id, artifact_path=artifact_path)) pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME) saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]) assert os.path.exists(saved_conda_env_path) assert saved_conda_env_path != onnx_custom_env with open(onnx_custom_env, "r") as f: onnx_custom_env_parsed = yaml.safe_load(f) with open(saved_conda_env_path, "r") as f: saved_conda_env_parsed = yaml.safe_load(f) assert saved_conda_env_parsed == onnx_custom_env_parsed
def test_model_log_without_specified_conda_env_uses_default_env_with_expected_dependencies( sklearn_knn_model): artifact_path = "model" knn_model = sklearn_knn_model.model with kiwi.start_run(): kiwi.sklearn.log_model( sk_model=knn_model, artifact_path=artifact_path, conda_env=None, serialization_format=kiwi.sklearn.SERIALIZATION_FORMAT_PICKLE) model_uri = "runs:/{run_id}/{artifact_path}".format( run_id=kiwi.active_run().info.run_id, artifact_path=artifact_path) model_path = _download_artifact_from_uri(artifact_uri=model_uri) pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME) conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]) with open(conda_env_path, "r") as f: conda_env = yaml.safe_load(f) assert conda_env == kiwi.sklearn.get_default_conda_env()
def test_save_model_without_specified_conda_env_uses_default_env_with_expected_dependencies( sklearn_logreg_model, main_scoped_model_class, tmpdir): sklearn_model_path = os.path.join(str(tmpdir), "sklearn_model") kiwi.sklearn.save_model(sk_model=sklearn_logreg_model, path=sklearn_model_path) pyfunc_model_path = os.path.join(str(tmpdir), "pyfunc_model") kiwi.pyfunc.save_model( path=pyfunc_model_path, artifacts={"sk_model": sklearn_model_path}, python_model=main_scoped_model_class(predict_fn=None), conda_env=_conda_env()) pyfunc_conf = _get_flavor_configuration( model_path=pyfunc_model_path, flavor_name=kiwi.pyfunc.FLAVOR_NAME) conda_env_path = os.path.join(pyfunc_model_path, pyfunc_conf[kiwi.pyfunc.ENV]) with open(conda_env_path, "r") as f: conda_env = yaml.safe_load(f) assert conda_env == _conda_env()
def load_model(model_uri): """ Load an XGBoost model from a local file or a run. :param model_uri: The location, in URI format, of the MLflow model. For example: - ``/Users/me/path/to/local/model`` - ``relative/path/to/local/model`` - ``s3://my_bucket/path/to/model`` - ``runs:/<mlflow_run_id>/run-relative/path/to/model`` For more information about supported URI schemes, see `Referencing Artifacts <https://www.mlflow.org/docs/latest/tracking.html# artifact-locations>`_. :return: An XGBoost model (an instance of `xgboost.Booster`_) """ local_model_path = _download_artifact_from_uri(artifact_uri=model_uri) flavor_conf = _get_flavor_configuration(model_path=local_model_path, flavor_name=FLAVOR_NAME) xgb_model_file_path = os.path.join(local_model_path, flavor_conf.get("data", "model.xgb")) return _load_model(path=xgb_model_file_path)
def test_log_model_without_specified_conda_env_uses_default_env_with_expected_dependencies( saved_tf_iris_model): artifact_path = "model" with kiwi.start_run(): kiwi.tensorflow.log_model( tf_saved_model_dir=saved_tf_iris_model.path, tf_meta_graph_tags=saved_tf_iris_model.meta_graph_tags, tf_signature_def_key=saved_tf_iris_model.signature_def_key, artifact_path=artifact_path, conda_env=None) model_uri = "runs:/{run_id}/{artifact_path}".format( run_id=kiwi.active_run().info.run_id, artifact_path=artifact_path) model_path = _download_artifact_from_uri(artifact_uri=model_uri) pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME) conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]) with open(conda_env_path, "r") as f: conda_env = yaml.safe_load(f) assert conda_env == kiwi.tensorflow.get_default_conda_env()
def load_model(model_uri): """ Load a scikit-learn model from a local file or a run. :param model_uri: The location, in URI format, of the MLflow model, for example: - ``/Users/me/path/to/local/model`` - ``relative/path/to/local/model`` - ``s3://my_bucket/path/to/model`` - ``runs:/<mlflow_run_id>/run-relative/path/to/model`` - ``models:/<model_name>/<model_version>`` - ``models:/<model_name>/<stage>`` For more information about supported URI schemes, see `Referencing Artifacts <https://www.mlflow.org/docs/latest/concepts.html# artifact-locations>`_. :return: A scikit-learn model. .. code-block:: python :caption: Example import mlflow.sklearn sk_model = mlflow.sklearn.load_model("runs:/96771d893a5e46159d9f3b49bf9013e2/sk_models") # use Pandas DataFrame to make predictions pandas_df = ... predictions = sk_model.predict(pandas_df) """ local_model_path = _download_artifact_from_uri(artifact_uri=model_uri) flavor_conf = _get_flavor_configuration(model_path=local_model_path, flavor_name=FLAVOR_NAME) sklearn_model_artifacts_path = os.path.join(local_model_path, flavor_conf['pickled_model']) serialization_format = flavor_conf.get('serialization_format', SERIALIZATION_FORMAT_PICKLE) return _load_model_from_local_file( path=sklearn_model_artifacts_path, serialization_format=serialization_format)
def test_load_model_raises_exception_when_pickle_module_cannot_be_imported( main_scoped_subclassed_model, model_path): kiwi.pytorch.save_model( path=model_path, pytorch_model=main_scoped_subclassed_model, conda_env=None) bad_pickle_module_name = "not.a.real.module" pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME) model_data_path = os.path.join(model_path, pyfunc_conf[pyfunc.DATA]) assert os.path.exists(model_data_path) assert kiwi.pytorch._PICKLE_MODULE_INFO_FILE_NAME in os.listdir(model_data_path) with open( os.path.join(model_data_path, kiwi.pytorch._PICKLE_MODULE_INFO_FILE_NAME), "w") as f: f.write(bad_pickle_module_name) with pytest.raises(MlflowException) as exc_info: kiwi.pytorch.load_model(model_uri=model_path) assert "Failed to import the pickle module" in str(exc_info) assert bad_pickle_module_name in str(exc_info)
def load_model(model_uri, **kwargs): """ Load a Keras model from a local file or a run. Extra arguments are passed through to keras.load_model. :param model_uri: The location, in URI format, of the MLflow model. For example: - ``/Users/me/path/to/local/model`` - ``relative/path/to/local/model`` - ``s3://my_bucket/path/to/model`` - ``runs:/<mlflow_run_id>/run-relative/path/to/model`` - ``models:/<model_name>/<model_version>`` - ``models:/<model_name>/<stage>`` For more information about supported URI schemes, see `Referencing Artifacts <https://www.mlflow.org/docs/latest/concepts.html# artifact-locations>`_. :return: A Keras model instance. .. code-block:: python :caption: Example # Load persisted model as a Keras model or as a PyFunc, call predict() on a pandas DataFrame keras_model = mlflow.keras.load_model("runs:/96771d893a5e46159d9f3b49bf9013e2" + "/models") predictions = keras_model.predict(x_test) """ local_model_path = _download_artifact_from_uri(artifact_uri=model_uri) flavor_conf = _get_flavor_configuration(model_path=local_model_path, flavor_name=FLAVOR_NAME) keras_module = importlib.import_module( flavor_conf.get("keras_module", "keras")) keras_model_artifacts_path = os.path.join( local_model_path, flavor_conf.get("data", _MODEL_SAVE_PATH)) return _load_model(model_path=keras_model_artifacts_path, keras_module=keras_module, **kwargs)
def test_sparkml_model_log_persists_specified_conda_env_in_mlflow_model_directory( spark_model_iris, model_path, spark_custom_env): artifact_path = "model" with kiwi.start_run(): sparkm.log_model( spark_model=spark_model_iris.model, artifact_path=artifact_path, conda_env=spark_custom_env) model_uri = "runs:/{run_id}/{artifact_path}".format( run_id=kiwi.active_run().info.run_id, artifact_path=artifact_path) model_path = _download_artifact_from_uri(artifact_uri=model_uri) pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME) saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]) assert os.path.exists(saved_conda_env_path) assert saved_conda_env_path != spark_custom_env with open(spark_custom_env, "r") as f: spark_custom_env_parsed = yaml.safe_load(f) with open(saved_conda_env_path, "r") as f: saved_conda_env_parsed = yaml.safe_load(f) assert saved_conda_env_parsed == spark_custom_env_parsed
def _load_pyfunc(model_path): pyfunc_config = _get_flavor_configuration( model_path=model_path, flavor_name=kiwi.pyfunc.FLAVOR_NAME) python_model_cloudpickle_version = pyfunc_config.get( CONFIG_KEY_CLOUDPICKLE_VERSION, None) if python_model_cloudpickle_version is None: kiwi.pyfunc._logger.warning( "The version of CloudPickle used to save the model could not be found in the MLmodel" " configuration") elif python_model_cloudpickle_version != cloudpickle.__version__: # CloudPickle does not have a well-defined cross-version compatibility policy. Micro version # releases have been known to cause incompatibilities. Therefore, we match on the full # library version kiwi.pyfunc._logger.warning( "The version of CloudPickle that was used to save the model, `CloudPickle %s`, differs" " from the version of CloudPickle that is currently running, `CloudPickle %s`, and may" " be incompatible", python_model_cloudpickle_version, cloudpickle.__version__) python_model_subpath = pyfunc_config.get(CONFIG_KEY_PYTHON_MODEL, None) if python_model_subpath is None: raise MlflowException( "Python model path was not specified in the model configuration") with open(os.path.join(model_path, python_model_subpath), "rb") as f: python_model = cloudpickle.load(f) artifacts = {} for saved_artifact_name, saved_artifact_info in\ pyfunc_config.get(CONFIG_KEY_ARTIFACTS, {}).items(): artifacts[saved_artifact_name] = os.path.join( model_path, saved_artifact_info[CONFIG_KEY_ARTIFACT_RELATIVE_PATH]) context = PythonModelContext(artifacts=artifacts) python_model.load_context(context=context) return _PythonModelPyfuncWrapper(python_model=python_model, context=context)
def test_log_model_persists_specified_conda_env_in_mlflow_model_directory( sklearn_knn_model, main_scoped_model_class, pyfunc_custom_env): sklearn_artifact_path = "sk_model" with kiwi.start_run(): kiwi.sklearn.log_model(sk_model=sklearn_knn_model, artifact_path=sklearn_artifact_path) sklearn_run_id = kiwi.active_run().info.run_id pyfunc_artifact_path = "pyfunc_model" with kiwi.start_run(): kiwi.pyfunc.log_model( artifact_path=pyfunc_artifact_path, artifacts={ "sk_model": utils_get_artifact_uri(artifact_path=sklearn_artifact_path, run_id=sklearn_run_id) }, python_model=main_scoped_model_class(predict_fn=None), conda_env=pyfunc_custom_env) 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)) pyfunc_conf = _get_flavor_configuration( model_path=pyfunc_model_path, flavor_name=kiwi.pyfunc.FLAVOR_NAME) saved_conda_env_path = os.path.join(pyfunc_model_path, pyfunc_conf[kiwi.pyfunc.ENV]) assert os.path.exists(saved_conda_env_path) assert saved_conda_env_path != pyfunc_custom_env with open(pyfunc_custom_env, "r") as f: pyfunc_custom_env_parsed = yaml.safe_load(f) with open(saved_conda_env_path, "r") as f: saved_conda_env_parsed = yaml.safe_load(f) assert saved_conda_env_parsed == pyfunc_custom_env_parsed
def test_get_flavor_configuration_throws_exception_when_model_configuration_does_not_exist( model_path): with pytest.raises(MlflowException) as exc: mlflow_model_utils._get_flavor_configuration( model_path=model_path, flavor_name=kiwi.mleap.FLAVOR_NAME) assert exc.error_code == RESOURCE_DOES_NOT_EXIST