import json
import math
import os
from collections import Counter
from collections.abc import Iterator
from typing import Any, Dict
import torch
import torch.distributed as dist
from loguru import logger
from torch.utils.data import IterableDataset, get_worker_info
from sequifier.config.depth_layout import DepthLayoutRegistryModel
from sequifier.helpers import (
configured_window_stride,
normalize_path,
resolve_window_sampling_plan,
stored_window_layout_from_metadata,
validate_stored_window_width,
)
from sequifier.io.batch import SequifierBatch
from sequifier.io.config import global_training
from sequifier.io.iteration_state import (
read_shared_int,
resolve_resume_worker,
shared_int,
skip_samples_for_batches,
write_shared_int,
)
from sequifier.io.pt_payload import load_pt_payload
from sequifier.io.window_sampling import build_window_batch
from sequifier.typechecking import beartype
[docs]class SequifierDatasetFromFolderPtLazy(IterableDataset):
"""Streams PT chunks into rank/worker-aligned batches."""
@beartype
def __init__(self, data_path: str, config: Any, shuffle: bool = True):
super().__init__()
self.data_dir = normalize_path(data_path, config.project_root)
self.config = config
self.batch_size = global_training(config).batch_size
self.shuffle = shuffle
self._epoch_state = shared_int(0)
self._start_batch_state = shared_int(0)
metadata_path = os.path.join(self.data_dir, "metadata.json")
if not os.path.exists(metadata_path):
raise FileNotFoundError(
f"metadata.json not found in '{self.data_dir}'. "
"Ensure data is pre-processed with write_format: pt."
)
with open(metadata_path, "r") as f:
metadata = json.load(f)
self.payload_n_classes = metadata.get("n_classes") or config.n_classes
self.depth_layouts = DepthLayoutRegistryModel.model_validate(
metadata.get("depth_layouts", {})
)
selected_layouts: DepthLayoutRegistryModel | None = getattr(
config, "depth_layouts", None
)
if selected_layouts is None:
selected_layouts = DepthLayoutRegistryModel()
if self.depth_layouts.compatibility_signature(
config.input_columns
) != selected_layouts.compatibility_signature(config.input_columns):
raise ValueError(
"PT folder depth layouts are incompatible with the selected interface"
)
self.folder_layout = stored_window_layout_from_metadata(metadata)
self.sampling_plan = resolve_window_sampling_plan(
self.folder_layout,
config.window_view,
configured_window_stride(config),
)
self.batch_files_info = []
for raw_file_info in metadata["batch_files"]:
file_info = dict(raw_file_info)
file_info["stored_samples"] = int(raw_file_info["samples"])
histogram = raw_file_info.get("left_pad_length_histogram")
if histogram is None and not self.sampling_plan.legacy_single_window:
file_path = os.path.join(self.data_dir, file_info["path"])
left_pad_lengths = load_pt_payload(
file_path,
layouts=self.depth_layouts,
n_classes=self.payload_n_classes,
).left_pad_lengths
histogram = {
str(value): count
for value, count in Counter(left_pad_lengths.tolist()).items()
}
if self.sampling_plan.legacy_single_window:
file_info["samples"] = file_info["stored_samples"]
else:
assert histogram is not None
file_info["samples"] = self.sampling_plan.sample_count_from_histogram(
histogram
)
if file_info["samples"] > 0:
self.batch_files_info.append(file_info)
self.total_samples = sum(info["samples"] for info in self.batch_files_info)
if self.total_samples == 0:
raise ValueError("No usable model windows were found in the dataset.")
self.target_samples = self._get_target_samples()
self.total_batches = self._calculate_total_batches(self.target_samples)
logger.info(
f"Lazy Dataset loaded into RAM with {self.target_samples} samples and {self.total_batches} batches."
)
@beartype
def _calculate_total_batches(self, target_samples: int) -> int:
num_workers = global_training(self.config).num_workers
num_workers_to_use = num_workers if num_workers > 0 else 1
total_batches = 0
for worker_id in range(num_workers_to_use):
worker_samples = target_samples // num_workers_to_use + (
1 if worker_id < target_samples % num_workers_to_use else 0
)
total_batches += math.ceil(worker_samples / self.batch_size)
return total_batches
[docs] @beartype
def set_epoch(self, epoch: int):
"""Set the shuffle epoch."""
write_shared_int(self._epoch_state, epoch)
[docs] @beartype
def set_start_batch(self, start_batch: int):
"""Set the first global batch to yield on the next iteration."""
write_shared_int(self._start_batch_state, start_batch)
@beartype
def _get_target_samples(self) -> int:
"""Return the padded per-rank sample count for aligned distributed steps."""
world_size = dist.get_world_size() if dist.is_initialized() else 1
num_files = len(self.batch_files_info)
samples_per_rank = []
for r in range(world_size):
f_r = list(range(r, num_files, world_size))
samples_per_rank.append(
sum(self.batch_files_info[i]["samples"] for i in f_r) if f_r else 0
)
return max(samples_per_rank)
@beartype
def __len__(self) -> int:
return self.total_batches
@beartype
def __iter__(
self,
) -> Iterator[SequifierBatch]:
world_size = dist.get_world_size() if dist.is_initialized() else 1
rank = dist.get_rank() if dist.is_initialized() else 0
worker_info = get_worker_info()
physical_worker_id = worker_info.id if worker_info is not None else 0
num_workers = worker_info.num_workers if worker_info is not None else 1
epoch = read_shared_int(self._epoch_state)
start_batch = read_shared_int(self._start_batch_state)
num_files = len(self.batch_files_info)
original_files_for_this_rank = list(range(rank, num_files, world_size))
rank_real_samples = sum(
self.batch_files_info[i]["samples"] for i in original_files_for_this_rank
)
files_for_this_rank = original_files_for_this_rank.copy()
if not files_for_this_rank:
if self.target_samples == 0:
return
files_for_this_rank = [rank % num_files]
base_samples_per_worker = self.target_samples // num_workers
remainder = self.target_samples % num_workers
worker_sample_counts = [
base_samples_per_worker + (1 if i < remainder else 0)
for i in range(num_workers)
]
worker_batch_counts = [
math.ceil(sample_count / self.batch_size)
for sample_count in worker_sample_counts
]
worker_id, skip_batches = resolve_resume_worker(
start_batch,
physical_worker_id,
num_workers,
worker_batch_counts,
)
worker_start_sample = 0
for i in range(worker_id):
worker_start_sample += worker_sample_counts[i]
worker_target_samples = worker_sample_counts[worker_id]
worker_end_sample = worker_start_sample + worker_target_samples
skipped_samples = skip_samples_for_batches(
skip_batches, self.batch_size, worker_target_samples
)
worker_start_sample += skipped_samples
worker_target_samples -= skipped_samples
if worker_target_samples <= 0:
return
g = torch.Generator()
g.manual_seed(self.config.seed + epoch)
if self.shuffle:
file_order = torch.randperm(len(files_for_this_rank), generator=g).tolist()
ordered_files = [files_for_this_rank[i] for i in file_order]
else:
ordered_files = files_for_this_rank.copy()
extended_files = []
current_samples = 0
file_idx = 0
while current_samples < self.target_samples:
f_id = ordered_files[file_idx % len(ordered_files)]
extended_files.append(f_id)
current_samples += self.batch_files_info[f_id]["samples"]
file_idx += 1
yielded_samples = 0
global_file_start_sample = 0
seq_buffer: Dict[str, torch.Tensor] = {}
tgt_buffer: Dict[str, torch.Tensor] = {}
meta_buffer: Dict[str, torch.Tensor] = {}
buffer_len = 0
for f_id in extended_files:
if yielded_samples >= worker_target_samples:
break
file_samples = self.batch_files_info[f_id]["samples"]
file_start = global_file_start_sample
file_end = global_file_start_sample + file_samples
global_file_start_sample += file_samples
if file_end <= worker_start_sample or file_start >= worker_end_sample:
continue
file_path = os.path.join(self.data_dir, self.batch_files_info[f_id]["path"])
payload = load_pt_payload(
file_path,
layouts=self.depth_layouts,
n_classes=self.payload_n_classes,
)
(
sequences_batch,
_,
_,
_,
left_pad_lengths_batch,
) = payload
for tensor in sequences_batch.values():
validate_stored_window_width(tensor, self.folder_layout.window_length)
sample_index = self.sampling_plan.build_index(left_pad_lengths_batch)
if len(sample_index) != file_samples:
raise RuntimeError(
f"Expanded sample count mismatch for {file_path}: "
f"metadata={file_samples}, loaded={len(sample_index)}."
)
indices = torch.arange(file_samples)
if self.shuffle:
g_file = torch.Generator()
g_file.manual_seed(self.config.seed + epoch + f_id + rank)
indices = indices[torch.randperm(file_samples, generator=g_file)]
worker_file_start_idx = max(0, worker_start_sample - file_start)
worker_file_end_idx = min(file_samples, worker_end_sample - file_start)
worker_indices = indices[worker_file_start_idx:worker_file_end_idx]
logical_positions = torch.arange(
file_start + worker_file_start_idx,
file_start + worker_file_end_idx,
dtype=torch.int64,
)
sample_is_real = logical_positions < rank_real_samples
num_new_samples = len(worker_indices)
if num_new_samples == 0:
del sequences_batch
continue
new_batch = build_window_batch(
sequences_batch,
self.config.input_columns,
self.config.target_columns,
sample_index,
worker_indices,
sample_is_real,
depth_valid_masks={
name: payload.depth_valid_masks[name]
for name in self.config.depth_layouts.root
},
)
new_seq = new_batch.inputs
new_tgt = new_batch.targets
new_meta = new_batch.metadata
del sequences_batch, left_pad_lengths_batch
if buffer_len == 0:
seq_buffer = new_seq
tgt_buffer = new_tgt
meta_buffer = new_meta
else:
seq_buffer = {
k: torch.cat([seq_buffer[k], new_seq[k]], dim=0) for k in seq_buffer
}
tgt_buffer = {
k: torch.cat([tgt_buffer[k], new_tgt[k]], dim=0) for k in tgt_buffer
}
if set(meta_buffer) != set(new_meta):
raise RuntimeError(
"Inconsistent leftPadLength metadata across PT chunks."
)
meta_buffer = {
k: torch.cat([meta_buffer[k], new_meta[k]], dim=0)
for k in meta_buffer
}
buffer_len += num_new_samples
while buffer_len >= self.batch_size:
if yielded_samples >= worker_target_samples:
break
batch_seq = {k: v[: self.batch_size] for k, v in seq_buffer.items()}
batch_tgt = {k: v[: self.batch_size] for k, v in tgt_buffer.items()}
batch_meta = {k: v[: self.batch_size] for k, v in meta_buffer.items()}
yield SequifierBatch(
inputs=batch_seq,
targets=batch_tgt,
metadata=batch_meta,
)
yielded_samples += self.batch_size
seq_buffer = {k: v[self.batch_size :] for k, v in seq_buffer.items()}
tgt_buffer = {k: v[self.batch_size :] for k, v in tgt_buffer.items()}
meta_buffer = {k: v[self.batch_size :] for k, v in meta_buffer.items()}
buffer_len -= self.batch_size
if buffer_len > 0 and yielded_samples < worker_target_samples:
remaining_needed = worker_target_samples - yielded_samples
final_yield_size = min(buffer_len, remaining_needed)
batch_seq = {k: v[:final_yield_size] for k, v in seq_buffer.items()}
batch_tgt = {k: v[:final_yield_size] for k, v in tgt_buffer.items()}
batch_meta = {k: v[:final_yield_size] for k, v in meta_buffer.items()}
yield SequifierBatch(
inputs=batch_seq,
targets=batch_tgt,
metadata=batch_meta,
)