| |
| """Packed, read-only map from TAR path/member name to byte offset and size. |
| |
| The corpus stays in the existing large TAR containers. A single compressed |
| record per TAR replaces repeated full TAR-directory scans on every DDP rank. |
| """ |
| from collections import OrderedDict |
| import json |
| import os |
| from pathlib import Path |
| import threading |
| import zlib |
|
|
|
|
| class TarMemberMap: |
| def __init__(self, manifest, capacity=256): |
| self.manifest_path = Path(manifest) |
| spec = json.loads(self.manifest_path.read_text()) |
| assert spec['status'] == 'complete' and spec['format'] == 'zlib-json-tar-member-map-v1' |
| root = self.manifest_path.parent |
| index_path = root / spec['index'] |
| data_path = root / spec['data'] |
| assert index_path.stat().st_size == spec['index_bytes'] |
| assert data_path.stat().st_size == spec['data_bytes'] |
| index = json.loads(index_path.read_text()) |
| assert len(index['paths']) == spec['tars'] |
| assert all(len(index[k]) == spec['tars'] for k in ('offsets', 'sizes', 'members', 'tar_sizes')) |
| self.rows = {path: (int(index['offsets'][i]), int(index['sizes'][i]), |
| int(index['members'][i]), int(index['tar_sizes'][i])) |
| for i, path in enumerate(index['paths'])} |
| self.fd = os.open(data_path, os.O_RDONLY) |
| self.capacity = int(capacity) |
| self.local = threading.local() |
|
|
| def _cache(self): |
| if not hasattr(self.local, 'members'): |
| self.local.members = OrderedDict() |
| return self.local.members |
|
|
| def locate(self, path, member): |
| cache = self._cache() |
| if path not in cache: |
| assert path in self.rows, f'TAR absent from packed member map: {path}' |
| offset, size, expected_members, expected_tar_size = self.rows[path] |
| assert os.path.getsize(path) == expected_tar_size, f'TAR size changed after index build: {path}' |
| compressed = os.pread(self.fd, size, offset) |
| assert len(compressed) == size |
| decoded = json.loads(zlib.decompress(compressed)) |
| assert len(decoded) == expected_members |
| cache[path] = decoded |
| while len(cache) > self.capacity: |
| cache.popitem(last=False) |
| cache.move_to_end(path) |
| value = cache[path].get(member) |
| assert value is not None, f'Member absent from packed map: {path}::{member}' |
| return int(value[0]), int(value[1]) |
|
|
|
|
