mlem.api.load: Optional types break model loading

[igordertigor]

I'm saving a dummy model function like this:

from typing import Optional
import mlem.api

def model(inp: float, state: Optional[float]):
    if state is None:
        state = 0.
    return inp + state

mlem.api.save(model, 'model')

There is no error message when writing the model. However, if I try to load the model

import mlem.api

model = mlem.api.load('model')

I get the following error message (but no model):

Traceback (most recent call last):
  File "dummy_load.py", line 3, in <module>
    model = mlem.api.load('model')
  File "/xxx/.venv/lib/python3.8/site-packages/mlem/core/metadata.py", line 122, in load
    meta = load_meta(
  File "/xxx/.venv/lib/python3.8/site-packages/mlem/core/metadata.py", line 202, in load_meta
    meta.load_value()
  File "/xxx/.venv/lib/python3.8/site-packages/mlem/core/objects.py", line 634, in load_value
    self.model_type.load(self.relative_artifacts)
  File "/xxx/.venv/lib/python3.8/site-packages/mlem/core/model.py", line 249, in load
    self.model = self.io.load(artifacts)
  File "/xxx/.venv/lib/python3.8/site-packages/mlem/contrib/callable.py", line 88, in load
    return self._deserialize_model(f, refs)
  File "/xxx/.venv/lib/python3.8/site-packages/mlem/contrib/callable.py", line 111, in _deserialize_model
    return _ModelUnpickler(refs, in_file).load()
  File "/usr/lib/python3.8/pickle.py", line 1212, in load
    dispatch[key[0]](self)
  File "/usr/lib/python3.8/pickle.py", line 1528, in load_global
    klass = self.find_class(module, name)
  File "/usr/lib/python3.8/pickle.py", line 1579, in find_class
    __import__(module, level=0)
  File "/home/ingo/tmp/try_sound_processor_layout/.venv/lib/python3.8/site-packages/mlem/ext.py", line 231, in load_module
    module = importlib.import_module(fullname)
  File "/usr/lib/python3.8/importlib/__init__.py", line 127, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
ModuleNotFoundError: No module named '__builtin__'

When I omit the type annotations like this:

import mlem.api

def model(inp, state):
    if state is None:
        state = 0.
    return inp + state

mlem.api.save(model, 'model')

The model loads correctly.

[mike0sv]

Ok, this was because we used dill to pickle and regular pickle to unpickle. This PR should fix your problem, can you confirm by installing from it's branch?

[igordertigor]

Sorry for the late reply. Yes, I can confirm that it's working.