| """Bounded random access to indexed JSONL metadata and existing codec shards. |
| |
| Only offsets are memory-mapped. No corpus-sized Python list or concatenated |
| audio/code tensor is copied into each GPU process. TAR member indices and |
| Parquet reference tables are kept in small per-thread LRU caches. |
| """ |
| from collections import OrderedDict |
| import hashlib |
| import io |
| import json |
| import os |
| from pathlib import Path |
| import tarfile |
| import threading |
| import numpy as np |
|
|
|
|
| class ShardReader: |
| def __init__(self, capacity=16, file_capacity=None): |
| import resource |
| self.capacity = capacity |
| soft_limit = resource.getrlimit(resource.RLIMIT_NOFILE)[0] |
| self.file_capacity = file_capacity or max(32, min(2048, soft_limit // 2)) |
| self.local = threading.local() |
|
|
| def cache(self): |
| if not hasattr(self.local, 'files'): |
| self.local.files = OrderedDict() |
| self.local.tables = OrderedDict() |
| self.local.members = OrderedDict() |
| return self.local |
|
|
| def bytes_at(self, path, offset, size): |
| state = self.cache() |
| if path not in state.files: |
| state.files[path] = os.open(path, os.O_RDONLY) |
| while len(state.files) > self.file_capacity: |
| _, fd = state.files.popitem(last=False) |
| os.close(fd) |
| fd = state.files[path] |
| state.files.move_to_end(path) |
| value = os.pread(fd, int(size), int(offset)) |
| if len(value) != size: |
| raise IOError(f'Short shard read: {path}@{offset}+{size}') |
| return value |
|
|
| def member(self, path, name): |
| state = self.cache() |
| if path not in state.members: |
| with tarfile.open(path, 'r:') as archive: |
| state.members[path] = {m.name: (m.offset_data, m.size) for m in archive if m.isfile()} |
| while len(state.members) > self.capacity: |
| state.members.popitem(last=False) |
| state.members.move_to_end(path) |
| offset, size = state.members[path][name] |
| return self.bytes_at(path, offset, size) |
|
|
| def codes(self, locator): |
| kind = locator.get('kind', 'tar_moss_npy') |
| if kind == 'tar_moss_npy': |
| blob = (self.bytes_at(locator['path'], locator['offset'], locator['size']) |
| if 'offset' in locator else self.member(locator['path'], locator['member'])) |
| try: |
| result = np.load(io.BytesIO(blob), allow_pickle=False) |
| except ValueError: |
| |
| |
| result = np.frombuffer(blob, dtype=np.uint16).reshape(-1, 12) |
| assert result.dtype == np.uint16 |
| elif kind == 'local_moss_parquet': |
| import pyarrow.parquet as pq |
| state = self.cache() |
| key = (locator['path'], locator['codes_field'], locator['frames_field']) |
| if key not in state.tables: |
| state.tables[key] = pq.read_table(locator['path'], columns=[locator['codes_field'], locator['frames_field']]) |
| while len(state.tables) > self.capacity: |
| state.tables.popitem(last=False) |
| state.tables.move_to_end(key) |
| row = state.tables[key].slice(locator['row_index'], 1).to_pylist()[0] |
| frames = int(row[locator['frames_field']]) |
| result = np.frombuffer(row[locator['codes_field']], dtype='<i2').reshape(frames, 12) |
| else: |
| raise ValueError(f'Unsupported verified codec locator: {kind}') |
| assert result.ndim == 2 and result.shape[1] == 12 and len(result) > 0 |
| assert int(result.min()) >= 0 and int(result.max()) < 1024 |
| expected = locator.get('full_frames') |
| if expected is not None: |
| assert len(result) == expected |
| if 'crop_frames' in locator: |
| start, length = int(locator.get('crop_start_frame', 0)), int(locator['crop_frames']) |
| assert start >= 0 and length > 0 and start + length <= len(result) |
| result = result[start:start + length] |
| return result |
|
|
| def record(self, locator): |
| blob = (self.bytes_at(locator['path'], locator['offset'], locator['size']) |
| if 'offset' in locator else self.member(locator['path'], locator['member'])) |
| return json.loads(blob) |
|
|
|
|
| class IndexedRecords: |
| def __init__(self, manifest): |
| self.manifest_path = Path(manifest) |
| self.spec = json.loads(self.manifest_path.read_text()) |
| root = self.manifest_path.parent |
| self.offsets = np.load(root / self.spec['offsets'], mmap_mode='r', allow_pickle=False) |
| self.path = root / self.spec['records'] |
| self.fd = os.open(self.path, os.O_RDONLY) |
| assert len(self.offsets) == self.spec['samples'] + 1 |
| assert int(self.offsets[0]) == 0 and int(self.offsets[-1]) == self.path.stat().st_size |
| self.reader = ShardReader() |
|
|
| def __len__(self): |
| return len(self.offsets) - 1 |
|
|
| def row(self, index): |
| start, end = int(self.offsets[index]), int(self.offsets[index + 1]) |
| return json.loads(os.pread(self.fd, end - start, start)) |
|
|
| def example(self, index, packer): |
| from modes import verify_record_mode |
| row = self.row(index) |
| verify_record_mode(row) |
| mode = row['mode'] |
| meta = self.reader.record(row['record_locator']) |
| assert meta['uid'] == row['uid'] |
| meta['frames'] = int(meta['moss_frames']) |
| codes = self.reader.codes(row['target_locator']) |
| assert len(codes) == meta['frames'] |
| reference = None |
| if mode == 'reference': |
| assert row['reference_uid'] != row['uid'] and row['reference_evidence'] |
| reference = self.reader.codes(row['reference_locator']) |
| assert 0 < len(reference) <= 37 |
| elif mode == 'score': |
| meta['conditioning_measurements'] = row['conditioning_measurements'] |
| result = packer.pack_mode(meta, codes, mode, reference) |
| result['accounting'] = {'uid': row['uid'], 'mode': mode, 'frames': len(codes), |
| 'target_audio_hours': float(meta['dur_s']) / 3600, |
| 'reference_frames': len(reference) if reference is not None else 0} |
| return result |
|
|
|
|
| def write_manifest(directory, rows, contract): |
| directory = Path(directory) |
| directory.mkdir(parents=True, exist_ok=True) |
| offsets, n, modes, duration = [0], 0, {}, 0. |
| path = directory / 'records.jsonl' |
| digest = hashlib.sha256() |
| seen = set() |
| with path.with_suffix('.jsonl.tmp').open('wb') as stream: |
| for row in rows: |
| assert row['uid'] not in seen, row['uid'] |
| seen.add(row['uid']) |
| assert row['mode'] in ('instruction', 'reference', 'score') |
| blob = (json.dumps(row, ensure_ascii=False, separators=(',', ':')) + '\n').encode() |
| stream.write(blob); digest.update(blob); offsets.append(offsets[-1] + len(blob)); n += 1 |
| modes[row['mode']] = modes.get(row['mode'], 0) + 1 |
| duration += row.get('dur_s', 0.) / 3600 |
| path.with_suffix('.jsonl.tmp').replace(path) |
| np.save(directory / 'offsets.npy', np.asarray(offsets, dtype=np.int64), allow_pickle=False) |
| result = dict(contract, records=path.name, offsets='offsets.npy', samples=n, modes=modes, |
| target_audio_hours=duration, records_sha256=digest.hexdigest()) |
| temporary = directory / 'manifest.json.tmp' |
| temporary.write_text(json.dumps(result, indent=2) + '\n') |
| temporary.replace(directory / 'manifest.json') |
| return directory / 'manifest.json' |
|
|
