Advice Needed: handling significant amount of streams

Hiya!

I have approximately 1k data streams, each containing pickled numpy arrays. When data is loaded, I need to sample a subsequence from it, so my dataloader looks like this:

class WatermazeDatasetNew(StreamingDataset):
    def __init__(
        self,
        streams,
        seq_len,
        batch_size,
    ):
        super().__init__(
            streams=streams,
            batch_size=batch_size,
        )
        self.seq_len = seq_len

    def __getitem__(self, idx: int):
        obj = super().__getitem__(idx)
        idx_traj = np.random.randint(0, len(obj["img"]) - self.seq_len)
        state = obj["img"][idx_traj : idx_traj + self.seq_len] / 255
        action = obj["action"][idx_traj : idx_traj + self.seq_len]
        reward = obj["reward"][idx_traj : idx_traj + self.seq_len]
        return state, action, reward

It does what it should, but data loading takes too much time. The data cannot be put into RAM altogether, since its capacity is around 100GiB, while data itself is ~150GiB. I wonder if there is a way to keep the .mds file open in RAM and close it only when RAM is close to full. Something like cache_limit, but for RAM. I tried to enlarge predownload and epoch_size, but it didn't help much. Timings I got were like this:

GETTING DATA TOOK: 66.773770
forward took 2.23513
GETTING DATA TOOK: 1.681198
forward took 0.01434
GETTING DATA TOOK: 1.234387
forward took 0.00848
GETTING DATA TOOK: 1.397633
forward took 0.00872
GETTING DATA TOOK: 1.306774
forward took 0.00877
GETTING DATA TOOK: 1.275940
forward took 0.01117
GETTING DATA TOOK: 1.560843
forward took 0.00788
GETTING DATA TOOK: 1.619301
forward took 0.01183
GETTING DATA TOOK: 42.097745
forward took 0.00698
GETTING DATA TOOK: 10.102559
forward took 0.01008
GETTING DATA TOOK: 1.982731
forward took 0.01147
GETTING DATA TOOK: 1.585924
forward took 0.00956
GETTING DATA TOOK: 1.719182
forward took 0.01002
GETTING DATA TOOK: 1.663528
forward took 0.00875
GETTING DATA TOOK: 1.777168
forward took 0.00906
GETTING DATA TOOK: 1.621303
forward took 0.01053
GETTING DATA TOOK: 49.404249

So clearly every now and then the data files are closed and reopened again.

Is there any simple solution, am I missing something?

Begin with StreamingDataset.getitem and inspect how cache_limit, predownload, and epoch_size affect file reopening; use the reported GETTING DATA timings as a baseline. Done would be a confirmed RAM-cache approach or a clear explanation of the available behavior for this workload.