Repository: SeldonIO/tempo (language: python)

Issue:
Reinvestigate dill
given https://github.com/uqfoundation/dill/issues/383
We should look if we can use dill with `recurse=True`


Discussion / hints:

Fix patch:
diff --git a/setup.py b/setup.py
index 39635fa3..9cfa03b6 100644
--- a/setup.py
+++ b/setup.py
@@ -48,7 +48,7 @@ def _load_description() -> str:
         "packaging",
         "requests",
         "pydantic",
-        "cloudpickle",
+        "dill",
         "python-rclone",
         "seldon-deploy-sdk",
         "conda-pack",
diff --git a/tempo/serve/loader/__init__.py b/tempo/serve/loader/__init__.py
new file mode 100644
index 00000000..804a4ef7
--- /dev/null
+++ b/tempo/serve/loader/__init__.py
@@ -0,0 +1,14 @@
+from .artifact import load, load_custom, load_remote, save, save_custom
+from .env import save_environment
+from .upload import download, upload
+
+__all__ = [
+    "save",
+    "save_custom",
+    "load",
+    "load_custom",
+    "load_remote",
+    "save_environment",
+    "upload",
+    "download",
+]
diff --git a/tempo/serve/loader/artifact.py b/tempo/serve/loader/artifact.py
new file mode 100644
index 00000000..471cdf6a
--- /dev/null
+++ b/tempo/serve/loader/artifact.py
@@ -0,0 +1,91 @@
+import os
+from types import FunctionType
+from typing import Any
+
+import dill
+
+from ..constants import DefaultModelFilename, DefaultRemoteFilename
+
+
+def _fixed_create_function(
+    fcode,
+    fglobals,
+    fname=None,
+    fdefaults=None,
+    fclosure=None,
+    fdict=None,
+    fkwdefaults=None,
+):
+    """
+    Patched version of `dill`'s `_create_function`, to make sure __builtins__
+    is serialised.
+
+    Patched version from:
+        https://github.com/uqfoundation/dill/issues/219#issuecomment-522068603
+    Original source code can be found at:
+        https://github.com/uqfoundation/dill/blob/846b888ef957dc1991464b562df425aa35142cb8/dill/_dill.py#L594-L603
+
+    """
+    # same as FunctionType, but enable passing __dict__ to new function,
+    # __dict__ is the storehouse for attributes added after function creation
+    if fdict is None:
+        fdict = dict()
+    func = FunctionType(fcode, fglobals or dict(), fname, fdefaults, fclosure)
+    func.__dict__.update(fdict)  # XXX: better copy? option to copy?
+    if fkwdefaults is not None:
+        func.__kwdefaults__ = fkwdefaults
+
+    # THE WORKAROUND:
+    # if the function was serialized without recurse, fglobals would actually contain
+    # __builtins__, but because of recurse only the referenced modules/objects
+    # end up in fglobals and we are missing the important __builtins__
+    if "__builtins__" not in func.__globals__:
+        func.__globals__["__builtins__"] = globals()["__builtins__"]
+
+    return func
+
+
+def _patch_dill(f):
+    def _patched(*args, **kwargs):
+        _dill = dill._dill
+        # Save original before patching
+        original_create_function = _dill._create_function
+        _dill._create_function = _fixed_create_function
+
+        ret = f(*args, **kwargs)
+
+        _dill._create_function = original_create_function
+
+        return ret
+
+    return _patched
+
+
+@_patch_dill
+def save(tempo_artifact: Any, save_env=True):
+    model = tempo_artifact.get_tempo()
+    model.save(save_env=save_env)
+
+
+@_patch_dill
+def save_custom(pipeline, file_path: str) -> str:
+    with open(file_path, "wb") as file:
+        dill.dump(pipeline, file, recurse=True)
+
+    return file_path
+
+
+def load(folder: str):
+    file_path_pkl = os.path.join(folder, DefaultModelFilename)
+    return load_custom(file_path_pkl)
+
+
+def load_custom(file_path: str):
+    with open(file_path, "rb") as file:
+        return dill.load(file)
+
+
+def load_remote(folder: str):
+    file_path = os.path.join(folder, DefaultRemoteFilename)
+    with open(file_path, "rb") as file:
+        return dill.load(file)
diff --git a/tempo/serve/loader.py b/tempo/serve/loader/env.py
similarity index 68%
rename from tempo/serve/loader.py
rename to tempo/serve/loader/env.py
index 7fdab4c7..625fc3f5 100644
--- a/tempo/serve/loader.py
+++ b/tempo/serve/loader/env.py
@@ -1,46 +1,14 @@
-import os
 import re
 import uuid
 from subprocess import run
 from tempfile import NamedTemporaryFile
-from typing import Any, Optional
+from typing import Optional
 
-import cloudpickle
 import conda_pack
-import rclone
 import yaml
 
-from tempo.conf import settings
-from tempo.serve.constants import DefaultModelFilename, DefaultRemoteFilename, MLServerEnvDeps
-from tempo.utils import logger
-
-
-def save(tempo_artifact: Any, save_env=True):
-    model = tempo_artifact.get_tempo()
-    model.save(save_env=save_env)
-
-
-def save_custom(pipeline, file_path: str) -> str:
-    with open(file_path, "wb") as file:
-        cloudpickle.dump(pipeline, file)
-
-    return file_path
-
-
-def load(folder: str):
-    file_path_pkl = os.path.join(folder, DefaultModelFilename)
-    return load_custom(file_path_pkl)
-
-
-def load_custom(file_path: str):
-    with open(file_path, "rb") as file:
-        return cloudpickle.load(file)
-
-
-def load_remote(folder: str):
-    file_path = os.path.join(folder, DefaultRemoteFilename)
-    with open(file_path, "rb") as file:
-        return cloudpickle.load(file)
+from ...utils import logger
+from ..constants import MLServerEnvDeps
 
 
 def _get_env(conda_env_file_path: str = None, env_name: str = None) -> dict:
@@ -157,26 +125,3 @@ def _create_and_pack_environment(env: dict, file_path: str):
             cmd = f"conda remove --name {tmp_env_name} --all --yes"
             logger.info("Removing conda env with: %s", cmd)
             run(cmd, shell=True, check=True)
-
-
-def _load_rclone_cfg() -> str:
-    with open(settings.rclone_cfg, "r") as f:
-        return f.read()
-
-
-def upload(tempo_artifact: Any):
-    "Upload local to remote using rclone"
-    model = tempo_artifact.get_tempo()
-    local_path = model.details.local_folder
-    remote_uri = model.details.uri
-    logger.info("Uploading %s to %s", local_path, remote_uri)
-    rclone.with_config(_load_rclone_cfg()).copy(local_path, remote_uri, flags=["-P"])
-
-
-def download(tempo_artifact: Any):
-    "Download remote to local using rclone"
-    model = tempo_artifact.get_tempo()
-    local_path = model.details.local_folder
-    remote_uri = model.details.uri
-    logger.info("Downloading %s to %s", remote_uri, local_path)
-    rclone.with_config(_load_rclone_cfg()).copy(remote_uri, local_path, flags=["-P"])
diff --git a/tempo/serve/loader/upload.py b/tempo/serve/loader/upload.py
new file mode 100644
index 00000000..91375031
--- /dev/null
+++ b/tempo/serve/loader/upload.py
@@ -0,0 +1,29 @@
+from typing import Any
+
+import rclone
+
+from ...conf import settings
+from ...utils import logger
+
+
+def _load_rclone_cfg() -> str:
+    with open(settings.rclone_cfg, "r") as f:
+        return f.read()
+
+
+def upload(tempo_artifact: Any):
+    "Upload local to remote using rclone"
+    model = tempo_artifact.get_tempo()
+    local_path = model.details.local_folder
+    remote_uri = model.details.uri
+    logger.info("Uploading %s to %s", local_path, remote_uri)
+    rclone.with_config(_load_rclone_cfg()).copy(local_path, remote_uri, flags=["-P"])
+
+
+def download(tempo_artifact: Any):
+    "Download remote to local using rclone"
+    model = tempo_artifact.get_tempo()
+    local_path = model.details.local_folder
+    remote_uri = model.details.uri
+    logger.info("Downloading %s to %s", remote_uri, local_path)
+    rclone.with_config(_load_rclone_cfg()).copy(remote_uri, local_path, flags=["-P"])


Test patch:
diff --git a/tests/serve/test_loader.py b/tests/serve/test_loader.py
index 5003ab1c..af5bfeff 100644
--- a/tests/serve/test_loader.py
+++ b/tests/serve/test_loader.py
@@ -3,7 +3,7 @@
 import pytest
 
 from tempo.serve.constants import DefaultEnvFilename, MLServerEnvDeps
-from tempo.serve.loader import (
+from tempo.serve.loader.env import (
     _add_required_deps,
     _create_and_pack_environment,
     _get_env,