Add _get_default_pip_requirements to all flavors (mlflow#4575)

* Add _get_default_pip_requirements to all flavors Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * Fix lightgbm Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * Fix statsmodels and pytorch flavors Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * revert change on cloudpickle Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * Address comments Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * remove docstrings of get_default_conda_env Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * add missing docs Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * Fix shap doc Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * remove underscore Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * rm mlflow/spark.py.a35ff0bf51ede0b51ed1e135b0eb55e2.tmp Signed-off-by: harupy <hkawamura0130@gmail.com>
BLACKBUCK-LABS · Jul 21, 2021 · 3182883 · 3182883
1 parent cfed457
commit 3182883
Show file tree

Hide file tree

Showing 18 changed files with 231 additions and 97 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -29,16 +29,23 @@
 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
              :func:`save_model() <mlflow.pyfunc.save_model>`
              and :func:`log_model() <mlflow.pyfunc.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
@@ -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
@@ -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):