Remove memory hog reserialization of block data (#1329)

Author: ml-evs

pydatalab/src/pydatalab/blocks/base.py

@@ -1,5 +1,4 @@
 import functools
-import json
 import random
 import warnings
 from collections.abc import Callable, Sequence
@@ -237,9 +236,7 @@ def to_web(self) -> dict[str, Any]:
         else:
             self.data.pop("warnings", None)
 
-        model = self.block_db_model(**self.data)
-        serialized = json.loads(model.json(exclude_unset=True, exclude_none=True))
-        return serialized
+        return self.block_db_model(**self.data).dict(exclude_unset=True, exclude_none=True)
 
     def process_events(self, events: list[dict] | dict):
         """Handle any supported events passed to the block."""

pydatalab/src/pydatalab/utils.py

@@ -8,7 +8,7 @@
 from math import ceil
 
 import pandas as pd
-from bson import json_util
+from bson import ObjectId
 from flask.json.provider import DefaultJSONProvider
 
 
@@ -43,7 +43,10 @@ def default(o):
         if isinstance(o, (datetime.date, datetime.datetime)):
             return o.isoformat()
 
-        return json_util.default(o)
+        elif isinstance(o, ObjectId):
+            return str(o)
+
+        raise RuntimeError(f"Type {type(o)} not serializable")
 
 
 class BSONProvider(DefaultJSONProvider):

pydatalab/tests/server/test_blocks.py

@@ -256,6 +256,7 @@ def test_uvvis_block_lifecycle(admin_client, default_sample_dict, example_data_d
 
     assert response.status_code == 200
     web_block = response.json["new_block_data"]
+    assert web_block["file_id"] == example_file_ids[0]
     assert "bokeh_plot_data" in web_block
     assert web_block.get("errors") is None
 
@@ -573,7 +574,7 @@ def create_large_xye_file(tmpdir):
     yield fname
 
 
-@pytest.mark.limit_memory("130MB")
+@pytest.mark.limit_memory("110MB")
 def test_large_fake_xrd_data_block_serialization(
     admin_client, default_sample_dict, tmpdir, create_large_xye_file
 ):