Add method to remove state when serialising model

tempo/serve/base.py

@@ -112,6 +112,18 @@ def loadmethod(self, load_func: LoadMethodSignature) -> LoadMethodSignature:
 
         return load_func
 
+    def __getstate__(self) -> dict:
+        """
+        __getstate__ gets called by pickle before serialising an object to get
+        its internal representation.
+        We override __getstate__ to make sure that the model's internal context
+        is not pickled with the object.
+        """
+        state = self.__dict__.copy()
+        state["context"] = SimpleNamespace()
+
+        return state
+
     @classmethod
     def load(cls, folder: str) -> "BaseModel":
         file_path_pkl = os.path.join(folder, DefaultModelFilename)

tests/serve/test_model.py

@@ -56,7 +56,9 @@ def custom_model(a: np.ndarray) -> np.ndarray:
 #
 # Test lambda function
 #
-@pytest.mark.parametrize("input, expected", [(np.array([[0, 0, 0, 1]]), np.array([[0, 0, 1]]))])
+@pytest.mark.parametrize(
+    "input, expected", [(np.array([[0, 0, 0, 1]]), np.array([[0, 0, 1]]))]
+)
 def test_lambda(input, expected):
     model = Model(
         name="test-iris-sklearn",
@@ -158,7 +160,9 @@ def test_custom_multiheaded_model_tuple(v2_input, expected):
         name="multi-headed",
         platform=ModelFramework.Custom,
     )
-    def custom_multiheaded_model_tuple(a: np.ndarray, b: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    def custom_multiheaded_model_tuple(
+        a: np.ndarray, b: np.ndarray
+    ) -> Tuple[np.ndarray, np.ndarray]:
         return a, b
 
     response = custom_multiheaded_model_tuple.request(v2_input)
@@ -236,3 +240,11 @@ def predict(self, X: str) -> str:
 def test_custom_loadmethod(custom_model):
     pred = custom_model(payload=np.array([1, 2, 3]))
     assert pred == np.array([6])
+
+
+def test_model_save(custom_model: Model):
+    custom_model.save(save_env=False)
+    loaded = Model.load(custom_model.details.local_folder)
+
+    assert len(custom_model.context.__dict__) > 0
+    assert len(loaded.context.__dict__) == 0