Files
digital-drosophila/models/male-cns-lif/conditioning-acceleration-v1/benchmark.py
T
2026-10-01 14:35:32 +08:00

258 lines
15 KiB
Python

"""Bounded replay benchmark. No brain simulation, training, or weight installation."""
from __future__ import annotations
import argparse
from datetime import datetime, timezone
import hashlib
import json
import math
import os
from pathlib import Path
import statistics
import subprocess
import sys
import time
import traceback
import urllib.request
import numpy as np
from gpu_backend import CudaEligibility, prepare_events
REG_SHA = 'b63e6281b6505ca3b60acb467dd0b3fae274cf513dcbfd653649ba52e9ba9740'
IDENTITY_KEYS = ('kc', 'candidate_edge_id', 'candidate_sources', 'candidate_targets',
'candidate_compartments', 'base_candidate')
OUTPUTS = ('sampled_trace', 'final_trace', 'snapshots', 'delta_after')
def sha(path):
h = hashlib.sha256()
with Path(path).open('rb') as stream:
for chunk in iter(lambda: stream.read(8 << 20), b''):
h.update(chunk)
return h.hexdigest()
def save(path, value):
with Path(path).open('x', encoding='utf-8') as stream:
json.dump(value, stream, indent=2, sort_keys=True, allow_nan=False)
stream.write('\n')
def read(path):
return json.loads(Path(path).read_text())
def evidence(a):
return dict(dtype=a.dtype.str, shape=list(a.shape), sha256=hashlib.sha256(a.tobytes()).hexdigest())
def contained(root, relative):
p = root / relative
assert not p.is_symlink() and p.resolve().is_relative_to(root.resolve()), 'Escaping fixture path'
return p
def idle_queues():
result = {}
for port in (8188, 8189, 8190):
with urllib.request.urlopen(f'http://192.168.10.51:{port}/queue', timeout=5) as response:
q = json.load(response)
assert set(('queue_running', 'queue_pending')).issubset(q)
result[str(port)] = {key: len(q[key]) for key in ('queue_running', 'queue_pending')}
assert not any(result[str(port)].values()), 'ComfyUI has work; benchmark will not compete'
return result
def gpu_snapshot(uuid):
text = subprocess.check_output(['nvidia-smi', '--id=' + uuid,
'--query-gpu=uuid,name,pci.bus_id,memory.total,memory.used,utilization.gpu,driver_version',
'--format=csv,noheader,nounits'], text=True).strip()
fields = [s.strip() for s in text.split(',')]
assert fields[0] == uuid
return dict(zip(('uuid', 'name', 'pci_bus_id', 'memory_total_mib', 'memory_used_mib',
'utilization_percent', 'driver_version'), fields))
def stats(values):
return dict(n=len(values), min=min(values), median=statistics.median(values),
mean=statistics.mean(values), max=max(values))
def compare(result, arrays):
report = {}
for key in OUTPUTS:
actual, expected = getattr(result, key), arrays['reference_' + key]
exact = actual.dtype == expected.dtype and actual.shape == expected.shape and actual.tobytes() == expected.tobytes()
diff = np.abs(actual.astype(np.float64) - expected.astype(np.float64))
report[key] = dict(bitwise_exact=exact, actual=evidence(actual), expected=evidence(expected),
max_abs_error=float(np.max(diff)) if diff.size else 0.,
mean_abs_error=float(np.mean(diff)) if diff.size else 0.)
return report
def load_fixtures(manifests):
fixtures = []
common_identity = None
for manifest_path in manifests:
m = read(manifest_path)
assert m.get('complete') is True and m['registration_sha256'] == REG_SHA
root = manifest_path.parent
for f in m['fixtures']:
npz = contained(root, f['npz'])
command_path = contained(root, f['command'])
assert sha(npz) == f['npz_sha256'] and sha(command_path) == f['command_sha256']
with np.load(npz, allow_pickle=False) as z:
a = {k: z[k].copy() for k in z.files}
identity = {k: evidence(a[k]) for k in IDENTITY_KEYS}
if common_identity is None:
common_identity = identity
assert common_identity == identity, 'Full KC/mask/base identity differs across fixtures'
assert len(a['kc']) == 4064 and len(a['candidate_edge_id']) == 14232
sources_path = npz.parent / 'reference-sources.json'
assert sha(sources_path) == m['files'][sources_path.relative_to(root).as_posix()]['sha256']
source_pins = read(sources_path)
expected_sources = dict(
conditioning_update_source='96efff678cff910654eb08e86f1a17789e4e5589e340345f419c7027a7d547e2',
eligibility_bridge_source='174fb8d14619189e74b095e77a6db18b14c93046dc431c0b55710a5aa2dd2e61',
phase4j_source='abb4574da85f00bd9ef61024c65993ef81d42651ff8d8d63c15dd69c6222471e')
assert set(source_pins) == set(expected_sources)
for role, info in source_pins.items():
assert info['sha256'] == expected_sources[role] and sha(info['path']) == info['sha256'], 'Reference source changed'
# Extract constants from the authenticated original bridge, without import/exec.
import ast
bridge = source_pins['eligibility_bridge_source']
assert bridge['sha256'] == '174fb8d14619189e74b095e77a6db18b14c93046dc431c0b55710a5aa2dd2e61'
constants = {}
for node in ast.parse(Path(bridge['path']).read_text()).body:
if isinstance(node, ast.Assign):
for target in node.targets:
if isinstance(target, ast.Name) and target.id in ('DT_MS', 'TAU_MS', 'ETA_MV', 'FLOOR_FRACTION', 'POLARITY'):
constants[target.id] = ast.literal_eval(node.value)
assert constants['POLARITY'] == -1
rule = dict(eta_mv=constants['ETA_MV'], floor_fraction=constants['FLOOR_FRACTION'],
decay=math.exp(-constants['DT_MS'] / constants['TAU_MS']))
fixtures.append(dict(group=root.name, record=f, arrays=a, command=read(command_path),
rule=rule, input_path=str(npz), manifest_sha256=sha(manifest_path)))
assert fixtures
return fixtures, common_identity
def check_rejections(manifests, backend, rule):
reports = []
for path in manifests:
for case in read(path).get('invalid_cases', []):
root = path.parent
npz, command = contained(root, case['npz']), contained(root, case['command'])
assert sha(npz) == case['npz_sha256'] and sha(command) == case['command_sha256']
with np.load(npz, allow_pickle=False) as z:
a = {k: z[k].copy() for k in z.files}
c = read(command)
commands = [] if c['us_tick'] is None else [c['us_tick']]
stage = 'prepare'
try:
prepared = prepare_events(a['step'], a['neuron_idx'], a['kc'], c['steps'], commands,
n_neurons=c['n_neurons'])
stage = 'run'
backend.run(prepared, a['candidate_sources'], a['candidate_compartments'],
c['compartment'], a['delta_before'], a['base_candidate'],
d=c['d_command'], frozen=c['frozen'], **rule)
except ValueError as error:
assert stage == case['stage'], 'Rejected at a different validation boundary'
reports.append(dict(case_name=case['case_name'], rejected=True,
stage=stage, expected_exception=case['expected_exception'], message=str(error)))
else:
raise AssertionError('Candidate accepted invalid case: ' + case['case_name'])
return reports
def main(args):
assert os.environ.get('CUDA_VISIBLE_DEVICES') == args.uuid, 'Use one explicit physical UUID per process'
assert args.output.is_absolute() and not args.output.exists()
args.output.mkdir()
started = time.perf_counter()
try:
queues_before = idle_queues()
gpu_before = gpu_snapshot(args.uuid)
initial_manifest_pins = {str(p): sha(p) for p in args.manifest}
fixtures, common_identity = load_fixtures(args.manifest)
backend = CudaEligibility(device_index=0)
cp = backend.cp
props = cp.cuda.runtime.getDeviceProperties(0)
bus = cp.cuda.runtime.deviceGetPCIBusId(0)
if isinstance(bus, bytes): bus = bus.decode()
assert bus.lower()[-10:] == gpu_before['pci_bus_id'].lower()[-10:], 'Physical GPU binding differs'
device = dict(**gpu_before, runtime_pci_bus_id=bus, compute_capability=backend.device.compute_capability,
initialization_seconds=backend.initialization_seconds, cupy_version=cp.__version__,
numpy_version=np.__version__, python_version=sys.version,
runtime_version=cp.cuda.runtime.runtimeGetVersion(), driver_version_runtime=cp.cuda.runtime.driverGetVersion(),
runtime_name=props['name'].decode() if isinstance(props['name'], bytes) else props['name'])
source = Path(__file__).resolve().parent
commit = subprocess.check_output(['git', '-C', str(source), 'rev-parse', 'HEAD'], text=True).strip()
assert subprocess.check_output(['git', '-C', str(source), 'status', '--porcelain'], text=True) == '', 'Candidate source worktree dirty'
save(args.output / 'authority.json', dict(device=device, queues_before=queues_before,
candidate_commit=commit, common_full_mapping_identity=common_identity,
fixture_manifests=initial_manifest_pins, repeats=20, warmups=3,
all_source_files={p.name: sha(p) for p in sorted(source.iterdir()) if p.is_file()},
scope='eligibility/delta replay only; no brain, training, or installation',
original_experiment_resumed=False, created_utc=datetime.now(timezone.utc).isoformat()))
reports = []
peak_reserved = 0
for i, fixture in enumerate(fixtures):
assert time.perf_counter() - started < 900, 'Benchmark wall budget exhausted'
idle_queues()
a, c, rule = fixture['arrays'], fixture['command'], fixture['rule']
arrays_before = {k: evidence(v) for k, v in a.items()}
f_id = fixture['group'] + '--' + fixture['record']['fixture_id']
def run_once(resident=False):
begin = time.perf_counter()
prepared = prepare_events(a['step'], a['neuron_idx'], a['kc'], c['steps'],
a['reference_command_steps'], n_neurons=c['n_neurons'])
result = backend.run(prepared, a['candidate_sources'], a['candidate_compartments'],
c['compartment'], a['delta_before'], a['base_candidate'],
d=c['d_command'], frozen=c['frozen'], **rule,
kernel_warmup=3 if resident else 0, kernel_repeats=20 if resident else 1)
comparison = compare(result, a)
elapsed = time.perf_counter() - begin
if not all(r['bitwise_exact'] for r in comparison.values()):
save(args.output / (f_id + '-mismatch.json'), comparison)
raise AssertionError('GPU bitwise mismatch: ' + f_id)
return result, comparison, elapsed
cold, comp, cold_wall = run_once()
for _ in range(3):
run_once()
samples = []
for _ in range(20):
result, comp, elapsed = run_once()
samples.append(dict(**result.timing, replay_wall_including_preparation_and_comparison_seconds=elapsed))
resident, _, resident_wall = run_once(resident=True)
peak_reserved = max(peak_reserved, cp.get_default_memory_pool().total_bytes())
assert arrays_before == {k: evidence(v) for k, v in a.items()}, 'Input arrays mutated'
report = dict(fixture_id=f_id, input_path=fixture['input_path'], fixture_manifest_sha256=fixture['manifest_sha256'],
command=c, rule=rule, first_call_timing=cold.timing, first_call_wall_seconds=cold_wall,
raw_warm_samples=samples, resident_timing=resident.timing,
summaries={key: stats([s[key] for s in samples]) for key in
('preparation_seconds', 'kernel_seconds', 'end_to_end_seconds',
'replay_wall_including_preparation_and_comparison_seconds')},
comparison=comp, input_unchanged=True, all_repeats_bitwise_exact=True,
reference_timing=fixture['record'].get('reference_timing'), audit=result.audit)
save(args.output / (f_id + '.json'), report)
reports.append(report)
print(json.dumps(dict(fixture=f_id, exact=True,
median_ms=1000 * report['summaries']['replay_wall_including_preparation_and_comparison_seconds']['median'])), flush=True)
# Recheck original exported files at the end, not just at read time.
assert initial_manifest_pins == {str(p): sha(p) for p in args.manifest}, 'Fixture manifest changed'
load_fixtures(args.manifest)
rejections = check_rejections(args.manifest, backend, fixtures[0]['rule'])
save(args.output / 'invalid-input-rejections.json', rejections)
assert commit == subprocess.check_output(['git', '-C', str(source), 'rev-parse', 'HEAD'], text=True).strip()
assert subprocess.check_output(['git', '-C', str(source), 'status', '--porcelain'], text=True) == ''
result = dict(status='PASS-ELIGIBILITY-UPDATE-REPLAY-EQUIVALENCE', fixtures=len(reports),
device=device, candidate_commit=commit, all_fixture_repeats_bitwise_exact=True,
observed_cuda_pool_reserved_peak_bytes=peak_reserved, gpu_after=gpu_snapshot(args.uuid),
queues_after=idle_queues(), wall_seconds=time.perf_counter()-started,
original_experiment_resumed=False, brain_simulation_tested=False,
rejected_invalid_cases=len(rejections),
full_training_speedup_measured=False, learning_success_claimed=False,
coverage_limit='available prefix + fixed synthetic edge cases, not all registered conditions',
files={p.name: sha(p) for p in sorted(args.output.iterdir()) if p.is_file()})
save(args.output / 'summary.json', result)
print(json.dumps(result), flush=True)
except Exception:
save(args.output / 'failure.json', dict(exit_code=1, traceback=traceback.format_exc(),
wall_seconds=time.perf_counter()-started, original_experiment_resumed=False))
raise
if __name__ == '__main__':
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--manifest', type=Path, action='append', required=True)
parser.add_argument('--output', type=Path, required=True)
parser.add_argument('--uuid', required=True)
main(parser.parse_args())