#!/usr/bin/env python3 """Bounded model-build commands and integrity checks (also shipped in d4r/tools).""" import argparse import hashlib import json import os from pathlib import Path import signal import subprocess import sys import time MANIFEST = 'files.sha256.json' MODELS = ('rr-d', 'rr-e', 'rr-f', 'rr-d-gfx11', 'rr-e-gfx11', 'k', 'k-ldr', 'm', 'm-ldr', 'l', 'l-ldr') def digest(path): with path.open('rb') as stream: return hashlib.file_digest(stream, 'sha256').hexdigest() def manifest(directory, write=False): directory = Path(directory) files = {str(p.relative_to(directory)): digest(p) for p in sorted(directory.rglob('*')) if p.is_file() and p.name != MANIFEST} if not files or not any(n in files for n in ('model.json', 'rr.txt')): raise ValueError(f'{directory}: missing model metadata') if write: (directory / MANIFEST).write_text(json.dumps(files, indent=2, sort_keys=True) + '\n') elif files != json.loads((directory / MANIFEST).read_text()): raise ValueError(f'{directory}: missing, changed or unexpected model files; rebuild the models') def capture(directory, name): """A completed, finite frame and all files referenced by each network-layer capture.""" import numpy as np directory = Path(directory) output = directory / 'out.raw' if output.stat().st_size != 1280 * 720 * 8: raise ValueError('capture did not produce a complete RGBA16F frame') pixels = np.fromfile(output, dtype=np.float16).reshape(-1, 4) if not np.isfinite(pixels).all() or not np.any(pixels[:, :3]): raise ValueError('capture produced a nonfinite or empty frame') if name == 'k': kernels = ['dltss_pwin_' + s + '_layer' for s in ('enc0', 'enc1', 'enc2', 'enc3', 'enc4', 'dec5', 'dec4', 'dec3', 'dec2', 'dec1', 'dec0')] else: kernels = ['rrlite_' + s + '_4x4' for s in ('enc1', 'enc2', 'enc3', 'enc4', 'enc5', 'dec4', 'dec3', 'dec2', 'dec1')] kernels += ['rrlite_enc0_4x4_mvlo_hdr' + ('_folded' if name == 'm' else ''), 'rrlite_dec0_4x4' + ('_folded' if name == 'm' else '')] for kernel in kernels: matches = sorted((directory / 'replay').glob('replay-*-' + kernel)) if not matches: raise ValueError(f'missing captured layer {kernel}') launch = matches[0] rows = [line.split() for line in (launch / 'manifest.txt').read_text().splitlines()] if ['kernel', kernel] not in rows: raise ValueError(f'{kernel}: incorrect capture metadata') args = next(int(r[1]) for r in rows if r[0] == 'args') if args <= 0 or (launch / 'args.bin').stat().st_size != args: raise ValueError(f'{kernel}: truncated arguments') allocations = {r[1]: int(r[3]) for r in rows if r[0] == 'alloc'} if not allocations: raise ValueError(f'{kernel}: no weight allocations captured') for index, size in allocations.items(): if size <= 0 or (launch / f'alloc-{index}.bin').stat().st_size != size: raise ValueError(f'{kernel}: truncated allocation {index}') for r in rows: if r[0] == 'pointer' and (r[2] not in allocations or not 0 <= int(r[3]) < allocations[r[2]]): raise ValueError(f'{kernel}: invalid captured pointer') def run(log, timeout, command): if command and command[0] == '--': command = command[1:] if not command or not 0 < timeout <= 86400: raise ValueError('command required; timeout must be between 0 and 86400 seconds') log = Path(log) log.parent.mkdir(parents=True, exist_ok=True) started = time.monotonic() print(f' starting {log.stem} (deadline {timeout:g}s; log {log})', flush=True) with log.open('wb') as output: child = subprocess.Popen(command, stdout=output, stderr=subprocess.STDOUT, start_new_session=True) def stop(): for sig in (signal.SIGTERM, signal.SIGKILL): try: os.killpg(child.pid, sig) except ProcessLookupError: break try: child.wait(timeout=3) except subprocess.TimeoutExpired: continue # Kill remaining members too, even when the group leader has exited. if sig == signal.SIGTERM: try: os.killpg(child.pid, signal.SIGKILL) except ProcessLookupError: pass break previous = {} def interrupted(sig, frame): raise KeyboardInterrupt for sig in (signal.SIGTERM, signal.SIGINT): previous[sig] = signal.signal(sig, interrupted) result = None try: heartbeat = 0 while child.poll() is None: elapsed = time.monotonic() - started if elapsed >= timeout: stop() print(f' {log.stem}: timed out after {elapsed:.0f}s', file=sys.stderr) result = 124 break if elapsed >= heartbeat + 15: print(f' {log.stem}: still running ({elapsed:.0f}s, pid {child.pid}); log {log}', flush=True) heartbeat = elapsed time.sleep(min(.25, max(.01, timeout - elapsed))) if result is None: result = child.returncode except KeyboardInterrupt: stop() result = 130 finally: for sig, handler in previous.items(): signal.signal(sig, handler) if result: print('\n'.join(log.read_text(errors='replace').splitlines()[-30:]), file=sys.stderr) print(f' {log.stem}: exit={result}, elapsed={time.monotonic() - started:.1f}s; log {log}', flush=True) return result if result >= 0 else 128 - result def main(): parser = argparse.ArgumentParser(description=__doc__) sub = parser.add_subparsers(dest='action', required=True) for action in ('manifest', 'validate'): sub.add_parser(action).add_argument('directory', type=Path) p = sub.add_parser('capture') p.add_argument('directory', type=Path) p.add_argument('name', choices=('k', 'm', 'l')) p = sub.add_parser('run') p.add_argument('--log', required=True) p.add_argument('--timeout', type=float, default=900) p.add_argument('command', nargs=argparse.REMAINDER) a = parser.parse_args() try: if a.action == 'run': return run(a.log, a.timeout, a.command) if a.action == 'capture': capture(a.directory, a.name) elif a.action == 'manifest': manifest(a.directory, write=True) else: for name in MODELS: manifest(a.directory / name) except (OSError, ValueError, KeyError, StopIteration) as error: print(f'model-tools: {error}', file=sys.stderr) return 1 return 0 if __name__ == '__main__': sys.exit(main())