zanj
ZANJ
Overview
The ZANJ
format is meant to be a way of saving arbitrary objects to disk, in a way that is flexible, allows keeping configuration and data together, and is human readable. It is very loosely inspired by HDF5 and the derived exdir
format, and the implementation is inspired by npz
files.
- You can take any
SerializableDataclass
from the muutils library and save it to disk -- any large arrays or lists will be stored efficiently as external files in the zip archive, while the basic structure and metadata will be stored in readable JSON files. - You can also specify a special
ConfiguredModel
, which inherits from atorch.nn.Module
which will let you save not just your model weights, but all required configuration information, plus any other metadata (like training logs) in a single file.
This library was originally a module in muutils
Installation
Available on PyPI as zanj
pip install zanj
Usage
You can find a runnable example of this in demo.ipynb
Saving a basic object
Any SerializableDataclass
of basic types can be saved as zanj:
import numpy as np
import pandas as pd
from muutils.json_serialize import SerializableDataclass, serializable_dataclass, serializable_field
from zanj import ZANJ
@serializable_dataclass
class BasicZanj(SerializableDataclass):
a: str
q: int = 42
c: list[int] = serializable_field(default_factory=list)
# initialize a zanj reader/writer
zj = ZANJ()
# create an instance
instance: BasicZanj = BasicZanj("hello", 42, [1, 2, 3])
path: str = "tests/junk_data/path_to_save_instancezanj.zanj"
zj.save(instance, path)
recovered: BasicZanj = zj.read(path)
ZANJ will intelligently handle nested serializable dataclasses, numpy arrays, pytorch tensors, and pandas dataframes:
import torch
import pandas as pd
@serializable_dataclass
class Complicated(SerializableDataclass):
name: str
arr1: np.ndarray
arr2: np.ndarray
iris_data: pd.DataFrame
brain_data: pd.DataFrame
container: list[BasicZanj]
torch_tensor: torch.Tensor
For custom classes, you can specify a serialization_fn
and loading_fn
to handle the logic of converting to and from a json-serializable format:
@serializable_dataclass
class Complicated(SerializableDataclass):
name: str
device: torch.device = serializable_field(
serialization_fn=lambda self: str(self.device),
loading_fn=lambda data: torch.device(data["device"]),
)
Note that loading_fn
takes the dictionary of the whole class -- this is in case you've stored data in multiple fields of the dict which are needed to reconstruct the object.
Saving Models
First, define a configuration class for your model. This class will hold the parameters for your model and any associated objects (like losses and optimizers). The configuration class should be a subclass of SerializableDataclass
and use the serializable_field
function to define fields that need special serialization.
Here's an example that defines a GPT-like model configuration:
from zanj.torchutil import ConfiguredModel, set_config_class
@serializable_dataclass
class MyNNConfig(SerializableDataclass):
input_dim: int
hidden_dim: int
output_dim: int
# store the activation function by name, reconstruct it by looking it up in torch.nn
act_fn: torch.nn.Module = serializable_field(
serialization_fn=lambda x: x.__name__,
loading_fn=lambda x: getattr(torch.nn, x["act_fn"]),
)
# same for the loss function
loss_kwargs: dict = serializable_field(default_factory=dict)
loss_factory: torch.nn.modules.loss._Loss = serializable_field(
default_factory=lambda: torch.nn.CrossEntropyLoss,
serialization_fn=lambda x: x.__name__,
loading_fn=lambda x: getattr(torch.nn, x["loss_factory"]),
)
loss = property(lambda self: self.loss_factory(**self.loss_kwargs))
Then, define your model class. It should be a subclass of ConfiguredModel
, and use the set_config_class
decorator to associate it with your configuration class. The __init__
method should take a single argument, which is an instance of your configuration class. You must also call the superclass __init__
method with the configuration instance.
@set_config_class(MyNNConfig)
class MyNN(ConfiguredModel[MyNNConfig]):
def __init__(self, config: MyNNConfig):
# call the superclass init!
# this will store the model in the zanj_model_config field
super().__init__(config)
# whatever you want here
self.net = torch.nn.Sequential(
torch.nn.Linear(config.input_dim, config.hidden_dim),
config.act_fn(),
torch.nn.Linear(config.hidden_dim, config.output_dim),
)
def forward(self, x):
return self.net(x)
You can now create instances of your model, save them to disk, and load them back into memory:
config = MyNNConfig(
input_dim=10,
hidden_dim=20,
output_dim=2,
act_fn=torch.nn.ReLU,
loss_kwargs=dict(reduction="mean"),
)
# create your model from the config, and save
model = MyNN(config)
fname = "tests/junk_data/path_to_save_modelzanj.zanj"
ZANJ().save(model, fname)
# load by calling the class method `read()`
loaded_model = MyNN.read(fname)
# zanj will actually infer the type of the object in the file
# -- and will warn you if you don't have the correct package installed
loaded_another_way = ZANJ().read(fname)
Configuration
When initializing a ZANJ
object, you can specify some configuration info about saving, such as:
- thresholds for how big an array/table has to be before moving to external file
- compression settings
- error modes
- additional handlers for serialization
# how big an array or list (including pandas DataFrame) can be before moving it from the core JSON file
external_array_threshold: int = ZANJ_GLOBAL_DEFAULTS.external_array_threshold
external_list_threshold: int = ZANJ_GLOBAL_DEFAULTS.external_list_threshold
# compression settings passed to `zipfile` package
compress: bool | int = ZANJ_GLOBAL_DEFAULTS.compress
# for doing very cursed things in your own custom loading or serialization functions
custom_settings: dict[str, Any] | None = ZANJ_GLOBAL_DEFAULTS.custom_settings
# specify additional serialization handlers
handlers_pre: MonoTuple[SerializerHandler] = tuple()
handlers_default: MonoTuple[SerializerHandler] = DEFAULT_SERIALIZER_HANDLERS_ZANJ,
Implementation
The on-disk format is a file <filename>zanj.zanj
is a zip file containing:
__zanj_meta__.json
: a file containing zanj-specific metadata including:- system information
- installed packages
- information about external files
__zanj__.json
: a file containing user-specified data- when an element is too big, it can be moved to an external file
.npy
for numpy arrays or torch tensors.jsonl
for pandas dataframes or large sequences
- list of external files stored in
__zanj_meta__.json
- "$ref" key, specified in
_REF_KEY
in muutils, will have value pointing to external file _FORMAT_KEY
key will detail an external format type
- when an element is too big, it can be moved to an external file
Comparison to other formats
Format | Safe | Zero-copy | Lazy loading | No file size limit | Layout control | Flexibility | Bfloat16 |
---|---|---|---|---|---|---|---|
pickle (PyTorch) | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ |
H5 (Tensorflow) | ✅ | ❌ | ✅ | ✅ | ~ | ~ | ❌ |
HDF5 | ✅ | ? | ✅ | ✅ | ~ | ✅ | ❌ |
SavedModel (Tensorflow) | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ |
MsgPack (flax) | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ |
Protobuf (ONNX) | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ |
Cap'n'Proto | ✅ | ✅ | ~ | ✅ | ✅ | ~ | ❌ |
Numpy (npy,npz) | ✅ | ? | ? | ❌ | ✅ | ❌ | ❌ |
SafeTensors | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ |
exdir | ✅ | ? | ? | ? | ? | ✅ | ❌ |
ZANJ | ✅ | ❌ | ❌* | ✅ | ✅ | ✅ | ❌* |
- Safe: Can I use a file randomly downloaded and expect not to run arbitrary code ?
- Zero-copy: Does reading the file require more memory than the original file ?
- Lazy loading: Can I inspect the file without loading everything ? And loading only some tensors in it without scanning the whole file (distributed setting) ?
- Layout control: Lazy loading, is not necessarily enough since if the information about tensors is spread out in your file, then even if the information is lazily accessible you might have to access most of your file to read the available tensors (incurring many DISK -> RAM copies). Controlling the layout to keep fast access to single tensors is important.
- No file size limit: Is there a limit to the file size ?
- Flexibility: Can I save custom code in the format and be able to use it later with zero extra code ? (~ means we can store more than pure tensors, but no custom code)
- Bfloat16: Does the format support native bfloat16 (meaning no weird workarounds are necessary)? This is becoming increasingly important in the ML world.
*
denotes this feature may be coming at a future date :)
(This table was stolen from safetensors)
1""" 2.. include:: ../README.md 3""" 4 5from __future__ import annotations 6 7from zanj.loading import register_loader_handler 8from zanj.zanj import ZANJ 9 10__all__ = [ 11 "register_loader_handler", 12 "ZANJ", 13 # modules 14 "externals", 15 "loading", 16 "serializing", 17 "torchutil", 18 "zanj", 19]
236def register_loader_handler(handler: LoaderHandler): 237 """register a custom loader handler""" 238 global LOADER_MAP, LOADER_MAP_LOCK 239 with LOADER_MAP_LOCK: 240 LOADER_MAP[handler.uid] = handler
register a custom loader handler
67class ZANJ(JsonSerializer): 68 """Zip up: Arrays in Numpy, JSON for everything else 69 70 given an arbitrary object, throw into a zip file, with arrays stored in .npy files, and everything else stored in a json file 71 72 (basically npz file with json) 73 74 - numpy (or pytorch) arrays are stored in paths according to their name and structure in the object 75 - everything else about the object is stored in a json file `zanj.json` in the root of the archive, via `muutils.json_serialize.JsonSerializer` 76 - metadata about ZANJ configuration, and optionally packages and versions, is stored in a `__zanj_meta__.json` file in the root of the archive 77 78 create a ZANJ-class via `z_cls = ZANJ().create(obj)`, and save/read instances of the object via `z_cls.save(obj, path)`, `z_cls.load(path)`. be sure to pass an **instance** of the object, to make sure that the attributes of the class can be correctly recognized 79 80 """ 81 82 def __init__( 83 self, 84 error_mode: ErrorMode = ZANJ_GLOBAL_DEFAULTS.error_mode, 85 internal_array_mode: ArrayMode = ZANJ_GLOBAL_DEFAULTS.internal_array_mode, 86 external_array_threshold: int = ZANJ_GLOBAL_DEFAULTS.external_array_threshold, 87 external_list_threshold: int = ZANJ_GLOBAL_DEFAULTS.external_list_threshold, 88 compress: bool | int = ZANJ_GLOBAL_DEFAULTS.compress, 89 custom_settings: dict[str, Any] | None = ZANJ_GLOBAL_DEFAULTS.custom_settings, 90 handlers_pre: MonoTuple[SerializerHandler] = tuple(), 91 handlers_default: MonoTuple[ 92 SerializerHandler 93 ] = DEFAULT_SERIALIZER_HANDLERS_ZANJ, 94 ) -> None: 95 super().__init__( 96 array_mode=internal_array_mode, 97 error_mode=error_mode, 98 handlers_pre=handlers_pre, 99 handlers_default=handlers_default, 100 ) 101 102 self.external_array_threshold: int = external_array_threshold 103 self.external_list_threshold: int = external_list_threshold 104 self.custom_settings: dict = ( 105 custom_settings if custom_settings is not None else dict() 106 ) 107 108 # process compression to int if bool given 109 self.compress = compress 110 if isinstance(compress, bool): 111 if compress: 112 self.compress = zipfile.ZIP_DEFLATED 113 else: 114 self.compress = zipfile.ZIP_STORED 115 116 # create the externals, leave it empty 117 self._externals: dict[str, ExternalItem] = dict() 118 119 def externals_info(self) -> dict[str, dict[str, str | int | list[int]]]: 120 """return information about the current externals""" 121 output: dict[str, dict] = dict() 122 123 key: str 124 item: ExternalItem 125 for key, item in self._externals.items(): 126 data = item.data 127 output[key] = { 128 "item_type": item.item_type, 129 "path": item.path, 130 "type(data)": str(type(data)), 131 "len(data)": len(data), 132 } 133 134 if item.item_type == "ndarray": 135 output[key].update(arr_metadata(data)) 136 elif item.item_type.startswith("jsonl"): 137 output[key]["data[0]"] = data[0] 138 139 return { 140 key: val 141 for key, val in sorted(output.items(), key=lambda x: len(x[1]["path"])) 142 } 143 144 def meta(self) -> JSONitem: 145 """return the metadata of the ZANJ archive""" 146 147 serialization_handlers = {h.uid: h.serialize() for h in self.handlers} 148 load_handlers = {h.uid: h.serialize() for h in LOADER_MAP.values()} 149 150 return dict( 151 # configuration of this ZANJ instance 152 zanj_cfg=dict( 153 error_mode=str(self.error_mode), 154 array_mode=str(self.array_mode), 155 external_array_threshold=self.external_array_threshold, 156 external_list_threshold=self.external_list_threshold, 157 compress=self.compress, 158 serialization_handlers=serialization_handlers, 159 load_handlers=load_handlers, 160 ), 161 # system info (python, pip packages, torch & cuda, platform info, git info) 162 sysinfo=json_serialize(SysInfo.get_all(include=("python", "pytorch"))), 163 externals_info=self.externals_info(), 164 timestamp=time.time(), 165 ) 166 167 def save(self, obj: Any, file_path: str | Path) -> str: 168 """save the object to a ZANJ archive. returns the path to the archive""" 169 170 # adjust extension 171 file_path = str(file_path) 172 if not file_path.endswith(".zanj"): 173 file_path += ".zanj" 174 175 # make directory 176 dir_path: str = os.path.dirname(file_path) 177 if dir_path != "": 178 if not os.path.exists(dir_path): 179 os.makedirs(dir_path, exist_ok=False) 180 181 # clear the externals! 182 self._externals = dict() 183 184 # serialize the object -- this will populate self._externals 185 # TODO: calling self.json_serialize again here might be slow 186 json_data: JSONitem = self.json_serialize(self.json_serialize(obj)) 187 188 # open the zip file 189 zipf: zipfile.ZipFile = zipfile.ZipFile( 190 file=file_path, mode="w", compression=self.compress 191 ) 192 193 # store base json data and metadata 194 zipf.writestr( 195 ZANJ_META, 196 json.dumps( 197 self.json_serialize(self.meta()), 198 indent="\t", 199 ), 200 ) 201 zipf.writestr( 202 ZANJ_MAIN, 203 json.dumps( 204 json_data, 205 indent="\t", 206 ), 207 ) 208 209 # store externals 210 for key, (ext_type, ext_data, ext_path) in self._externals.items(): 211 # why force zip64? numpy.savez does it 212 with zipf.open(key, "w", force_zip64=True) as fp: 213 EXTERNAL_STORE_FUNCS[ext_type](self, fp, ext_data) 214 215 zipf.close() 216 217 # clear the externals, again 218 self._externals = dict() 219 220 return file_path 221 222 def read( 223 self, 224 file_path: Union[str, Path], 225 ) -> Any: 226 """load the object from a ZANJ archive 227 # TODO: load only some part of the zanj file by passing an ObjectPath 228 """ 229 file_path = Path(file_path) 230 if not file_path.exists(): 231 raise FileNotFoundError(f"file not found: {file_path}") 232 if not file_path.is_file(): 233 raise FileNotFoundError(f"not a file: {file_path}") 234 235 loaded_zanj: LoadedZANJ = LoadedZANJ( 236 path=file_path, 237 zanj=self, 238 ) 239 240 loaded_zanj.populate_externals() 241 242 return load_item_recursive( 243 loaded_zanj._json_data, 244 path=tuple(), 245 zanj=self, 246 error_mode=self.error_mode, 247 # lh_map=loader_handlers, 248 )
Zip up: Arrays in Numpy, JSON for everything else
given an arbitrary object, throw into a zip file, with arrays stored in .npy files, and everything else stored in a json file
(basically npz file with json)
- numpy (or pytorch) arrays are stored in paths according to their name and structure in the object
- everything else about the object is stored in a json file
zanj.json
in the root of the archive, viamuutils.json_serialize.JsonSerializer
- metadata about ZANJ configuration, and optionally packages and versions, is stored in a
__zanj_meta__.json
file in the root of the archive
create a ZANJ-class via z_cls = ZANJ().create(obj)
, and save/read instances of the object via z_cls.save(obj, path)
, z_cls.load(path)
. be sure to pass an instance of the object, to make sure that the attributes of the class can be correctly recognized
82 def __init__( 83 self, 84 error_mode: ErrorMode = ZANJ_GLOBAL_DEFAULTS.error_mode, 85 internal_array_mode: ArrayMode = ZANJ_GLOBAL_DEFAULTS.internal_array_mode, 86 external_array_threshold: int = ZANJ_GLOBAL_DEFAULTS.external_array_threshold, 87 external_list_threshold: int = ZANJ_GLOBAL_DEFAULTS.external_list_threshold, 88 compress: bool | int = ZANJ_GLOBAL_DEFAULTS.compress, 89 custom_settings: dict[str, Any] | None = ZANJ_GLOBAL_DEFAULTS.custom_settings, 90 handlers_pre: MonoTuple[SerializerHandler] = tuple(), 91 handlers_default: MonoTuple[ 92 SerializerHandler 93 ] = DEFAULT_SERIALIZER_HANDLERS_ZANJ, 94 ) -> None: 95 super().__init__( 96 array_mode=internal_array_mode, 97 error_mode=error_mode, 98 handlers_pre=handlers_pre, 99 handlers_default=handlers_default, 100 ) 101 102 self.external_array_threshold: int = external_array_threshold 103 self.external_list_threshold: int = external_list_threshold 104 self.custom_settings: dict = ( 105 custom_settings if custom_settings is not None else dict() 106 ) 107 108 # process compression to int if bool given 109 self.compress = compress 110 if isinstance(compress, bool): 111 if compress: 112 self.compress = zipfile.ZIP_DEFLATED 113 else: 114 self.compress = zipfile.ZIP_STORED 115 116 # create the externals, leave it empty 117 self._externals: dict[str, ExternalItem] = dict()
119 def externals_info(self) -> dict[str, dict[str, str | int | list[int]]]: 120 """return information about the current externals""" 121 output: dict[str, dict] = dict() 122 123 key: str 124 item: ExternalItem 125 for key, item in self._externals.items(): 126 data = item.data 127 output[key] = { 128 "item_type": item.item_type, 129 "path": item.path, 130 "type(data)": str(type(data)), 131 "len(data)": len(data), 132 } 133 134 if item.item_type == "ndarray": 135 output[key].update(arr_metadata(data)) 136 elif item.item_type.startswith("jsonl"): 137 output[key]["data[0]"] = data[0] 138 139 return { 140 key: val 141 for key, val in sorted(output.items(), key=lambda x: len(x[1]["path"])) 142 }
return information about the current externals
144 def meta(self) -> JSONitem: 145 """return the metadata of the ZANJ archive""" 146 147 serialization_handlers = {h.uid: h.serialize() for h in self.handlers} 148 load_handlers = {h.uid: h.serialize() for h in LOADER_MAP.values()} 149 150 return dict( 151 # configuration of this ZANJ instance 152 zanj_cfg=dict( 153 error_mode=str(self.error_mode), 154 array_mode=str(self.array_mode), 155 external_array_threshold=self.external_array_threshold, 156 external_list_threshold=self.external_list_threshold, 157 compress=self.compress, 158 serialization_handlers=serialization_handlers, 159 load_handlers=load_handlers, 160 ), 161 # system info (python, pip packages, torch & cuda, platform info, git info) 162 sysinfo=json_serialize(SysInfo.get_all(include=("python", "pytorch"))), 163 externals_info=self.externals_info(), 164 timestamp=time.time(), 165 )
return the metadata of the ZANJ archive
167 def save(self, obj: Any, file_path: str | Path) -> str: 168 """save the object to a ZANJ archive. returns the path to the archive""" 169 170 # adjust extension 171 file_path = str(file_path) 172 if not file_path.endswith(".zanj"): 173 file_path += ".zanj" 174 175 # make directory 176 dir_path: str = os.path.dirname(file_path) 177 if dir_path != "": 178 if not os.path.exists(dir_path): 179 os.makedirs(dir_path, exist_ok=False) 180 181 # clear the externals! 182 self._externals = dict() 183 184 # serialize the object -- this will populate self._externals 185 # TODO: calling self.json_serialize again here might be slow 186 json_data: JSONitem = self.json_serialize(self.json_serialize(obj)) 187 188 # open the zip file 189 zipf: zipfile.ZipFile = zipfile.ZipFile( 190 file=file_path, mode="w", compression=self.compress 191 ) 192 193 # store base json data and metadata 194 zipf.writestr( 195 ZANJ_META, 196 json.dumps( 197 self.json_serialize(self.meta()), 198 indent="\t", 199 ), 200 ) 201 zipf.writestr( 202 ZANJ_MAIN, 203 json.dumps( 204 json_data, 205 indent="\t", 206 ), 207 ) 208 209 # store externals 210 for key, (ext_type, ext_data, ext_path) in self._externals.items(): 211 # why force zip64? numpy.savez does it 212 with zipf.open(key, "w", force_zip64=True) as fp: 213 EXTERNAL_STORE_FUNCS[ext_type](self, fp, ext_data) 214 215 zipf.close() 216 217 # clear the externals, again 218 self._externals = dict() 219 220 return file_path
save the object to a ZANJ archive. returns the path to the archive
222 def read( 223 self, 224 file_path: Union[str, Path], 225 ) -> Any: 226 """load the object from a ZANJ archive 227 # TODO: load only some part of the zanj file by passing an ObjectPath 228 """ 229 file_path = Path(file_path) 230 if not file_path.exists(): 231 raise FileNotFoundError(f"file not found: {file_path}") 232 if not file_path.is_file(): 233 raise FileNotFoundError(f"not a file: {file_path}") 234 235 loaded_zanj: LoadedZANJ = LoadedZANJ( 236 path=file_path, 237 zanj=self, 238 ) 239 240 loaded_zanj.populate_externals() 241 242 return load_item_recursive( 243 loaded_zanj._json_data, 244 path=tuple(), 245 zanj=self, 246 error_mode=self.error_mode, 247 # lh_map=loader_handlers, 248 )
load the object from a ZANJ archive
TODO: load only some part of the zanj file by passing an ObjectPath
Inherited Members
- muutils.json_serialize.json_serialize.JsonSerializer
- array_mode
- error_mode
- write_only_format
- handlers
- json_serialize
- hashify