[solarsys]

Problem

On the disk-backed path (set_task, create_sample_dataset(in_memory=False)), litdata's BinaryWriter.serialize flattens each sample and infers one serializer per leaf from the first sample only, then zips later samples against it. Any non-tensor field whose length, shape or type varies between samples breaks:

• Later sample longer than the first: the extra items are silently dropped.
• First sample longer than later ones (e.g. floats): the write succeeds, then the read fails with UnicodeDecodeError or returns wrong values. This silent-corruption case was not tested before.
• None in one sample, a list in another: the write crashes.

This hits RawProcessor fields (notes per admission, measurement times), tuple_time_text without a tokenizer, StageNet's optional times, text with None, and any key that is not in a schema. For example, the multimodal mortality task (MultimodalMortalityPredictionMIMIC4: raw discharge, radiology, lab_times) fails on master with 1 and 2 workers ('str' object has no attribute 'device').

Fix

Every non-tensor field is stored as one opaque value per sample. Its layout can no longer vary, so litdata's first-sample inference holds for all samples.

• SampleBuilder.encode: fields whose processor declares stores_tensor = True stay as tensors. Everything else, including non-schema keys, is pickled together into one bytes item.
• SampleDataset.__getitem__ decodes that item. Index access, iteration (shuffled or not), subset, split parts and get_dataloader all go through it, and key order is preserved.
• FeatureProcessor.stores_tensor (default False) is set on processors whose process() always returns a tensor. ImageProcessor makes it a property that follows to_tensor. A tensor-declared processor that returns a non-tensor raises a ValueError that names the field.
• Cache invalidation: storage_format is added to the processed-cache key and to metadata, so caches in the old layout are rebuilt rather than misread.
• Worker failures: set_task sends the task to pool workers pickled. A task that a worker can't unpickle (e.g. a class defined under if __name__ == "__main__") now raises a RuntimeError with guidance instead of hanging the pool.
• RawProcessor docstring: documents the contract, the trust boundary (the cache already holds pickles in schema.pkl) and the ≤2.0.2 note. The float warning is removed, since floats now round-trip.
• Docs: new section "How Processed Samples Are Stored" in docs/api/processors.rst. Example: examples/raw_fields_disk_cache.py.

Processor audit

Stored as tensor (stores_tensor)	Stored opaque
binary / multiclass / multilabel / regression labels, sequence, nested_sequence, nested_floats, deep_nested_*, nested_multi_hot, multi_hot, tensor, timeseries, cehr, audio, image (when to_tensor)	raw, ignore, text, tuple_time_text, stagenet, stagenet_tensor, temporal_timeseries, time_image, graph, plus non-schema keys

Why the litdata cap (>=0.2.64,<0.2.68) is kept, and why it's now safe

The fix doesn't depend on version-specific litdata behavior. Every sample now has the same leaf layout (fixed tensor fields plus one bytes item), which is the case litdata's first-sample inference handles correctly. I verified the round-trip on 0.2.64, 0.2.67 and 0.2.76, so lifting the cap later is a separate decision this PR doesn't block.

Testing

• New tests/core/test_raw_storage.py:
  • 16 shapes: scalar str/int/float, None, empty lists, varying flat str/int/float, nested, tuples, dicts with fixed and varying keys, lists with None, a 100 KB string, bytes, numpy arrays.
  • 4 first-sample orders each: as given, shortest first, longest first, empty first.
  • Checks exact equality between the disk and in-memory datasets, and exact types (tuples stay tuples).
  • Covers every read path, the misdeclared-processor error and the encode layout.
  • set_task round-trip with num_workers 1 and 2, and with split=PatientSplit(...) parts.
• test_split_set_task: the cache-key test is updated, and a test checks that caches in the old layout are not reused.
• Version matrix: the new tests and test_split_set_task pass on litdata 0.2.64, 0.2.67 and 0.2.76.
• Integration: a stand-in for the multimodal mortality task (real schema, synthetic notes) fails on master and round-trips exactly with this branch, with 1 and 2 workers.
• Performance: an 8k-sample epoch takes 0.11 s on master and on this branch, shuffled and unshuffled, and 0.11–0.12 s with variable-length raw fields.
• Full suite (unittest discover, fresh cache): 1437 tests. All pass except 3 modules that need pytest, which isn't installed in that venv. Run under pytest, those modules give 31 passed and 3 failed in test_tuple_time_text_tokenizer. The same 3 fail on master: the HF tokenizer can't load in this environment.
• tools/check_pr_rules.py passes.

Notes

• Caches written by a newer litdata. Caches written by litdata ≥0.2.7x use Arrow chunks and can't be read by 0.2.67. The cache key doesn't include the litdata version, so a machine that switched versions can hit a struct.error in deserialize whenever a cached task_df.ld is re-read. On this machine, 3 test_processor_schemas tests hit this with the shared default cache, because this branch's new cache key forces a re-read; they pass with a fresh cache. It's out of scope here, but adding the litdata version to the cache key might be worth a separate change.
• Collate quirk: collate_fn_dict_with_padding treats a list of 2-tuples as (time, value) pairs. That's existing behavior, now documented rather than changed.

🤖 Generated with Claude Code

[jhnwu3]

It appears this PR might break something in the CI here, that causes some part of the code to go from 6 minute unit tests to 6 hours here. Not entirely sure what or why.

[jhnwu3]

Hey professor, something here breaks the CI. Not sure if you can take a look later.

The runtime of our unit tests is usually ~10 minutes, so it taking upwards of 6 hours and timing out suggests something else is up.

[solarsys]

@jhnwu3 Thanks for flagging the 6-minute-to-6-hour CI slowdown. It was a hang, not slow tests, and this PR did trigger it. The cause is now found, and the fix is in 36b8c44.

Cause. create_sample_dataset(in_memory=False) wrote through litdata.optimize, which starts helper processes even with num_workers=0. When the sample transform raises, litdata never stops those helpers. They keep Python's multiprocessing resource-tracker pipe open, so the next tracker shutdown waits forever. That shutdown can come from a later dataset build or from interpreter exit, which is why one attempt printed Ran 1432 tests … OK and still never finished.

Master never hit this because no test raised inside optimize. This PR's test_misdeclared_tensor_processor_raises_with_field_name raises there on purpose. Users could hit the same hang after any transform error in create_sample_dataset(in_memory=False).

How it was confirmed. Runs on ubuntu-latest with a watchdog and py-spy:

• The main process was blocked in resource_tracker._stop_locked, with two orphaned litdata _upload_fn workers holding the pipe.
• The raising test followed by test_tsv_load hangs every time. All the other new tests followed by test_tsv_load pass.

Fix (36b8c44). create_sample_dataset(in_memory=False) now writes with BinaryWriter in the calling process, the way set_task's single-worker path already does. No processes start, and errors reach the caller directly. A new test fails if any process is started; it trips on the old optimize path.

CI on the fix (merged with current master): 1,473 core tests and 2 NLP tests pass, in a normal 8m20s, and contribution-rules passes. The earlier hang was in the test environment, not a problem with the storage change itself.

[jhnwu3]

After a long session of reading documentation and digging, lgtm so far. Will keep an eye out in future PRs.