"""Streaming full-dataset ingest for ATOM-native training. Resolve one or more shard files, read them in UTF-8-safe byte chunks, atomize online with ``Atomizer.encode_bytes(..., flush=False)`` (flush at shard ends), and yield consecutive ``(current, target)`` packet pairs without ever materializing the full packet list in RAM. Capacity thesis: ATOM targets GPT-3 *capability* via persistent structured matter while streaming complete corpora on a single box — not by matching 175B params. """ from __future__ import annotations import queue import threading from collections import deque from glob import glob from pathlib import Path from typing import Deque, Iterator, Sequence, TypeVar from src.io.atomizer import AtomPacket, Atomizer _T = TypeVar("_T") _PREFETCH_DONE = object() def resolve_shard_paths( data: str | Path | None = None, data_glob: str | None = None, *, root: str | Path | None = None, ) -> list[Path]: """Resolve ``--data`` and/or ``--data-glob`` into an ordered list of shards. Explicit ``data`` paths come first (deduped), then glob matches sorted lexicographically for stable epoch order. """ base = Path(root) if root is not None else Path.cwd() paths: list[Path] = [] seen: set[Path] = set() def _add(path: Path) -> None: resolved = path if path.is_absolute() else (base / path) resolved = resolved.resolve() if resolved in seen: return if not resolved.is_file(): raise FileNotFoundError(f"shard not found: {resolved}") seen.add(resolved) paths.append(resolved) if data: _add(Path(data)) if data_glob: pattern = data_glob matches = sorted(Path(p).resolve() for p in glob(pattern, root_dir=str(base))) if not matches: # Also try the pattern as absolute / cwd-relative without root_dir matches = sorted(Path(p).resolve() for p in glob(pattern)) if not matches: raise FileNotFoundError(f"no shards matched glob: {data_glob!r}") for match in matches: if match in seen: continue if not match.is_file(): continue seen.add(match) paths.append(match) if not paths: raise ValueError("provide --data and/or --data-glob resolving to at least one file") return paths def _utf8_safe_split(buffer: bytes) -> tuple[bytes, bytes]: """Split ``buffer`` into (emit, hold) so hold has no incomplete UTF-8 tail.""" if not buffer: return b"", b"" # Walk back over continuation bytes / incomplete lead. i = len(buffer) while i > 0 and 0x80 <= buffer[i - 1] <= 0xBF: i -= 1 if i == 0: # Entire buffer is continuations — hold everything until more arrives. return b"", buffer lead = buffer[i - 1] if lead < 0x80: return buffer, b"" if 0xC2 <= lead <= 0xDF: need = 2 elif 0xE0 <= lead <= 0xEF: need = 3 elif 0xF0 <= lead <= 0xF4: need = 4 else: # Invalid lead; emit as-is (atomizer is byte-level). return buffer, b"" have = len(buffer) - (i - 1) if have < need: return buffer[: i - 1], buffer[i - 1 :] return buffer, b"" def one_ahead_prefetch(iterable: Iterator[_T]) -> Iterator[_T]: """Yield items from ``iterable`` while reading one ahead on a daemon thread. Same order and values as iterating ``iterable`` directly. The worker blocks when the one-slot buffer is full, so at most one unread item is buffered. On generator close / StopIteration the worker is joined. """ source = iter(iterable) buf: queue.Queue = queue.Queue(maxsize=1) error_box: list[BaseException] = [] def _worker() -> None: try: for item in source: buf.put(item) except BaseException as exc: # noqa: BLE001 — forward to consumer error_box.append(exc) finally: buf.put(_PREFETCH_DONE) thread = threading.Thread(target=_worker, name="atom-stream-prefetch", daemon=True) thread.start() try: while True: item = buf.get() if item is _PREFETCH_DONE: break yield item # type: ignore[misc] if error_box: raise error_box[0] finally: # Drain so a blocked put can finish, then join. try: while True: buf.get_nowait() except queue.Empty: pass thread.join(timeout=5.0) def iter_byte_chunks( paths: Sequence[str | Path], chunk_bytes: int = 1 << 20, start_offset: int = 0, ) -> Iterator[tuple[int, bytes, bool]]: """Yield ``(shard_index, chunk, shard_ended)`` UTF-8-safe byte chunks. Chunks never split mid-codepoint. When a shard ends, the final yield for that shard has ``shard_ended=True`` (possibly with empty ``chunk`` if only a flush is needed after a prior partial). ``start_offset`` (opt-in, default 0 = unchanged): begin at that byte of the concatenated shard sequence (whole shards before it are skipped without reading; inside the target shard the handle seeks and then moves forward past UTF-8 continuation bytes to the next codepoint start). """ if chunk_bytes < 1: raise ValueError("chunk_bytes must be >= 1") path_list = [Path(p) for p in paths] remaining = max(0, int(start_offset)) for shard_index, path in enumerate(path_list): seek_to = 0 if remaining: size = path.stat().st_size if remaining >= size: remaining -= size continue seek_to, remaining = remaining, 0 hold = b"" emitted_any = False with path.open("rb") as handle: if seek_to: handle.seek(seek_to) while True: b = handle.read(1) if not b: break if not (0x80 <= b[0] <= 0xBF): handle.seek(-1, 1) break while True: block = handle.read(chunk_bytes) if not block: break combined = hold + block emit, hold = _utf8_safe_split(combined) if emit: emitted_any = True yield shard_index, emit, False if hold: emitted_any = True yield shard_index, hold, True elif emitted_any: yield shard_index, b"", True else: # Empty shard: still signal end so callers can flush atomizer state. yield shard_index, b"", True class StreamingPacketSource: """Online ``(current, target)`` transitions from corpus shards. Parameters ---------- atomizer: Stateful ``Atomizer`` used for incremental encode (owned by caller; this source does not deep-copy it). paths: Ordered shard paths. chunk_bytes: Read size before UTF-8-safe trimming (default 1 MiB). loop: If True, restart from the first shard forever (infinite continuum). buffer_size: Bounded ring of recent packets for validation samples (default 2048). prefetch: If True (default), read the next byte chunk on a background thread while the current chunk is consumed/atomized. Same bytes and order; disable with ``prefetch=False`` / ``--no-stream-prefetch``. """ def __init__( self, atomizer: Atomizer, paths: Sequence[str | Path], *, chunk_bytes: int = 1 << 20, loop: bool = True, buffer_size: int = 2048, reset_atomizer: bool = True, prefetch: bool = True, start_offset: int = 0, ) -> None: if not paths: raise ValueError("StreamingPacketSource requires at least one shard path") if buffer_size < 2: raise ValueError("buffer_size must be >= 2 to hold a transition pair") self.atomizer = atomizer self.paths = [Path(p) for p in paths] self.chunk_bytes = int(chunk_bytes) self.loop = bool(loop) self.buffer_size = int(buffer_size) self.prefetch = bool(prefetch) # Opt-in: first pass starts at this byte of the concatenated shards (later loops from 0). self.start_offset = max(0, int(start_offset)) self.bytes_seen = 0 self.packets_seen = 0 self.shard_index = 0 self.epoch = 0 self._ring: Deque[AtomPacket] = deque(maxlen=self.buffer_size) self._prev: AtomPacket | None = None self._pending: Deque[tuple[AtomPacket, AtomPacket]] = deque() self._exhausted = False if reset_atomizer: self.atomizer.reset() def stats(self) -> dict: return { "bytes_seen": self.bytes_seen, "packets_seen": self.packets_seen, "shard_index": self.shard_index, "epoch": self.epoch, "ring_size": len(self._ring), "n_shards": len(self.paths), "loop": self.loop, "chunk_bytes": self.chunk_bytes, "prefetch": self.prefetch, "paths": [str(p) for p in self.paths], } def validation_packets(self) -> list[AtomPacket]: """Snapshot of the ring buffer (oldest → newest).""" return list(self._ring) def validation_transitions(self) -> list[tuple[AtomPacket, AtomPacket]]: packets = self.validation_packets() if len(packets) < 2: return [] return list(zip(packets[:-1], packets[1:])) def _ingest_packets(self, packets: list[AtomPacket]) -> None: for packet in packets: self.packets_seen += 1 self._ring.append(packet) if self._prev is not None: self._pending.append((self._prev, packet)) self._prev = packet def _feed_chunk(self, chunk: bytes, *, flush: bool) -> None: if chunk: self.bytes_seen += len(chunk) packets = self.atomizer.encode_bytes(chunk, reset=False, flush=False) self._ingest_packets(packets) if flush: # Empty encode with flush=True drains the atomizer pending span. flushed = self.atomizer.encode_bytes(b"", reset=False, flush=True) self._ingest_packets(flushed) def _chunk_stream(self) -> Iterator[tuple[int, bytes, bool]]: first = True while True: if first and self.start_offset: chunks: Iterator[tuple[int, bytes, bool]] = iter_byte_chunks( self.paths, chunk_bytes=self.chunk_bytes, start_offset=self.start_offset ) else: chunks = iter_byte_chunks(self.paths, chunk_bytes=self.chunk_bytes) first = False if self.prefetch: chunks = one_ahead_prefetch(chunks) for shard_index, chunk, shard_ended in chunks: yield shard_index, chunk, shard_ended if not self.loop: return self.epoch += 1 def __iter__(self) -> Iterator[tuple[AtomPacket, AtomPacket]]: return self def __next__(self) -> tuple[AtomPacket, AtomPacket]: if self._pending: return self._pending.popleft() if self._exhausted: raise StopIteration if not hasattr(self, "_chunk_iter"): self._chunk_iter = self._chunk_stream() while not self._pending: try: shard_index, chunk, shard_ended = next(self._chunk_iter) except StopIteration: self._exhausted = True raise self.shard_index = shard_index self._feed_chunk(chunk, flush=shard_ended) return self._pending.popleft() def next_transition(self) -> tuple[AtomPacket, AtomPacket]: """Pull the next ``(current, target)`` pair (same as ``next(self)``).""" return next(self)