Add bounded cold and warm preparation-only comparison command

2026-11-04 19:43:49 +08:00 · 2026-09-19 14:50:44 -07:00
parent 6995c688db
commit a7551f9df4
1 changed files with 81 additions and 0 deletions
@@ -0,0 +1,81 @@
+"""Exactly two real semantic preparations; no diffusion, cache hits, or device access."""
+import argparse
+from dataclasses import replace
+import json
+from pathlib import Path
+import resource
+import sys
+import time
+
+import numpy as np
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path[:0] = [str(ROOT / 'prototype'), str(ROOT / 'experiments/ace_chestnut_20260916')]
+from hook_planning import PlanRequest
+from host_hook_adapter import HostHookAdapter, PreparationOnlyDecoder
+
+
+def compare(reference, actual):
+    result = {}
+    for name in ('encoder_hidden_states', 'encoder_attention_mask', 'context_latents'):
+        before = np.load(reference / (name + '.npy'), allow_pickle=False)
+        after = np.load(actual / (name + '.npy'), allow_pickle=False)
+        if before.shape != after.shape:
+            raise ValueError('Conditioning shape changed: ' + name)
+        delta = before.astype(np.float64) - after.astype(np.float64)
+        result[name] = {'shape': list(after.shape), 'exact': bool(np.array_equal(before, after)),
+                        'max_absolute_error': float(np.max(np.abs(delta))),
+                        'rms_error': float(np.sqrt(np.mean(delta * delta)))}
+    old = json.loads((reference / 'semantic_plan.json').read_text())
+    new = json.loads((actual / 'semantic_plan.json').read_text())
+    result['semantic_codes_equal'] = old['audio_codes'] == new['audio_codes']
+    return result
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--reference', type=Path, required=True)
+    parser.add_argument('--assets-root', type=Path, required=True)
+    parser.add_argument('--out', type=Path, required=True)
+    args = parser.parse_args()
+    original = PlanRequest(**json.loads((args.reference / 'prepared.json').read_text())['request'])
+    if original.plan_index != 0 or original.prefix_seconds or original.hook_reference_sha256:
+        raise ValueError('This bounded check only accepts the saved initial request')
+    args.out.mkdir(parents=True, exist_ok=False)
+    report = {'reference': str(args.reference), 'runs': [], 'audio_diffusion_permitted': False}
+    try:
+        adapter = HostHookAdapter(args.assets_root, preparation_only=True)
+        identity = adapter.fingerprints()
+        if identity['model_fingerprint'] != original.model_fingerprint:
+            raise ValueError('Model assets differ from the original request')
+        request = replace(original, **identity)
+        report.update(original_request=original.identity(), request=request.identity())
+        for label in ('cold', 'warm'):
+            output = args.out / label
+            output.mkdir()
+            started = time.monotonic()
+            adapter(request, {}, output)  # Bypass PlanCache: both calls perform real preparation.
+            elapsed = time.monotonic() - started
+            if not isinstance(adapter.handler.mlx_decoder, PreparationOnlyDecoder):
+                raise RuntimeError('Preparation-only decoder was not retained')
+            prepared = json.loads((output / 'prepared.json').read_text())
+            if prepared['audio_diffusion_called'] or not prepared['host_preparation_only']:
+                raise RuntimeError('Unexpected preparation provenance')
+            import mlx.core as mx
+            result = {'phase': label, 'seconds': elapsed,
+                      'process_peak_rss_bytes_macos': resource.getrusage(resource.RUSAGE_SELF).ru_maxrss,
+                      'mlx_active_bytes_after': mx.get_active_memory(),
+                      'comparison_to_original': compare(args.reference, output)}
+            if label == 'warm': result['comparison_to_cold'] = compare(args.out / 'cold', output)
+            report['runs'].append(result)
+            (args.out / 'comparison.json').write_text(json.dumps(report, indent=2))
+            print(label, elapsed, 'seconds; comparisons saved', flush=True)
+    except BaseException as error:
+        report['error'] = repr(error)
+        raise
+    finally:
+        (args.out / 'comparison.json').write_text(json.dumps(report, indent=2))
+
+
+if __name__ == '__main__':
+    main()