fix: allow weights_only=False when loading --cache-dataset torch cache

tinyml-tinyverse/tests/test_cache_dataset_weights_only.py

@@ -0,0 +1,145 @@
+"""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
+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
+
+
+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', args=None):
+        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_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)
+        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", side_effect=get_cache_path):
+            dataset, dataset_test, train_sampler, test_sampler = utils.load_data(
+                "/some/datadir", args, dataset_loader_dict={}
+            )
+
+    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)

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):
+# 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(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,11 +228,17 @@ 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', 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))
-        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)
@@ -212,19 +247,19 @@ 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)
+    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))
-        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,
@@ -236,11 +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))
-        # 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 and is_main_process():
+            logger.info("Saving dataset_test to {}".format(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: