import lmdb
import random
import asyncio
import torch.utils.data as data
import time
import threading
import warnings
import torch
import pickle
import numpy as np

class LMDBDataset(data.Dataset):
    def __init__(self, root, num_samples, buffer_size=8, num_shards=0, shuffle_shards=False):
        self.root = root
        self.num_samples = num_samples
        self.num_shards = num_shards
        self.shuffle_shards = shuffle_shards
        self.buffer_size = buffer_size

        self.queue = []
        self.shard_idxs = list(range(num_shards))
        if self.shuffle_shards:
            random.shuffle(self.shard_idxs)
        self.shard_cursor = 0
        for i in range(self.buffer_size):
            self.read_shard()
        threading.Thread(target=self.ensure_buffer, daemon=True).start()
        print("START")

    def ensure_buffer(self):
        print("ensure_buffer")
        while True:
            if len(self.queue) != self.buffer_size:
                self.read_shard()
            time.sleep(0.01)

    def read_shard(self):
        # print(self.shard_cursor)
        self.shard_cursor += 1
        if self.shard_cursor == self.num_shards:
            self.shard_cursor = 0
            if self.shuffle_shards:
                random.shuffle(self.shard_idxs)
        # time.sleep(0.5)
        env = lmdb.open('{}/{:08d}'.format(self.root, self.shard_idxs[self.shard_cursor]))
        txn = env.begin()
        self.queue.append(txn.cursor())


    def __len__(self):
        return self.num_samples

    # async def query_data(self, idx):

    def __getitem__(self, idx):
        print(self.shard_cursor)

        row = self.queue[0].next()
        if row is False:
            self.queue.pop(0)
            assert(len(self.queue))
            row = self.queue[0].next()
        time.sleep(0.005)
        return 1

train_dataset = LMDBDataset(root='/home/xxx/test_lmdb/imagenet', num_samples=1000000, num_shards=10, shuffle_shards=True)


train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=256, shuffle=False,
        num_workers=4, pin_memory=True)

for i, data in enumerate(train_loader):
    print(i, data.shape)

I am trying to create a Dataset that is able to constantly fetch new LMDB shard when the size of the queue falls below the specified buffer_size. It works when it is not wrapped in a DataLoader. However, the function ensure_buffer() does not run when it is wrapped in a DataLoader. Why? How should I make this work?