fix: unsafe non-blocking H2D transfers and MPS fallback timing

tinyml-modelmaker/tinyml_modelmaker/run_tinyml_modelmaker.py

@@ -34,6 +34,14 @@
 import os
 import sys
 
+# PyTorch's MPS-fallback flag is read once when its MPS backend is registered
+# (during `import torch`, pulled in transitively below by `import tinyml_modelmaker`)
+# -- setting it later via os.environ from within already-running Python code has no
+# effect. It must be in the process environment before torch is ever imported, so it
+# is set here, at the top of this script, ahead of any project import. Harmless on
+# CUDA/CPU since it only changes MPS dispatch behavior.
+os.environ.setdefault('PYTORCH_ENABLE_MPS_FALLBACK', '1')
+
 import yaml
 
 logger = logging.getLogger(__name__)

tinyml-tinyverse/tests/test_anomalydetection_train_device_crash.py

@@ -0,0 +1,156 @@
+"""Regression test for a crash introduced by this session's own H2D/MPS fix.
+
+timeseries_anomalydetection/train.py's get_reconstruction_errors_stats()
+gates non_blocking transfers with `device.type == 'cuda'` (the fix applied
+across the codebase to stop unsafe non_blocking=True on non-CUDA devices).
+Before that fix, non_blocking was hardcoded True and never touched `device`
+beyond passing it straight to `.to(device, ...)`, which tolerates a plain
+device string ('cuda') just fine.
+
+The fix exposed a latent bug in the caller: main() called this function
+with `args.device` -- the raw argparse string ('cuda' by default, never
+converted to a torch.device anywhere in this file) -- instead of `device`,
+the torch.device already constructed by setup_training_environment() and
+in scope in main(). `device.type` on a plain str raises AttributeError, so
+this crashed on every anomaly-detection training run, right after export,
+while calculating the detection threshold -- not an edge case, the normal
+path.
+
+Fixed by passing the existing local `device` (torch.device) instead of
+`args.device` (str) at the call site.
+
+Two tests:
+1. A function-contract test characterizing get_reconstruction_errors_stats()
+   directly: it works with a real torch.device and crashes with a raw
+   device string -- this is the exact defect surface and locks in the
+   function's contract against future regressions.
+2. A main()-level test that actually drives the real call site (heavily
+   mocking every other dependency, with args.start_epoch == args.epochs so
+   the training loop is skipped and get_reconstruction_errors_stats is
+   reached quickly) and asserts what main() actually passes as the device
+   argument -- this is the one that genuinely exercises the fixed line,
+   since test 1 alone would pass identically whether or not the call site
+   itself were fixed.
+"""
+from argparse import Namespace
+from contextlib import ExitStack
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pytest
+import torch
+
+from tinyml_tinyverse.references.timeseries_anomalydetection import train as anomaly_train
+
+
+class _FakeAnomalyDataset(torch.utils.data.Dataset):
+    """(raw, data, label) tuples matching what utils.collate_fn expects, with
+    data shaped (C, H, W) = (1, 4, 4) so the per-sample reconstruction-error
+    reduction (dim=(1, 2, 3) over an unsqueezed (1, C, H, W) tensor) has real
+    dimensions to reduce over."""
+
+    classes = ["a", "b"]
+    X = np.zeros((4, 3, 4), dtype=np.float32)
+
+    def __len__(self):
+        return 2
+
+    def __getitem__(self, idx):
+        return torch.tensor(idx), torch.zeros(1, 4, 4), torch.tensor(0)
+
+
+def _fake_ort_sess():
+    sess = MagicMock()
+    sess.run.return_value = [np.zeros((1, 1, 4, 4), dtype=np.float32)]
+    return sess
+
+
+def _fake_data_loader():
+    return torch.utils.data.DataLoader(
+        _FakeAnomalyDataset(), batch_size=2, collate_fn=anomaly_train.utils.collate_fn)
+
+
+def test_get_reconstruction_errors_stats_works_with_a_real_torch_device():
+    """The fixed call site now passes a real torch.device -- confirm the
+    function actually works with one (device.type resolves normally)."""
+    with patch.object(anomaly_train.ort, "InferenceSession", return_value=_fake_ort_sess()):
+        mean, std = anomaly_train.get_reconstruction_errors_stats(
+            generic_model=True, model_path="/fake/model.onnx",
+            device=torch.device("cpu"), data_loader=_fake_data_loader(),
+        )
+
+    assert torch.is_tensor(mean)
+    assert torch.is_tensor(std)
+
+
+def test_get_reconstruction_errors_stats_crashes_if_given_a_raw_device_string():
+    """Characterizes the exact regression the fix closes: this is what the
+    pre-fix call site (`get_reconstruction_errors_stats(..., args.device,
+    ...)`, args.device being the unconverted argparse string) actually
+    passed, and it crashes immediately on `device.type`."""
+    with patch.object(anomaly_train.ort, "InferenceSession", return_value=_fake_ort_sess()):
+        with pytest.raises(AttributeError, match="'str' object has no attribute 'type'"):
+            anomaly_train.get_reconstruction_errors_stats(
+                generic_model=True, model_path="/fake/model.onnx",
+                device="cpu", data_loader=_fake_data_loader(),
+            )
+
+
+def test_main_passes_a_torch_device_not_the_raw_args_device_string():
+    """Drives the real main() (heavily mocked elsewhere) far enough to reach
+    the actual call site and captures what it passes. This is the test that
+    genuinely exercises the fixed line -- the two tests above only
+    characterize the callee's contract and would pass unchanged whether or
+    not main()'s call site itself were fixed."""
+    args = Namespace(
+        quantization=True, ondevice_training=False, model='dummy', model_config=None,
+        model_spec=None, dual_op=False, output_int=True, auto_quantization=False,
+        weight_bitwidth=8, activation_bitwidth=8, epochs=1, start_epoch=1,  # zero iterations
+        quantization_method='QAT', distributed=False, apex=False, print_freq=10,
+        output_dir='/tmp/fake-output', autoquant_tolerance_anomaly=0.1,
+        opset_version=17, generic_model=True, gen_golden_vectors=False,
+        device='cuda',  # the raw argparse string that must NOT reach get_reconstruction_errors_stats
+    )
+    dataset = _FakeAnomalyDataset()
+    anomaly_train.dataset_load_state['dataset'] = dataset
+    anomaly_train.dataset_load_state['dataset_test'] = dataset
+    anomaly_train.dataset_load_state['train_sampler'] = None
+    anomaly_train.dataset_load_state['test_sampler'] = None
+
+    real_device = torch.device("cpu")
+    fake_loaders = ([1, 2], [1, 2])
+
+    with ExitStack() as stack:
+        stack.enter_context(patch.object(
+            anomaly_train, "setup_training_environment", return_value=(anomaly_train.getLogger("test"), real_device)))
+        stack.enter_context(patch.object(anomaly_train, "prepare_transforms"))
+        stack.enter_context(patch.object(anomaly_train, "create_data_loaders", return_value=fake_loaders))
+        stack.enter_context(patch.object(anomaly_train.models, "get_model", return_value=torch.nn.Identity()))
+        stack.enter_context(patch.object(anomaly_train, "log_model_summary"))
+        stack.enter_context(patch.object(anomaly_train, "load_pretrained_weights", side_effect=lambda model, a, l: model))
+        stack.enter_context(patch.object(anomaly_train, "handle_export_only", return_value=False))
+        stack.enter_context(patch.object(anomaly_train, "move_model_to_device"))
+        stack.enter_context(patch.object(anomaly_train, "compile_model_if_enabled", side_effect=lambda model, a, l: model))
+        stack.enter_context(patch.object(anomaly_train.utils, "quantization_wrapped_model", side_effect=lambda model, *a, **kw: model))
+        stack.enter_context(patch.object(anomaly_train, "setup_optimizer_and_scheduler", return_value=(MagicMock(), MagicMock())))
+        stack.enter_context(patch.object(
+            anomaly_train, "setup_distributed_model", side_effect=lambda model, a, d: (model, model, None)))
+        stack.enter_context(patch.object(anomaly_train, "resume_from_checkpoint"))
+        stack.enter_context(patch.object(anomaly_train, "get_amp_context", return_value=MagicMock()))
+        stack.enter_context(patch.object(anomaly_train, "get_grad_scaler", return_value=None))
+        stack.enter_context(patch.object(anomaly_train.utils, "export_model"))
+        stack.enter_context(patch.object(anomaly_train, "log_training_time"))
+        stack.enter_context(patch.object(anomaly_train, "shutdown_data_loaders"))
+        mock_get_stats = stack.enter_context(patch.object(
+            anomaly_train, "get_reconstruction_errors_stats",
+            return_value=(torch.tensor(0.0), torch.tensor(0.0))))
+
+        anomaly_train.main(0, args)
+
+    mock_get_stats.assert_called_once()
+    passed_device = mock_get_stats.call_args[0][2]
+    assert passed_device is real_device, (
+        f"main() passed {passed_device!r} (type {type(passed_device).__name__}) as the device "
+        "argument -- expected the real torch.device from setup_training_environment(), not "
+        "args.device (a raw string)."
+    )

tinyml-tinyverse/tests/test_onnx_robustness_bugs.py

@@ -0,0 +1,102 @@
+"""Regression tests for two test_onnx.py robustness bugs.
+
+1. timeseries_anomalydetection/test_onnx.py's get_reconstruction_errors_stats()
+   built its DataLoader with `pin_memory=True if args.gpu > 0 else False`, but
+   the script's argparser (common/test_onnx_base.py) only ever defines
+   `--gpus` (plural), never `--gpu`. Every call raised:
+     AttributeError: 'Namespace' object has no attribute 'gpu'
+   before any model loading or data processing even started. The same
+   file's main() had a related (non-crashing but still wrong) variant:
+   `pin_memory=True if gpu > 0 else False`, gating on the DDP rank
+   parameter rather than device type. Both are now unconditional
+   `pin_memory=True`, matching every sibling test_onnx.py.
+
+2. audio_classification/test_onnx.py never imported or called
+   shutdown_data_loaders() on its DataLoader, unlike every one of its five
+   sibling test_onnx.py scripts (image_classification, timeseries_
+   classification, timeseries_forecasting, timeseries_regression,
+   timeseries_anomalydetection), leaking DataLoader worker processes /
+   POSIX semaphores whenever --workers > 0.
+"""
+import tempfile
+from argparse import Namespace
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pytest
+import torch
+
+from tinyml_tinyverse.references.audio_classification import test_onnx as audio_test_onnx
+from tinyml_tinyverse.references.timeseries_anomalydetection import test_onnx as anomaly_test_onnx
+
+
+class _FakeAnomalyDataset(torch.utils.data.Dataset):
+    """(raw, data, label) tuples matching what utils.collate_fn expects, with
+    data shaped (C, H, W) = (1, 4, 4) so the per-sample reconstruction-error
+    reduction (dim=(1, 2, 3) over an unsqueezed (1, C, H, W) tensor) has real
+    dimensions to reduce over."""
+
+    def __len__(self):
+        return 2
+
+    def __getitem__(self, idx):
+        return torch.tensor(idx), torch.zeros(1, 4, 4), torch.tensor(0)
+
+
+class _FakeAudioDataset(torch.utils.data.Dataset):
+    classes = ["a", "b"]
+
+    def __len__(self):
+        return 2
+
+    def __getitem__(self, idx):
+        return torch.tensor(idx), torch.zeros(1, 4), torch.tensor(0)
+
+
+def _fake_ort_sess():
+    sess = MagicMock()
+    sess.run.return_value = [np.zeros((1, 1, 4, 4), dtype=np.float32)]
+    return sess
+
+
+def test_get_reconstruction_errors_stats_does_not_crash_without_args_gpu():
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        args = Namespace(
+            output_dir=tmp_dir, lis=None, DEBUG=False, seed=0, device="cpu",
+            data_path="/fake/data", batch_size=2, workers=0,
+            gpus=1,  # note: no `gpu` attribute -- matches the real argparser
+            model_path="/fake/model.onnx", generic_model=True,
+        )
+        fake_dataset = _FakeAnomalyDataset()
+
+        with patch.object(anomaly_test_onnx, "prepare_transforms"), \
+             patch.object(anomaly_test_onnx.utils, "load_data",
+                          return_value=(fake_dataset, fake_dataset, None, None)), \
+             patch.object(anomaly_test_onnx, "load_onnx_model",
+                          return_value=(_fake_ort_sess(), "input", "output")):
+            mean, std = anomaly_test_onnx.get_reconstruction_errors_stats(args)
+
+    assert torch.is_tensor(mean)
+    assert torch.is_tensor(std)
+
+
+def test_audio_test_onnx_shuts_down_data_loader_even_when_model_load_fails():
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        args = Namespace(
+            output_dir=tmp_dir, lis=None, DEBUG=False, seed=0, device="cpu",
+            data_path="/fake/data", batch_size=2, workers=0,
+            model_path="/fake/model.onnx", generic_model=True,
+            distributed=False, nn_for_feature_extraction=False,
+        )
+        fake_dataset = _FakeAudioDataset()
+
+        with patch.object(audio_test_onnx.utils, "init_distributed_mode"), \
+             patch.object(audio_test_onnx, "prepare_transforms"), \
+             patch.object(audio_test_onnx.utils, "load_data",
+                          return_value=(fake_dataset, fake_dataset, None, None)), \
+             patch.object(audio_test_onnx, "load_onnx_model", side_effect=RuntimeError("boom")), \
+             patch.object(audio_test_onnx, "shutdown_data_loaders") as mock_shutdown:
+            with pytest.raises(RuntimeError):
+                audio_test_onnx.main(0, args)
+
+        mock_shutdown.assert_called_once()

tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py

@@ -1294,11 +1294,13 @@ def evaluate_forecasting(model, criterion, data_loader, device, transform=None,
     targets=[]
     outputs=[]
 
+    # See evaluate_classification for why non_blocking must be gated on CUDA.
+    non_blocking = (device.type == 'cuda')
     with torch.no_grad():
         for _, data, target in metric_logger.log_every(data_loader, print_freq, header):
             # Move data and target to the specified device
-            data = data.float().to(device, non_blocking=True)
-            target = target.float().to(device, non_blocking=True)
+            data = data.float().to(device, non_blocking=non_blocking)
+            target = target.float().to(device, non_blocking=non_blocking)
 
             # Apply transformation if provided
             if transform:
@@ -1365,14 +1367,16 @@ def evaluate_regression(model, criterion, data_loader, device, transform, log_su
     print_freq = print_freq if print_freq else len(data_loader)
     header = f'Test: {log_suffix}'
 
+    # See evaluate_classification for why non_blocking must be gated on CUDA.
+    non_blocking = (device.type == 'cuda')
     with torch.no_grad():
         val_loss = 0
         target_list = []
         predictions_list = []
         # for _, data, target in metric_logger.log_every(data_loader, print_freq, header):
         for _, data, target in data_loader:
-            data = data.float().to(device, non_blocking=True)
-            target = target.float().to(device, non_blocking=True)
+            data = data.float().to(device, non_blocking=non_blocking)
+            target = target.float().to(device, non_blocking=non_blocking)
 
             if transform:
                 data = transform(data)
@@ -1472,11 +1476,13 @@ def evaluate_anomalydetection(
     print_freq = print_freq if print_freq else len(data_loader)
     header = f'Validation{log_suffix} - Epoch[{epoch}]: '
 
+    # See evaluate_classification for why non_blocking must be gated on CUDA.
+    non_blocking = (device.type == 'cuda')
     with torch.no_grad():
         for _, data, labels in metric_logger.log_every(data_loader, print_freq, header):
             # for data, target in data_loader:
-            data = data.float().to(device, non_blocking=True)
-            #In anomlay detection with auto encoder, the target and the input data both are same. 
+            data = data.float().to(device, non_blocking=non_blocking)
+            #In anomlay detection with auto encoder, the target and the input data both are same.
             target = data
             if transform:
                 data = transform(data)
@@ -1564,14 +1570,20 @@ def evaluate_classification(model, criterion, data_loader, device, transform, lo
     target_list = []
     predictions_list = []
 
+    # non_blocking H2D transfers are only safe/beneficial with pinned source memory.
+    # create_data_loaders() only pins memory for CUDA (pin_memory=False for MPS/CPU),
+    # so non_blocking must be disabled on those backends -- otherwise the async copy
+    # can race with reuse of the source buffer, corrupting the transferred tensor
+    # (observed on MPS as NaN activations reaching the quantization observers).
+    non_blocking = (device.type == 'cuda')
     with torch.no_grad():
         for data_raw, data_feat_ext, target in metric_logger.log_every(data_loader, print_freq, header):
             if nn_for_feature_extraction:
-                data = data_raw.float().to(device, non_blocking=True)
+                data = data_raw.float().to(device, non_blocking=non_blocking)
             else:
-                data = data_feat_ext.float().to(device, non_blocking=True)
+                data = data_feat_ext.float().to(device, non_blocking=non_blocking)
 
-            target = target.long().to(device, non_blocking=True)
+            target = target.long().to(device, non_blocking=non_blocking)
             if transform:
                 data = transform(data)