Skip to content

Commit

Permalink
[GraphBolt][CUDA] Use deque instead of Queue for Bufferer.
Browse files Browse the repository at this point in the history
  • Loading branch information
mfbalin committed Jan 31, 2024
1 parent 7a97609 commit ab21a36
Showing 1 changed file with 8 additions and 8 deletions.
16 changes: 8 additions & 8 deletions python/dgl/graphbolt/dataloader.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
"""Graph Bolt DataLoaders"""

from queue import Queue
from collections import deque

import torch
import torch.utils.data
Expand Down Expand Up @@ -69,18 +69,18 @@ def __init__(self, datapipe, buffer_size=1):
raise ValueError(
"'buffer_size' is required to be a positive integer."
)
self.buffer = Queue(buffer_size)
self.buffer = deque(maxlen=buffer_size)

def __iter__(self):
for data in self.datapipe:
if not self.buffer.full():
self.buffer.put(data)
if len(self.buffer) < self.buffer.maxlen:
self.buffer.append(data)
else:
return_data = self.buffer.get()
self.buffer.put(data)
return_data = self.buffer.popleft()
self.buffer.append(data)
yield return_data
while not self.buffer.empty():
yield self.buffer.get()
while len(self.buffer) > 0:
yield self.buffer.popleft()


class Awaiter(dp.iter.IterDataPipe):
Expand Down

0 comments on commit ab21a36

Please sign in to comment.