Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / runtime / data_iterator.py: 80%
107 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates
2# Copyright 2026 Huawei Technologies Co., Ltd
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8# http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15# ============================================================================
17"""Trainer-side dataloader iteration helpers.
19``BackgroundPrefetcher`` and ``HyperIter`` are split out of the former
20``auto_models/trainer/base.py`` in stage 7 (05 §15.11 step 3); class names and
21checkpointing semantics are retained.
22"""
24__all__ = [
25 "BackgroundPrefetcher",
26 "HyperIter",
27]
29import logging
30import queue
31import threading
32from typing import Any, Dict, Optional
35logger = logging.getLogger(__name__)
38class BackgroundPrefetcher:
39 """
40 Prefetches batches from a dataloader in a background thread to overlap data loading
41 with GPU computation. Synchronizes dataloader state for correct checkpointing.
42 """
44 def __init__(self, dataloader: Any, maxsize: int = 1) -> None:
45 """Start the background worker prefetching from ``dataloader``."""
46 self.dataloader = dataloader
47 self.iterator = iter(dataloader)
48 self.queue = queue.Queue(maxsize=maxsize)
49 self.stop_event = threading.Event()
50 self.original_state_dict = getattr(dataloader, "state_dict", None)
51 self.current_state = None
52 self.thread = threading.Thread(target=self._worker)
53 self.thread.daemon = True
54 self.thread.start()
56 def _worker(self) -> None:
57 """Prefetch data and capture dataloader state in a background thread."""
58 try:
59 while not self.stop_event.is_set():
60 try:
61 item = next(self.iterator)
62 except StopIteration:
63 self._put_result((StopIteration, None))
64 break
66 # A stop request cannot interrupt next(), so discard the item once
67 # that call returns instead of retaining another batch and state.
68 if self.stop_event.is_set():
69 break
71 # Ensure we capture the state so that subsequent dataloader advances
72 # don't mutate the captured state in-place. The underlying dataloader's
73 # state_dict() should handle deepcopying if necessary.
74 state = self.original_state_dict() if self.original_state_dict else None
75 if not self._put_result((item, state)):
76 break
77 # The worker must transfer any producer failure back to the training thread.
78 except Exception as exc: # pylint: disable=broad-exception-caught
79 self._put_result((exc, None))
80 finally:
81 # A timed-out stop cannot cancel next(), so the worker must release
82 # its own references when that call eventually returns.
83 if self.stop_event.is_set():
84 self._cleanup_after_stop()
86 def _put_result(self, result: tuple[Any, Any]) -> bool:
87 """Put a worker result while allowing a stop request to cancel the write."""
88 while not self.stop_event.is_set():
89 try:
90 self.queue.put(result, timeout=0.1)
91 return True
92 except queue.Full:
93 continue
94 return False
96 def _drain_queue(self) -> None:
97 """Release every result currently retained by the prefetch queue."""
98 while True:
99 try:
100 self.queue.get_nowait()
101 except queue.Empty:
102 return
104 def _cleanup_after_stop(self) -> None:
105 """Release queued results and live dataloader references after stopping."""
106 self._drain_queue()
107 self.iterator = None
108 self.dataloader = None
109 self.original_state_dict = None
110 # The final checkpoint may be collected after stop(), so keep the
111 # consumed-batch snapshot while releasing live dataloader references.
113 def __iter__(self) -> "BackgroundPrefetcher":
114 """Return this prefetcher as its own iterator."""
115 return self
117 def __next__(self) -> Any:
118 """Return the next prefetched item, re-raising worker failures."""
119 res = self.queue.get()
120 if isinstance(res, tuple) and len(res) == 2:
121 item, state = res
122 if item is StopIteration:
123 raise StopIteration
124 if isinstance(item, Exception):
125 raise item
126 self.current_state = state
127 return item
128 if res is StopIteration:
129 raise StopIteration
130 if isinstance(res, Exception):
131 raise res
132 return res
134 def state_dict(self) -> Dict[str, Any]:
135 """Return the dataloader state captured alongside the current item."""
136 if self.current_state is not None:
137 return self.current_state
138 if self.original_state_dict:
139 return self.original_state_dict()
140 return {}
142 def stop(self, timeout: Optional[float] = 5.0) -> bool:
143 """Stop the background worker and release its dataloader references.
145 Args:
146 timeout: Maximum seconds to wait. ``None`` waits until the worker exits.
148 Returns:
149 Whether the worker terminated within the requested timeout.
150 """
151 self.stop_event.set()
152 self._drain_queue()
153 if self.thread.is_alive():
154 self.thread.join(timeout=timeout)
155 if self.thread.is_alive():
156 logger.warning("BackgroundPrefetcher worker thread did not terminate within timeout.")
157 return False
159 self._cleanup_after_stop()
160 return True
163class HyperIter:
164 """
165 A unified iterator wrapper that handles both standard iteration and background prefetching.
166 """
168 def __init__(self, dataloader: Any, use_background_prefetcher: bool = False, maxsize: int = 1) -> None:
169 """Wrap ``dataloader`` in either direct iteration or background prefetching."""
170 self.dataloader = dataloader
171 self.use_background_prefetcher = use_background_prefetcher
172 if use_background_prefetcher:
173 self.iterator = BackgroundPrefetcher(dataloader, maxsize=maxsize)
174 else:
175 self.iterator = iter(dataloader)
177 def __iter__(self) -> "HyperIter":
178 """Return this wrapper as its own iterator."""
179 return self
181 def __next__(self) -> Any:
182 """Return the next batch from the underlying iterator."""
183 return next(self.iterator)
185 def stop(self, timeout: Optional[float] = 5.0) -> bool:
186 """Stop the background prefetch worker when one is active.
188 Args:
189 timeout: Maximum seconds to wait. ``None`` waits until the worker exits.
191 Returns:
192 Whether the worker is stopped, or ``True`` when prefetching is disabled.
193 """
194 if self.use_background_prefetcher and hasattr(self.iterator, "stop"):
195 return self.iterator.stop(timeout=timeout)
196 return True
198 def state_dict(self) -> Dict[str, Any]:
199 """Return the underlying dataloader or prefetcher state for checkpointing."""
200 if self.use_background_prefetcher and hasattr(self.iterator, "state_dict"):
201 return self.iterator.state_dict()
202 if hasattr(self.dataloader, "state_dict"):
203 return self.dataloader.state_dict()
204 return {}