diff --git a/mlflow/catboost.py b/mlflow/catboost.py index 60041445e3ae9..7a6954d3b7d44 100644 --- a/mlflow/catboost.py +++ b/mlflow/catboost.py @@ -39,17 +39,23 @@ _MODEL_BINARY_FILE_NAME = "model.cb" +def get_default_pip_requirements(): + """ + :return: A list of default pip requirements for MLflow Models produced by this flavor. + Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment + that, at minimum, contains these requirements. + """ + import catboost as cb + + return ["catboost=={}".format(cb.__version__)] + + def get_default_conda_env(): """ :return: The default Conda environment for MLflow Models produced by calls to :func:`save_model()` and :func:`log_model()`. """ - import catboost as cb - - return _mlflow_conda_env( - # CatBoost is not yet available via the default conda channels, so we install it via pip - additional_pip_deps=["catboost=={}".format(cb.__version__)] - ) + return _mlflow_conda_env(additional_pip_deps=get_default_pip_requirements()) def save_model( diff --git a/mlflow/fastai.py b/mlflow/fastai.py index 158e565d6653c..0160217400889 100644 --- a/mlflow/fastai.py +++ b/mlflow/fastai.py @@ -43,6 +43,23 @@ FLAVOR_NAME = "fastai" +def get_default_pip_requirements(include_cloudpickle=False): + """ + :return: A list of default pip requirements for MLflow Models produced by this flavor. + Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment + that, at minimum, contains these requirements. + """ + import fastai + + pip_deps = ["fastai=={}".format(fastai.__version__)] + if include_cloudpickle: + import cloudpickle + + pip_deps.append("cloudpickle=={}".format(cloudpickle.__version__)) + + return pip_deps + + def get_default_conda_env(include_cloudpickle=False): """ :return: The default Conda environment as a dictionary for MLflow Models produced by calls to @@ -71,15 +88,7 @@ def get_default_conda_env(include_cloudpickle=False): 'dependencies': ['python=3.7.5', 'fastai=1.0.61', 'pip', {'pip': ['mlflow']}]} """ - - import fastai - - pip_deps = ["fastai=={}".format(fastai.__version__)] - if include_cloudpickle: - import cloudpickle - - pip_deps.append("cloudpickle=={}".format(cloudpickle.__version__)) - return _mlflow_conda_env(additional_pip_deps=pip_deps) + return _mlflow_conda_env(additional_pip_deps=get_default_pip_requirements(include_cloudpickle)) def save_model( diff --git a/mlflow/gluon.py b/mlflow/gluon.py index 4dc874caa4b81..a41f15aac5bf9 100644 --- a/mlflow/gluon.py +++ b/mlflow/gluon.py @@ -232,16 +232,23 @@ def save_model( mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME)) -def get_default_conda_env(): +def get_default_pip_requirements(): """ - :return: The default Conda environment for MLflow Models produced by calls to - :func:`save_model()` and :func:`log_model()`. + :return: A list of default pip requirements for MLflow Models produced by this flavor. + Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment + that, at minimum, contains these requirements. """ import mxnet as mx - pip_deps = ["mxnet=={}".format(mx.__version__)] + return ["mxnet=={}".format(mx.__version__)] - return _mlflow_conda_env(additional_pip_deps=pip_deps) + +def get_default_conda_env(): + """ + :return: The default Conda environment for MLflow Models produced by calls to + :func:`save_model()` and :func:`log_model()`. + """ + return _mlflow_conda_env(additional_pip_deps=get_default_pip_requirements()) @experimental diff --git a/mlflow/h2o.py b/mlflow/h2o.py index d0347bbe5b2b3..b39933c0742b6 100644 --- a/mlflow/h2o.py +++ b/mlflow/h2o.py @@ -24,14 +24,23 @@ FLAVOR_NAME = "h2o" +def get_default_pip_requirements(): + """ + :return: A list of default pip requirements for MLflow Models produced by this flavor. + Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment + that, at minimum, contains these requirements. + """ + import h2o + + return ["h2o=={}".format(h2o.__version__)] + + def get_default_conda_env(): """ :return: The default Conda environment for MLflow Models produced by calls to :func:`save_model()` and :func:`log_model()`. """ - import h2o - - return _mlflow_conda_env(additional_pip_deps=["h2o=={}".format(h2o.__version__)]) + return _mlflow_conda_env(additional_pip_deps=get_default_pip_requirements()) def save_model( diff --git a/mlflow/keras.py b/mlflow/keras.py index 2c28d8ea46c21..28ba1c333b65b 100644 --- a/mlflow/keras.py +++ b/mlflow/keras.py @@ -51,10 +51,11 @@ _PIP_ENV_SUBPATH = "requirements.txt" -def get_default_conda_env(include_cloudpickle=False, keras_module=None): +def get_default_pip_requirements(include_cloudpickle=False, keras_module=None): """ - :return: The default Conda environment for MLflow Models produced by calls to - :func:`save_model()` and :func:`log_model()`. + :return: A list of default pip requirements for MLflow Models produced by this flavor. + Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment + that, at minimum, contains these requirements. """ import tensorflow as tf @@ -76,7 +77,18 @@ def get_default_conda_env(include_cloudpickle=False, keras_module=None): # see https://github.com/tensorflow/tensorflow/issues/44467 if Version(tf.__version__) < Version("2.4"): pip_deps.append("h5py<3.0.0") - return _mlflow_conda_env(additional_pip_deps=pip_deps) + + return pip_deps + + +def get_default_conda_env(include_cloudpickle=False, keras_module=None): + """ + :return: The default Conda environment for MLflow Models produced by calls to + :func:`save_model()` and :func:`log_model()`. + """ + return _mlflow_conda_env( + additional_pip_deps=get_default_pip_requirements(include_cloudpickle, keras_module) + ) def save_model( diff --git a/mlflow/lightgbm.py b/mlflow/lightgbm.py index 0a161e2cf601b..dba0b3d98a729 100644 --- a/mlflow/lightgbm.py +++ b/mlflow/lightgbm.py @@ -56,17 +56,23 @@ _logger = logging.getLogger(__name__) +def get_default_pip_requirements(): + """ + :return: A list of default pip requirements for MLflow Models produced by this flavor. + Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment + that, at minimum, contains these requirements. + """ + import lightgbm as lgb + + return ["lightgbm=={}".format(lgb.__version__)] + + def get_default_conda_env(): """ :return: The default Conda environment for MLflow Models produced by calls to :func:`save_model()` and :func:`log_model()`. """ - import lightgbm as lgb - - return _mlflow_conda_env( - # LightGBM is not yet available via the default conda channels, so we install it via pip - additional_pip_deps=["lightgbm=={}".format(lgb.__version__)], - ) + return _mlflow_conda_env(additional_pip_deps=get_default_pip_requirements()) def save_model( diff --git a/mlflow/onnx.py b/mlflow/onnx.py index 4a3868d44373a..7259c9066fc56 100644 --- a/mlflow/onnx.py +++ b/mlflow/onnx.py @@ -31,24 +31,31 @@ # TEMPORARY CHANGE to trigger cross version tests +def get_default_pip_requirements(): + """ + :return: A list of default pip requirements for MLflow Models produced by this flavor. + Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment + that, at minimum, contains these requirements. + """ + import onnx + import onnxruntime + + return [ + "onnx=={}".format(onnx.__version__), + # The ONNX pyfunc representation requires the OnnxRuntime + # inference engine. Therefore, the conda environment must + # include OnnxRuntime + "onnxruntime=={}".format(onnxruntime.__version__), + ] + + @experimental def get_default_conda_env(): """ :return: The default Conda environment for MLflow Models produced by calls to :func:`save_model()` and :func:`log_model()`. """ - import onnx - import onnxruntime - - return _mlflow_conda_env( - additional_pip_deps=[ - "onnx=={}".format(onnx.__version__), - # The ONNX pyfunc representation requires the OnnxRuntime - # inference engine. Therefore, the conda environment must - # include OnnxRuntime - "onnxruntime=={}".format(onnxruntime.__version__), - ], - ) + return _mlflow_conda_env(additional_pip_deps=get_default_pip_requirements()) @experimental diff --git a/mlflow/paddle/__init__.py b/mlflow/paddle/__init__.py index f46c792b5d1a0..5d5b5f027b3d8 100644 --- a/mlflow/paddle/__init__.py +++ b/mlflow/paddle/__init__.py @@ -33,15 +33,23 @@ _logger = logging.getLogger(__name__) +def get_default_pip_requirements(): + """ + :return: A list of default pip requirements for MLflow Models produced by this flavor. + Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment + that, at minimum, contains these requirements. + """ + import paddle + + return ["paddlepaddle=={}".format(paddle.__version__)] + + def get_default_conda_env(): """ :return: The default Conda environment for MLflow Models produced by calls to :func:`save_model()` and :func:`log_model()`. """ - import paddle - - pip_deps = ["paddlepaddle=={}".format(paddle.__version__)] - return _mlflow_conda_env(additional_pip_deps=pip_deps, additional_conda_channels=None) + return _mlflow_conda_env(additional_pip_deps=get_default_pip_requirements()) def save_model( diff --git a/mlflow/pyfunc/model.py b/mlflow/pyfunc/model.py index 3be6cc1ad0ab9..2b360b639cc7a 100644 --- a/mlflow/pyfunc/model.py +++ b/mlflow/pyfunc/model.py @@ -29,6 +29,15 @@ CONFIG_KEY_CLOUDPICKLE_VERSION = "cloudpickle_version" +def get_default_pip_requirements(): + """ + :return: A list of default pip requirements for MLflow Models produced by this flavor. + Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment + that, at minimum, contains these requirements. + """ + return ["cloudpickle=={}".format(cloudpickle.__version__)] + + def get_default_conda_env(): """ :return: The default Conda environment for MLflow Models produced by calls to @@ -36,9 +45,7 @@ def get_default_conda_env(): and :func:`log_model() ` when a user-defined subclass of :class:`PythonModel` is provided. """ - return _mlflow_conda_env( - additional_pip_deps=["cloudpickle=={}".format(cloudpickle.__version__)], - ) + return _mlflow_conda_env(additional_pip_deps=get_default_pip_requirements()) class PythonModel(object): diff --git a/mlflow/pytorch/__init__.py b/mlflow/pytorch/__init__.py index 61b35528bfe4f..6d7931e0af62f 100644 --- a/mlflow/pytorch/__init__.py +++ b/mlflow/pytorch/__init__.py @@ -47,6 +47,25 @@ _logger = logging.getLogger(__name__) +def get_default_pip_requirements(): + """ + :return: A list of default pip requirements for MLflow Models produced by this flavor. + Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment + that, at minimum, contains these requirements. + """ + import torch + import torchvision + + return [ + "torch=={}".format(torch.__version__), + "torchvision=={}".format(torchvision.__version__), + # We include CloudPickle in the default environment because + # it's required by the default pickle module used by `save_model()` + # and `log_model()`: `mlflow.pytorch.pickle_module`. + "cloudpickle=={}".format(cloudpickle.__version__), + ] + + def get_default_conda_env(): """ :return: The default Conda environment as a dictionary for MLflow Models produced by calls to @@ -76,19 +95,7 @@ def get_default_conda_env(): 'mlflow', 'cloudpickle==1.6.0']}]} """ - import torch - import torchvision - - return _mlflow_conda_env( - additional_pip_deps=[ - "torch=={}".format(torch.__version__), - "torchvision=={}".format(torchvision.__version__), - # We include CloudPickle in the default environment because - # it's required by the default pickle module used by `save_model()` - # and `log_model()`: `mlflow.pytorch.pickle_module`. - "cloudpickle=={}".format(cloudpickle.__version__), - ] - ) + return _mlflow_conda_env(additional_pip_deps=get_default_pip_requirements()) def log_model( diff --git a/mlflow/shap.py b/mlflow/shap.py index 19a6611724271..7a9af8bed6b4b 100644 --- a/mlflow/shap.py +++ b/mlflow/shap.py @@ -72,17 +72,24 @@ def get_underlying_model_flavor(model): return _UNKNOWN_MODEL_FLAVOR +def get_default_pip_requirements(): + """ + :return: A list of default pip requirements for MLflow Models produced by this flavor. + Calls to :func:`save_explainer()` and :func:`log_explainer()` produce a pip environment + that, at minimum, contains these requirements. + """ + import shap + + return ["shap=={}".format(shap.__version__)] + + def get_default_conda_env(): """ :return: The default Conda environment for MLflow Models produced by calls to :func:`save_explainer()` and :func:`log_explainer()`. """ - import shap - - pip_deps = ["shap=={}".format(shap.__version__)] - - return _mlflow_conda_env(additional_pip_deps=pip_deps) + return _mlflow_conda_env(additional_pip_deps=get_default_pip_requirements()) def _load_pyfunc(path): diff --git a/mlflow/sklearn/__init__.py b/mlflow/sklearn/__init__.py index 64091723ae704..8c0633422961a 100644 --- a/mlflow/sklearn/__init__.py +++ b/mlflow/sklearn/__init__.py @@ -51,10 +51,11 @@ _SklearnTrainingSession = _get_new_training_session_class() -def get_default_conda_env(include_cloudpickle=False): +def get_default_pip_requirements(include_cloudpickle=False): """ - :return: The default Conda environment for MLflow Models produced by calls to - :func:`save_model()` and :func:`log_model()`. + :return: A list of default pip requirements for MLflow Models produced by this flavor. + Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment + that, at minimum, contains these requirements. """ import sklearn @@ -63,7 +64,16 @@ def get_default_conda_env(include_cloudpickle=False): import cloudpickle pip_deps += ["cloudpickle=={}".format(cloudpickle.__version__)] - return _mlflow_conda_env(additional_pip_deps=pip_deps) + + return pip_deps + + +def get_default_conda_env(include_cloudpickle=False): + """ + :return: The default Conda environment for MLflow Models produced by calls to + :func:`save_model()` and :func:`log_model()`. + """ + return _mlflow_conda_env(additional_pip_deps=get_default_pip_requirements(include_cloudpickle)) def save_model( diff --git a/mlflow/spacy.py b/mlflow/spacy.py index ff486403846ad..f091256fc63fb 100644 --- a/mlflow/spacy.py +++ b/mlflow/spacy.py @@ -30,14 +30,23 @@ _logger = logging.getLogger(__name__) +def get_default_pip_requirements(): + """ + :return: A list of default pip requirements for MLflow Models produced by this flavor. + Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment + that, at minimum, contains these requirements. + """ + import spacy + + return ["spacy=={}".format(spacy.__version__)] + + def get_default_conda_env(): """ :return: The default Conda environment for MLflow Models produced by calls to :func:`save_model()` and :func:`log_model()`. """ - import spacy - - return _mlflow_conda_env(additional_pip_deps=["spacy=={}".format(spacy.__version__)]) + return _mlflow_conda_env(additional_pip_deps=get_default_pip_requirements()) def save_model( diff --git a/mlflow/spark.py b/mlflow/spark.py index c8912ff3ed265..875520e0a54d7 100644 --- a/mlflow/spark.py +++ b/mlflow/spark.py @@ -66,6 +66,21 @@ def _format_exception(ex): return "".join(traceback.format_exception(type(ex), ex, ex.__traceback__)) +def get_default_pip_requirements(): + """ + :return: A list of default pip requirements for MLflow Models produced by this flavor. + Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment + that, at minimum, contains these requirements. + """ + import pyspark + + # Strip the suffix from `dev` versions of PySpark, which are not + # available for installation from Anaconda or PyPI + pyspark_version = re.sub(r"(\.?)dev.*", "", pyspark.__version__) + + return ["pyspark=={}".format(pyspark_version)] + + def get_default_conda_env(): """ :return: The default Conda environment for MLflow Models produced by calls to @@ -76,13 +91,7 @@ def get_default_conda_env(): ``2.4.5.dev0``, invoking this method produces a Conda environment with a dependency on PySpark version ``2.4.5``). """ - import pyspark - - # Strip the suffix from `dev` versions of PySpark, which are not - # available for installation from Anaconda or PyPI - pyspark_version = re.sub(r"(\.?)dev.*", "", pyspark.__version__) - - return _mlflow_conda_env(additional_pip_deps=["pyspark=={}".format(pyspark_version)]) + return _mlflow_conda_env(additional_pip_deps=get_default_pip_requirements()) def log_model( diff --git a/mlflow/statsmodels.py b/mlflow/statsmodels.py index a9d85971ff178..f68c52558c938 100644 --- a/mlflow/statsmodels.py +++ b/mlflow/statsmodels.py @@ -48,16 +48,23 @@ _logger = logging.getLogger(__name__) +def get_default_pip_requirements(): + """ + :return: A list of default pip requirements for MLflow Models produced by this flavor. + Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment + that, at minimum, contains these requirements. + """ + import statsmodels + + return ["statsmodels=={}".format(statsmodels.__version__)] + + def get_default_conda_env(): """ :return: The default Conda environment for MLflow Models produced by calls to :func:`save_model()` and :func:`log_model()`. """ - import statsmodels - - return _mlflow_conda_env( - additional_pip_deps=["statsmodels=={}".format(statsmodels.__version__)] - ) + return _mlflow_conda_env(additional_pip_deps=get_default_pip_requirements()) def save_model( diff --git a/mlflow/tensorflow.py b/mlflow/tensorflow.py index b90930704f788..54b0f178f3795 100644 --- a/mlflow/tensorflow.py +++ b/mlflow/tensorflow.py @@ -67,14 +67,23 @@ _AUTOLOG_RUN_ID = None +def get_default_pip_requirements(): + """ + :return: A list of default pip requirements for MLflow Models produced by this flavor. + Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment + that, at minimum, contains these requirements. + """ + import tensorflow + + return ["tensorflow=={}".format(tensorflow.__version__)] + + def get_default_conda_env(): """ :return: The default Conda environment for MLflow Models produced by calls to :func:`save_model()` and :func:`log_model()`. """ - import tensorflow - - return _mlflow_conda_env(additional_pip_deps=["tensorflow=={}".format(tensorflow.__version__)]) + return _mlflow_conda_env(additional_pip_deps=get_default_pip_requirements()) @keyword_only diff --git a/mlflow/xgboost.py b/mlflow/xgboost.py index 21bf238136718..411c811dd482b 100644 --- a/mlflow/xgboost.py +++ b/mlflow/xgboost.py @@ -72,7 +72,12 @@ _logger = logging.getLogger(__name__) -def _get_default_pip_requirements(): +def get_default_pip_requirements(): + """ + :return: A list of default pip requirements for MLflow Models produced by this flavor. + Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment + that, at minimum, contains these requirements. + """ import xgboost as xgb return ["xgboost=={}".format(xgb.__version__)] @@ -83,7 +88,7 @@ def get_default_conda_env(): :return: The default Conda environment for MLflow Models produced by calls to :func:`save_model()` and :func:`log_model()`. """ - return _mlflow_conda_env(additional_pip_deps=_get_default_pip_requirements()) + return _mlflow_conda_env(additional_pip_deps=get_default_pip_requirements()) def save_model( @@ -190,7 +195,7 @@ def save_model( conda_env, pip_requirements, pip_constraints = ( _process_pip_requirements( - _get_default_pip_requirements(), pip_requirements, extra_pip_requirements, + get_default_pip_requirements(), pip_requirements, extra_pip_requirements, ) if conda_env is None else _process_conda_env(conda_env) diff --git a/tests/xgboost/test_xgboost_model_export.py b/tests/xgboost/test_xgboost_model_export.py index ed0e160c38dda..5154051df6602 100644 --- a/tests/xgboost/test_xgboost_model_export.py +++ b/tests/xgboost/test_xgboost_model_export.py @@ -239,7 +239,7 @@ def test_save_model_with_pip_requirements(xgb_model, tmpdir): @pytest.mark.large def test_save_model_with_extra_pip_requirements(xgb_model, tmpdir): - default_reqs = mlflow.xgboost._get_default_pip_requirements() + default_reqs = mlflow.xgboost.get_default_pip_requirements() # Path to a requirements file tmpdir1 = tmpdir.join("1") @@ -295,7 +295,7 @@ def test_log_model_with_pip_requirements(xgb_model, tmpdir): @pytest.mark.large def test_log_model_with_extra_pip_requirements(xgb_model, tmpdir): - default_reqs = mlflow.xgboost._get_default_pip_requirements() + default_reqs = mlflow.xgboost.get_default_pip_requirements() # Path to a requirements file req_file = tmpdir.join("requirements.txt")