update pickling with joblib (#87)

* update pickling with `joblib`
* usage

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: Borda

.github/workflows/ci-testing.yml

@@ -21,10 +21,10 @@ jobs:
         requires: ["latest"]
         dependency: ["lightning"]
         include:
-          - { os: "ubuntu-22.04", python-version: "3.9", requires: "oldest", dependency: "lightning" }
-          - { os: "ubuntu-24.04", python-version: "3.10", requires: "latest", dependency: "pytorch_lightning" }
-          - { os: "ubuntu-24.04", python-version: "3.12", requires: "latest", dependency: "pytorch_lightning" }
-          - { os: "macOS-13", python-version: "3.12", requires: "latest", dependency: "pytorch_lightning" }
+          - { requires: "oldest", dependency: "lightning", os: "ubuntu-22.04", python-version: "3.9" }
+          - { requires: "latest", dependency: "pytorch_lightning", os: "ubuntu-24.04", python-version: "3.12" }
+          - { requires: "latest", dependency: "pytorch_lightning", os: "windows-2022", python-version: "3.12" }
+          - { requires: "latest", dependency: "pytorch_lightning", os: "macOS-13", python-version: "3.12" }
 
     # Timeout: https://stackoverflow.com/a/59076067/4521646
     timeout-minutes: 35

_requirements/extra.txt

@@ -1,2 +1,3 @@
 lightning >= 2.0.0
 numpy <2.0.0 ; platform_system == "Darwin" # compatibility fix for Torch
+joblib >= 1.0.0

requirements.txt

@@ -1,5 +1,2 @@
-# NOTE: once we add more dependencies, consider update dependabot to check for updates
-
 lightning-sdk >=0.2.9
 lightning-utilities
-joblib

src/litmodels/integrations/mixins.py

@@ -1,6 +1,5 @@
 import inspect
 import json
-import pickle
 import tempfile
 import warnings
 from abc import ABC
@@ -10,6 +9,7 @@
 from lightning_utilities.core.rank_zero import rank_zero_warn
 
 from litmodels.io.cloud import download_model_files, upload_model_files
+from litmodels.io.utils import dump_pickle, load_pickle
 
 if TYPE_CHECKING:
     import torch
@@ -89,8 +89,7 @@ def upload_model(
         """
         name, model_name, temp_folder = self._setup(name, temp_folder)
         pickle_path = Path(temp_folder) / f"{model_name}.pkl"
-        with open(pickle_path, "wb") as fp:
-            pickle.dump(self, fp, protocol=pickle.HIGHEST_PROTOCOL)
+        dump_pickle(model=self, path=pickle_path)
         if version:
             name = f"{name}:{version}"
         self._upload_model_files(name=name, path=pickle_path, metadata=metadata)
@@ -116,8 +115,7 @@ def download_model(
         if len(pkl_files) > 1:
             raise RuntimeError(f"Multiple pickle files found for model: {model_registry} with {pkl_files}")
         pkl_path = Path(temp_folder) / pkl_files[0]
-        with open(pkl_path, "rb") as fp:
-            obj = pickle.load(fp)
+        obj = load_pickle(path=pkl_path)
         if not isinstance(obj, cls):
             raise RuntimeError(f"Unpickled object is not of type {cls.__name__}: {type(obj)}")
         return obj

src/litmodels/io/gateway.py

@@ -3,10 +3,10 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
-import joblib
 from lightning_utilities import module_available
 
 from litmodels.io.cloud import download_model_files, upload_model_files
+from litmodels.io.utils import dump_pickle, load_pickle
 
 if module_available("torch"):
     import torch
@@ -56,7 +56,7 @@ def upload_model(
         torch.save(model.state_dict(), path)
     else:
         path = os.path.join(staging_dir, f"{model.__class__.__name__}.pkl")
-        joblib.dump(model, path)
+        dump_pickle(model=model, path=path)
 
     return upload_model_files(
         path=path,
@@ -111,7 +111,7 @@ def load_model(name: str, download_dir: str = ".") -> Any:
         raise NotImplementedError("Downloaded model with multiple files is not supported yet.")
     model_path = Path(download_dir) / download_paths[0]
     if model_path.suffix.lower() == ".pkl":
-        return joblib.load(model_path)
+        return load_pickle(path=model_path)
     if model_path.suffix.lower() == ".ts":
         return torch.jit.load(model_path)
     raise NotImplementedError(f"Loading model from {model_path.suffix} is not supported yet.")

src/litmodels/io/utils.py

@@ -0,0 +1,39 @@
+import pickle
+from pathlib import Path
+from typing import Any, Union
+
+from lightning_utilities import module_available
+
+if module_available("joblib"):
+    import joblib
+else:
+    joblib = None
+
+
+def dump_pickle(model: Any, path: Union[str, Path]) -> None:
+    """Dump a model to a pickle file.
+
+    Args:
+        model: The model to be pickled.
+        path: The path where the model will be saved.
+    """
+    if joblib is not None:
+        joblib.dump(model, filename=path, compress=7)
+    else:
+        with open(path, "wb") as fp:
+            pickle.dump(model, fp, protocol=pickle.HIGHEST_PROTOCOL)
+
+
+def load_pickle(path: Union[str, Path]) -> Any:
+    """Load a model from a pickle file.
+
+    Args:
+        path: The path to the pickle file.
+
+    Returns:
+        The unpickled model.
+    """
+    if joblib is not None:
+        return joblib.load(path)
+    with open(path, "rb") as fp:
+        return pickle.load(fp)