fix: remove the use of BufferedIterator (#4737)

This pull request refactors the data loading mechanism in the training
module by replacing the `BufferedIterator` class with a simpler
`cycle_iterator` function.

I was able to reproduce this error on a system of Ni6Fe10 provided by
@iProzd . The error log shows there is something strange when running
garbage collection. I've tried the following approaches:
- Not using GPU for training, but CPU: failed
- Set `NUM_WORKERS`=0: worked
- Manually run garbage collection `gc.collect()` when a DataLoader
finishes its epoch: worked, but ~10% slower

<details><summary>Error log</summary>
<p>

```
Fatal Python error: Aborted

Thread 0x00007f68b9289700 (most recent call first):
  File "/root/miniconda3/lib/python3.10/threading.py", line 320 in wait
  File "/root/miniconda3/lib/python3.10/multiprocessing/queues.py", line 231 in _feed
  File "/root/miniconda3/lib/python3.10/threading.py", line 953 in run
  File "/root/miniconda3/lib/python3.10/threading.py", line 1016 in _bootstrap_inner
  File "/root/miniconda3/lib/python3.10/threading.py", line 973 in _bootstrap

Current thread 0x00007f6c595d7280 (most recent call first):
  Garbage-collecting
  File "/root/miniconda3/lib/python3.10/ast.py", line 99 in _convert
  File "/root/miniconda3/lib/python3.10/ast.py", line 110 in literal_eval
  File "/root/miniconda3/lib/python3.10/site-packages/numpy/lib/utils.py", line 1078 in safe_eval
  File "/root/miniconda3/lib/python3.10/site-packages/numpy/lib/format.py", line 623 in _read_array_header
  File "/root/miniconda3/lib/python3.10/site-packages/numpy/lib/format.py", line 784 in read_array
  File "/root/miniconda3/lib/python3.10/site-packages/numpy/lib/npyio.py", line 456 in load
  File "/aisi/cc/deepmd-kit/deepmd/utils/path.py", line 187 in load_numpy
  File "/aisi/cc/deepmd-kit/deepmd/utils/data.py", line 634 in _load_data
  File "/aisi/cc/deepmd-kit/deepmd/utils/data.py", line 526 in _load_set
  File "/aisi/cc/deepmd-kit/deepmd/utils/data.py", line 251 in get_item_torch
  File "/aisi/cc/deepmd-kit/deepmd/pt/utils/dataset.py", line 39 in __getitem__
  File "/root/miniconda3/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 52 in <listcomp>
  File "/root/miniconda3/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 52 in fetch
  File "/root/miniconda3/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 789 in _next_data
  File "/root/miniconda3/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 733 in __next__
  File "/aisi/cc/deepmd-kit/deepmd/pt/utils/dataloader.py", line 204 in __getitem__
  File "/root/miniconda3/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 54 in fetch
  File "/root/miniconda3/lib/python3.10/site-packages/torch/utils/data/_utils/worker.py", line 349 in _worker_loop
  File "/root/miniconda3/lib/python3.10/multiprocessing/process.py", line 108 in run
  File "/root/miniconda3/lib/python3.10/multiprocessing/process.py", line 314 in _bootstrap
  File "/root/miniconda3/lib/python3.10/multiprocessing/popen_fork.py", line 71 in _launch
  File "/root/miniconda3/lib/python3.10/multiprocessing/popen_fork.py", line 19 in __init__
  File "/root/miniconda3/lib/python3.10/multiprocessing/context.py", line 281 in _Popen
  File "/root/miniconda3/lib/python3.10/multiprocessing/context.py", line 224 in _Popen
  File "/root/miniconda3/lib/python3.10/multiprocessing/process.py", line 121 in start
  File "/root/miniconda3/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1171 in __init__
  File "/root/miniconda3/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 424 in _get_iterator
  File "/root/miniconda3/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 493 in __iter__
  File "/aisi/cc/deepmd-kit/deepmd/pt/train/training.py", line 1075 in get_data
  File "/aisi/cc/deepmd-kit/deepmd/pt/train/training.py", line 689 in step
  File "/aisi/cc/deepmd-kit/deepmd/pt/train/training.py", line 960 in run
  File "/aisi/cc/deepmd-kit/deepmd/pt/entrypoints/main.py", line 361 in train
  File "/aisi/cc/deepmd-kit/deepmd/pt/entrypoints/main.py", line 530 in main
  File "/root/miniconda3/lib/python3.10/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 355 in wrapper
  File "/aisi/cc/deepmd-kit/deepmd/main.py", line 930 in main
  File "/root/miniconda3/bin/dp", line 8 in <module>

Extension modules: numpy.core._multiarray_umath, numpy.core._multiarray_tests, numpy.linalg._umath_linalg, numpy.fft._pocketfft_internal, numpy.random._common, numpy.random.bit_generator, numpy.random._bounded_integers, numpy.random._mt19937, numpy.random.mtrand, numpy.random._philox, numpy.random._pcg64, numpy.random._sfc64, numpy.random._generator, torch._C, torch._C._dynamo.autograd_compiler, torch._C._dynamo.eval_frame, torch._C._dynamo.guards, torch._C._dynamo.utils, torch._C._fft, torch._C._linalg, torch._C._nested, torch._C._nn, torch._C._sparse, torch._C._special, h5py._errors, h5py.defs, h5py._objects, h5py.h5, h5py.utils, h5py.h5t, h5py.h5s, h5py.h5ac, h5py.h5p, h5py.h5r, h5py._proxy, h5py._conv, h5py.h5z, h5py.h5a, h5py.h5d, h5py.h5ds, h5py.h5g, h5py.h5i, h5py.h5o, h5py.h5f, h5py.h5fd, h5py.h5pl, h5py.h5l, h5py._selector, yaml._yaml, scipy._lib._ccallback_c, scipy.special._ufuncs_cxx, scipy.special._ufuncs, scipy.special._specfun, scipy.special._comb, scipy.linalg._fblas, scipy.linalg._flapack, scipy.linalg.cython_lapack, scipy.linalg._cythonized_array_utils, scipy.linalg._solve_toeplitz, scipy.linalg._decomp_lu_cython, scipy.linalg._matfuncs_sqrtm_triu, scipy.linalg._matfuncs_expm, scipy.linalg._linalg_pythran, scipy.linalg.cython_blas, scipy.linalg._decomp_update, scipy.sparse._sparsetools, _csparsetools, scipy.sparse._csparsetools, scipy.sparse.linalg._dsolve._superlu, scipy.sparse.linalg._eigen.arpack._arpack, scipy.sparse.linalg._propack._spropack, scipy.sparse.linalg._propack._dpropack, scipy.sparse.linalg._propack._cpropack, scipy.sparse.linalg._propack._zpropack, scipy.sparse.csgraph._tools, scipy.sparse.csgraph._shortest_path, scipy.sparse.csgraph._traversal, scipy.sparse.csgraph._min_spanning_tree, scipy.sparse.csgraph._flow, scipy.sparse.csgraph._matching, scipy.sparse.csgraph._reordering, scipy.special._ellip_harm_2, scipy.interpolate._fitpack, scipy.interpolate._dfitpack, scipy.optimize._group_columns, scipy._lib.messagestream, scipy.optimize._trlib._trlib, scipy.optimize._lbfgsb, _moduleTNC, scipy.optimize._moduleTNC, scipy.optimize._cobyla, scipy.optimize._slsqp, scipy.optimize._minpack, scipy.optimize._lsq.givens_elimination, scipy.optimize._zeros, scipy.optimize._cython_nnls, scipy._lib._uarray._uarray, scipy.linalg._decomp_interpolative, scipy.optimize._bglu_dense, scipy.optimize._lsap, scipy.spatial._ckdtree, scipy.spatial._qhull, scipy.spatial._voronoi, scipy.spatial._distance_wrap, scipy.spatial._hausdorff, scipy.spatial.transform._rotation, scipy.optimize._direct, scipy.interpolate._dierckx, scipy.interpolate._ppoly, scipy.interpolate._interpnd, scipy.interpolate._rbfinterp_pythran, scipy.interpolate._rgi_cython, scipy.interpolate._bspl (total: 113)
Traceback (most recent call last):
  File "/root/miniconda3/bin/dp", line 8, in <module>
    sys.exit(main())
  File "/aisi/cc/deepmd-kit/deepmd/main.py", line 930, in main
    deepmd_main(args)
  File "/root/miniconda3/lib/python3.10/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 355, in wrapper
    return f(*args, **kwargs)
  File "/aisi/cc/deepmd-kit/deepmd/pt/entrypoints/main.py", line 530, in main
    train(
  File "/aisi/cc/deepmd-kit/deepmd/pt/entrypoints/main.py", line 361, in train
    trainer.run()
  File "/aisi/cc/deepmd-kit/deepmd/pt/train/training.py", line 960, in run
    step(step_id)
  File "/aisi/cc/deepmd-kit/deepmd/pt/train/training.py", line 705, in step
    loss.backward()
  File "/root/miniconda3/lib/python3.10/site-packages/torch/_tensor.py", line 648, in backward
    torch.autograd.backward(
  File "/root/miniconda3/lib/python3.10/site-packages/torch/autograd/__init__.py", line 353, in backward
    _engine_run_backward(
  File "/root/miniconda3/lib/python3.10/site-packages/torch/autograd/graph.py", line 824, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/root/miniconda3/lib/python3.10/site-packages/torch/utils/data/_utils/signal_handling.py", line 73, in handler
    _error_if_any_worker_fails()
RuntimeError: DataLoader worker (pid 434731) is killed by signal: Aborted. 
```

</p>
</details> 

The problem couples with pytorch tensor and python threads and pipes,
which is hard to locate the root cause.

I tested this PR on single-task training and multi-task training, and
the training speed (in s/1000 steps) is almost the same:

|data|before|after|
|---|---|---|
|omat|235.52|235.89|
|multi-task pretraining|290.63|290.00|

Fix #4586


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **Refactor**
- Enhanced data loading with an infinite cycling iterator for
uninterrupted batch retrieval during training and validation.
- Removed background prefetching and threading to simplify data loading
utilities.
- **Style**
- Added a clarifying comment about shuffling behavior when distributed
sampling is active.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Chun Cai <amoycaic@gmail.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

deepmd/pt/train/training.py

@@ -2,6 +2,9 @@
 import functools
 import logging
 import time
+from collections.abc import (
+    Iterable,
+)
 from copy import (
     deepcopy,
 )
@@ -47,7 +50,6 @@
     dp_random,
 )
 from deepmd.pt.utils.dataloader import (
-    BufferedIterator,
     get_sampler_from_params,
 )
 from deepmd.pt.utils.env import (
@@ -159,8 +161,24 @@ def get_opt_param(params):
             }
             return opt_type, opt_param
 
+        def cycle_iterator(iterable: Iterable):
+            """
+            Produces an infinite iterator by repeatedly cycling through the given iterable.
+
+            Args:
+                iterable (Iterable): The iterable to cycle through.
+
+            Yields
+            ------
+            Any: The next item from the iterable, cycling back to the beginning when the end is reached.
+            """
+            while True:
+                with torch.device("cpu"):
+                    it = iter(iterable)
+                yield from it
+
         def get_data_loader(_training_data, _validation_data, _training_params):
-            def get_dataloader_and_buffer(_data, _params):
+            def get_dataloader_and_iter(_data, _params):
                 _sampler = get_sampler_from_params(_data, _params)
                 if _sampler is None:
                     log.warning(
@@ -177,33 +195,32 @@ def get_dataloader_and_buffer(_data, _params):
                     collate_fn=lambda batch: batch,  # prevent extra conversion
                     pin_memory=True,
                 )
-                with torch.device("cpu"):
-                    _data_buffered = BufferedIterator(iter(_dataloader))
-                return _dataloader, _data_buffered
+                _data_iter = cycle_iterator(_dataloader)
+                return _dataloader, _data_iter
 
-            training_dataloader, training_data_buffered = get_dataloader_and_buffer(
+            training_dataloader, training_data_iter = get_dataloader_and_iter(
                 _training_data, _training_params["training_data"]
             )
 
             if _validation_data is not None:
                 (
                     validation_dataloader,
-                    validation_data_buffered,
-                ) = get_dataloader_and_buffer(
+                    validation_data_iter,
+                ) = get_dataloader_and_iter(
                     _validation_data, _training_params["validation_data"]
                 )
                 valid_numb_batch = _training_params["validation_data"].get(
                     "numb_btch", 1
                 )
             else:
                 validation_dataloader = None
-                validation_data_buffered = None
+                validation_data_iter = None
                 valid_numb_batch = 1
             return (
                 training_dataloader,
-                training_data_buffered,
+                training_data_iter,
                 validation_dataloader,
-                validation_data_buffered,
+                validation_data_iter,
                 valid_numb_batch,
             )
 
@@ -1064,48 +1081,15 @@ def save_model(self, save_path, lr=0.0, step=0) -> None:
             checkpoint_files[0].unlink()
 
     def get_data(self, is_train=True, task_key="Default"):
-        if not self.multi_task:
-            if is_train:
-                try:
-                    batch_data = next(iter(self.training_data))
-                except StopIteration:
-                    # Refresh the status of the dataloader to start from a new epoch
-                    with torch.device("cpu"):
-                        self.training_data = BufferedIterator(
-                            iter(self.training_dataloader)
-                        )
-                    batch_data = next(iter(self.training_data))
-            else:
-                if self.validation_data is None:
-                    return {}, {}, {}
-                try:
-                    batch_data = next(iter(self.validation_data))
-                except StopIteration:
-                    self.validation_data = BufferedIterator(
-                        iter(self.validation_dataloader)
-                    )
-                    batch_data = next(iter(self.validation_data))
+        if is_train:
+            iterator = self.training_data
         else:
-            if is_train:
-                try:
-                    batch_data = next(iter(self.training_data[task_key]))
-                except StopIteration:
-                    # Refresh the status of the dataloader to start from a new epoch
-                    self.training_data[task_key] = BufferedIterator(
-                        iter(self.training_dataloader[task_key])
-                    )
-                    batch_data = next(iter(self.training_data[task_key]))
-            else:
-                if self.validation_data[task_key] is None:
-                    return {}, {}, {}
-                try:
-                    batch_data = next(iter(self.validation_data[task_key]))
-                except StopIteration:
-                    self.validation_data[task_key] = BufferedIterator(
-                        iter(self.validation_dataloader[task_key])
-                    )
-                    batch_data = next(iter(self.validation_data[task_key]))
-
+            iterator = self.validation_data
+        if self.multi_task:
+            iterator = iterator[task_key]
+        if iterator is None:
+            return {}, {}, {}
+        batch_data = next(iterator)
         for key in batch_data.keys():
             if key == "sid" or key == "fid" or key == "box" or "find_" in key:
                 continue

deepmd/pt/utils/dataloader.py

@@ -1,16 +1,9 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import logging
 import os
-import time
 from multiprocessing.dummy import (
     Pool,
 )
-from queue import (
-    Queue,
-)
-from threading import (
-    Thread,
-)
 
 import h5py
 import numpy as np
@@ -173,7 +166,9 @@ def construct_dataset(system):
                 num_workers=0,  # Should be 0 to avoid too many threads forked
                 sampler=system_sampler,
                 collate_fn=collate_batch,
-                shuffle=(not (dist.is_available() and dist.is_initialized()))
+                shuffle=(
+                    not (dist.is_available() and dist.is_initialized())
+                )  # distributed sampler will do the shuffling by default
                 and shuffle,
             )
             self.dataloaders.append(system_dataloader)
@@ -200,11 +195,12 @@ def __len__(self) -> int:
 
     def __getitem__(self, idx):
         # log.warning(str(torch.distributed.get_rank())+" idx: "+str(idx)+" index: "+str(self.index[idx]))
-        try:
-            batch = next(self.iters[idx])
-        except StopIteration:
-            self.iters[idx] = iter(self.dataloaders[idx])
-            batch = next(self.iters[idx])
+        with torch.device("cpu"):
+            try:
+                batch = next(self.iters[idx])
+            except StopIteration:
+                self.iters[idx] = iter(self.dataloaders[idx])
+                batch = next(self.iters[idx])
         batch["sid"] = idx
         return batch
 
@@ -235,54 +231,6 @@ def print_summary(
             )
 
 
-class BackgroundConsumer(Thread):
-    def __init__(self, queue, source) -> None:
-        super().__init__()
-        self.daemon = True
-        self._queue = queue
-        self._source = source  # Main DL iterator
-
-    def run(self) -> None:
-        for item in self._source:
-            self._queue.put(item)  # Blocking if the queue is full
-
-        # Signal the consumer we are done; this should not happen for DataLoader
-        self._queue.put(StopIteration())
-
-
-QUEUESIZE = 32
-
-
-class BufferedIterator:
-    def __init__(self, iterable) -> None:
-        self._queue = Queue(QUEUESIZE)
-        self._iterable = iterable
-        self._consumer = BackgroundConsumer(self._queue, self._iterable)
-        self._consumer.start()
-        self.last_warning_time = time.time()
-
-    def __iter__(self):
-        return self
-
-    def __len__(self) -> int:
-        return len(self._iterable)
-
-    def __next__(self):
-        start_wait = time.time()
-        item = self._queue.get()
-        wait_time = time.time() - start_wait
-        if (
-            wait_time > 1.0 and start_wait - self.last_warning_time > 15 * 60
-        ):  # Even for Multi-Task training, each step usually takes < 1s
-            log.warning(
-                f"Data loading is slow, waited {wait_time:.2f} seconds. Ignoring this warning for 15 minutes."
-            )
-            self.last_warning_time = start_wait
-        if isinstance(item, Exception):
-            raise item
-        return item
-
-
 def collate_batch(batch):
     example = batch[0]
     result = {}
@@ -320,7 +268,11 @@ def get_weighted_sampler(training_data, prob_style, sys_prob=False):
     # training_data.total_batch is the size of one epoch, you can increase it to avoid too many  rebuilding of iterators
     len_sampler = training_data.total_batch * max(env.NUM_WORKERS, 1)
     with torch.device("cpu"):
-        sampler = WeightedRandomSampler(probs, len_sampler, replacement=True)
+        sampler = WeightedRandomSampler(
+            probs,
+            len_sampler,
+            replacement=True,
+        )
     return sampler
 
 

source/tests/pt/model/test_saveload_dpa1.py

@@ -25,7 +25,6 @@
     env,
 )
 from deepmd.pt.utils.dataloader import (
-    BufferedIterator,
     DpLoaderSet,
 )
 from deepmd.pt.utils.stat import (
@@ -72,8 +71,13 @@ def setUp(self) -> None:
             drop_last=False,
             pin_memory=True,
         )
+
+        def cycle_iterator(iterable):
+            while True:
+                yield from iterable
+
         with torch.device("cpu"):
-            self.training_data = BufferedIterator(iter(self.training_dataloader))
+            self.training_data = cycle_iterator(self.training_dataloader)
         self.loss = EnergyStdLoss(**self.config["loss"])
         self.cur_lr = 1
         self.task_key = "Default"
@@ -111,12 +115,8 @@ def create_wrapper(self, read: bool):
         return ModelWrapper(model, self.loss)
 
     def get_data(self):
-        try:
-            batch_data = next(iter(self.training_data))
-        except StopIteration:
-            # Refresh the status of the dataloader to start from a new epoch
-            self.training_data = BufferedIterator(iter(self.training_dataloader))
-            batch_data = next(iter(self.training_data))
+        with torch.device("cpu"):
+            batch_data = next(self.training_data)
         input_dict = {}
         for item in ["coord", "atype", "box"]:
             if item in batch_data:

source/tests/pt/model/test_saveload_se_e2_a.py

@@ -25,7 +25,6 @@
     env,
 )
 from deepmd.pt.utils.dataloader import (
-    BufferedIterator,
     DpLoaderSet,
 )
 from deepmd.pt.utils.stat import (
@@ -72,8 +71,13 @@ def setUp(self) -> None:
             drop_last=False,
             pin_memory=True,
         )
+
+        def cycle_iterator(iterable):
+            while True:
+                yield from iterable
+
         with torch.device("cpu"):
-            self.training_data = BufferedIterator(iter(self.training_dataloader))
+            self.training_data = cycle_iterator(self.training_dataloader)
         self.loss = EnergyStdLoss(**self.config["loss"])
         self.cur_lr = 1
         self.task_key = "Default"
@@ -105,12 +109,8 @@ def create_wrapper(self):
         return ModelWrapper(model, self.loss)
 
     def get_data(self):
-        try:
-            batch_data = next(iter(self.training_data))
-        except StopIteration:
-            # Refresh the status of the dataloader to start from a new epoch
-            self.training_data = BufferedIterator(iter(self.training_dataloader))
-            batch_data = next(iter(self.training_data))
+        with torch.device("cpu"):
+            batch_data = next(self.training_data)
         input_dict = {}
         for item in ["coord", "atype", "box"]:
             if item in batch_data: