Coverage for dynasor/trajectory/prefetch.py: 95%

1"""Two-thread reader/GPU-worker pipeline used by the GPU backends

2(``backend='torch'``/``'cupy'``) in :mod:`dynasor.correlation_functions` to

3overlap trajectory reading with GPU computation.

5Internal implementation detail, not part of the public API.

6"""

8import queue

9import threading

12def _drop_frames_outside_windows(raw_iter, width, window_step):

13 """Yield only the items of *raw_iter* that fall inside a window.

15 Windows of *width* items start every *window_step* items, so for

16 ``window_step > width`` the items in between belong to no window at all

17 and nothing is ever computed from them. Dropping them here, ahead of the

18 batched GPU processor, keeps the amount of work equal to the per-frame

19 path, which applies its processor after the window iterator has already

20 skipped them.

22 The items that remain form consecutive groups of *width*, one per

23 window, so the caller has to walk the result with a window step of

24 *width* rather than *window_step*.

25 """

26 for index, item in enumerate(raw_iter):

27 if index % window_step < width:

28 yield item

31def _prefetch_batched_iter(raw_iter, fn_batch, batch_size, n_prefetch=2):

32 """Two-thread pipeline that overlaps disk I/O with GPU computation.

34 * **Reader thread** pulls raw items from *raw_iter* (typically NetCDF/XYZ

35 reads) and places them into a bounded raw queue.

36 * **GPU thread** drains *batch_size* raw items at a time, calls

37 *fn_batch* (which does H2D transfer + GPU kernel + D2H), and places the

38 processed batch into a result queue.

40 While the GPU thread is running the kernel on batch N, the reader thread

41 is simultaneously reading batch N+1 from disk, achieving true pipeline

42 overlap.

44 Parameters

45 ----------

46 raw_iter:

47 Source iterator of raw items (e.g. a :class:`Trajectory`).

48 fn_batch:

49 Callable that accepts a list of up to *batch_size* raw items and

50 returns a list of the same length with processed results.

51 batch_size:

52 Number of frames per GPU kernel call.

53 n_prefetch:

54 Number of processed batches to buffer ahead of the caller.

56 Both queues are bounded, so both threads can block indefinitely waiting

57 for the other end of the pipeline. Two events make every such wait

58 abandonable, so that neither thread outlives the generator:

60 * ``stop_event`` is set when the caller stops consuming the generator,

61 either by closing it or by abandoning it after an exception further

62 down the pipeline. It releases the GPU thread from waiting for raw

63 items and from waiting for room in the result queue.

64 * ``drain_event`` is set when the GPU thread stops draining the raw

65 queue, whether because the trajectory is exhausted or because

66 *fn_batch* raised. It releases the reader thread from waiting for room

67 in the raw queue that will never free up.

68 """

69 _DONE = object()

70 _STOPPED = object()

71 exc_holder: list = [None]

72 stop_event = threading.Event()

73 drain_event = threading.Event()

75 # raw_q: disk → GPU thread. Allow the reader to stay up to one full

76 # batch ahead of the GPU thread so it is never waiting for work.

77 raw_q: queue.Queue = queue.Queue(maxsize=batch_size * (n_prefetch + 1))

78 # proc_q: GPU thread → caller.

79 proc_q: queue.Queue = queue.Queue(maxsize=n_prefetch)

81 def _put_or_stop(q, item, *events):

82 """Put *item* on *q*, retrying until there is room or one of *events*

83 fires. Returns False, without placing *item*, if an event fired

84 first."""

85 while not any(event.is_set() for event in events):

86 try:

87 q.put(item, timeout=0.1)

88 return True

89 except queue.Full:

90 continue

91 return False

93 def _get_or_stop(q, *events):

94 """Get the next item from *q*, retrying until one arrives or one of

95 *events* fires. Returns :data:`_STOPPED` if an event fired first."""

96 while not any(event.is_set() for event in events):

97 try:

98 return q.get(timeout=0.1)

99 except queue.Empty:

100 continue

101 return _STOPPED

102

103 def _reader():

104 try:

105 for item in raw_iter:

106 if not _put_or_stop(raw_q, item, stop_event, drain_event):

107 return

108 except Exception as exc: # noqa: BLE001

109 exc_holder[0] = exc

110 finally:

111 _put_or_stop(raw_q, _DONE, stop_event, drain_event)

112

113 def _gpu_worker():

114 try:

115 batch = []

116 while True:

117 item = _get_or_stop(raw_q, stop_event)

118 if item is _STOPPED:

119 return

120 if item is _DONE:

121 if batch:

122 _put_or_stop(proc_q, fn_batch(batch), stop_event)

123 break

124 batch.append(item)

125 if len(batch) == batch_size:

126 if not _put_or_stop(proc_q, fn_batch(batch), stop_event): 126 ↛ 127line 126 didn't jump to line 127 because the condition on line 126 was never true

127 return

128 batch = []

129 except Exception as exc: # noqa: BLE001

130 if exc_holder[0] is None: 130 ↛ 133line 130 didn't jump to line 133 because the condition on line 130 was always true

131 exc_holder[0] = exc

132 finally:

133 drain_event.set()

134 # the caller is still waiting for the sentinel unless it has

135 # abandoned the generator, in which case stop_event releases us

136 _put_or_stop(proc_q, _DONE, stop_event)

137

138 reader_t = threading.Thread(target=_reader, daemon=True)

139 gpu_t = threading.Thread(target=_gpu_worker, daemon=True)

140 reader_t.start()

141 gpu_t.start()

142 try:

143 while True:

144 item = proc_q.get()

145 if item is _DONE:

146 break

147 yield from item

148 if exc_holder[0] is not None:

149 raise exc_holder[0]

150 finally:

151 stop_event.set()