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..cf0bc901 --- /dev/null +++ b/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) diff --git a/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py b/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py index a8fb82fa..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): +# 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: