Spaces:
Sleeping
Sleeping
File size: 4,890 Bytes
a42d9ad | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | """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)
|