Julian lechnerandClaude Opus 5.5
14b9185e86
Sandboxed clients may write only to data/ and cannot create it themselves. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
231 lines
10 KiB
Python
231 lines
10 KiB
Python
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
|
|
|
|
FILES = ('server.py', 'benchmark.py', 'requirements.txt')
|
|
CUDA_INDEX = 'https://download.pytorch.org/whl/cu126'
|
|
MODEL_PRELOAD_TIMEOUT = 900
|
|
|
|
|
|
def digest(path):
|
|
return hashlib.sha256(path.read_bytes()).hexdigest()
|
|
|
|
|
|
def python_path(target):
|
|
return target / ('Scripts/python.exe' if os.name == 'nt' else 'bin/python')
|
|
|
|
|
|
def ensure_runtime(target, dry_run):
|
|
if dry_run:
|
|
return python_path(target / '.venv')
|
|
venv = target / '.venv'
|
|
if not venv.exists():
|
|
subprocess.run([sys.executable, '-m', 'venv', str(venv)], check=True, timeout=180)
|
|
python = python_path(venv)
|
|
if not python.is_file():
|
|
raise ValueError('Semantic retrieval virtual environment is incomplete.')
|
|
arguments = [str(python), '-m', 'pip', 'install', '--disable-pip-version-check', '-r', str(target / 'requirements.txt')]
|
|
if shutil.which('nvidia-smi'):
|
|
try:
|
|
subprocess.run([*arguments, '--extra-index-url', CUDA_INDEX], check=True, timeout=900)
|
|
except subprocess.CalledProcessError:
|
|
print('Semantic retrieval: CUDA runtime installation failed; falling back to CPU.', file=sys.stderr)
|
|
subprocess.run(arguments, check=True, timeout=900)
|
|
else:
|
|
subprocess.run(arguments, check=True, timeout=900)
|
|
return python
|
|
|
|
|
|
def ensure_models(python, target):
|
|
subprocess.run([str(python), str(target / 'server.py'), 'preload', '--model-cache', str(target / 'model-cache')],
|
|
check=True, timeout=MODEL_PRELOAD_TIMEOUT)
|
|
|
|
|
|
def load_state(path):
|
|
if not path.exists():
|
|
return {'version': 1, 'clients': [], 'files': {}}
|
|
state = json.loads(path.read_text(encoding='utf-8'))
|
|
if state.get('version') != 1 or not isinstance(state.get('clients'), list) or not isinstance(state.get('files'), dict):
|
|
raise ValueError('Invalid semantic retrieval ownership state.')
|
|
if any(client not in ('codex', 'claude') for client in state['clients']):
|
|
raise ValueError('Invalid semantic retrieval client state.')
|
|
if any(name not in FILES or not isinstance(value, str) or len(value) != 64 or any(character not in '0123456789abcdef' for character in value) for name, value in state['files'].items()):
|
|
raise ValueError('Invalid semantic retrieval file state.')
|
|
return state
|
|
|
|
|
|
def save_state(path, state, dry_run):
|
|
if not dry_run:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = None
|
|
try:
|
|
with tempfile.NamedTemporaryFile(mode='w', encoding='utf-8', dir=path.parent, delete=False) as stream:
|
|
temporary = Path(stream.name)
|
|
json.dump(state, stream, indent=2)
|
|
stream.write('\n')
|
|
stream.flush()
|
|
os.fsync(stream.fileno())
|
|
os.replace(temporary, path)
|
|
finally:
|
|
if temporary is not None:
|
|
temporary.unlink(missing_ok=True)
|
|
|
|
|
|
def check_runtime(target, state):
|
|
if target.is_symlink():
|
|
raise ValueError('Semantic retrieval directory must not be a symbolic link.')
|
|
if not target.exists():
|
|
return
|
|
if not target.is_dir():
|
|
raise ValueError('Semantic retrieval target is not a directory.')
|
|
for name, expected in state['files'].items():
|
|
path = target / name
|
|
if path.is_symlink() or not path.is_file() or digest(path) != expected:
|
|
raise ValueError(f'Owned semantic retrieval file changed: {name}')
|
|
allowed = set(state['files']) | {'.venv', 'model-cache', '__pycache__', 'data'}
|
|
for path in target.iterdir():
|
|
if path.is_symlink() or path.name not in allowed:
|
|
raise ValueError(f'Foreign or linked semantic retrieval path: {path.name}')
|
|
if path.name == 'data' and (not path.is_dir() or any(
|
|
child.is_symlink() or not (
|
|
child.is_file() and (child.name in ('daemon.key', 'daemon.starting') or re.fullmatch(
|
|
r'[0-9a-f]{64}\.sqlite3(?:-(?:wal|shm|journal))?|daemon-[0-9a-f]{12}\.starting', child.name))
|
|
or child.is_socket() and re.fullmatch(r'daemon-[0-9a-f]{16}\.sock', child.name))
|
|
for child in path.iterdir())):
|
|
raise ValueError('Foreign or linked semantic retrieval path: data')
|
|
if path.name == '__pycache__' and (not path.is_dir() or any(
|
|
child.is_symlink() or not child.is_file() or not re.fullmatch(
|
|
r'(?:server|benchmark)\.[A-Za-z0-9_-]+(?:\.opt-\d+)?\.pyc', child.name)
|
|
for child in path.iterdir())):
|
|
raise ValueError('Foreign or linked semantic retrieval path: __pycache__')
|
|
|
|
|
|
def sync(root, home, clients, enabled, dry_run, update, summary=False):
|
|
base = home / '.my-ai-configuration'
|
|
target = base / 'semantic-retrieval'
|
|
state_path = base / 'semantic-retrieval.json'
|
|
state = load_state(state_path)
|
|
clients = list(dict.fromkeys(clients))
|
|
if base.exists() and base.is_symlink():
|
|
raise ValueError('Semantic retrieval state directory must not be a symbolic link.')
|
|
if not enabled:
|
|
owned = [client for client in state['clients'] if client in clients]
|
|
if dry_run:
|
|
if not summary:
|
|
print('DRYRUN disable semantic retrieval: ' + (', '.join(owned) if owned else 'nothing owned'))
|
|
return
|
|
state['clients'] = [client for client in state['clients'] if client not in clients]
|
|
if state['clients']:
|
|
save_state(state_path, state, False)
|
|
return
|
|
if not state['files']:
|
|
state_path.unlink(missing_ok=True)
|
|
return
|
|
check_runtime(target, state)
|
|
if target.exists() and target.is_dir():
|
|
shutil.rmtree(target)
|
|
state_path.unlink(missing_ok=True)
|
|
return
|
|
|
|
check_runtime(target, state)
|
|
if target.exists() and not state['files']:
|
|
raise ValueError('Semantic retrieval directory exists without setup ownership; refusing to overwrite it.')
|
|
if dry_run:
|
|
if not summary:
|
|
print('DRYRUN enable semantic retrieval: Qwen3-Embedding-0.6B + Qwen3-Reranker-0.6B')
|
|
return
|
|
target.mkdir(parents=True, exist_ok=True)
|
|
source = root / 'shared' / 'retrieval'
|
|
for name in FILES:
|
|
shutil.copy2(source / name, target / name)
|
|
state['files'][name] = digest(target / name)
|
|
save_state(state_path, state, False)
|
|
python = ensure_runtime(target, False)
|
|
ensure_models(python, target)
|
|
# Sandboxed clients may write only to data/; it cannot be created from inside the sandbox.
|
|
(target / 'data').mkdir(exist_ok=True)
|
|
state['clients'] = list(dict.fromkeys(state['clients']))
|
|
save_state(state_path, state, False)
|
|
for client in clients:
|
|
if client not in state['clients']:
|
|
state['clients'].append(client)
|
|
save_state(state_path, state, False)
|
|
save_state(state_path, state, False)
|
|
|
|
|
|
def benchmark(root, home, dry_run, summary=False):
|
|
base = home / '.my-ai-configuration'
|
|
target = base / 'semantic-retrieval'
|
|
state_path = base / 'semantic-retrieval.json'
|
|
if dry_run:
|
|
if not summary:
|
|
print('DRYRUN benchmark semantic retrieval: Qwen3-Embedding-0.6B + Qwen3-Reranker-0.6B')
|
|
return True
|
|
state = load_state(state_path)
|
|
source = root / 'shared' / 'retrieval'
|
|
if state['files']:
|
|
check_runtime(target, state)
|
|
elif target.is_symlink():
|
|
raise ValueError('Semantic retrieval directory must not be a symbolic link.')
|
|
if target.exists() and not state['files']:
|
|
if not all((target / name).is_file() and digest(target / name) == digest(source / name) for name in FILES):
|
|
raise ValueError('Semantic retrieval directory exists without setup ownership; refusing to overwrite it.')
|
|
state['files'] = {name: digest(target / name) for name in FILES}
|
|
check_runtime(target, state)
|
|
target.mkdir(parents=True, exist_ok=True)
|
|
for name in FILES:
|
|
shutil.copy2(source / name, target / name)
|
|
state['files'][name] = digest(target / name)
|
|
python = ensure_runtime(target, False)
|
|
result = subprocess.run([str(python), str(target / 'benchmark.py'), '--model-cache', str(target / 'model-cache')], check=True,
|
|
stdout=subprocess.PIPE, text=True, timeout=300)
|
|
outcome = json.loads(result.stdout)
|
|
if not isinstance(outcome, dict) or not isinstance(outcome.get('recommended'), bool):
|
|
raise ValueError('Semantic retrieval benchmark returned an invalid result.')
|
|
if outcome['recommended'] or state['clients']:
|
|
save_state(state_path, state, False)
|
|
else:
|
|
check_runtime(target, state)
|
|
shutil.rmtree(target)
|
|
state_path.unlink(missing_ok=True)
|
|
if not summary:
|
|
print('Semantic retrieval warm-model check: ' + ('passed' if outcome['recommended'] else 'failed') + f" (embedding {outcome['embeddingSeconds']}s, reranking {outcome['rerankingSeconds']}s; excludes model startup and indexing)")
|
|
return outcome['recommended']
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('action', choices=('sync', 'benchmark'))
|
|
parser.add_argument('--root', required=True, type=Path)
|
|
parser.add_argument('--home', required=True, type=Path)
|
|
parser.add_argument('--client', choices=('codex', 'claude', 'both'))
|
|
parser.add_argument('--enabled', choices=('true', 'false'))
|
|
parser.add_argument('--dry-run', action='store_true')
|
|
parser.add_argument('--update', action='store_true')
|
|
parser.add_argument('--summary', action='store_true')
|
|
args = parser.parse_args()
|
|
if args.action == 'benchmark':
|
|
recommended = benchmark(args.root.resolve(), args.home.resolve(), args.dry_run, args.summary)
|
|
raise SystemExit(0 if recommended else 2)
|
|
if args.client is None or args.enabled is None:
|
|
parser.error('sync requires --client and --enabled')
|
|
clients = ('codex', 'claude') if args.client == 'both' else (args.client,)
|
|
sync(args.root.resolve(), args.home.resolve(), clients, args.enabled == 'true', args.dry_run, args.update, args.summary)
|
|
print('SEMANTIC RETRIEVAL: reconciliation complete' + (' (dry run)' if args.dry_run else ''))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
try:
|
|
main()
|
|
except (OSError, ValueError, json.JSONDecodeError, subprocess.SubprocessError) as error:
|
|
print(f'Semantic retrieval: {error}', file=sys.stderr)
|
|
raise SystemExit(1)
|