343 lines
16 KiB
Python
343 lines
16 KiB
Python
"""Install, benchmark, migrate, and remove the setup-owned QMD local search."""
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
|
|
PACKAGE = '@tobilu/qmd'
|
|
VERSION = '2.8.3'
|
|
SCRIPTS = ('qmd-lib.mjs', 'qmd-config.mjs', 'qmd-daemon.mjs', 'qmd-search.mjs', 'qmd-warm.mjs', 'qmd-benchmark.mjs')
|
|
DEVICE_TIMEOUT = 600
|
|
WINDOWS = os.name == 'nt'
|
|
|
|
|
|
def run(arguments, env=None, timeout=1800, check=True):
|
|
result = subprocess.run(arguments, env=env, capture_output=True, text=True, encoding='utf-8', errors='replace',
|
|
timeout=timeout, shell=WINDOWS and arguments[0] in ('npm', 'qmd'))
|
|
if check and result.returncode != 0:
|
|
raise ValueError(f'{" ".join(arguments[:3])} failed: {(result.stderr or result.stdout).strip()[-300:]}')
|
|
return result.stdout
|
|
|
|
|
|
def process_alive(pid):
|
|
if WINDOWS: # os.kill(pid, 0) would send CTRL_C_EVENT on Windows
|
|
import ctypes
|
|
kernel32 = ctypes.windll.kernel32
|
|
handle = kernel32.OpenProcess(0x1000, False, pid) # PROCESS_QUERY_LIMITED_INFORMATION
|
|
if not handle:
|
|
return False
|
|
code = ctypes.c_ulong()
|
|
try:
|
|
return bool(kernel32.GetExitCodeProcess(handle, ctypes.byref(code))) and code.value == 259 # STILL_ACTIVE
|
|
finally:
|
|
kernel32.CloseHandle(handle)
|
|
try:
|
|
os.kill(pid, 0)
|
|
except ProcessLookupError:
|
|
return False
|
|
except PermissionError:
|
|
return True
|
|
return True
|
|
|
|
|
|
def decide(results):
|
|
for device in ('gpu', 'cpu'):
|
|
if any(r.get('device') == device and r.get('passed') is True for r in results):
|
|
return device
|
|
return None
|
|
|
|
|
|
def node_version_ok(text):
|
|
match = re.fullmatch(r'v(\d+)\.\d+\.\d+', text.strip())
|
|
return bool(match) and int(match.group(1)) >= 22
|
|
|
|
|
|
def digest(path):
|
|
return hashlib.sha256(path.read_bytes()).hexdigest()
|
|
|
|
|
|
def atomic_json(path, value):
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with tempfile.NamedTemporaryFile('w', encoding='utf-8', dir=path.parent, delete=False) as stream:
|
|
json.dump(value, stream, indent=2)
|
|
stream.write('\n')
|
|
os.replace(stream.name, path)
|
|
|
|
|
|
class Setup:
|
|
def __init__(self, home, root, runner=run, dry_run=False, summary=False):
|
|
self.home, self.root, self.run, self.dry_run, self.summary = Path(home), Path(root), runner, dry_run, summary
|
|
self.base = self.home / '.my-ai-configuration'
|
|
self.state_path = self.base / 'qmd.json'
|
|
self.scripts = self.base / 'qmd'
|
|
self.models = Path(os.environ.get('XDG_CACHE_HOME') or self.home / '.cache') / 'qmd' / 'models'
|
|
|
|
def say(self, text):
|
|
if not self.summary:
|
|
print(text)
|
|
|
|
def load(self):
|
|
if not self.state_path.exists():
|
|
return {'version': 1, 'clients': [], 'device': None, 'owned_package': False, 'owned_models': [], 'models_block': False, 'previous_models': None, 'files': {}}
|
|
state = json.loads(self.state_path.read_text(encoding='utf-8'))
|
|
if state.get('version') != 1 or state.get('device') not in (None, 'gpu', 'cpu'):
|
|
raise ValueError('Invalid QMD setup state.')
|
|
return state
|
|
|
|
def node(self, script, *arguments, env=None, timeout=1800):
|
|
return self.run(['node', str(self.scripts / script), *arguments], env=env, timeout=timeout)
|
|
|
|
def package_installed(self):
|
|
# npm list exits non-zero on tree problems (ELSPROBLEMS) but still prints valid JSON.
|
|
output = self.run(['npm', 'list', '--global', '--depth=0', '--json'], check=False)
|
|
try:
|
|
listing = json.loads(output)
|
|
except json.JSONDecodeError:
|
|
raise ValueError('npm list did not return JSON.')
|
|
if not isinstance(listing, dict) or not isinstance(listing.get('dependencies'), dict):
|
|
raise ValueError('npm list did not report the global packages.')
|
|
return PACKAGE in listing['dependencies']
|
|
|
|
def require_node(self):
|
|
if not node_version_ok(self.run(['node', '--version'])):
|
|
raise ValueError('QMD local search needs Node.js 22 or newer.')
|
|
|
|
def stop_legacy_daemon(self, legacy):
|
|
"""Ask the old Python daemon to exit so its files can be deleted (Windows keeps open files locked)."""
|
|
python = legacy / '.venv' / ('Scripts/python.exe' if WINDOWS else 'bin/python')
|
|
server = legacy / 'server.py'
|
|
if python.is_file() and server.is_file():
|
|
try:
|
|
self.run([str(python), str(server), 'stop', '--data-dir', str(legacy / 'data'),
|
|
'--model-cache', str(legacy / 'model-cache')], timeout=60, check=False)
|
|
except (ValueError, OSError, subprocess.SubprocessError):
|
|
pass
|
|
|
|
def migrate(self, state=None):
|
|
state = self.load() if state is None else state
|
|
legacy_state = self.base / 'semantic-retrieval.json'
|
|
legacy = self.base / 'semantic-retrieval'
|
|
if legacy_state.exists() or legacy.exists():
|
|
if legacy.is_symlink():
|
|
raise ValueError('Legacy semantic retrieval directory must not be a symbolic link.')
|
|
if legacy.exists():
|
|
self.stop_legacy_daemon(legacy)
|
|
shutil.rmtree(legacy)
|
|
legacy_state.unlink(missing_ok=True)
|
|
self.say('Removed the previous Python semantic retrieval installation.')
|
|
ledger_path = self.base / 'extensions.json'
|
|
if ledger_path.exists():
|
|
ledger = json.loads(ledger_path.read_text(encoding='utf-8'))
|
|
owned = [r for r in ledger.get('resources', []) if r.get('kind') == 'qmd']
|
|
if owned:
|
|
state['owned_package'] = True
|
|
atomic_json(self.state_path, state) # record ownership before the ledger forgets it
|
|
ledger['resources'] = [r for r in ledger['resources'] if r.get('kind') != 'qmd']
|
|
atomic_json(ledger_path, ledger)
|
|
|
|
def install_scripts(self, state):
|
|
if self.scripts.is_symlink():
|
|
raise ValueError('QMD script directory must not be a symbolic link.')
|
|
self.scripts.mkdir(parents=True, exist_ok=True)
|
|
for name in SCRIPTS:
|
|
target = self.scripts / name
|
|
if target.is_symlink():
|
|
raise ValueError(f'QMD script must not be a symbolic link: {name}')
|
|
expected = state['files'].get(name)
|
|
if target.exists() and expected and digest(target) != expected:
|
|
raise ValueError(f'Setup-owned QMD script changed: {name}')
|
|
shutil.copy2(self.root / 'shared' / 'qmd' / name, target)
|
|
state['files'][name] = digest(target)
|
|
|
|
def install(self, state):
|
|
"""Steps shared by Yes and Auto: package, scripts, models block, model download."""
|
|
self.require_node()
|
|
self.migrate(state)
|
|
if not self.package_installed():
|
|
self.run(['npm', 'install', '--global', f'{PACKAGE}@{VERSION}'])
|
|
state['owned_package'] = True
|
|
self.install_scripts(state)
|
|
atomic_json(self.state_path, state)
|
|
if not state['models_block']:
|
|
output = self.node('qmd-config.mjs', 'set-models')
|
|
try:
|
|
previous = json.loads(output.strip().splitlines()[-1])['previous']
|
|
except (IndexError, KeyError, TypeError, json.JSONDecodeError):
|
|
raise ValueError('qmd-config set-models returned no previous models.')
|
|
state['previous_models'] = previous
|
|
state['models_block'] = True
|
|
atomic_json(self.state_path, state)
|
|
before = {p.name for p in self.models.iterdir()} if self.models.is_dir() else set()
|
|
self.run(['qmd', 'pull'], timeout=3600)
|
|
after = {p.name for p in self.models.iterdir()} if self.models.is_dir() else set()
|
|
state['owned_models'] = sorted(set(state['owned_models']) | (after - before))
|
|
atomic_json(self.state_path, state)
|
|
|
|
def detect_gpu(self):
|
|
try:
|
|
found = json.loads(self.node('qmd-benchmark.mjs', '--detect', timeout=DEVICE_TIMEOUT) or '{}')
|
|
gpu = found.get('gpu') or False
|
|
except Exception as error: # any probe failure means: use the CPU
|
|
self.say(f'QMD GPU detection failed, using the CPU: {str(error)[:200]}')
|
|
return False
|
|
names = ', '.join(found.get('devices') or []) or 'unnamed device'
|
|
self.say(f"QMD devices: {f'{str(gpu).upper()}: {names}' if gpu else 'no GPU'}; CPU threads: {found.get('threads', '?')}")
|
|
return gpu
|
|
|
|
def enable(self, clients):
|
|
state = self.load()
|
|
if self.dry_run:
|
|
self.say('DRYRUN enable QMD local search: Qwen3-Embedding-0.6B, Qwen3-Reranker-0.6B, QMD query expansion 1.7B')
|
|
return
|
|
self.install(state)
|
|
if state['device'] is None:
|
|
state['device'] = 'gpu' if self.detect_gpu() else 'cpu'
|
|
state['clients'] = sorted(set(state['clients']) | set(clients))
|
|
atomic_json(self.state_path, state)
|
|
self.say(f'QMD local search enabled on {state["device"].upper()}.')
|
|
|
|
def remove_all(self, state):
|
|
daemon = self.scripts / 'daemon.json'
|
|
try:
|
|
pid = json.loads(daemon.read_text(encoding='utf-8')).get('pid') if daemon.is_file() else None
|
|
except (OSError, ValueError, AttributeError):
|
|
pid = None
|
|
try:
|
|
if self.scripts.joinpath('qmd-search.mjs').exists():
|
|
self.node('qmd-search.mjs', '--stop', timeout=60)
|
|
except (ValueError, OSError, subprocess.SubprocessError):
|
|
pass
|
|
if isinstance(pid, int) and pid > 0:
|
|
deadline = time.monotonic() + 10 # the daemon must release the model files and the package
|
|
while process_alive(pid) and time.monotonic() < deadline:
|
|
time.sleep(0.2)
|
|
collections = self.scripts / 'collections.json'
|
|
names = json.loads(collections.read_text(encoding='utf-8')) if collections.exists() else []
|
|
for name in names:
|
|
if re.fullmatch(r'repo-[0-9a-f]{12}', name):
|
|
try:
|
|
self.run(['qmd', 'collection', 'remove', name])
|
|
except (ValueError, subprocess.SubprocessError):
|
|
pass
|
|
if state['models_block'] and self.scripts.joinpath('qmd-config.mjs').exists():
|
|
previous = state.get('previous_models')
|
|
try:
|
|
if previous is not None:
|
|
self.node('qmd-config.mjs', 'restore-models', json.dumps(previous))
|
|
else:
|
|
self.node('qmd-config.mjs', 'unset-models')
|
|
except (ValueError, OSError, subprocess.SubprocessError) as error:
|
|
print(f'QMD local search: could not restore the models block in index.yml: {error}', file=sys.stderr)
|
|
kept = []
|
|
for name in state['owned_models']:
|
|
path = self.models / name
|
|
if name == Path(name).name and path.is_file() and not path.is_symlink():
|
|
try:
|
|
path.unlink()
|
|
except OSError as error:
|
|
print(f'QMD local search: could not delete model {name}; the next removal retries it: {error}', file=sys.stderr)
|
|
kept.append(name)
|
|
if state['owned_package']:
|
|
self.run(['npm', 'uninstall', '--global', PACKAGE])
|
|
if self.scripts.exists() and not self.scripts.is_symlink():
|
|
shutil.rmtree(self.scripts)
|
|
if kept:
|
|
state.update(clients=[], owned_models=kept, owned_package=False, models_block=False, previous_models=None, files={})
|
|
atomic_json(self.state_path, state)
|
|
else:
|
|
self.state_path.unlink(missing_ok=True)
|
|
|
|
def disable(self, clients):
|
|
state = self.load()
|
|
if not self.state_path.exists():
|
|
if not self.dry_run:
|
|
self.migrate(state)
|
|
if state['owned_package']:
|
|
self.run(['npm', 'uninstall', '--global', PACKAGE])
|
|
self.state_path.unlink(missing_ok=True)
|
|
return
|
|
if self.dry_run:
|
|
self.say('DRYRUN disable QMD local search')
|
|
return
|
|
state['clients'] = [c for c in state['clients'] if c not in clients]
|
|
if state['clients']:
|
|
atomic_json(self.state_path, state)
|
|
return
|
|
self.remove_all(state)
|
|
self.say('QMD local search removed.')
|
|
|
|
def benchmark(self):
|
|
state = self.load()
|
|
if self.dry_run:
|
|
self.say('DRYRUN benchmark QMD local search on GPU, then CPU')
|
|
return True
|
|
try:
|
|
self.install(state)
|
|
except (ValueError, OSError, subprocess.SubprocessError, json.JSONDecodeError) as error:
|
|
print(f'QMD local search is not suitable for this computer: {error}')
|
|
try:
|
|
self.remove_all(state)
|
|
except (ValueError, OSError, subprocess.SubprocessError, json.JSONDecodeError) as cleanup:
|
|
print(f'QMD cleanup was incomplete: {cleanup}')
|
|
return False
|
|
results = []
|
|
devices = (['gpu'] if self.detect_gpu() else []) + ['cpu']
|
|
for device in devices:
|
|
env = dict(os.environ)
|
|
env.pop('QMD_FORCE_CPU' if device == 'gpu' else 'QMD_LLAMA_GPU', None)
|
|
env['QMD_LLAMA_GPU' if device == 'gpu' else 'QMD_FORCE_CPU'] = 'auto' if device == 'gpu' else '1'
|
|
try:
|
|
result = json.loads(self.node('qmd-benchmark.mjs', '--device', device, env=env, timeout=DEVICE_TIMEOUT))
|
|
if not isinstance(result, dict):
|
|
raise ValueError('benchmark returned a non-object result')
|
|
except (ValueError, subprocess.SubprocessError) as error:
|
|
result = {'device': device, 'passed': False, 'error': str(error)[:200]}
|
|
results.append(result)
|
|
print(f'QMD benchmark {device.upper()}: ' + ('passed' if result.get('passed') else 'failed')
|
|
+ f" (search {result.get('querySeconds', '-')} s, embedding {result.get('chunksPerSecond', '-')} chunks/s)")
|
|
if result.get('passed'):
|
|
break
|
|
device = decide(results)
|
|
if device is None:
|
|
self.remove_all(state)
|
|
print('QMD local search is not suitable for this computer; removed what the setup installed.')
|
|
return False
|
|
state['device'] = device
|
|
atomic_json(self.state_path, state)
|
|
return True
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
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()
|
|
setup = Setup(args.home.resolve(), args.root.resolve(), run, args.dry_run, args.summary)
|
|
if args.action == 'benchmark':
|
|
raise SystemExit(0 if setup.benchmark() 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]
|
|
(setup.enable if args.enabled == 'true' else setup.disable)(clients)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
try:
|
|
main()
|
|
except (OSError, ValueError, json.JSONDecodeError, subprocess.SubprocessError) as error:
|
|
print(f'QMD local search: {error}', file=sys.stderr)
|
|
raise SystemExit(1)
|