Repository: bentoml/BentoML
Pull request: feat: support picklable_model
<!--- Thanks for sending a pull request!

Please make sure to read the contribution guidelines, then fill out the blanks below before requesting a code review.

Name your Pull Request with one of the following prefixes, e.g. "feat: add support for PyTorch", to indicate the type of changes proposed. This is based on the [Conventional Commits specification](https://www.conventionalcommits.org/en/v1.0.0/#summary).

- feat: (new feature for the user, not a new feature for build script)
- fix: (bug fix for the user, not a fix to a build script)
- docs: (changes to the documentation)
- style: (formatting, missing semicolons, etc; no production code change)
- refactor: (refactoring production code, eg. renaming a variable)
- perf: (code changes that improve performance)
- test: (adding missing tests, refactoring tests; no production code change)
- chore: (updating grunt tasks etc; no production code change)
- build: (changes that affect the build system or external dependencies)
- ci: (changes to configuration files and scripts)
- revert: (reverts a previous commit)
-->

## Description
<!--- Describe your changes in detail -->
<!--- Attach screenshots here if appropriate. -->

## Motivation and Context
<!--- Why is this change required? What problem does it solve? -->
<!--- If it fixes an open issue, please link to the issue here. -->
<!--- If it is based on a conversation in slack channel, pls quote related messages here -->

## How Has This Been Tested?
<!--- Please describe in detail how you tested your changes. -->
<!--- Include details of your testing environment, and the tests you ran to -->
<!--- see how your change affects other areas of the code, etc. -->

## Checklist:
<!--- Go over all the following points, and put an `x` in all the boxes that apply. -->
<!--- If you're unsure about any of these, don't hesitate to ask. We're here to help! -->
<!--- If you plan to update documentation or tests in follow-up, please note -->
- [ ] My code follows the bentoml code style, both `make format` and
  `make lint` script have passed
  ([instructions](https://github.com/bentoml/BentoML/blob/master/DEVELOPMENT.md#style-check-and-auto-formatting-your-code)).
- [ ] My change reduces project test coverage and requires unit tests to be added
- [ ] I have added unit tests covering my code change
- [ ] My change requires a change to the documentation
- [ ] I have updated the documentation accordingly

Problem statement:
Support PickableModel as a new model type
Save a model or python object, that is pickable, in the model store:
```python
bentoml.pickable_model.save('my_model', obj, method="__call__")
```

Load the obj back:
```python
obj = bentoml.pickable_model.load('my_model:latest')
```

Load the obj back as a runner:
```python
simple_runner = bentoml.pickable_model.load_runner('my_model', batch=False)
```
Returns a simple runner when `batch=false`, which should work in most cases, as long as the input to the function is pickable.

If the function being pickled supports processing batch input, user may use `batch=True`, :
```python
runner = bentoml.pickable_model.load_runner('my_model', batch=True)
```

Hints from issue discussion:

Commit message: Apply suggestions from code review

Co-authored-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>

Patch to review:
diff --git a/.github/workflows/frameworks.yml b/.github/workflows/frameworks.yml
index 1730913fe0f..bf670de40c7 100644
--- a/.github/workflows/frameworks.yml
+++ b/.github/workflows/frameworks.yml
@@ -30,6 +30,7 @@ jobs:
       onnx: ${{ steps.filter.outputs.onnx }}
       onnxmlir: ${{ steps.filter.outputs.onnxmlir }}
       paddle: ${{ steps.filter.outputs.paddle }}
+      pickle: ${{ steps.filter.outputs.pickle }}
       pycaret: ${{ steps.filter.outputs.pycaret }}
       pyspark_mllib: ${{ steps.filter.outputs.pyspark_mllib }}
       pytorch: ${{ steps.filter.outputs.pytorch }}
@@ -126,6 +127,11 @@ jobs:
               - bentoml/paddle.py
               - bentoml/_internal/frameworks/paddle.py
               - tests/integration/frameworks/paddle
+            pickle:
+              - *runner
+              - bentoml/pickle.py
+              - bentoml/_internal/frameworks/pickle.py
+              - tests/integration/frameworks/pickle
             pycaret:
               - *runner
               - bentoml/pycaret.py
@@ -717,6 +723,50 @@ jobs:
           flags: paddle
           verbose: true
 
+  pickle_integration_tests:
+    needs: diff
+    if: ${{ (github.event_name == 'pull_request' && needs.diff.outputs.pickle == 'true') || github.event_name == 'push' }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          fetch-depth: 0 # fetch all tags and branches
+      - name: Setup python
+        uses: actions/setup-python@v2
+        with:
+          python-version: 3.8
+
+      - name: Get pip cache dir
+        id: cache-dir
+        run: |
+          echo ::set-output name=dir::$(pip cache dir)
+
+      - name: Cache pip dependencies
+        uses: actions/cache@v2
+        id: cache-pip
+        with:
+          path: ${{ steps.cache-dir.outputs.dir }}
+          key: ${{ runner.os }}-tests-${{ hashFiles('requirements/tests-requirements.txt') }}
+
+
+      - name: Install dependencies
+        run: |
+          pip install .
+          pip install -r requirements/tests-requirements.txt
+
+      - name: Run tests and generate coverage report
+        run: make tests-pickle
+        shell: bash
+
+      - name: Upload test coverage to Codecov
+        if: ${{ needs.diff.outputs.paddle == 'true' }}
+        uses: codecov/codecov-action@v2
+        with:
+          directory: ./
+          files: ./pickle.xml
+          flags: pickle
+          verbose: true
+
   pycaret_integration_tests:
     needs: diff
     if: ${{ (github.event_name == 'pull_request' && needs.diff.outputs.pycaret == 'true') || github.event_name == 'push' }}
diff --git a/bentoml/__init__.py b/bentoml/__init__.py
index eb678dd0ad1..8871f2aecc5 100644
--- a/bentoml/__init__.py
+++ b/bentoml/__init__.py
@@ -42,6 +42,7 @@
     from bentoml import easyocr
     from bentoml import pycaret
     from bentoml import pytorch
+    from bentoml import pickle
     from bentoml import sklearn
     from bentoml import xgboost
     from bentoml import catboost
@@ -71,6 +72,7 @@
         "bentoml.pytorch_lightning", globals(), "bentoml.pytorch_lightning"
     )
     sklearn = _LazyLoader("bentoml.sklearn", globals(), "bentoml.sklearn")
+    pickle = _LazyLoader("bentoml.pickle", globals(), "bentoml.pickle")
     spacy = _LazyLoader("bentoml.spacy", globals(), "bentoml.spacy")
     statsmodels = _LazyLoader("bentoml.statsmodels", globals(), "bentoml.statsmodels")
     tensorflow = _LazyLoader("bentoml.tensorflow", globals(), "bentoml.tensorflow")
@@ -110,6 +112,7 @@
     "onnx",
     "onnxmlir",
     "paddle",
+    "pickle",
     "pycaret",
     "pytorch",
     "pytorch_lightning",
diff --git a/bentoml/_internal/frameworks/pickle.py b/bentoml/_internal/frameworks/pickle.py
new file mode 100644
index 00000000000..fb31ee23fe5
--- /dev/null
+++ b/bentoml/_internal/frameworks/pickle.py
@@ -0,0 +1,287 @@
+import typing as t
+from typing import TYPE_CHECKING
+
+from simple_di import inject
+from simple_di import Provide
+
+from bentoml import Tag
+from bentoml import Model
+from bentoml import Runner
+from bentoml import SimpleRunner
+from bentoml.exceptions import BentoMLException
+from bentoml.exceptions import MissingDependencyException
+
+from ..types import PathType
+from ..utils import LazyLoader
+from ..models import PKL_EXT
+from ..models import SAVE_NAMESPACE
+from ..utils.pkg import get_pkg_version
+from ..configuration.containers import BentoMLContainer
+
+if TYPE_CHECKING:
+    import numpy as np
+    from .. import ext_typing as ext
+
+    from ..models import ModelStore
+
+import pickle
+
+try:
+    import cloudpickle
+except ImportError:  # pragma: no cover
+    raise MissingDependencyException(
+        """cloudpickle is required in order to use the module `bentoml.pickle`, install
+         cloudpickle with `pip install cloudpickle`.
+        """
+    )
+
+MODULE_NAME = "bentoml.pickle"
+
+
+np = LazyLoader("np", globals(), "numpy")  # noqa: F811
+pd = LazyLoader("pd", globals(), "pandas")
+
+
+def _get_model_info(tag: Tag, model_store: "ModelStore") -> t.Tuple["Model", PathType]:
+    model = model_store.get(tag)
+    if model.info.module not in (MODULE_NAME, __name__):
+        raise BentoMLException(
+            f"Model {tag} was saved with module {model.info.module}, failed loading with {MODULE_NAME}."
+        )
+    model_file = model.path_of(f"{SAVE_NAMESPACE}{PKL_EXT}")
+
+    return model, model_file
+
+
+@inject
+def load(
+    tag: t.Union[str, Tag],
+    model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
+) -> t.Union["BaseEstimator", "Pipeline"]:
+    """
+    Load a model from BentoML local modelstore with given name.
+
+    Args:
+        tag (:code:`Union[str, Tag]`):
+            Tag of a saved model in BentoML local modelstore.
+        model_store (:mod:`~bentoml._internal.models.store.ModelStore`, default to :mod:`BentoMLContainer.model_store`):
+            BentoML modelstore, provided by DI Container.
+
+    Returns:
+        :obj:`Union[BaseEstimator, Pipeline]`: an instance of :obj: model from BentoML modelstore.
+
+    Examples:
+
+    .. code-block:: python
+
+        import bentoml
+
+        unpickled_model = bentoml.pickle.load('my_model:latest')
+    """  # noqa
+    _, model_file = _get_model_info(tag, model_store)
+    with open(model_file, 'rb') as f:
+        # The protocol version used is detected automatically, so we do not
+        # have to specify it.
+        return pickle.load(f)
+
+
+@inject
+def save(
+    name: str,
+    model: t.Union["BaseEstimator", "Pipeline"],
+    *,
+    metadata: t.Optional[t.Dict[str, t.Any]] = None,
+    model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
+) -> Tag:
+    """
+    Save a model instance to BentoML modelstore.
+
+    Args:
+        name (:code:`str`):
+            Name for given model instance. This should pass Python identifier check.
+        model (:code:`Union[BaseEstimator, Pipeline]`):
+            Instance of model to be saved.
+        metadata (:code:`Dict[str, Any]`, `optional`,  default to :code:`None`):
+            Custom metadata for given model.
+        model_store (:mod:`~bentoml._internal.models.store.ModelStore`, default to :mod:`BentoMLContainer.model_store`):
+            BentoML modelstore, provided by DI Container.
+
+    Returns:
+        :obj:`~bentoml._internal.types.Tag`: A :obj:`tag` with a format `name:version` where `name` is the user-defined model's name, and a generated `version` by BentoML.
+
+    Examples:
+
+    .. code-block:: python
+
+        import bentoml
+        
+        class MyCoolModel:
+            def predict(self, some_integer: int):
+                return some_integer**2
+
+        model_to_save = MyCoolModel();
+        tag_info = bentoml.pickle.save("test_pickle_model", model_to_save)
+        runner = bentoml.pickle.load_runner(tag_info)
+        runner.run(3)
+        
+    """  # noqa
+    context = {
+        "framework_name": "pickle"
+    }
+
+    _model = Model.create(
+        name,
+        module=MODULE_NAME,
+        metadata=metadata,
+        context=context,
+    )
+
+    with open(_model.path_of(f"{SAVE_NAMESPACE}{PKL_EXT}"), 'wb') as f:
+        cloudpickle.dump(model, f)
+
+    _model.save(model_store)
+    return _model.tag
+
+
+class _PickleRunner(Runner):
+    @inject
+    def __init__(
+        self,
+        tag: Tag,
+        function_name: str,
+        name: str,
+        resource_quota: t.Optional[t.Dict[str, t.Any]],
+        batch_options: t.Optional[t.Dict[str, t.Any]],
+        model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
+    ):
+        super().__init__(name, resource_quota, batch_options)
+        model_info, model_file = _get_model_info(tag, model_store)
+        self._model_store = model_store
+        self._model_info = model_info
+        self._model_file = model_file
+        self._backend = "loky"
+        self._function_name = function_name
+
+    @property
+    def num_replica(self) -> int:
+        return 1
+
+    @property
+    def required_models(self) -> t.List[Tag]:
+        return [self._model_info.tag]
+
+    # pylint: disable=attribute-defined-outside-init
+    def _setup(self) -> None:
+        with open(self._model_file, 'rb') as f:
+            self._model = pickle.load(f)
+        self._infer_func = getattr(self._model, self._function_name)
+
+    # pylint: disable=arguments-differ
+    def _run_batch(  # type: ignore[reportIncompatibleMethodOverride]
+        self,
+        inputs: t.Union["ext.NpNDArray", "ext.PdDataFrame"],
+    ) -> "ext.NpNDArray":
+        return self._infer_func(inputs)
+
+
+class _PickleSimpleRunner(SimpleRunner):
+    @inject
+    def __init__(
+        self,
+        tag: Tag,
+        function_name: str,
+        name: str,
+        resource_quota: t.Optional[t.Dict[str, t.Any]],
+        batch_options: t.Optional[t.Dict[str, t.Any]],
+        model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
+    ):
+        super().__init__(name, resource_quota, batch_options)
+        model_info, model_file = _get_model_info(tag, model_store)
+        self._model_store = model_store
+        self._model_info = model_info
+        self._model_file = model_file
+        self._backend = "loky"
+        self._function_name = function_name
+
+    @property
+    def num_replica(self) -> int:
+        return 1
+
+    @property
+    def required_models(self) -> t.List[Tag]:
+        return [self._model_info.tag]
+
+    # pylint: disable=attribute-defined-outside-init
+    def _setup(self) -> None:
+        with open(self._model_file, 'rb') as f:
+            self._model = pickle.load(f)
+        self._infer_func = getattr(self._model, self._function_name)
+
+    def _run(  # type: ignore[reportIncompatibleMethodOverride]
+        self,
+        inputs: t.Union["ext.NpNDArray", "ext.PdDataFrame"],
+    ) -> "ext.NpNDArray":
+        return self._infer_func(inputs)
+
+
+@inject
+def load_runner(
+    tag: t.Union[str, Tag],
+    function_name: str = "predict",
+    *,
+    name: t.Optional[str] = None,
+    resource_quota: t.Union[None, t.Dict[str, t.Any]] = None,
+    batch_options: t.Union[None, t.Dict[str, t.Any]] = None,
+    model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
+) -> "BaseRunner":
+    """
+    Runner represents a unit of serving logic that can be scaled horizontally to
+    maximize throughput. :func:`bentoml.pickle.load_runner` implements a Runner class that
+    wraps the commands that dump and load a pickled object, which optimizes it for the BentoML runtime.
+
+    Args:
+        tag (:code:`Union[str, Tag]`):
+            Tag of a saved model in BentoML local modelstore..
+        function_name (:code:`str`, `optional`, default to :code:`predict`):
+            Predict function used by a given pickled model.
+        resource_quota (:code:`Dict[str, Any]`, default to :code:`None`):
+            Dictionary to configure resources allocation for runner.
+        batch_options (:code:`Dict[str, Any]`, default to :code:`None`):
+            Dictionary to configure batch options for runner in a service context.
+        model_store (:mod:`~bentoml._internal.models.store.ModelStore`, default to :mod:`BentoMLContainer.model_store`):
+            BentoML modelstore, provided by DI Container.
+
+    Returns:
+        :obj:`~bentoml._internal.runner.Runner`: Runner instances for the target :mod:`bentoml.pickle` model
+
+    Examples:
+
+    .. code-block:: python
+
+        import bentoml
+
+        runner = bentoml.pickle.load_runner("my_model:latest")
+        runner.run([[1,2,3,4]])
+    """  # noqa
+    tag = Tag.from_taglike(tag)
+    if name is None:
+        name = tag.name
+
+    if batch_options and batch_options.get('enabled'):
+        return _PickleRunner(
+            tag=tag,
+            function_name=function_name,
+            name=name,
+            resource_quota=resource_quota,
+            batch_options=batch_options,
+            model_store=model_store,
+        )
+    else:
+        return _PickleSimpleRunner(
+            tag=tag,
+            function_name=function_name,
+            name=name,
+            resource_quota=resource_quota,
+            batch_options=batch_options,
+            model_store=model_store,
+        )
diff --git a/bentoml/pickle.py b/bentoml/pickle.py
new file mode 100644
index 00000000000..06f2792165c
--- /dev/null
+++ b/bentoml/pickle.py
@@ -0,0 +1,5 @@
+from ._internal.frameworks.pickle import load
+from ._internal.frameworks.pickle import save
+from ._internal.frameworks.pickle import load_runner
+
+__all__ = ["load", "load_runner", "save"]
diff --git a/scripts/ci/config.yml b/scripts/ci/config.yml
index 3c52916ca86..7cd84ad98fd 100644
--- a/scripts/ci/config.yml
+++ b/scripts/ci/config.yml
@@ -139,6 +139,9 @@ paddle:
   external_scripts: |
     export SETUPTOOLS_USE_DISTUTILS=stdlib
 
+pickle:
+  <<: *tmpl
+
 pycaret:
   <<: *tmpl
   dependencies:
diff --git a/tests/integration/frameworks/test_pickle_impl.py b/tests/integration/frameworks/test_pickle_impl.py
new file mode 100644
index 00000000000..4224a82cf93
--- /dev/null
+++ b/tests/integration/frameworks/test_pickle_impl.py
@@ -0,0 +1,93 @@
+import typing as t
+from typing import TYPE_CHECKING
+
+import numpy as np
+import pytest
+
+import bentoml
+import bentoml.models
+
+if TYPE_CHECKING:
+    from bentoml._internal.store import Tag
+    from bentoml._internal.models import ModelStore
+
+
+class MyCoolModel:
+    def predict(self, some_integer: int):
+        return some_integer**2
+
+
+class MyCoolBatchModel:
+    def predict(self, some_integer: t.List):
+        return list(map(lambda x: x**2, some_integer))
+
+
+class MockBatchOptions:
+    enabled = True
+
+
+def save_procedure(
+    metadata: t.Dict[str, t.Any],
+    _modelstore: "ModelStore",
+) -> "Tag":
+    model_to_save = MyCoolModel()
+    tag_info = bentoml.pickle.save(
+        "test_pickle_model",
+        model_to_save,
+        metadata=metadata,
+        model_store=_modelstore,
+    )
+    return tag_info
+
+
+def save_batch_procedure(
+    metadata: t.Dict[str, t.Any],
+    _modelstore: "ModelStore",
+) -> "Tag":
+    model_to_save = MyCoolBatchModel()
+    tag_info = bentoml.pickle.save(
+        "test_pickle_model",
+        model_to_save,
+        metadata=metadata,
+        model_store=_modelstore,
+    )
+    return tag_info
+
+@pytest.mark.parametrize(
+    "metadata",
+    [
+        ({"model": "Pickle", "test": True}),
+        ({"acc": 0.876}),
+    ],
+)
+def test_pickle_save_load(
+    metadata: t.Dict[str, t.Any],
+    modelstore: "ModelStore",
+) -> None:
+    tag = save_procedure(metadata, _modelstore=modelstore)
+    _model = bentoml.models.get(tag, _model_store=modelstore)
+    assert _model.info.metadata is not None
+
+    loaded_model = bentoml.pickle.load(
+        _model.tag, model_store=modelstore
+    )
+    assert isinstance(loaded_model, MyCoolModel)
+    assert loaded_model.predict(4) == np.array([16])
+
+
+def test_pickle_runner_setup_run(modelstore: "ModelStore") -> None:
+
+    tag = save_procedure({}, _modelstore=modelstore)
+    runner = bentoml.pickle.load_runner(tag, model_store=modelstore)
+
+    assert tag in runner.required_models
+    assert runner.run(3) == np.array([9])
+
+
+def test_pickle_runner_setup_run_batch(modelstore: "ModelStore") -> None:
+
+    tag = save_batch_procedure({}, _modelstore=modelstore)
+    runner = bentoml.pickle.load_runner(tag, model_store=modelstore, batch_options={'enabled': True})
+
+    assert tag in runner.required_models
+    assert runner.run_batch([3, 9]) == [9, 81]


Code review comments:
diff_hunk: @@ -0,0 +1,287 @@
+import typing as t
+from typing import TYPE_CHECKING
+
+from simple_di import inject
+from simple_di import Provide
+
+from bentoml import Tag
+from bentoml import Model
+from bentoml import Runner
+from bentoml import SimpleRunner
+from bentoml.exceptions import BentoMLException
+from bentoml.exceptions import MissingDependencyException
+
+from ..types import PathType
+from ..utils import LazyLoader
+from ..models import PKL_EXT
+from ..models import SAVE_NAMESPACE
+from ..utils.pkg import get_pkg_version
+from ..configuration.containers import BentoMLContainer
+
+if TYPE_CHECKING:
+    import numpy as np
+    from .. import ext_typing as ext
+
+    from ..models import ModelStore
+
+import pickle
+
+try:
+    import cloudpickle
+except ImportError:  # pragma: no cover
+    raise MissingDependencyException(
+        """cloudpickle is required in order to use the module `bentoml.pickle`, install
+         cloudpickle with `pip install cloudpickle`.
+        """
path: bentoml/_internal/frameworks/pickle.py
text: @user1:
I think this will add the spaces from the indent into the string?
```suggestion
        "cloudpickle is required in order to use the module `bentoml.pickle`, install cloudpickle with `pip install cloudpickle`."
```

diff_hunk: @@ -0,0 +1,287 @@
+import typing as t
+from typing import TYPE_CHECKING
+
+from simple_di import inject
+from simple_di import Provide
+
+from bentoml import Tag
+from bentoml import Model
+from bentoml import Runner
+from bentoml import SimpleRunner
+from bentoml.exceptions import BentoMLException
+from bentoml.exceptions import MissingDependencyException
+
+from ..types import PathType
+from ..utils import LazyLoader
+from ..models import PKL_EXT
+from ..models import SAVE_NAMESPACE
+from ..utils.pkg import get_pkg_version
+from ..configuration.containers import BentoMLContainer
+
+if TYPE_CHECKING:
+    import numpy as np
+    from .. import ext_typing as ext
+
+    from ..models import ModelStore
+
+import pickle
+
+try:
+    import cloudpickle
+except ImportError:  # pragma: no cover
+    raise MissingDependencyException(
+        """cloudpickle is required in order to use the module `bentoml.pickle`, install
+         cloudpickle with `pip install cloudpickle`.
+        """
+    )
+
+MODULE_NAME = "bentoml.pickle"
+
+
+np = LazyLoader("np", globals(), "numpy")  # noqa: F811
+pd = LazyLoader("pd", globals(), "pandas")
+
+
+def _get_model_info(tag: Tag, model_store: "ModelStore") -> t.Tuple["Model", PathType]:
+    model = model_store.get(tag)
+    if model.info.module not in (MODULE_NAME, __name__):
+        raise BentoMLException(
+            f"Model {tag} was saved with module {model.info.module}, failed loading with {MODULE_NAME}."
path: bentoml/_internal/frameworks/pickle.py
text: @user1:
(We should change this in all the other frameworks too)

```suggestion
            f"Model {tag} was saved with module {model.info.module}, cannot load with {MODULE_NAME}."
```