Bentoml: Add Pytorch Lightning support

[IskanderNakipov]

I'm trying to use BentoML to serve my LightningModule. However, BentoML is unable to save my module, and I get this error:

File "/media/iskander/TB/Ubuntu/Programs/anaconda3/envs/stocks/lib/python3.8/site-packages/bentoml/utils/cloudpickle.py", line 847, in save_global
    if obj.__module__ == "__main__":
AttributeError: 'torch.dtype' object has no attribute '__module__'

Here's minimal code example to reproduce error (lightning is needed):

import bentoml
from bentoml.adapters import JsonInput
from bentoml.frameworks.pytorch import PytorchModelArtifact

from pytorch_lightning.core import LightningModule


class Net(LightningModule):
    def __init__(self):
        super(Net, self).__init__()

    def forward(self, *args, **kwargs):
        pass


@bentoml.env(infer_pip_packages=True)
@bentoml.artifacts([PytorchModelArtifact('net')])
class PytorchModelService(bentoml.BentoService):
    @bentoml.api(input=JsonInput(), batch=True)
    def predict(self, data):
        return data


svc = PytorchModelService()
svc.pack('net', Net())
svc.save()

I've tried debugging, and figured out, that during pickling in python's pickle module method
def save_reduce(self, func, args, state=None, listitems=None, dictitems=None, state_setter=None, obj=None):
(line 619) receives OrderedDict, where one of values is torch.dtype. As i understand, it comes from __reduce_ex__ method of LightningModule

Describe the solution you'd like
It would be awesome to be able to use BentoML with Lightning

Describe alternatives you've considered
One of alternatives is to rewrite my code, so that I have separate class derived from torch.nn.Module, however it is not the best solution

Additional context
python==3.8 (Anaconda)
BentoML==0.9.0
pytorch==1.6
pytorch-lightning==0.9.0

Everithing is installed with pip

[parano]

Thank you @IskanderNakipov - we've been hearing about PyTorch-Lightning quite a lot from BentoML users lately and will prioritize supporting it! We are releasing version 0.9.1 today, and we will target supporting it in the next release 0.9.2, coming out in about 2 weeks from now.

And let me know if you'd be interested in helping to build support for Pytorch-lightening in BentoML.

[IskanderNakipov]

@parano I'll be happy to help, what can i do?

[parano]

hi @IskanderNakipov - sorry about the delay in getting back to you, could you ping me in the bentoml slack channel? Let's chat about it!

[yubozhao]

While we are in the process of adding support for Pytorch Lighting, you can export the pytorch lighting model to ONNX and bundle with BentoML's ONNXModelArtifact.

You can find how to export to ONNX model here: https://github.com/PyTorchLightning/pytorch-lightning#and-even-export-for-production-via-onnx-or-torchscript