[BugFix] JSON/orjson compatibility (#1373)

Author: vmoens

pyproject.toml

@@ -40,6 +40,7 @@ dependencies = [
     "importlib_metadata",
     # orjson fails to be installed in python 3.13t
     'orjson ; python_version < "3.13"',
+    "pyvers (>=0.1.0,<0.2.0)",
 ]
 
 [project.urls]

tensordict/_lazy.py

@@ -81,11 +81,6 @@
 from torch import Tensor
 from torch.nn.utils.rnn import pad_sequence
 
-try:
-    import orjson as json
-except ImportError:
-    # Fallback
-    import json
 
 try:
     from functorch import dim as ftdim
@@ -2876,11 +2871,16 @@ def save_metadata(prefix=prefix, self=self):
                 if not prefix.exists():
                     os.makedirs(prefix, exist_ok=True)
                 with open(prefix / "meta.json", "wb") as f:
-                    f.write(
-                        json.dumps(
-                            {"_type": str(type(self)), "stack_dim": self.stack_dim}
-                        )
+                    from tensordict.utils import json_dumps
+
+                    json_str = json_dumps(
+                        {"_type": str(type(self)), "stack_dim": self.stack_dim}
                     )
+                    # Ensure we write bytes to the binary file
+                    if isinstance(json_str, str):
+                        f.write(json_str.encode("utf-8"))
+                    else:
+                        f.write(json_str)
 
             if executor is None:
                 save_metadata()
@@ -4201,7 +4201,14 @@ def save_metadata(data: TensorDictBase, filepath, metadata=None):
                 }
             )
             with open(filepath, "wb") as json_metadata:
-                json_metadata.write(json.dumps(metadata))
+                from tensordict.utils import json_dumps
+
+                json_str = json_dumps(metadata)
+                # Ensure we write bytes to the binary file
+                if isinstance(json_str, str):
+                    json_metadata.write(json_str.encode("utf-8"))
+                else:
+                    json_metadata.write(json_str)
 
         if prefix is not None:
             prefix = Path(prefix)

tensordict/_td.py

@@ -93,11 +93,6 @@
 from torch.nn.utils._named_member_accessor import swap_tensor
 from torch.utils._pytree import tree_map
 
-try:
-    import orjson as json
-except ImportError:
-    # Fallback
-    import json
 try:
     from functorch import dim as ftdim
 
@@ -4175,14 +4170,17 @@ def save_metadata(prefix=prefix, self=self):
                 if not prefix.exists():
                     os.makedirs(prefix, exist_ok=True)
                 with open(prefix / "meta.json", "wb") as f:
-                    f.write(
-                        json.dumps(
-                            {
-                                "_type": str(type(self)),
-                                "index": _index_to_str(self.idx),
-                            }
-                        )
+                    from tensordict.utils import json_dumps
+
+                    metadata_json = json_dumps(
+                        {
+                            "_type": str(type(self)),
+                            "index": _index_to_str(self.idx),
+                        }
                     )
+                    if isinstance(metadata_json, str):
+                        metadata_json = metadata_json.encode("utf-8")
+                    f.write(metadata_json)
 
             if executor is None:
                 save_metadata()
@@ -4682,7 +4680,14 @@ def _save_metadata(data: TensorDictBase, prefix: Path, metadata=None):
         }
     )
     with open(filepath, "wb") as json_metadata:
-        json_metadata.write(json.dumps(metadata))
+        from tensordict.utils import json_dumps
+
+        json_str = json_dumps(metadata)
+        # Ensure we write bytes to the binary file
+        if isinstance(json_str, str):
+            json_metadata.write(json_str.encode("utf-8"))
+        else:
+            json_metadata.write(json_str)
 
 
 # user did specify location and memmap is in wrong place, so we copy

tensordict/base.py

@@ -13,6 +13,9 @@
 import gc
 import importlib
 import importlib.util
+
+# JSON backend is now handled by utils.json_dumps
+import json
 import os.path
 import queue
 import uuid
@@ -111,12 +114,6 @@
 from torch.nn.parameter import Parameter, UninitializedTensorMixin
 from torch.utils._pytree import tree_map
 
-try:
-    import orjson as json
-except ImportError:
-    # Fallback for 3.13
-    import json
-
 try:
     from torch.compiler import is_compiling
 except ImportError:  # torch 2.0
@@ -5355,7 +5352,9 @@ def consolidate(
         else:
             # Convert the dict to json
             try:
-                metadata_dict_json = json.dumps(metadata_dict)
+                from tensordict.utils import json_dumps
+
+                metadata_dict_json = json_dumps(metadata_dict)
             except TypeError as e:
                 raise RuntimeError(
                     "Failed to convert the metatdata to json. "
@@ -5364,6 +5363,8 @@ def consolidate(
                     "If you encounter this error, please file an issue on github."
                 ) from e
             # Represent as a tensor
+            if isinstance(metadata_dict_json, str):
+                metadata_dict_json = metadata_dict_json.encode("utf-8")
             metadata_dict_json = torch.as_tensor(
                 bytearray(metadata_dict_json), dtype=torch.uint8
             )

tensordict/persistent.py

@@ -56,11 +56,6 @@
 )
 from torch import multiprocessing as mp
 
-try:
-    import orjson as json
-except ImportError:
-    # Fallback for 3.13
-    import json
 
 _has_h5 = importlib.util.find_spec("h5py", None) is not None
 
@@ -768,7 +763,14 @@ def save_metadata(data: TensorDictBase, filepath, metadata=None):
                 }
             )
             with open(filepath, "wb") as json_metadata:
-                json_metadata.write(json.dumps(metadata))
+                from tensordict.utils import json_dumps
+
+                json_str = json_dumps(metadata)
+                # Ensure we write bytes to the binary file
+                if isinstance(json_str, str):
+                    json_metadata.write(json_str.encode("utf-8"))
+                else:
+                    json_metadata.write(json_str)
 
         if prefix is not None:
             prefix = Path(prefix)

tensordict/tensorclass.py

@@ -12,6 +12,7 @@
 import dataclasses
 import functools
 import inspect
+
 import multiprocessing.managers
 import multiprocessing.sharedctypes
 import numbers
@@ -80,11 +81,6 @@
 from torch.multiprocessing import Manager
 from torch.utils._pytree import tree_map
 
-try:
-    import orjson as json
-except ImportError:
-    # Fallback for 3.13
-    import json
 try:
     from torch.compiler import is_compiling
 except ImportError:  # torch 2.0
@@ -1532,7 +1528,14 @@ def save_metadata(cls=cls, _non_tensordict=_non_tensordict, prefix=prefix):
                         metadata[key] = value
                     else:
                         to_pickle[key] = value
-                f.write(json.dumps(metadata))
+                from tensordict.utils import json_dumps
+
+                json_str = json_dumps(metadata)
+                # Ensure we write bytes to the binary file
+                if isinstance(json_str, str):
+                    f.write(json_str.encode("utf-8"))
+                else:
+                    f.write(json_str)
                 if to_pickle:
                     with open(prefix / "other.pickle", "wb") as pickle_file:
                         pickle.dump(to_pickle, pickle_file)
@@ -4094,7 +4097,14 @@ def save_metadata(prefix=prefix, self=self):
                     with open(prefix / "pickle.pkl", "wb") as f:
                         pickle.dump(data, f)
                 with open(prefix / "meta.json", "wb") as f:
-                    f.write(json.dumps(jsondict))
+                    from tensordict.utils import json_dumps
+
+                    json_str = json_dumps(jsondict, separators=(",", ":"))
+                    # Ensure we write bytes to the binary file
+                    if isinstance(json_str, str):
+                        f.write(json_str.encode("utf-8"))
+                    else:
+                        f.write(json_str)
 
             if executor is None:
                 save_metadata()