prepare_generation_worker.py (10521B)
1 #!/usr/bin/env python3 2 """Prepare an explicit artifact-only transition worker for an explicitly selected generation backend. 3 The recovery frontend is never invoked. Compilation APIs are retained self-built 4 exports; current map/graph/transport policy is compiled by the selected portable 5 compiler. Generation backend selection is separate from worker compilation. 6 """ 7 import argparse 8 import hashlib 9 import json 10 from pathlib import Path 11 import re 12 import subprocess 13 import time 14 from process_rss import run_guarded 15 from generation_backend import BACKENDS, DEFAULT_BACKEND 16 from generation_io import cli_version, compiler_flags, versions 17 18 ROOT = Path(__file__).resolve().parent.parent 19 parser = argparse.ArgumentParser(description=__doc__) 20 parser.add_argument('tricu', type=Path) 21 parser.add_argument('previous', type=Path) 22 parser.add_argument('output', type=Path) 23 parser.add_argument('--compiler-bundle', type=Path, required=True) 24 parser.add_argument('--compiler-io-version', type=cli_version, default=1) 25 parser.add_argument('--backend', choices=BACKENDS, default=DEFAULT_BACKEND) 26 parser.add_argument('--plan-only', action='store_true', help='report selected support and flat size without writes or execution') 27 parser.add_argument('--support-flat-bytes', type=int, default=64 * 1024 * 1024, help='explicit flat transport budget for retained support') 28 parser.add_argument('--legacy-compiler-header', action='store_true', 29 help='explicit one-time bridge header for a retained pre-import compiler') 30 parser.add_argument('--memory-mib', type=int, default=32768) 31 parser.add_argument('--seconds', type=float, default=1800, help='whole preparation budget, maximum 7200') 32 args = parser.parse_args() 33 if not 0 < args.memory_mib <= 65536 or not 0 < args.seconds <= 7200 or args.support_flat_bytes <= 0: 34 parser.error('preparation limits must be within 64 GiB and two hours') 35 if versions(json.loads((args.previous / 'inputs.json').read_text()))[1] != 1: 36 parser.error('legacy flat bridge requires v1 support; use prepare_compiler_kit.py for source-mapped v2 preparation') 37 start = time.monotonic() 38 peak_kib = 0 39 tricu = str(args.tricu.resolve()) 40 old = args.previous.resolve() / 'cas' 41 root = args.output.resolve() 42 store, tools = root / 'store', root / 'tools' 43 44 def remaining(): 45 budget = args.seconds - (time.monotonic() - start) 46 if budget <= 0: 47 raise TimeoutError('worker preparation wall-time guard') 48 return budget 49 50 def run(command): 51 global peak_kib 52 try: 53 reply, peak = run_guarded([tricu, *map(str, command), '+RTS', '-N1', 54 '-M' + str(args.memory_mib) + 'm', '-RTS'], 55 memory_mib=args.memory_mib, seconds=remaining()) 56 except subprocess.CalledProcessError as error: 57 print(error.output, flush=True) 58 raise 59 peak_kib = max(peak_kib, peak) 60 return reply 61 62 def imported(bundle, destination, indexed=False): 63 reply = run(['arboricx', 'import', '--file', bundle, '--store', destination, *(['--indexed'] if indexed else ['--flat', '--max-flat-bytes', str(args.support_flat_bytes)])]) 64 return re.search(r'Created module alias .* -> ([0-9a-f]{64})', reply)[1] 65 66 parts = [] 67 for name in ['object', 'module', 'lock', 'build', 'map']: 68 text = (ROOT / 'lib/source' / (name + '.tri')).read_text() 69 # This explicit bridge flattens only unqualified dependencies into its 70 # selected support facade. Never discard an authored namespace silently. 71 if name == 'object': 72 # Preserve the qualified arithmetic binding under a distinct facade name. 73 if '!import native_math as N' not in text or set(re.findall(r'\bN\.([A-Za-z_?]+)', text)) != {'lowerHex'}: 74 raise ValueError('unexpected qualified object dependencies') 75 text = text.replace('!import native_math as N', '').replace('N.lowerHex', 'workerNativeLowerHex') 76 if any(line.startswith('!import ') and line.split()[-1:] != ['*'] 77 for line in text.splitlines()): 78 raise ValueError('transition flattening requires unqualified native uses: ' + name) 79 text = '\n'.join(line for line in text.splitlines() 80 if not line.startswith(('!import', '!source', '!use', '!export'))) 81 if name == 'build': 82 # Native packing is retained; this unused bridge bundle entrypoint is 83 # omitted, not the canonical units or either generation backend. 84 text = re.sub(r'sourceBuildGraphBundle workspace.*?(?=-- Same graph/IO)', '', text, flags=re.S) 85 parts.append(text) 86 body = '\n\n'.join(parts) 87 # Conservative local reachability for the selected entrypoint. Token mentions 88 # can retain extra definitions (including comments), never rewrite expressions. 89 # The portable frontend still parses/resolves every retained definition. This 90 # avoids transporting the unused other backend's nearly-GiB flat compiler root. 91 entrypoint = BACKENDS[args.backend][0] 92 matches = list(re.finditer(r'^([A-Za-z_][A-Za-z0-9_?\-]*)[^\n=]*=', body, re.M)) 93 blocks = {} 94 order = [] 95 for index, match in enumerate(matches): 96 name = match[1] 97 if name not in blocks: 98 order.append(name); blocks[name] = '' 99 blocks[name] += body[match.start():matches[index + 1].start() if index + 1 < len(matches) else len(body)] 100 needed, pending = set(), [entrypoint] 101 while pending: 102 name = pending.pop() 103 if name in needed: continue 104 if name not in blocks: raise ValueError('missing worker definition: ' + name) 105 needed.add(name) 106 pending.extend(set(re.findall(r'[A-Za-z_][A-Za-z0-9_?\-]*', blocks[name])) & blocks.keys() - needed) 107 body = '\n'.join(blocks[name] for name in order if name in needed) 108 local = set(re.findall(r'^([A-Za-z_][A-Za-z0-9_?\-]*)[^\n=]*=', body, re.M)) 109 # These constructors are defined below even when the previous IO module already 110 # exports them. Do not also import them and trigger canonical shadow rejection. 111 constructor_names = ('readTreeTerm', 'hashTreeTerm', 'writeTreeTermIfAbsent', 112 'readSourceFile', 'encodeSourceText', 'readIndexedTree', 113 'hashIndexedTree', 'writeIndexedTreeIfAbsent') 114 local.update(constructor_names) 115 words = set(re.findall(r'[A-Za-z_][A-Za-z0-9_?\-]*', body)) 116 selected = {} 117 for module in ['prelude', 'io', 'source_unit', 'source_resolve', 'source_lexer', 'parser', 'native_math']: 118 pin = (old / 'cache/source-build' / module).read_text().split('\t')[2] 119 assert re.fullmatch('[a-f0-9]{64}', pin), 'invalid retained manifest pin' 120 raw = (old / 'objects' / pin[:3] / pin).read_bytes() 121 assert hashlib.sha256(b'arboricx.module-manifest.v1\0' + raw).hexdigest() == pin 122 for row in raw.decode().splitlines()[1:]: 123 _, name, kind, tree, abi = row.split('\t') 124 assert kind == 'arboricx.tree-term.v1' and abi == 'arboricx.abi.tree.v1' 125 if module == 'native_math': 126 if name != 'lowerHex': continue 127 name = 'workerNativeLowerHex' 128 if name in words and name not in local: 129 assert re.fullmatch('[a-f0-9]{64}', tree), 'invalid retained tree pin' 130 selected[name] = tree 131 # New IO constructors are portable data, not Haskell-compiled glue. 132 constructor_lines = [line for line in (ROOT / 'lib/io.tri').read_text().splitlines() 133 if any(line.startswith(name + ' =') for name in constructor_names)] 134 if sorted(line.split()[0] for line in constructor_lines) != sorted(constructor_names): 135 raise ValueError('missing or duplicate required IO constructor: ' + ', '.join(constructor_names)) 136 constructors = '\n'.join(constructor_lines) 137 expanded = sum((old / 'objects' / pin[:3] / pin).stat().st_size for pin in selected.values()) 138 if args.plan_only: 139 print(json.dumps(dict(selected=selected, expanded_support_bytes=expanded, 140 support_flat_bytes=args.support_flat_bytes, backend=args.backend, 141 entrypoint=entrypoint, local_definitions=len(needed)), indent=2)) 142 raise SystemExit(0) 143 if expanded > args.support_flat_bytes: 144 raise ValueError('selected support exceeds explicit flat transport budget; inspect --plan-only') 145 root.mkdir(parents=True, exist_ok=False) 146 command = ['arboricx', 'export', '--store', old, '--output', root / 'support.arboricx'] 147 for name, tree in selected.items(): 148 command += ['--target', tree, '--name', name] 149 print(run(command), end='') 150 support = imported(root / 'support.arboricx', store) 151 manifest_bytes = (store / 'objects' / support[:3] / support).read_bytes() 152 assert hashlib.sha256(b'arboricx.module-manifest.v1\0' + manifest_bytes).hexdigest() == support 153 rows = [row.split('\t') for row in manifest_bytes.decode().splitlines()[1:]] 154 assert len(rows) == len(selected) and {row[1]: row[3] for row in rows} == selected 155 for row in rows: 156 assert (row[0], row[2], row[4]) == ('export', 'arboricx.tree-term.v1', 'arboricx.abi.tree.v1') 157 for tree in set(selected.values()): 158 hasher = hashlib.sha256(b'arboricx.tree-term.v1\0') 159 with (store / 'objects' / tree[:3] / tree).open('rb') as stream: 160 for chunk in iter(lambda: stream.read(1024 * 1024), b''): hasher.update(chunk) 161 assert hasher.hexdigest() == tree, 'transported support digest mismatch' 162 (root / 'support.tri').write_text(f'!export * from {support}\n') 163 header = ('!source support\n!use support *\n' if args.legacy_compiler_header else '!import support *\n') 164 source = header + '!export ' + entrypoint + '\n' + constructors + '\n' + body 165 (root / 'worker.tri').write_text(source) 166 compiler = imported(args.compiler_bundle.resolve(), tools, True) 167 print(run(['arboricx', 'compile', '--file', root / 'worker.tri', '--store', store, 168 '--toolchain', compiler, *compiler_flags(args.compiler_io_version), '--tool-store', tools, '--output', root / 'worker.arboricx', 169 '--compile-timeout', str(max(1, min(900, int(remaining()))))]), end='') 170 pin = imported(root / 'worker.arboricx', store, indexed=True) 171 (root / 'worker.pin').write_text(pin) 172 (root / 'provenance.json').write_text(json.dumps(dict( 173 support=support, compiler=compiler, worker=pin, selected=selected, 174 compiler_io_version=args.compiler_io_version, produced_io_version=1, 175 backend=args.backend, entrypoint=entrypoint, local_definitions=sorted(needed), 176 elapsed_seconds=time.monotonic()-start, peak_kib=peak_kib, 177 memory_mib=args.memory_mib, seconds=args.seconds, 178 source_sha256=hashlib.sha256(source.encode()).hexdigest(), 179 artifact_sha256=hashlib.sha256((root / 'worker.arboricx').read_bytes()).hexdigest()), indent=2)+'\n') 180 print('Worker:', root / 'worker.arboricx')