From ec838bc0a328321b371e0f6af40986be5a5765dd Mon Sep 17 00:00:00 2001 From: M Platypus Date: Tue, 4 Aug 2026 18:39:44 -0400 Subject: [PATCH 1/3] fix: allow weights_only=False when loading --cache-dataset torch cache Both cache-hit branches in load_data() (training and validation/test data) called torch.load(cache_path) with no weights_only argument. The cached object is a (dataset, datadir) tuple where dataset is one of this project's own Dataset subclasses -- not a type PyTorch 2.6+'s new weights_only=True default allowlists. Any run that reused a --cache-dataset cache written by an earlier run failed with: _pickle.UnpicklingError: Weights only load failed... WeightsUnpickler error: Unsupported global: GLOBAL was not an allowed global by default... This is not a security fix: cache_path is derived from a hash of datadir and is a purely local cache this same process wrote moments earlier, never fetched from a URL or otherwise externally supplied. Explicit weights_only=False is the correct, safe choice here, unlike genuinely untrusted/external checkpoint loads elsewhere in the codebase which should keep the safe default or use safe_globals(). Adds a regression test that reproduces the crash end-to-end through load_data() itself (not just a direct torch.load call), verified to fail pre-fix with the exact UnpicklingError above and pass post-fix. Co-Authored-By: Claude Sonnet 5 --- .../tests/test_cache_dataset_weights_only.py | 49 +++++++++++++++++++ .../tinyml_tinyverse/common/utils/utils.py | 11 ++++- 2 files changed, 58 insertions(+), 2 deletions(-) create mode 100644 tinyml-tinyverse/tests/test_cache_dataset_weights_only.py diff --git a/tinyml-tinyverse/tests/test_cache_dataset_weights_only.py b/tinyml-tinyverse/tests/test_cache_dataset_weights_only.py new file mode 100644 index 00000000..8d0be4c6 --- /dev/null +++ b/tinyml-tinyverse/tests/test_cache_dataset_weights_only.py @@ -0,0 +1,49 @@ +"""Regression test for load_data()'s --cache-dataset torch.load() crash. + +Both cache-hit branches in load_data() (training data and validation data) +called torch.load(cache_path) with no weights_only argument. The cached +object is a (dataset, datadir) tuple where dataset is one of this project's +own Dataset subclasses -- not the kind of type torch's weights_only=True +default (PyTorch 2.6+) allowlists, so any second run using --cache-dataset +(after the first run wrote the cache) failed with +UnpicklingError: Weights only load failed. +""" +import os +import tempfile +from argparse import Namespace +from unittest.mock import patch + +import torch + +from tinyml_tinyverse.common.utils import utils + + +class _FakeDataset: + """Stands in for a real Dataset subclass -- the point is only that it's a + custom class instance, not the specific dataset implementation. Needs + __len__ since load_data() wraps the loaded dataset in a + RandomSampler/SequentialSampler afterward.""" + + def __init__(self, tag): + self.tag = tag + + def __len__(self): + return 4 + + +def test_load_data_reads_a_cached_dataset_without_weights_only_error(): + with tempfile.TemporaryDirectory() as tmp_dir: + cache_path = os.path.join(tmp_dir, "cache.pt") + torch.save((_FakeDataset("cached"), "/some/datadir"), cache_path) + + args = Namespace(cache_dataset=True, dataset_loader="unused", + distributed=False, loader_type="regression") + + with patch.object(utils, "_get_cache_path", return_value=cache_path): + dataset, dataset_test, train_sampler, test_sampler = utils.load_data( + "/some/datadir", args, dataset_loader_dict={} + ) + + assert isinstance(dataset, _FakeDataset) + assert dataset.tag == "cached" + assert isinstance(dataset_test, _FakeDataset) diff --git a/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py b/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py index a8fb82fa..4497d96c 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py +++ b/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py @@ -203,7 +203,13 @@ def load_data(datadir, args, dataset_loader_dict, test_only=False): if args.cache_dataset and os.path.exists(cache_path): # Attention, as the transforms are also cached! logger.info("Loading dataset_train from {}".format(cache_path)) - dataset, _ = torch.load(cache_path) + # weights_only=False is safe here: cache_path is a purely local cache this + # same process wrote moments earlier (derived from a hash of datadir, never + # fetched from a URL or otherwise externally supplied), and the cached object + # is one of this project's own Dataset subclasses -- not the kind of type + # torch's weights_only=True default (PyTorch 2.6+) allowlists, so loading + # this cache without the override raises UnpicklingError on every run. + dataset, _ = torch.load(cache_path, weights_only=False) else: if args.dataset == 'modelmaker': train_folders = os.path.normpath(datadir).split(os.sep) @@ -224,7 +230,8 @@ def load_data(datadir, args, dataset_loader_dict, test_only=False): if args.cache_dataset and os.path.exists(cache_path): # Attention, as the transforms are also cached! logger.info("Loading dataset_test from {}".format(cache_path)) - dataset_test, _ = torch.load(cache_path) + # See the training-cache load above for why weights_only=False is safe here. + dataset_test, _ = torch.load(cache_path, weights_only=False) else: # val_transform = presets.ClassificationPresetEval(crop_size=crop_size, resize_size=resize_size, # interpolation=interpolation, From de8ec33a731f4ebf0ed87d87fe648ca89f731884 Mon Sep 17 00:00:00 2001 From: M Platypus Date: Tue, 4 Aug 2026 20:05:11 -0400 Subject: [PATCH 2/3] fix: --cache-dataset silently made validation data equal training data An independent peer review of this PR's weights_only fix caught a deeper, more serious pre-existing bug it had unblocked: _get_cache_path(datadir) computed the SAME cache path for the training dataset and the validation dataset (both derived only from datadir), and only the training branch ever wrote its cache -- the validation write-back was dead code, commented out with "TODO: Add utils and uncomment the if block". Consequence, verified by reproduction: on the very first run with --cache-dataset enabled, load_data() writes the training dataset to cache_path, then immediately re-derives the identical cache_path for validation and finds it already exists (since it was just written) -- so dataset_test silently becomes a torch.load() of the TRAINING dataset. Not a crash: --cache-dataset appeared to work, but every validation metric from that point on was computed against training data instead of held-out data, at every run using this now-crash-free flag. Fixed by: - Giving _get_cache_path a `tag` parameter ('train'/'val') folded into the hash, so the two datasets get distinct cache paths. - Enabling the validation cache write-back, which was dead code -- without it, fixing only the path collision would mean the validation dataset simply never got cached (safe, but silently defeats the point of --cache-dataset for validation). Rewrites tests/test_cache_dataset_weights_only.py: the previous test patched _get_cache_path with a single fixed return_value for both the training and validation calls, which (ironically) meant it would have passed even under this exact bug -- both calls resolving to the same path is precisely the defect. New tests patch _get_cache_path with a tag-aware side_effect and assert the training and validation datasets come back distinct, and that a missing validation cache actually gets written. Verified all three fail on the pre-this-commit code with the exact symptom (dataset_test.tag == 'cached-train' instead of the expected validation value) and pass post-fix. Co-Authored-By: Claude Sonnet 5 --- .../tests/test_cache_dataset_weights_only.py | 98 ++++++++++++++++--- .../tinyml_tinyverse/common/utils/utils.py | 17 ++-- 2 files changed, 90 insertions(+), 25 deletions(-) diff --git a/tinyml-tinyverse/tests/test_cache_dataset_weights_only.py b/tinyml-tinyverse/tests/test_cache_dataset_weights_only.py index 8d0be4c6..71be3d14 100644 --- a/tinyml-tinyverse/tests/test_cache_dataset_weights_only.py +++ b/tinyml-tinyverse/tests/test_cache_dataset_weights_only.py @@ -1,12 +1,26 @@ -"""Regression test for load_data()'s --cache-dataset torch.load() crash. - -Both cache-hit branches in load_data() (training data and validation data) -called torch.load(cache_path) with no weights_only argument. The cached -object is a (dataset, datadir) tuple where dataset is one of this project's -own Dataset subclasses -- not the kind of type torch's weights_only=True -default (PyTorch 2.6+) allowlists, so any second run using --cache-dataset -(after the first run wrote the cache) failed with -UnpicklingError: Weights only load failed. +"""Regression tests for two --cache-dataset bugs in load_data(). + +1. Both cache-hit branches called torch.load(cache_path) with no + weights_only argument. The cached object is a (dataset, datadir) tuple + where dataset is one of this project's own Dataset subclasses -- not the + kind of type torch's weights_only=True default (PyTorch 2.6+) + allowlists, so any run reusing a --cache-dataset cache failed with + UnpicklingError: Weights only load failed. + +2. An independent peer review of the fix for (1) caught a second, more + serious bug _get_cache_path(datadir) exposed: it computed the SAME + cache path for the training dataset and the validation dataset (both + derived only from datadir), and only the training branch ever wrote its + cache -- the validation write-back was dead code, commented out with + "TODO: Add utils and uncomment the if block". On the very first run + with --cache-dataset enabled, the validation cache-hit check + (os.path.exists(cache_path)) became true immediately after the training + branch's write (same path), so dataset_test silently became a + torch.load() of the TRAINING dataset -- not a crash, just meaningless + validation metrics (train==val) from that point on. Fixed by giving + _get_cache_path a `tag` parameter ('train'/'val') folded into the hash, + producing distinct paths, and by enabling the validation cache + write-back that was dead code. """ import os import tempfile @@ -31,19 +45,71 @@ def __len__(self): return 4 -def test_load_data_reads_a_cached_dataset_without_weights_only_error(): +class _StubPreparable: + """Stands in for what dataset_loader(subset, dataset_dir=..., **kwargs) + returns -- load_data() immediately calls .prepare(**kwargs) on it.""" + + def __init__(self, dataset): + self._dataset = dataset + + def prepare(self, **kwargs): + return self._dataset + + +def _fake_cache_path(tmp_dir): + def _get(datadir, tag='train'): + return os.path.join(tmp_dir, f"cache_{tag}.pt") + return _get + + +def test_get_cache_path_differs_between_train_and_val_for_the_same_datadir(): + train_path = utils._get_cache_path("/some/datadir", tag='train') + val_path = utils._get_cache_path("/some/datadir", tag='val') + assert train_path != val_path + + +def test_load_data_reads_distinct_cached_train_and_val_datasets_without_weights_only_error(): with tempfile.TemporaryDirectory() as tmp_dir: - cache_path = os.path.join(tmp_dir, "cache.pt") - torch.save((_FakeDataset("cached"), "/some/datadir"), cache_path) + get_cache_path = _fake_cache_path(tmp_dir) + torch.save((_FakeDataset("cached-train"), "/some/datadir"), + get_cache_path("/some/datadir", tag='train')) + torch.save((_FakeDataset("cached-val"), "/some/datadir"), + get_cache_path("/some/datadir", tag='val')) args = Namespace(cache_dataset=True, dataset_loader="unused", distributed=False, loader_type="regression") - with patch.object(utils, "_get_cache_path", return_value=cache_path): + with patch.object(utils, "_get_cache_path", side_effect=get_cache_path): dataset, dataset_test, train_sampler, test_sampler = utils.load_data( "/some/datadir", args, dataset_loader_dict={} ) - assert isinstance(dataset, _FakeDataset) - assert dataset.tag == "cached" - assert isinstance(dataset_test, _FakeDataset) + assert dataset.tag == "cached-train" + assert dataset_test.tag == "cached-val" + + +def test_load_data_writes_a_validation_cache_when_missing(): + """The validation cache write-back was dead code, so --cache-dataset + never actually cached the validation dataset even after a successful + training-cache write. Confirm the val cache file gets written now.""" + with tempfile.TemporaryDirectory() as tmp_dir: + get_cache_path = _fake_cache_path(tmp_dir) + train_cache_path = get_cache_path("/some/datadir", tag='train') + val_cache_path = get_cache_path("/some/datadir", tag='val') + torch.save((_FakeDataset("cached-train"), "/some/datadir"), train_cache_path) + assert not os.path.exists(val_cache_path) + + def _fake_loader(subset, dataset_dir, **kwargs): + return _StubPreparable(_FakeDataset(f"loaded-{subset}")) + + args = Namespace(cache_dataset=True, dataset_loader="fake", dataset="local", + data_path="/some/datadir", distributed=False, loader_type="regression") + + with patch.object(utils, "_get_cache_path", side_effect=get_cache_path): + dataset, dataset_test, _, _ = utils.load_data( + "/some/datadir", args, dataset_loader_dict={"fake": _fake_loader} + ) + + assert dataset.tag == "cached-train" + assert dataset_test.tag == "loaded-val" + assert os.path.exists(val_cache_path) diff --git a/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py b/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py index 4497d96c..21a5b6bf 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py +++ b/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py @@ -162,9 +162,9 @@ def collate_fn(batch): return raw_tensors, tensors, targets -def _get_cache_path(filepath): +def _get_cache_path(filepath, tag='train'): import hashlib - h = hashlib.sha1(filepath.encode()).hexdigest() + h = hashlib.sha1(f'{tag}:{filepath}'.encode()).hexdigest() cache_path = os.path.join("~", ".torch", "audio_classification", "datasets", "audiofolder", h[:10] + ".pt") cache_path = os.path.expanduser(cache_path) return cache_path @@ -199,7 +199,7 @@ def load_data(datadir, args, dataset_loader_dict, test_only=False): logger.info("Loading training data") st = timeit.default_timer() - cache_path = _get_cache_path(datadir) + cache_path = _get_cache_path(datadir, tag='train') if args.cache_dataset and os.path.exists(cache_path): # Attention, as the transforms are also cached! logger.info("Loading dataset_train from {}".format(cache_path)) @@ -226,7 +226,7 @@ def load_data(datadir, args, dataset_loader_dict, test_only=False): logger.info("Loading validation data") st = timeit.default_timer() - cache_path = _get_cache_path(datadir) + cache_path = _get_cache_path(datadir, tag='val') if args.cache_dataset and os.path.exists(cache_path): # Attention, as the transforms are also cached! logger.info("Loading dataset_test from {}".format(cache_path)) @@ -243,11 +243,10 @@ def load_data(datadir, args, dataset_loader_dict, test_only=False): dataset_test = dataset_loader("val", dataset_dir=args.data_path, validation_list=val_list, **vars(args)).prepare(**vars(args)) else: dataset_test = dataset_loader("val", dataset_dir=args.data_path, **vars(args)).prepare(**vars(args)) - # TODO: Add utils and uncomment the if block - # if args.cache_dataset: - # logger.info("Saving dataset_test to {}".format(cache_path)) - # utils.mkdir(os.path.dirname(cache_path)) - # utils.save_on_master((dataset_test, datadir), cache_path) + if args.cache_dataset: + logger.info("Saving dataset_test to {}".format(cache_path)) + mkdir(os.path.dirname(cache_path)) + save_on_master((dataset_test, datadir), cache_path) logger.info("Took {:.2f} seconds".format(timeit.default_timer() - st)) logger.info("\nCreating data loaders") if args.distributed: From d0d1e2b276eb152ec3a4e91e33fc0cb7ae7bd045 Mon Sep 17 00:00:00 2001 From: M Platypus Date: Tue, 4 Aug 2026 21:49:42 -0400 Subject: [PATCH 3/3] fix: cache key ignores dataset config; cache writes were not atomic Two major CodeRabbit findings on this PR's cache handling, both extending pre-existing gaps the earlier reviews had also noted: 1. _get_cache_path keyed only on tag+datadir, but the cached object is the fully PREPARED dataset -- the code's own comment warns "Attention, as the transforms are also cached!". A run changing transforms, loader, frame size, etc. against the same datadir silently reused a dataset prepared under the OLD configuration. Now a documented tuple of dataset-shaping args (_CACHE_KEY_ARGS: loader, annotation, transform, and framing settings) is folded into the hash via getattr-with-None-default reprs -- stable for equivalent configs, different for any change. Deliberately not exhaustive (an exotic loader could consume an arg not listed; extend as needed) and biased toward over-invalidation, which costs a cache miss, never wrong data. Side effect: existing caches are invalidated once more (key changes); they were already invalidated by the tag change in the previous commit, so no additional migration impact. 2. Cache writes used a plain torch.save straight to the final path -- another rank (or a later run, after a mid-save kill) could see the file exist and read it partially written. New _save_cache_atomically helper writes to a .tmp sibling and publishes via os.replace (atomic on POSIX within a filesystem); both write sites now use it, wrapped in is_main_process() to preserve save_on_master's rank-0-only semantics. New tests pin both: config change produces a different key (same config produces the same key), and the atomic helper leaves no .tmp file and a loadable cache. Both verified to fail pre-fix and pass post-fix. Co-Authored-By: Claude Sonnet 5 --- .../tests/test_cache_dataset_weights_only.py | 32 ++++++++++++- .../tinyml_tinyverse/common/utils/utils.py | 47 +++++++++++++++---- 2 files changed, 68 insertions(+), 11 deletions(-) diff --git a/tinyml-tinyverse/tests/test_cache_dataset_weights_only.py b/tinyml-tinyverse/tests/test_cache_dataset_weights_only.py index 71be3d14..cf0bc901 100644 --- a/tinyml-tinyverse/tests/test_cache_dataset_weights_only.py +++ b/tinyml-tinyverse/tests/test_cache_dataset_weights_only.py @@ -57,7 +57,7 @@ def prepare(self, **kwargs): def _fake_cache_path(tmp_dir): - def _get(datadir, tag='train'): + def _get(datadir, tag='train', args=None): return os.path.join(tmp_dir, f"cache_{tag}.pt") return _get @@ -68,6 +68,36 @@ def test_get_cache_path_differs_between_train_and_val_for_the_same_datadir(): assert train_path != val_path +def test_get_cache_path_differs_when_dataset_config_changes(): + """The cached object is the fully PREPARED dataset -- transforms and + loader settings are baked in ("Attention, as the transforms are also + cached!"). A run with a different transform pipeline against the same + datadir must therefore miss the cache, not silently reuse a dataset + prepared under the old configuration.""" + args_a = Namespace(feat_ext_transform=["MFCC"], dataset_loader="GenericTSDataset") + args_b = Namespace(feat_ext_transform=["RAW"], dataset_loader="GenericTSDataset") + + path_a = utils._get_cache_path("/some/datadir", tag='train', args=args_a) + path_a_again = utils._get_cache_path("/some/datadir", tag='train', args=args_a) + path_b = utils._get_cache_path("/some/datadir", tag='train', args=args_b) + + assert path_a == path_a_again # stable for an equivalent configuration + assert path_a != path_b # any config change invalidates + + +def test_save_cache_atomically_publishes_only_a_complete_file(tmp_path): + """A reader that sees cache_path exist must never read a partial file: + the write goes to a .tmp sibling and is published via os.replace.""" + cache_path = str(tmp_path / "sub" / "cache_train.pt") + + utils._save_cache_atomically((_FakeDataset("payload"), "/some/datadir"), cache_path) + + assert os.path.exists(cache_path) + assert not os.path.exists(cache_path + ".tmp") + dataset, _ = torch.load(cache_path, weights_only=False) + assert dataset.tag == "payload" + + def test_load_data_reads_distinct_cached_train_and_val_datasets_without_weights_only_error(): with tempfile.TemporaryDirectory() as tmp_dir: get_cache_path = _fake_cache_path(tmp_dir) diff --git a/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py b/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py index 21a5b6bf..8c2e7dee 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py +++ b/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py @@ -162,14 +162,43 @@ def collate_fn(batch): return raw_tensors, tensors, targets -def _get_cache_path(filepath, tag='train'): +# Dataset-shaping args folded into the cache key: the cached object is the +# fully PREPARED dataset ("Attention, as the transforms are also cached!"), so +# a run that changes any of these against the same datadir must miss the cache +# rather than silently reuse a dataset prepared under the old configuration. +# Not exhaustive -- prepare(**vars(args)) passes everything through, so an +# exotic loader could consume an arg not listed here; extend as needed. +# Over-invalidation (a needless cache miss) is safe; under-invalidation is not. +_CACHE_KEY_ARGS = ( + 'dataset', 'dataset_loader', 'annotation_prefix', 'loader_type', + 'data_proc_transforms', 'feat_ext_transform', 'augmentation_transform', + 'transforms', 'frame_size', 'stride_size', 'sampling_rate', 'new_sr', + 'variables', 'resampling_factor', +) + + +def _get_cache_path(filepath, tag='train', args=None): import hashlib - h = hashlib.sha1(f'{tag}:{filepath}'.encode()).hexdigest() + key = f'{tag}:{filepath}' + if args is not None: + config = [(k, repr(getattr(args, k, None))) for k in _CACHE_KEY_ARGS] + key += f':{config!r}' + h = hashlib.sha1(key.encode()).hexdigest() cache_path = os.path.join("~", ".torch", "audio_classification", "datasets", "audiofolder", h[:10] + ".pt") cache_path = os.path.expanduser(cache_path) return cache_path +def _save_cache_atomically(payload, cache_path): + """Write to a temp file in the target directory, then atomically publish + via os.replace -- a reader that sees cache_path exist never reads a + partially-written file (e.g. another rank, or a run killed mid-save).""" + mkdir(os.path.dirname(cache_path)) + tmp_path = cache_path + '.tmp' + torch.save(payload, tmp_path) + os.replace(tmp_path, cache_path) + + def load_data(datadir, args, dataset_loader_dict, test_only=False): # Data loading code logger = getLogger("root.load_data") @@ -199,7 +228,7 @@ def load_data(datadir, args, dataset_loader_dict, test_only=False): logger.info("Loading training data") st = timeit.default_timer() - cache_path = _get_cache_path(datadir, tag='train') + cache_path = _get_cache_path(datadir, tag='train', args=args) if args.cache_dataset and os.path.exists(cache_path): # Attention, as the transforms are also cached! logger.info("Loading dataset_train from {}".format(cache_path)) @@ -218,15 +247,14 @@ def load_data(datadir, args, dataset_loader_dict, test_only=False): dataset = dataset_loader("training", dataset_dir=args.data_path, training_list=training_list, **vars(args)).prepare(**vars(args)) else: dataset = dataset_loader("training", dataset_dir=args.data_path, **vars(args)).prepare(**vars(args)) - if args.cache_dataset: + if args.cache_dataset and is_main_process(): logger.info("Saving dataset_train to {}".format(cache_path)) - mkdir(os.path.dirname(cache_path)) - save_on_master((dataset, datadir), cache_path) + _save_cache_atomically((dataset, datadir), cache_path) logger.info("Took {0:.2f} seconds".format(timeit.default_timer() - st)) logger.info("Loading validation data") st = timeit.default_timer() - cache_path = _get_cache_path(datadir, tag='val') + cache_path = _get_cache_path(datadir, tag='val', args=args) if args.cache_dataset and os.path.exists(cache_path): # Attention, as the transforms are also cached! logger.info("Loading dataset_test from {}".format(cache_path)) @@ -243,10 +271,9 @@ def load_data(datadir, args, dataset_loader_dict, test_only=False): dataset_test = dataset_loader("val", dataset_dir=args.data_path, validation_list=val_list, **vars(args)).prepare(**vars(args)) else: dataset_test = dataset_loader("val", dataset_dir=args.data_path, **vars(args)).prepare(**vars(args)) - if args.cache_dataset: + if args.cache_dataset and is_main_process(): logger.info("Saving dataset_test to {}".format(cache_path)) - mkdir(os.path.dirname(cache_path)) - save_on_master((dataset_test, datadir), cache_path) + _save_cache_atomically((dataset_test, datadir), cache_path) logger.info("Took {:.2f} seconds".format(timeit.default_timer() - st)) logger.info("\nCreating data loaders") if args.distributed: