Make data shard downloads more robust

[joellidin]

Description

The current implementation of shard downloads is not fully reliable. In some cases, downloads fail or stall, and retries do not always resolve the issue. Since shards are swapped only occasionally, one possible alternative would be to directly await the download of a new shard before proceeding, instead of relying on the current background/preparation logic.

Tasks

• Investigate current shard download logic and failure modes.
• Improve robustness of shard downloads:
  • Add stronger retry and validation handling, or
  • Switch to an approach where the system awaits the shard download directly when needed.
• Ensure that failures are logged clearly and that fallback behavior (e.g., retry, skip, or fail fast) is consistent.

Acceptance Criteria:

• Shard downloads succeed reliably without leaving the system in a broken state.
• If a shard cannot be downloaded within a reasonable timeframe, the system fails gracefully with clear logs.
• Optionally support an await-based direct download mode for shards, given the low frequency of shard swaps.

[joellidin]

We recently got this error as well:

                    ValueError: mmap length is greater than file size
                    ↩  received SIGTERM – shutting down …
                    [DCP] rank 1/4 save finished (window=58181)
[10/19/25 04:33:53] Unhandled exception in run()
                    ╭─────────────────────────────────────────────────────────────────────── Traceback (most recent call last) ────────────────────────────────────────────────────────────────────────╮
                    │ /app/neurons/base_node.py:99 in main                                                                                                                                             │
                    │                                                                                                                                                                                  │
                    │    96 │   │                                                                                                                                                                      │
                    │    97 │   │   # subclasses do their normal work here ----------------------------                                                                                                │
                    │    98 │   │   try:                                                                                                                                                               │
                    │ ❱  99 │   │   │   await self.run()                                                                                                                                               │
                    │   100 │   │   except Exception:                                                                                                                                                  │
                    │   101 │   │   │   # Catch and log the unhandled exception before shutting down.                                                                                                  │
                    │   102 │   │   │   tplr.logger.error("Unhandled exception in run()", exc_info=True)                                                                                               │
                    │                                                                                                                                                                                  │
                    │ /app/neurons/validator.py:1195 in run                                                                                                                                            │
                    │                                                                                                                                                                                  │
                    │   1192 │   │   │   │   and self.global_step % self.outer_steps_per_shard == 0                                                                                                    │
                    │   1193 │   │   │   ):                                                                                                                                                            │
                    │   1194 │   │   │   │   tplr.logger.info(f"Swapping dataset at window {self.current_window}")                                                                                     │
                    │ ❱ 1195 │   │   │   │   await self.dataset_manager.swap_datasets()                                                                                                                │
                    │   1196 │   │   │   │   self.set_dataloader(validator=True)                                                                                                                       │
                    │   1197 │   │   │   │   dist_helper.safe_barrier("sync_shard_switch", self.local_rank)                                                                                            │
                    │   1198                                                                                                                                                                           │
                    │                                                                                                                                                                                  │
                    │ /app/src/tplr/sharded_dataset.py:339 in swap_datasets                                                                                                                            │
                    │                                                                                                                                                                                  │
                    │   336 │   │   │   await self.upcoming_dataset                                                                                                                                    │
                    │   337 │   │                                                                                                                                                                      │
                    │   338 │   │   old_dataset = self.active_dataset                                                                                                                                  │
                    │ ❱ 339 │   │   await self.initialize_datasets(self.shard_index)                                                                                                                   │
                    │   340 │   │   tplr.logger.info("successfully swapped datasets.")                                                                                                                 │
                    │   341 │   │                                                                                                                                                                      │
                    │   342 │   │   if old_dataset and self.rank == 0:                                                                                                                                 │
                    │                                                                                                                                                                                  │
                    │ /app/src/tplr/sharded_dataset.py:314 in initialize_datasets                                                                                                                      │
                    │                                                                                                                                                                                  │
                    │   311 │   │   Args:                                                                                                                                                              │
                    │   312 │   │   │   current_shard_index: The index of the shard to make active.                                                                                                    │
                    │   313 │   │   """                                                                                                                                                                │
                    │ ❱ 314 │   │   self.active_dataset = await self.create_dataset(current_shard_index)                                                                                               │
                    │   315 │   │   next_shard = (current_shard_index + 1) % self.max_dataset_idx                                                                                                      │
                    │   316 │   │                                                                                                                                                                      │
                    │   317 │   │   # Only rank 0 prepares the next shard to avoid duplicate downloads                                                                                                 │
                    │                                                                                                                                                                                  │
                    │ /app/src/tplr/sharded_dataset.py:295 in create_dataset                                                                                                                           │
                    │                                                                                                                                                                                  │
                    │   292 │   │   │   await download_task                                                                                                                                            │
                    │   293 │   │   # Non-master ranks will just check if files exist (downloaded by rank 0)                                                                                           │
                    │   294 │   │                                                                                                                                                                      │
                    │ ❱ 295 │   │   dataset = SharedShardedDataset(                                                                                                                                    │
                    │   296 │   │   │   shard_index=shard_index,                                                                                                                                       │
                    │   297 │   │   │   sequence_length=self.sequence_length,                                                                                                                          │
                    │   298 │   │   │   rank=self.rank,                                                                                                                                                │
                    │                                                                                                                                                                                  │
                    │ /app/src/tplr/sharded_dataset.py:65 in __init__                                                                                                                                  │
                    │                                                                                                                                                                                  │
                    │    62 │   │   │   shard_index, file_prefix=file_prefix                                                                                                                           │
                    │    63 │   │   )                                                                                                                                                                  │
                    │    64 │   │   _ = self.check_paths([self.tokens_file, self.ids_file])                                                                                                            │
                    │ ❱  65 │   │   _ = self.mmap_tokens_and_ids(token_dtype)                                                                                                                          │
                    │    66 │   │                                                                                                                                                                      │
                    │    67 │   │   # should wrap in a timer                                                                                                                                           │
                    │    68 │   │   tplr.logger.info(                                                                                                                                                  │
                    │                                                                                                                                                                                  │
                    │ /app/src/tplr/sharded_dataset.py:131 in mmap_tokens_and_ids                                                                                                                      │
                    │                                                                                                                                                                                  │
                    │   128 │   │   # Tokens: support .npy (NumPy format) and raw .bin                                                                                                                 │
                    │   129 │   │   if tokens_path.suffix == ".npy":                                                                                                                                   │
                    │   130 │   │   │   # Correct way to memory-map a NumPy .npy file                                                                                                                  │
                    │ ❱ 131 │   │   │   arr = np.load(tokens_path, mmap_mode="r", allow_pickle=False)                                                                                                  │
                    │   132 │   │   │   # Ensure dtype is uint32 (no copy if already correct)                                                                                                          │
                    │   133 │   │   │   if arr.dtype != np.uint32:                                                                                                                                     │
                    │   134 │   │   │   │   arr = arr.astype(np.uint32, copy=False)                                                                                                                    │
                    │                                                                                                                                                                                  │
                    │ /app/.venv/lib/python3.11/site-packages/numpy/lib/_npyio_impl.py:481 in load                                                                                                     │
                    │                                                                                                                                                                                  │
                    │    478 │   │   │   if mmap_mode:                                                                                                                                                 │
                    │    479 │   │   │   │   if allow_pickle:                                                                                                                                          │
                    │    480 │   │   │   │   │   max_header_size = 2**64                                                                                                                               │
                    │ ❱  481 │   │   │   │   return format.open_memmap(file, mode=mmap_mode,                                                                                                           │
                    │    482 │   │   │   │   │   │   │   │   │   │     max_header_size=max_header_size)                                                                                                │
                    │    483 │   │   │   else:                                                                                                                                                         │
                    │    484 │   │   │   │   return format.read_array(fid, allow_pickle=allow_pickle,                                                                                                  │
                    │                                                                                                                                                                                  │
                    │ /app/.venv/lib/python3.11/site-packages/numpy/lib/format.py:972 in open_memmap                                                                                                   │
                    │                                                                                                                                                                                  │
                    │    969 │   if mode == 'w+':                                                                                                                                                      │
                    │    970 │   │   mode = 'r+'                                                                                                                                                       │
                    │    971 │                                                                                                                                                                         │
                    │ ❱  972 │   marray = numpy.memmap(filename, dtype=dtype, shape=shape, order=order,                                                                                                │
                    │    973 │   │   mode=mode, offset=offset)                                                                                                                                         │
                    │    974 │                                                                                                                                                                         │
                    │    975 │   return marray                                                                                                                                                         │
                    │                                                                                                                                                                                  │
                    │ /app/.venv/lib/python3.11/site-packages/numpy/_core/memmap.py:279 in __new__                                                                                                     │
                    │                                                                                                                                                                                  │
                    │   276 │   │   │   start = offset - offset % mmap.ALLOCATIONGRANULARITY                                                                                                           │
                    │   277 │   │   │   bytes -= start                                                                                                                                                 │
                    │   278 │   │   │   array_offset = offset - start                                                                                                                                  │
                    │ ❱ 279 │   │   │   mm = mmap.mmap(fid.fileno(), bytes, access=acc, offset=start)                                                                                                  │
                    │   280 │   │   │                                                                                                                                                                  │
                    │   281 │   │   │   self = ndarray.__new__(subtype, shape, dtype=descr, buffer=mm,                                                                                                 │
                    │   282 │   │   │   │   │   │   │   │      offset=array_offset, order=order)                                                                                                       │
                    ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
                    ValueError: mmap length is greater than file size
                    ↩  received SIGTERM – shutting down …
[10/19/25 04:34:03] [DCP] rank 2/4 start (window=58181) | owned shards: 1
                    (~69378.35 MiB)
[10/19/25 04:34:03] [DCP] rank 3/4 start (window=58181) | owned shards: 1
                    (~69378.35 MiB)
[10/19/25 04:34:03] [DCP] rank 1/4 start (window=58181) | owned shards: 1
                    (~69378.35 MiB)

It also tries to download shard 1 and 2 for some reason when we should be at shard 4 and 5.