import json
import math
import os
from collections.abc import Iterator
from typing import Any, Dict
import polars as pl
import torch
import torch.distributed as dist
from loguru import logger
from torch.utils.data import IterableDataset, get_worker_info
from sequifier.helpers import (
PANDAS_TO_TORCH_TYPES,
columns_from_slice,
configured_window_stride,
get_left_pad_lengths_from_preprocessed_data,
normalize_path,
resolve_window_sampling_plan,
stored_window_layout_from_metadata,
)
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.window_sampling import build_window_batch
from sequifier.typechecking import beartype
[docs]class SequifierDatasetFromFolderParquet(IterableDataset):
"""Eager Parquet-folder dataset yielding 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 merge_output: False."
)
with open(metadata_path, "r") as f:
metadata = json.load(f)
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),
)
logger.info(
f"Loading Parquet folder dataset into memory from '{self.data_dir}'..."
)
column_torch_types = {
col: PANDAS_TO_TORCH_TYPES[config.column_data_types[col]]
for col in config.column_data_types
}
# Sequence formatting structures matching long-format schema boundaries
sequence_columns = columns_from_slice(
slice(0, self.folder_layout.window_length),
self.folder_layout.window_length,
)
all_sequences: Dict[str, list[torch.Tensor]] = {
col: [] for col in set(config.input_columns + config.target_columns)
}
all_left_pad_lengths: list[torch.Tensor] = []
# Step 1: Eager I/O reduction pass over all chunk allocations
for file_info in metadata["batch_files"]:
file_path = os.path.join(self.data_dir, file_info["path"])
df = pl.read_parquet(file_path)
left_pad_lengths = get_left_pad_lengths_from_preprocessed_data(df)
if left_pad_lengths is not None:
all_left_pad_lengths.append(left_pad_lengths)
for col in all_sequences.keys():
feature_df = df.filter(pl.col("inputCol") == col)
if not feature_df.is_empty():
tensor_seq = torch.tensor(
feature_df.sort(["sequenceId", "subsequenceId"])
.select(sequence_columns)
.to_numpy(),
dtype=column_torch_types[col],
)
all_sequences[col].append(tensor_seq)
del df
# Step 2: Consolidate data lists into contiguous blocks
self.sequences: Dict[str, torch.Tensor] = {
col: torch.cat(tensors, dim=0)
for col, tensors in all_sequences.items()
if tensors
}
self.left_pad_lengths = torch.cat(all_left_pad_lengths)
self.sample_index = self.sampling_plan.build_index(self.left_pad_lengths)
self.n_samples = len(self.sample_index)
if self.n_samples == 0:
raise ValueError("No usable model windows were found in the dataset.")
# Step 3: Prevent serialization duplications across worker forks via shared memory flags
for tensor in self.sequences.values():
tensor.share_memory_()
self.sample_index.share_memory_()
self.target_samples = self._get_target_samples()
self.total_batches = self._calculate_total_batches(self.target_samples)
logger.info(
f"Parquet 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
samples_per_rank = [
len(range(r, self.n_samples, world_size)) for r in range(world_size)
]
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)
indices = torch.arange(self.n_samples)
if self.shuffle:
g = torch.Generator()
g.manual_seed(self.config.seed + epoch)
indices = indices[torch.randperm(self.n_samples, generator=g)]
indices_for_rank = indices[rank::world_size].tolist()
sample_is_real = [True] * len(indices_for_rank)
real_count = len(indices_for_rank)
if real_count == 0:
fallback_indices = indices.tolist()
n = min(len(fallback_indices), self.target_samples)
indices_for_rank.extend(fallback_indices[:n])
sample_is_real.extend([False] * n)
else:
while len(indices_for_rank) < self.target_samples:
n = min(real_count, self.target_samples - len(indices_for_rank))
indices_for_rank.extend(indices_for_rank[:n])
sample_is_real.extend([False] * n)
worker_batch_counts = [
math.ceil(len(indices_for_rank[i::num_workers]) / self.batch_size)
for i in range(num_workers)
]
worker_id, skip_batches = resolve_resume_worker(
start_batch,
physical_worker_id,
num_workers,
worker_batch_counts,
)
indices_for_worker = indices_for_rank[worker_id::num_workers]
sample_is_real_for_worker = sample_is_real[worker_id::num_workers]
skipped_samples = skip_samples_for_batches(
skip_batches, self.batch_size, len(indices_for_worker)
)
indices_for_worker = indices_for_worker[skipped_samples:]
sample_is_real_for_worker = sample_is_real_for_worker[skipped_samples:]
for i in range(0, len(indices_for_worker), self.batch_size):
batch_indices = indices_for_worker[i : i + self.batch_size]
batch_sample_is_real = sample_is_real_for_worker[i : i + self.batch_size]
yield build_window_batch(
self.sequences,
self.config.input_columns,
self.config.target_columns,
self.sample_index,
batch_indices,
batch_sample_is_real,
)