Files
digital-drosophila/models/male-cns-lif/conditioning-interface-v2/isolated_workers.py
T
2026-09-30 18:19:04 +08:00

382 lines
22 KiB
Python

"""Isolated AP workers; invoke each job in a NEW ``python -B`` process.
Top-level imports are stdlib only. Brain loads the frozen ordinary runtime,
then its read-only checkpoint sealer; update loads ONLY the frozen epoch
adapter. Never import the old coordinator, share/patch loader caches, or use a
forked interpreter. The external job SHA binds the caller's authority; this
worker does not independently prove preregistration or anatomical provenance.
Success receipts are outside all inputs/checkpoints. A separate exclusive
``<receipt>.attempt.json`` prevents retries even after a pre-output failure.
The parent must additionally bind the real child exit code and stdout/stderr;
the receipt cannot certify a process exit that has not yet happened.
"""
from __future__ import annotations
import argparse
from datetime import datetime, timezone
import hashlib
import importlib.util
import json
import os
from pathlib import Path
import re
import subprocess
import sys
import time
import urllib.request
SCHEMA = 'phase4ap-isolated-worker-job-v1'
GPU = 'GPU-79f3a85d-adce-7833-7003-6488140ca596'
SCIENCE_BIN = '/data/digital-drosophila/models/male-cns-lif/.venv/bin'
COMFY_URLS = tuple(f'http://192.168.10.51:{port}/queue' for port in (8188, 8189, 8190))
SOURCE_PINS = {
'runtime_light': '522b8c31ea938890fe15364fccd25ccc5bded54101754909f255532e7ad01aae',
'runtime_continuation': '0b8e8e0a6a23a079dc62a88355ce355d53b3ae22f6665519a8f6aef2dcfc5aba',
'conditioning_epoch': '757935a4815d9f625640df685184ed3fe44a8845323562796da7aa43cbb194e1',
}
GRAPH_FILES = ('nodes.npy', 'edge_id.npy', 'src.npy', 'dst.npy', 'signed_raw.npy', 'ptr.npy', 'weights.npy')
APL_FILES = ('input_matrix.npy', 'output_matrix.npy', 'kc_globals.npy', 'apl_globals.npy')
COMMON = {'schema', 'job_id', 'kind', 'registration_sha256', 'worker_sha256',
'output_dir', 'receipt_file', 'sources', 'input_sha256'}
BRAIN_FIELDS = {'case_dir', 'build_root', 'overlay_file', 'overlay_manifest', 'diagnostic_mode', 'admission'}
UPDATE_FIELDS = {'authority_file', 'authority_sha256'}
def require(ok, message):
if not ok:
raise ValueError(message)
def digest(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 valid_sha(value):
return type(value) is str and re.fullmatch('[0-9a-f]{64}', value) is not None
def read_json(path):
def pairs(items):
result = {}
for key, value in items:
require(key not in result, 'Duplicate JSON key: ' + key)
result[key] = value
return result
return json.loads(Path(path).read_text(encoding='utf-8'), object_pairs_hook=pairs,
parse_constant=lambda x: (_ for _ in ()).throw(ValueError('Nonfinite JSON: ' + x)))
def save_json(path, value):
with Path(path).open('x', encoding='utf-8', newline='\n') as stream:
json.dump(value, stream, sort_keys=True, indent=2, allow_nan=False)
stream.write('\n')
def absolute_path(value, *, kind=None):
require(type(value) is str and Path(value).is_absolute(), 'Explicit absolute path required')
path = Path(value)
for part in (path, *path.parents):
require(not part.is_symlink() and not (hasattr(part, 'is_junction') and part.is_junction()),
'Linked/junction path is forbidden: ' + str(part))
path = path.resolve(strict=kind is not None)
if kind == 'file':
require(path.is_file(), 'Regular input file required: ' + str(path))
elif kind == 'dir':
require(path.is_dir(), 'Existing directory required: ' + str(path))
return path
def verify_pins(pins):
for path, expected in pins.items():
require(absolute_path(str(path), kind='file') == path and valid_sha(expected) and digest(path) == expected,
'Pinned file changed: ' + str(path))
def _put(pins, path, sha):
require(valid_sha(sha) and (path not in pins or pins[path] == sha), 'Invalid/conflicting file SHA: ' + str(path))
pins[path] = sha
def _required(pins, path, expected=None):
path = absolute_path(str(path), kind='file')
require(path in pins and (expected is None or pins[path] == expected), 'Required external input pin missing: ' + str(path))
return path
def validate_job(job_file, expected_job_sha256, mode):
"""Read/verify all bytes and paths before importing any scientific module."""
require(mode in ('brain', 'update'), 'Unknown worker mode')
job_file = absolute_path(str(job_file), kind='file')
require(valid_sha(expected_job_sha256) and digest(job_file) == expected_job_sha256, 'External job SHA mismatch')
job = read_json(job_file)
require(type(job) is dict and set(job) == COMMON | (BRAIN_FIELDS if mode == 'brain' else UPDATE_FIELDS),
'Unknown or missing job fields')
require(job['schema'] == SCHEMA and job['kind'] == mode, 'Job mode/schema mismatch')
require(type(job['job_id']) is str and re.fullmatch('[A-Za-z0-9][A-Za-z0-9_.-]*', job['job_id']) is not None and
'..' not in job['job_id'], 'Unsafe job identity')
require(valid_sha(job['registration_sha256']) and valid_sha(job['worker_sha256']) and
digest(__file__) == job['worker_sha256'], 'Worker/registration pin invalid')
require(type(job['input_sha256']) is dict and bool(job['input_sha256']), 'Nonempty external input pins required')
pins = {job_file: expected_job_sha256, Path(__file__).resolve(): job['worker_sha256']}
for name, sha in job['input_sha256'].items():
_put(pins, absolute_path(name, kind='file'), sha)
roles = {'runtime_light', 'runtime_continuation'} if mode == 'brain' else {'conditioning_epoch'}
require(type(job['sources']) is dict and set(job['sources']) == roles, 'Wrong worker source roles')
sources = {}
for role, entry in job['sources'].items():
require(type(entry) is dict and set(entry) == {'path', 'sha256'} and entry['sha256'] == SOURCE_PINS[role],
'Frozen worker dependency pin differs: ' + role)
sources[role] = absolute_path(entry['path'], kind='file')
_put(pins, sources[role], entry['sha256'])
output = absolute_path(job['output_dir'])
receipt = absolute_path(job['receipt_file'])
attempt = Path(str(receipt) + '.attempt.json')
require(output.parent.is_dir() and receipt.parent.is_dir(), 'Output/receipt parents must be created by coordinator')
require(not output.exists() and not receipt.exists() and not attempt.exists(), 'Existing output/receipt/attempt; retry forbidden')
paths = dict(job=job_file, output=output, receipt=receipt, attempt=attempt)
if mode == 'brain':
require(job['diagnostic_mode'] in ('light', 'reference'), 'Invalid recording mode')
require(job['admission'] == {'gpu_uuid': GPU, 'comfy_queue_urls': list(COMFY_URLS)}, 'GPU/Comfy admission differs')
case = absolute_path(job['case_dir'], kind='dir')
spec_path = _required(pins, case / 'spec.json')
_required(pins, case / 'requests.npz')
spec = read_json(spec_path)
overlay = _required(pins, job['overlay_file'])
manifest = _required(pins, job['overlay_manifest'])
meta = read_json(manifest)
for role in ('mask', 'base_candidate', 'am_overlay'):
_required(pins, meta[role + '_path'], meta[role + '_sha256'])
graph, apl = (absolute_path(spec[key], kind='dir') for key in ('graph_dir', 'apl_bundle_dir'))
for name in GRAPH_FILES:
_required(pins, graph / name)
for name in APL_FILES:
_required(pins, apl / name)
_required(pins, spec['coefficients_path'])
require(sources['runtime_continuation'].parent / 'runtime_light.py' == sources['runtime_light'],
'Checkpoint reader must use declared adjacent runtime')
build = absolute_path(job['build_root'])
require(not build.exists() or build.is_dir(), 'Build cache is not a directory')
require(build.parent.is_dir(), 'Build parent must already exist')
paths.update(case=case, overlay=overlay, manifest=manifest, build=build)
else:
authority = _required(pins, job['authority_file'], job['authority_sha256'])
value = read_json(authority)
require(type(value.get('files')) is dict and type(value.get('frozen_inputs')) is dict, 'Authority input inventory missing')
for entry in value['files'].values():
require(type(entry) is dict and set(entry) == {'path', 'sha256'}, 'Invalid authority input row')
_required(pins, entry['path'], entry['sha256'])
for name, sha in value['frozen_inputs'].items():
_required(pins, name, sha)
checkpoint = absolute_path(value['checkpoint_dir'], kind='dir')
_required(pins, checkpoint / 'checkpoint-manifest.json', value['checkpoint_manifest_sha256'])
paths.update(authority=authority, checkpoint=checkpoint)
protected = {path.parent for path in pins}
if mode == 'update':
protected.add(paths['checkpoint'])
for target in (output, receipt, attempt, *([paths['build']] if mode == 'brain' else [])):
require(not any(target.is_relative_to(root) for root in protected), 'Output/receipt/cache inside protected input/source directory')
require(not receipt.is_relative_to(output) and not output.is_relative_to(receipt), 'Receipt must remain outside sealed output')
if mode == 'brain':
build = paths['build']
require(not output.is_relative_to(build) and not build.is_relative_to(output) and not receipt.is_relative_to(build),
'Build cache/output/receipt scopes overlap')
verify_pins(pins)
return job, paths, pins, sources
def _load_source(name, path, expected):
"""Execute exactly verified bytes; never reuse/repair an existing binding."""
raw = path.read_bytes()
require(hashlib.sha256(raw).hexdigest() == expected and name not in sys.modules, 'Unverified source/module binding: ' + name)
spec = importlib.util.spec_from_file_location(name, path)
require(spec is not None and spec.loader is not None, 'No Python source loader')
module = importlib.util.module_from_spec(spec)
sys.modules[name] = module # Normal verified import registration, not cache repair.
exec(compile(raw, str(path), 'exec'), module.__dict__)
return module
def _clean_process(mode):
require(sys.dont_write_bytecode, 'Use Python -B')
forbidden = ('torch', 'runtime_light', 'runtime_continuation', 'conditioning_epoch',
'conditioning_preflight_workflow', 'conditioning_preflight_prepare', 'engineering_workflow',
'phase4an_continuation_pinned_light', 'phase4an_sealed_am_overlay',
'_conditioning_epoch_continuation', '_phase4ap_epoch', '_phase4ap_checkpoint')
require(not any(name in sys.modules for name in forbidden), 'Worker requires a fresh exec interpreter, not a shared/forked runtime')
os.environ['CUDA_VISIBLE_DEVICES'] = '' if mode == 'update' else GPU
if mode == 'brain':
os.environ['CUDA_HOME'] = '/usr/local/cuda'
os.environ['MAX_JOBS'] = '2'
os.environ['TORCH_CUDA_ARCH_LIST'] = '8.9'
os.environ['PATH'] = SCIENCE_BIN + os.pathsep + '/usr/local/cuda/bin' + os.pathsep + os.environ.get('PATH', '')
def _gpu_inventory():
result = subprocess.run(['nvidia-smi', '-i', GPU,
'--query-gpu=uuid,name,memory.total,driver_version', '--format=csv,noheader'],
check=True, capture_output=True, text=True, timeout=20).stdout.strip()
require(len(result.splitlines()) == 1 and result.startswith(GPU + ',') and '4090' in result, 'Wrong physical GPU inventory')
return result
def _admit_brain():
rows = []
for url in COMFY_URLS:
with urllib.request.urlopen(url, timeout=5) as response:
require(response.geturl() == url, 'Comfy admission redirect forbidden')
raw = response.read(2 * 1024 * 1024 + 1)
require(len(raw) <= 2 * 1024 * 1024, 'Oversized Comfy admission response')
queue = json.loads(raw)
require(type(queue) is dict and queue.get('queue_running') == [] and queue.get('queue_pending') == [],
'ComfyUI busy or malformed; never evict a service')
rows.append(dict(url=url, queue_running=0, queue_pending=0, response_sha256=hashlib.sha256(raw).hexdigest()))
return dict(gpu_uuid=GPU, nvidia_smi_inventory=_gpu_inventory(), comfy_queues=rows,
observed_utc=datetime.now(timezone.utc).isoformat(), check_is_not_a_service_lock=True)
def _environment(mode):
result = dict(python_executable=sys.executable, python_version=sys.version,
python_no_bytecode=sys.dont_write_bytecode, PATH=os.environ.get('PATH'),
environment_scope='selected execution variables and library/precision metadata; not a dump of credentials',
CUDA_VISIBLE_DEVICES=os.environ.get('CUDA_VISIBLE_DEVICES'),
CUBLAS_WORKSPACE_CONFIG=os.environ.get('CUBLAS_WORKSPACE_CONFIG'),
TORCH_CUDA_ARCH_LIST=os.environ.get('TORCH_CUDA_ARCH_LIST'),
CUDA_HOME=os.environ.get('CUDA_HOME'), MAX_JOBS=os.environ.get('MAX_JOBS'),
settings_changed_by_environment_recorder=False)
torch, numpy = sys.modules.get('torch'), sys.modules.get('numpy')
result.update(numpy_version=None if numpy is None else numpy.__version__,
torch_version=None if torch is None else str(torch.__version__))
if mode == 'update':
require(os.environ.get('CUDA_VISIBLE_DEVICES') == '' and (torch is None or not torch.cuda.is_initialized()),
'Update worker initialized or exposed CUDA')
result.update(gpu_uuid=None, CUDA_hidden=True, CUDA_initialized=False,
device_scope='CPU eligibility only; no brain execution')
return result
require(torch is not None and os.environ.get('CUDA_VISIBLE_DEVICES') == GPU and torch.cuda.device_count() == 1,
'Brain runtime is not using the single admitted GPU')
properties = torch.cuda.get_device_properties(0)
require('4090' in properties.name, 'Unexpected GPU device class')
result.update(torch_cuda_version=torch.version.cuda, device_name=properties.name,
total_memory_bytes=int(properties.total_memory), gpu_uuid=GPU, nvidia_smi_inventory=_gpu_inventory(),
matmul_allow_tf32=bool(torch.backends.cuda.matmul.allow_tf32), cudnn_allow_tf32=bool(torch.backends.cudnn.allow_tf32),
deterministic_algorithms=bool(torch.are_deterministic_algorithms_enabled()),
cudnn_deterministic=bool(torch.backends.cudnn.deterministic), cudnn_benchmark=bool(torch.backends.cudnn.benchmark),
float32_matmul_precision=torch.get_float32_matmul_precision(), environment_recorder_adds_no_timing=True)
return result
def _verify_precision(environment, component):
require(component.get('cuda_matmul_allow_tf32') is False and component.get('float32_matmul_precision') == 'highest' and
component.get('precision_configuration_changed_by_component') is False, 'Frozen precision configuration differs')
pairs = {'matmul_allow_tf32': 'cuda_matmul_allow_tf32', 'cudnn_allow_tf32': 'cudnn_allow_tf32',
'deterministic_algorithms': 'deterministic_algorithms', 'cudnn_deterministic': 'cudnn_deterministic',
'cudnn_benchmark': 'cudnn_benchmark', 'float32_matmul_precision': 'float32_matmul_precision',
'CUBLAS_WORKSPACE_CONFIG': 'CUBLAS_WORKSPACE_CONFIG', 'device_name': 'device_name',
'torch_version': 'torch', 'torch_cuda_version': 'cuda', 'numpy_version': 'numpy'}
require(all(left in environment and right in component and
type(environment[left]) is type(component[right]) and environment[left] == component[right]
for left, right in pairs.items()), 'Actual environment and APL component disagree')
def _run_brain(job, paths, pins, sources):
r = _load_source('runtime_light', sources['runtime_light'], SOURCE_PINS['runtime_light'])
meta = read_json(paths['manifest'])
for relative, sha in r.PINS.items():
_required(pins, Path(meta['aj_model_dir']).parent / relative, sha)
r.execution_paths(paths['case'], paths['output'], paths['build'], protected={p.parent for p in pins})
admission = _admit_brain()
execute = r.run_reference if job['diagnostic_mode'] == 'reference' else r.run_light
summary = execute(paths['case'], paths['output'], paths['build'], paths['overlay'], paths['manifest'])
# This reader uses a separate pinned R but does NOT load the shared helper.
# E must never be loaded in this process; no caches are patched or cleared.
c = _load_source('_phase4ap_checkpoint', sources['runtime_continuation'], SOURCE_PINS['runtime_continuation'])
seal = c.seal_checkpoint(paths['output'])
record = c.load_checkpoint(paths['output'], seal['sha256'])
require(record['fresh'] is True and record['global_tick_offset'] == 0 and
record['summary'] == summary and summary['diagnostic_mode'] == job['diagnostic_mode'], 'Sealed runtime scope differs')
environment = _environment('brain')
_verify_precision(environment, summary['apl_component']['environment'])
return dict(checkpoint_manifest_sha256=seal['sha256'], checkpoint_manifest_file=str(paths['output'] / 'checkpoint-manifest.json'),
complete_fresh_trial=True, steps=record['steps'], admission=admission, environment=environment,
within_trial_weight_updates=False, epoch_update_performed=False)
def _run_update(job, paths, pins, sources):
e = _load_source('_phase4ap_epoch', sources['conditioning_epoch'], SOURCE_PINS['conditioning_epoch'])
chain = e.postprocess_epoch(paths['authority'], job['authority_sha256'], paths['output'])
require(chain.get('complete') is True and chain.get('authority_sha256') == job['authority_sha256'] and
chain.get('weights_installed') is False and chain.get('brain_executed') is False and
chain.get('learning_pass_claimed') is False, 'Incomplete/incorrect epoch postprocessing scope')
require(chain == read_json(paths['output'] / 'chain-receipt.json'), 'Stored chain differs from actual result')
require('runtime_light' not in sys.modules and 'runtime_continuation' not in sys.modules and
'conditioning_preflight_workflow' not in sys.modules, 'Update imported the mixed-runtime coordinator')
return dict(chain_receipt_sha256=digest(paths['output'] / 'chain-receipt.json'),
source_neural_checkpoint_manifest_sha256=chain['source_neural_checkpoint_manifest_sha256'],
before_overlay_file_sha256=chain['before_overlay_file_sha256'], after_overlay_file_sha256=chain['after_overlay_file_sha256'],
before_overlay_manifest_sha256=chain['before_overlay_manifest_sha256'], after_overlay_manifest_sha256=chain['after_overlay_manifest_sha256'],
environment=_environment('update'), brain_executed=False, weights_installed=False,
registration_authority_independently_proven=False)
def tree_hashes(directory):
require(directory.is_dir(), 'Worker produced no output directory')
result = {}
for path in sorted(directory.rglob('*')):
require(not path.is_symlink() and path.resolve().is_relative_to(directory), 'Unsafe worker output tree')
if path.is_dir():
continue
require(path.is_file(), 'Nonregular worker output')
result[path.relative_to(directory).as_posix()] = digest(path)
require(bool(result), 'Empty worker output')
return result
def execute_job(mode, job_file, expected_job_sha256):
"""CLI implementation. Production never substitutes dependencies or arrays."""
_clean_process(mode)
job, paths, pins, sources = validate_job(job_file, expected_job_sha256, mode)
started = time.perf_counter()
save_json(paths['attempt'], dict(schema='phase4ap-worker-attempt-v1', complete=False,
job_sha256=expected_job_sha256, job_id=job['job_id'], kind=mode, pid=os.getpid(),
started_utc=datetime.now(timezone.utc).isoformat(), no_automatic_retry=True))
result = (_run_brain if mode == 'brain' else _run_update)(job, paths, pins, sources)
outputs = tree_hashes(paths['output'])
verify_pins(pins)
require(outputs == tree_hashes(paths['output']), 'Worker output changed during receipt creation')
receipt = dict(schema='phase4ap-isolated-worker-receipt-v1', complete=True,
kind=mode, job_id=job['job_id'], job_file=str(paths['job']), job_sha256=expected_job_sha256,
registration_sha256=job['registration_sha256'], worker_sha256=job['worker_sha256'],
pid=os.getpid(), argv=[sys.executable, *sys.argv], wall_seconds=time.perf_counter() - started,
attempt_sha256=digest(paths['attempt']), source_files=job['sources'],
verified_input_sha256={str(path): sha for path, sha in sorted(pins.items(), key=lambda item: str(item[0]))},
output_dir=str(paths['output']), receipt_file=str(paths['receipt']), output_files_sha256=outputs, result=result,
environment=result['environment'], completion_semantics='worker returned successfully before writing this receipt; parent must independently require actual process exit 0',
process_exit_independently_observed=False, registration_authority_independently_proven=False,
learning_success_claimed=False, scientific_gate_pass_claimed=False,
old_sources_modified=False, loader_caches_patched=False)
save_json(paths['receipt'], receipt)
return receipt
def main(argv=None):
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('mode', choices=('brain', 'update'))
parser.add_argument('--job', required=True)
parser.add_argument('--job-sha256', required=True)
args = parser.parse_args(argv)
receipt = execute_job(args.mode, args.job, args.job_sha256)
print(json.dumps(dict(complete=True, kind=args.mode, job_id=receipt['job_id'],
receipt_file=receipt['receipt_file'],
receipt_sha256=digest(receipt['receipt_file'])), sort_keys=True), flush=True)
return 0
if __name__ == '__main__':
raise SystemExit(main())