One warm dataset for every node, job, and cluster.

from streaming import MDSWriter, StreamingDataset
from torch.utils.data import DataLoader
 
# Once: write MDS shards onto the disk, uncompressed,
# so every node reads them in place
with MDSWriter(out="/data/mds/train",
               columns={"tokens": "ndarray"}) as out:
    for sample in samples():
        out.write(sample)
 
# Every node, every job: local only, no remote and no
# cache_limit, so nothing is downloaded or evicted
dataset = StreamingDataset(local="/data/mds/train",
                           shuffle=True, batch_size=8)
loader = DataLoader(dataset, batch_size=8, num_workers=8)

import webdataset as wds
 
# Tar shards on the disk, opened like local files.
# No cache_dir: the disk is already the shared cache.
urls = "/data/wds/train-{000000..000999}.tar"
 
dataset = (
    wds.WebDataset(urls, shardshuffle=1000,
                   nodesplitter=wds.split_by_node)
    .shuffle(1000)
    .decode("pil")
    .to_tuple("png", "json")
)

# One Archil disk over your training bucket
apiVersion: v1
kind: PersistentVolume
metadata:
  name: datasets
spec:
  capacity:
    storage: 1000Gi  # Archil disks are elastic
  accessModes:
    - ReadWriteMany
  persistentVolumeReclaimPolicy: Retain
  csi:
    driver: csi.archil.com
    volumeHandle: "dsk-XXXXXXXXXXXX"
    volumeAttributes:
      region: "aws-us-east-1"
---
apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: datasets
spec:
  accessModes:
    - ReadWriteMany
  resources:
    requests:
      storage: 1000Gi
  volumeName: datasets

import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset
 
 
class TokenWindows(Dataset):
    """Random windows from one raw token file on the mount."""
 
    def __init__(self, path="/data/tokens/train.bin", ctx=2048):
        self.tokens = np.memmap(path, dtype=np.uint16, mode="r")
        self.ctx = ctx
 
    def __len__(self):
        return (len(self.tokens) - 1) // self.ctx
 
    def __getitem__(self, i):
        start = i * self.ctx
        window = self.tokens[start : start + self.ctx + 1]
        window = torch.from_numpy(window.astype(np.int64))
        return window[:-1], window[1:]
 
 
loader = DataLoader(TokenWindows(), batch_size=8,
                    shuffle=True, num_workers=8)