Spaces:
Sleeping
Sleeping
| """Memory-mapped checkpoint manifests and atomic publication of masks.""" | |
| import fcntl | |
| import json | |
| import os | |
| from pathlib import Path | |
| import shutil | |
| import tempfile | |
| import numpy as np | |
| from .core import _inputs, _options, _run | |
| def _input_path(root, relative, suffix): | |
| if not isinstance(relative, str) or Path(relative).is_absolute(): | |
| raise ValueError("manifest paths must be relative to the input root") | |
| path = (root / relative).resolve(strict=True) | |
| if not path.is_relative_to(root) or not path.is_file() or path.suffix != suffix: | |
| raise ValueError(f"expected a {suffix} file inside the input root") | |
| return path | |
| def _open_npy(root, name): | |
| return np.load(_input_path(root, name, ".npy"), mmap_mode="r", allow_pickle=False) | |
| def _read_manifest(root, manifest, max_weights): | |
| path = _input_path(root, manifest, ".json") | |
| with path.open("rb") as file: | |
| raw = file.read(1_048_577) | |
| if len(raw) > 1_048_576: | |
| raise ValueError("manifest exceeds 1 MiB") | |
| spec = json.loads(raw) | |
| if not isinstance(spec, dict) or set(spec) != {"keep", "projections"}: | |
| raise ValueError("manifest must contain keep and projections") | |
| items = spec["projections"] | |
| if not isinstance(items, list) or not 1 <= len(items) <= 1024: | |
| raise ValueError("provide 1 to 1024 projection records") | |
| if type(spec["keep"]) is not int: | |
| raise ValueError("manifest keep must be an integer") | |
| names, weights, moments = [], [], [] | |
| seen = set() | |
| total = 0 | |
| for item in items: | |
| if (not isinstance(item, dict) or not {"name", "weights"} <= item.keys() | |
| or item.keys() - {"name", "weights", "input_second_moment"}): | |
| raise ValueError("projection records require name and weights") | |
| name = item["name"] | |
| if not isinstance(name, str) or not 1 <= len(name) <= 256 or name in seen: | |
| raise ValueError("projection names must be nonempty and unique") | |
| seen.add(name) | |
| w = _open_npy(root, item["weights"]) | |
| total += w.size | |
| if max_weights is not None and total > max_weights: | |
| raise ValueError("manifest exceeds the configured weight limit") | |
| names.append(name) | |
| weights.append(w) | |
| moments.append(_open_npy(root, item["input_second_moment"]) | |
| if "input_second_moment" in item else None) | |
| return names, weights, moments, spec["keep"] | |
| def prune_manifest(input_root, manifest, output_dir, *, scratch_dir=None, | |
| threads=1, seed=0, max_weights=None): | |
| """Publish disk-backed .npy masks and result.json to a new directory.""" | |
| root = Path(input_root).resolve(strict=True) | |
| names, weights, moments, keep = _read_manifest(root, manifest, max_weights) | |
| return _publish_masks(names, weights, moments, keep, output_dir, | |
| scratch_dir=scratch_dir, threads=threads, seed=seed) | |
| def _publish_masks(names, weights, moments, keep, output_dir, *, | |
| scratch_dir, threads, seed): | |
| output = Path(output_dir).absolute() | |
| arrays, moments = _inputs(weights, moments) | |
| total = sum(w.size for w in arrays) | |
| keep, threads, seed = _options(total, keep, threads, seed) | |
| output.parent.mkdir(parents=True, exist_ok=True) | |
| lock_path = output.parent / f".{output.name}.lock" | |
| with lock_path.open("a+b") as lock: | |
| fcntl.flock(lock, fcntl.LOCK_EX) | |
| if output.exists() or output.is_symlink(): | |
| raise FileExistsError(output) | |
| stage = Path(tempfile.mkdtemp(prefix=f".{output.name}-", dir=output.parent)) | |
| try: | |
| paths = [stage / f"mask-{j:04d}.npy" for j in range(len(arrays))] | |
| masks = [np.lib.format.open_memmap(path, mode="w+", dtype=bool, shape=w.shape) | |
| for path, w in zip(paths, arrays)] | |
| for path in paths: | |
| with path.open("r+b") as file: | |
| os.posix_fallocate(file.fileno(), 0, path.stat().st_size) | |
| scratch = Path(scratch_dir) if scratch_dir is not None else stage | |
| _run(arrays, moments, masks, total, keep, threads, seed, scratch) | |
| entries = [] | |
| for name, path, mask in zip(names, paths, masks): | |
| mask.flush() | |
| entries.append(dict(name=name, mask=path.name, shape=list(mask.shape), | |
| kept_slots=int(np.count_nonzero(mask)))) | |
| result = dict(weights=total, kept_slots=keep, projections=entries) | |
| if sum(item["kept_slots"] for item in entries) != keep: | |
| raise RuntimeError("native selector returned an incorrect retention budget") | |
| (stage / "result.json").write_text(json.dumps(result, indent=2) + "\n") | |
| stage.rename(output) | |
| return result | |
| finally: | |
| if stage.exists(): | |
| shutil.rmtree(stage) | |