From f7fcf114e3858157156fec24edf66e6761652533 Mon Sep 17 00:00:00 2001 From: Aleksandr Dremov Date: Tue, 19 May 2026 16:17:40 +0200 Subject: [PATCH 1/3] Change state_dict to use deepcopy Use deepcopy for source state_dict to avoid mutation. --- torchdata/nodes/_populate_queue.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/torchdata/nodes/_populate_queue.py b/torchdata/nodes/_populate_queue.py index cf22a4edc..19e90722f 100644 --- a/torchdata/nodes/_populate_queue.py +++ b/torchdata/nodes/_populate_queue.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import copy import queue import threading from typing import Any, Dict, Optional, Union @@ -75,7 +76,7 @@ def _put( yielded += 1 snapshot = None if snapshot_frequency > 0 and yielded % snapshot_frequency == 0: - snapshot = source.state_dict() + snapshot = copy.deepcopy(source.state_dict()) _put(item, block=False, snapshot=snapshot) except StopIteration as e: _put(e, block=False) From 4b36d9869876613a2a61ea6543280fed31247372 Mon Sep 17 00:00:00 2001 From: Aleksandr Dremov Date: Tue, 19 May 2026 16:38:55 +0200 Subject: [PATCH 2/3] Optimize snapshot creation in _populate_queue.py Refactor snapshot handling to avoid unnecessary deepcopy. --- torchdata/nodes/_populate_queue.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/torchdata/nodes/_populate_queue.py b/torchdata/nodes/_populate_queue.py index 19e90722f..02c6e7fc1 100644 --- a/torchdata/nodes/_populate_queue.py +++ b/torchdata/nodes/_populate_queue.py @@ -76,7 +76,9 @@ def _put( yielded += 1 snapshot = None if snapshot_frequency > 0 and yielded % snapshot_frequency == 0: - snapshot = copy.deepcopy(source.state_dict()) + snapshot = source.state_dict() + if snapshot is not None: + snapshot = copy.deepcopy(snapshot) _put(item, block=False, snapshot=snapshot) except StopIteration as e: _put(e, block=False) From 3f77ba256c7718d4e8843581777f06c31dfef6e0 Mon Sep 17 00:00:00 2001 From: Aleksandr Dremov Date: Tue, 19 May 2026 16:45:26 +0200 Subject: [PATCH 3/3] Add deepcopy for snapshot in _populate_queue --- torchdata/nodes/_populate_queue.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/torchdata/nodes/_populate_queue.py b/torchdata/nodes/_populate_queue.py index 02c6e7fc1..9e1779b59 100644 --- a/torchdata/nodes/_populate_queue.py +++ b/torchdata/nodes/_populate_queue.py @@ -61,7 +61,10 @@ def _put( assert ( isinstance(snapshot_frequency, int) and snapshot_frequency >= 0 ), f"snapshot_frequency must be non-negative integer! Got {snapshot_frequency}" - snapshot_store.append_initial_snapshot(snapshot=source.state_dict()) + snapshot = source.state_dict() + if snapshot is not None: + snapshot = copy.deepcopy(snapshot) + snapshot_store.append_initial_snapshot(snapshot=snapshot) except Exception: e = StartupExceptionWrapper(where="in _populate_queue startup for device") snapshot_store.append_initial_snapshot(snapshot=e)