Fix HDF5Writer id()-reuse bug corrupting class_index in multi-worker generation

tests/utils/test_writer.py

@@ -182,6 +182,57 @@ def test_writer_reproducibility(tmp_path, dataset_length, multithreading):
         assert m0 == m1
         assert np.allclose(x0, x1, rtol=1e-6)
 
+@pytest.mark.full
+def test_writer_no_parent_metadata_id_reuse_corruption(tmp_path):
+    """Regression test for an HDF5Writer bug where a signal's metadata
+    ``.parent`` chain (e.g. per-generator/class-level metadata) fell back to
+    ``str(id(obj))`` for its HDF5 key instead of a stamped, collision-free one.
+    Parent objects aren't guaranteed to stay alive for the whole write
+    session - a new one can be built per batch/worker and then garbage
+    collected - so CPython could recycle a freed parent's address for a
+    later, unrelated parent. Both then resolved to the same HDF5 key, and the
+    writer's "already written" dedup guard silently kept the first parent's
+    metadata while skipping the second, so every signal pointing at the
+    second parent read back the first parent's fields instead of its own
+    (most visibly, ``class_index`` coming back as a list mixing two labels
+    instead of a scalar).
+
+    Requires multiple DataLoader workers and enough samples to make the
+    address-recycling race likely; empirically this reproduced on ~1.4% of
+    3,000 samples with 4 workers before the fix, and 0% single-process.
+    """
+    seed = 987654321
+    num_iq_samples = 4096
+    small_meta = {
+        **TorchSigDefaults().default_dataset_metadata,
+        "num_iq_samples_dataset": num_iq_samples,
+        "signal_duration_in_samples_min": num_iq_samples * 0.8,
+        "signal_duration_in_samples_max": num_iq_samples * 1.0,
+        "num_signals_min": 1,
+        "num_signals_max": 1,
+    }
+    dataset_length = 3000
+
+    ds = TorchSigIterableDataset(
+        metadata=small_meta, target_labels=["class_index"],
+    )
+    dl = WorkerSeedingDataLoader(ds, seed=seed, batch_size=16, num_workers=4)
+    DatasetCreator(
+        dataloader=dl, dataset_length=dataset_length, root=tmp_path,
+        overwrite=True, multithreading=True,
+    ).create()
+
+    sds = StaticTorchSigDataset(root=str(tmp_path), target_labels=["class_index"])
+    assert len(sds) == dataset_length
+
+    bad = [i for i in range(len(sds)) if isinstance(sds[i][1], list)]
+    assert not bad, (
+        f"{len(bad)}/{dataset_length} samples had a list-valued class_index "
+        f"(first offending indices: {bad[:10]}) - parent metadata objects are "
+        "colliding on a recycled id()-based HDF5 key."
+    )
+
+
 @pytest.mark.full
 def test_writer_memory_growth(tmp_path):
     """Measure how RAM usage grows while writing datasets of increasing size.

torchsig/utils/file_handlers/hdf5.py

@@ -25,16 +25,21 @@
 def _hdf5_key(obj) -> str:
     """Return the HDF5 group key to use for *obj*.
 
-    Ephemeral objects (generated Signal instances) receive a short sequential
-    integer key that is stamped onto them by ``HDF5Writer._assign_hdf5_keys``
-    immediately before writing. Persistent objects (generators, datasets)
-    that are never garbage-collected within a write session fall back to
-    ``str(id(obj))``, which is stable for the lifetime of the writer.
-
-    Using a counter for signals avoids the id()-reuse hazard that arises when
-    CPython recycles the memory address of a freed signal and a later signal
-    lands at the same address, causing the "already exists" guard to skip the
-    write silently.
+    Signal instances and their metadata ``.parent`` chain all receive a short
+    sequential integer key, stamped on by ``HDF5Writer._assign_hdf5_keys``
+    (and ``_assign_hdf5_keys_to_parent_chain``) immediately before writing.
+    Any object reaching this function without one - which should no longer
+    happen for objects written via ``HDF5Writer`` - falls back to
+    ``str(id(obj))``.
+
+    Using a counter avoids the id()-reuse hazard that arises when CPython
+    recycles the memory address of a freed object (a signal, or a parent
+    metadata object rebuilt per batch/worker) and a later, unrelated object
+    lands at the same address: without stamped keys, both would resolve to
+    the same HDF5 key, and the "already exists" guard in
+    ``populate_hdf5_group_with_metadata``/``populate_hdf5_group_with_signal_data``
+    would silently skip the second write, leaving the second object reading
+    back the first's data.
     """
     try:
         return obj._hdf5_key
@@ -264,12 +269,41 @@ def _assign_hdf5_keys(self, signal) -> None:
         is used by the module-level populate helpers instead of ``str(id(signal))``,
         making the HDF5 layout independent of CPython memory addresses and
         allowing signals to be garbage-collected as soon as they leave scope.
+
+        Also stamps the signal's metadata ``.parent`` chain (e.g. per-class or
+        per-generator metadata objects walked by ``populate_hdf5_group_with_metadata``).
+        Those objects are not necessarily long-lived - a new one can be built for
+        each batch/worker and then garbage-collected - so leaving them to
+        ``_hdf5_key``'s ``str(id(obj))`` fallback lets CPython recycle a freed
+        parent's address for an unrelated later parent. Both then resolve to the
+        same HDF5 key, so ``populate_hdf5_group_with_metadata``'s
+        ``if key in group: return False`` guard silently keeps the first parent's
+        metadata and skips writing the second, and every signal pointing at the
+        second parent ends up reading back the first's fields (e.g. the wrong
+        ``class_index``). Stamping a counter-based key here - reusing it via
+        ``hasattr`` when the same parent object is genuinely shared across many
+        signals - makes parent identity independent of memory address too.
         """
         signal._hdf5_key = str(self._key_counter)
         self._key_counter += 1
+        self._assign_hdf5_keys_to_parent_chain(signal)
         for cs in signal.component_signals:
             self._assign_hdf5_keys(cs)
 
+    def _assign_hdf5_keys_to_parent_chain(self, metadata_obj) -> None:
+        """Stamp a stable ``_hdf5_key`` on *metadata_obj*'s ``.parent`` chain.
+
+        Skips objects that already carry a ``_hdf5_key`` (a genuinely shared
+        parent instance keeps its key, so it still dedupes to one HDF5 group).
+        See ``_assign_hdf5_keys`` for why this must not fall back to ``id()``.
+        """
+        parent = getattr(metadata_obj, "parent", None)
+        while parent is not None:
+            if not hasattr(parent, "_hdf5_key"):
+                parent._hdf5_key = str(self._key_counter)
+                self._key_counter += 1
+            parent = getattr(parent, "parent", None)
+
     def _write_batch_to_hdf5(self, data) -> None:
         """Writes a batch of signals (as List[Signal]) to the file.