   1: from __future__ import annotations
   2: 
   3: from dataclasses import dataclass
   4: from pathlib import Path
   5: 
   6: import numpy as np
   7: 
   8: from mintdim.cli.unit_build import log_event
   9: 
  10: from ....shared.errors import TokenDTypeValidationError
  11: 
  12: 
  13: _DTYPE_MAP = {
  14:     "uint16": np.uint16,
  15:     "uint32": np.uint32,
  16: }
  17: 
  18: _BYTES_PER_TOKEN = {
  19:     "uint16": 2,
  20:     "uint32": 4,
  21: }
  22: 
  23: _MAGIC_VALUE = {
  24:     "uint16": 0xFFFF,
  25:     "uint32": 0xFFFFFFFF,
  26: }
  27: 
  28: _RECOMMENDED_SHARD_MIB = 512
  29: _EMPTY_UNIT_NOTICE = "EMPTY_UNIT.md"
  30: _FAILED_UNIT_NOTICE = "UNIT_BUILD_FAILED.md"
  31: 
  32: 
  33: def select_token_dtype(vocab_size: int) -> str:
  34:     if vocab_size <= _MAGIC_VALUE["uint16"]:
  35:         return "uint16"
  36:     if vocab_size <= _MAGIC_VALUE["uint32"]:
  37:         return "uint32"
  38:     raise TokenDTypeValidationError()
  39: 
  40: 
  41: def bytes_per_token(dtype: str) -> int:
  42:     return _BYTES_PER_TOKEN[dtype]
  43: 
  44: 
  45: def magic_for_dtype(dtype: str) -> int:
  46:     return _MAGIC_VALUE[dtype]
  47: 
  48: 
  49: def assert_no_magic_collision(*, vocab_size: int, token_dtype: str) -> None:
  50:     from ....shared.errors import MagicVocabCollisionError
  51: 
  52:     magic = _MAGIC_VALUE[token_dtype]
  53:     if vocab_size > magic:
  54:         raise MagicVocabCollisionError(
  55:             vocab_size=vocab_size,
  56:             token_dtype=token_dtype,
  57:             magic_value=magic,
  58:         )
  59: 
  60: 
  61: @dataclass(frozen=True)
  62: class ShardRecord:
  63:     """One sample's payload as it will be written to a shard.
  64: 
  65:     `segment_lengths` is parallel to the session's `sequence_template`. Its
  66:     sum must be <= unit_size; the remainder of the payload area is padded.
  67:     `payload_tokens` is the concatenation of every segment's tokens, in
  68:     sequence_template order.
  69:     """
  70: 
  71:     segment_lengths: list[int]
  72:     payload_tokens: list[int]
  73: 
  74: 
  75: class ShardWriter:
  76:     """Append fixed-width records to `unit_<size>/shard_*.bin`.
  77: 
  78:     Each record on disk is laid out as:
  79:         [magic] [len_seg_0] ... [len_seg_{n-1}] [payload_0..unit_size-1]
  80:     All slots are `dtype` (uint16 or uint32). Total record size in slots
  81:     is `1 + n_segments + unit_size`.
  82:     """
  83: 
  84:     def __init__(
  85:         self,
  86:         *,
  87:         root: Path,
  88:         unit_size: int,
  89:         dtype: str,
  90:         pad_id: int,
  91:         samples_per_shard: int,
  92:         n_segments: int,
  93:     ) -> None:
  94:         self.root = Path(root)
  95:         self.unit_size = unit_size
  96:         self.dtype = dtype
  97:         self.np_dtype = _DTYPE_MAP[dtype]
  98:         self.pad_id = pad_id
  99:         self.samples_per_shard = samples_per_shard
 100:         self.n_segments = n_segments
 101:         self.magic = _MAGIC_VALUE[dtype]
 102:         self.record_slots = 1 + n_segments + unit_size
 103: 
 104:         self._shard_index = -1
 105:         self._samples_in_current = 0
 106:         self._total_samples = 0
 107:         self._bytes_in_current = 0
 108:         self._handle = None
 109:         self._current_shard_path: Path | None = None
 110: 
 111:         self.root.mkdir(parents=True, exist_ok=True)
 112: 
 113:     def write_batch(self, records: list[ShardRecord]) -> list[tuple[int, int]]:
 114:         """Append a batch and return (shard_index, position_in_shard) per record."""
 115:         locations: list[tuple[int, int]] = []
 116:         cursor = 0
 117:         while cursor < len(records):
 118:             if self._handle is None or self._samples_in_current >= self.samples_per_shard:
 119:                 self._roll()
 120: 
 121:             capacity = self.samples_per_shard - self._samples_in_current
 122:             chunk = records[cursor : cursor + capacity]
 123:             shard_index = self._shard_index
 124:             start_position = self._samples_in_current
 125: 
 126:             arr = np.full(
 127:                 (len(chunk), self.record_slots),
 128:                 self.pad_id,
 129:                 dtype=self.np_dtype,
 130:             )
 131:             payload_offset = 1 + self.n_segments
 132:             for row, rec in enumerate(chunk):
 133:                 arr[row, 0] = self.magic
 134:                 if len(rec.segment_lengths) != self.n_segments:
 135:                     raise ValueError(
 136:                         f"ShardRecord.segment_lengths must have length {self.n_segments}, "
 137:                         f"got {len(rec.segment_lengths)}"
 138:                     )
 139:                 for j, length in enumerate(rec.segment_lengths):
 140:                     arr[row, 1 + j] = length
 141:                 payload = rec.payload_tokens
 142:                 if len(payload) > self.unit_size:
 143:                     raise ValueError(
 144:                         f"payload_tokens length {len(payload)} exceeds unit_size {self.unit_size}"
 145:                     )
 146:                 arr[row, payload_offset : payload_offset + len(payload)] = payload
 147:             arr.tofile(self._handle)
 148: 
 149:             locations.extend(
 150:                 (shard_index, start_position + offset) for offset in range(len(chunk))
 151:             )
 152:             self._samples_in_current += len(chunk)
 153:             self._total_samples += len(chunk)
 154:             self._bytes_in_current += (
 155:                 len(chunk) * self.record_slots * _BYTES_PER_TOKEN[self.dtype]
 156:             )
 157:             cursor += len(chunk)
 158: 
 159:         return locations
 160: 
 161:     def close(self, *, success: bool = True, failure_reason: str | None = None) -> None:
 162:         if self._handle is not None:
 163:             self._handle.close()
 164:             self._handle = None
 165:             self._emit_shard_status()
 166: 
 167:         if success and self._total_samples == 0:
 168:             self._write_empty_unit_notice()
 169:         elif not success:
 170:             self._write_failed_unit_notice(failure_reason)
 171: 
 172:     def _roll(self) -> None:
 173:         if self._handle is not None:
 174:             self._handle.close()
 175:             self._emit_shard_status()
 176:         self._shard_index += 1
 177:         path = self.root / f"shard_{self._shard_index:06d}.bin"
 178:         self._handle = open(path, "wb")
 179:         self._current_shard_path = path
 180:         self._samples_in_current = 0
 181:         self._bytes_in_current = 0
 182: 
 183:     def _emit_shard_status(self) -> None:
 184:         mib = self._bytes_in_current / (1024 * 1024)
 185: 
 186:         fields: dict[str, object] = {
 187:             "shard": f"{self._shard_index:06d}",
 188:             "samples": self._samples_in_current,
 189:             "bytes": f"{mib:.1f} MiB",
 190:             "unit": self.unit_size,
 191:         }
 192:         if self._current_shard_path is not None:
 193:             fields["path"] = str(self._current_shard_path)
 194: 
 195:         log_event("shard", **fields)
 196: 
 197:         if mib > _RECOMMENDED_SHARD_MIB:
 198:             log_event(
 199:                 "warning",
 200:                 shard=f"{self._shard_index:06d}",
 201:                 size=f"{mib:.1f} MiB",
 202:                 recommended=f"{_RECOMMENDED_SHARD_MIB:.1f} MiB",
 203:             )
 204: 
 205:     def _write_empty_unit_notice(self) -> None:
 206:         path = self.root / _EMPTY_UNIT_NOTICE
 207:         path.write_text(
 208:             "\n".join(
 209:                 [
 210:                     "# EmptyUnitNotice",
 211:                     "",
 212:                     "Status:",
 213:                     "No samples were assigned to this unit.",
 214:                     "",
 215:                     "Unit:",
 216:                     str(self.unit_size),
 217:                     "",
 218:                     "Reason:",
 219:                     (
 220:                         "No rendered JSONL sample produced a token_count that "
 221:                         "the unit planner assigned to this unit."
 222:                     ),
 223:                     "",
 224:                     "This is not an error:",
 225:                     (
 226:                         "The build completed normally. The unit directory is kept "
 227:                         "so the declared output layout remains visible."
 228:                     ),
 229:                     "",
 230:                 ]
 231:             ),
 232:             encoding="utf-8",
 233:         )
 234: 
 235:     def _write_failed_unit_notice(self, failure_reason: str | None) -> None:
 236:         path = self.root / _FAILED_UNIT_NOTICE
 237:         path.write_text(
 238:             "\n".join(
 239:                 [
 240:                     "# UnitBuildFailedNotice",
 241:                     "",
 242:                     "Status:",
 243:                     "The unit-build run failed before this output was finalized.",
 244:                     "",
 245:                     "Unit:",
 246:                     str(self.unit_size),
 247:                     "",
 248:                     "Samples written before failure:",
 249:                     str(self._total_samples),
 250:                     "",
 251:                     "Failure:",
 252:                     failure_reason or "Unknown failure",
 253:                     "",
 254:                     "This is not an empty unit:",
 255:                     (
 256:                         "EMPTY_UNIT.md is written only after a successful build "
 257:                         "where no samples were assigned to this unit."
 258:                     ),
 259:                     "",
 260:                 ]
 261:             ),
 262:             encoding="utf-8",
 263:         )
