-
Notifications
You must be signed in to change notification settings - Fork 4.3k
I have a use case where I want to save some information, which can consist of numpy nd array of variable shape, numpy 1D arrays, objects like pytorch model state_dict, pytorch optimizer state_dict, scalars like floats, ints, strings, custom types like MetricsInfo etc. These formats can be encoded as various datatypes in arrow which is straight forward for primitive types, tricks like #48099 for variable shape tensors and binary type for other python objects.
I only want to use a single file for all these. The information will be generated periodically like after each epoch in training deep learning models. So at each period, say epoch end, I need to save this information to the file. This is important because if the run is interrupted, I don't want to lose all information till the current epoch. Nor do I want to pressure the memory by buffering all information for a single write.
Can this be done with files of arrow or parquet format?
EDIT: after adding the data, I would like to get random access reads to the saved data
All reactions
Replies: 3 comments 5 replies
yes, you can use either. you just need to be able to represent your data as a recordbatch as a precondition (or any type a write function accepts).
All reactions
Hey @drin ,
I forgot to mention in my top post but after saving the record batches I would like to get random access reads to the data at specific epoch. For eg: after training the model and doing further analysis, I'll have a list of epoch number with best performance. So, I would like to only get data at these best epoch numbers or single best epoch number
All reactions
you can make epoch a column of your record batch or put some index-like structure in the schema metadata that you can use to identify record batch index from epoch.
All reactions
Quick trade-off summary for "periodic writes + random access + crash tolerant": Parquet writes row groups per call, but the file is only valid once FileWriter::Close() writes the footer, so a crash mid-training can lose everything. Arrow IPC stream format is append-safe (no footer, batches self-contained with length prefix) but has no random access. Arrow IPC file format has random access via its footer but the same close-to-read issue. In practice the cleanest shape for your case is one IPC file per epoch plus an append-only JSONL manifest (fsync per line), random access goes through the manifest, crash loss is bounded to the in-flight file. Dataset API unions the per-epoch files for scans when you need to read everything.
All reactions
@MukundaKatta Please don't post random AI-generated responses. People can use AI by themselves to get answers if they want to. Thanks in advance.
All reactions
Got it, thanks for the feedback.
I used AI to help structure my thinking, but I went through the problem and verified the details myself before posting. I’ll make sure future responses are more concise and clearly reflect my own reasoning.
Appreciate the heads-up.
All reactions
Fair callout, @pitrou. You are right, I should not have posted the long AI-drafted version in the first place. Apologies for the noise on the thread.
All reactions
I have similar use-case. If you use IPC file and there is a crash you can recover most(flushed) data by skipping file header and reading it as if IPC stream. That is what I do to recover missing footer. For 1 batch=1epoch this should work for you.