searchi/tools/embeddinggemma2/benchmark.py

"""Explicit CPU-only measurements of the real verified Searchi Google runtime.

Native mode isolates worker IPC/inference. MCP mode measures the actual shared
service over stdio against an explicitly supplied, already indexed test corpus.
No models are downloaded, files indexed, or production defaults modified.
"""
import argparse
import hashlib
import json
import math
import os
from pathlib import Path
import platform
import queue
import statistics
import subprocess
import threading
import time

ROOT = Path(__file__).resolve().parents[2]
QUERIES = [
    'cancel the supplier partnership',
    'mountains reflected in calm water',
    'die Kündigungsfrist im Lieferantenvertrag',
    'source code that checks indexed file permissions',
    'music with piano and a slow melody',
    'a boat crossing a lake at sunset',
    'notes about returning unused materials',
    'find the recording discussing orange soil',
]

def elapsed(start):
    return (time.perf_counter() - start) * 1000

def distribution(values):
    ordered = sorted(values)
    return {'samples': len(values), 'mean_ms': statistics.mean(values),
            'p50_ms': statistics.median(values),
            'p95_ms': ordered[max(0, math.ceil(len(values) * .95) - 1)]}

def require(condition, message):
    if not condition:
        raise ValueError(message)


class Session:
    def __init__(self, command):
        environment = os.environ.copy()
        environment['SEARCHI_FORCE_CPU'] = '1'
        flags = subprocess.CREATE_NO_WINDOW if os.name == 'nt' else 0
        self.started = time.perf_counter()
        self.process = subprocess.Popen(command, stdin=subprocess.PIPE,
            stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True,
            encoding='utf-8', env=environment, creationflags=flags)
        self.lines = queue.Queue()
        self.errors = []
        threading.Thread(target=self.read, daemon=True).start()
        threading.Thread(target=self.read_errors, daemon=True).start()

    def read(self):
        try:
            for line in self.process.stdout:
                self.lines.put(json.loads(line))
            self.lines.put(RuntimeError('Runtime stdout closed'))
        except Exception as error:
            self.lines.put(error)

    def read_errors(self):
        for line in self.process.stderr:
            if len(self.errors) < 20:
                self.errors.append(line.rstrip())

    def next(self):
        result = self.lines.get(timeout=120)
        if isinstance(result, Exception):
            raise result
        return result

    def send(self, value):
        self.process.stdin.write(json.dumps(value, ensure_ascii=False) + '\n')
        self.process.stdin.flush()

    def worker(self, value):
        start = time.perf_counter()
        self.send(value)
        response = self.next()
        if not response.get('ok'):
            raise RuntimeError(response.get('error', 'Worker request failed'))
        require(response['device'] == 'cpu', 'Benchmark contract failed at benchmark.py:84')
        return response, elapsed(start)

    def mcp(self, identifier, method, params):
        start = time.perf_counter()
        self.send({'jsonrpc': '2.0', 'id': identifier, 'method': method, 'params': params})
        while True:
            response = self.next()
            if response.get('id') == identifier:
                require('error' not in response, response)
                return response['result'], elapsed(start)

    def close(self):
        self.process.stdin.close()
        try:
            self.process.wait(timeout=10)
        except subprocess.TimeoutExpired:
            self.process.kill()
            self.process.wait(timeout=10)

def vectors(response, count):
    result = response['vectors']
    require(len(result) == count, 'Benchmark contract failed at benchmark.py:106')
    for vector in result:
        require(len(vector) == 768 and all(math.isfinite(n) for n in vector), 'Benchmark contract failed at benchmark.py:108')
        require(abs(sum(n*n for n in vector) - 1) < .0001, 'Benchmark contract failed at benchmark.py:109')
    return result

def native(args):
    manifest = json.loads((ROOT / 'tools/embeddinggemma2/model-manifest.json').read_text())
    model = args.models / manifest['files'][0]['path']
    start = time.perf_counter()
    require(model.stat().st_size == manifest['files'][0]['bytes'], 'Benchmark contract failed at benchmark.py:116')
    with model.open('rb') as stream:
        require(hashlib.file_digest(stream, 'sha256').hexdigest() == manifest['files'][0]['sha256'], 'Benchmark contract failed at benchmark.py:118')
    verification_ms = elapsed(start)
    reports = []
    for threads in args.threads:
        session = Session([str(args.worker.resolve()), '--model', str(model.resolve()),
                           '--threads', str(threads), '--device', 'cpu'])
        try:
            hello = session.next()
            load_ms = elapsed(session.started)
            require(hello['device'] == 'cpu' and hello['model'] == manifest['model_id'], 'Benchmark contract failed at benchmark.py:127')
            first, first_ms = session.worker({'op':'query', 'text':'first unseen foreground query',
                                              'debug_performance':True})
            vectors(first, 1)
            samples = []
            for repeat in range(args.repeats):
                for query in QUERIES:
                    response, duration = session.worker({'op':'query','text':query,
                                                         'debug_performance':True})
                    vectors(response, 1)
                    samples.append({'query':query,'repeat':repeat,'round_trip_ms':duration,
                                    'worker':response['performance']})
            window_samples = []
            passage = ('Supplier cancellation requires sixty days written notice. '
                       'Kündigung und Rückgabe der Materialien. ') * 500
            for _ in range(3):
                response, duration = session.worker({'op':'windows','text':passage,'title':'Agreement'})
                encoded = passage.encode()
                for window in response['windows']:
                    require(encoded[window['start_byte']:window['end_byte']].decode() == window['text'], 'Benchmark contract failed at benchmark.py:146')
                require(response['windows'][-1]['end_byte'] == len(encoded), 'Benchmark contract failed at benchmark.py:147')
                window_samples.append({'bytes':len(encoded),'windows':len(response['windows']),
                                       'round_trip_ms':duration})
            documents = [(f'Document {n}. Supplier cancellation requires sixty days written notice. '
                         'Unused materials must be returned. ') * 5 for n in range(8)]
            singles = []
            single_times = []
            for document in documents:
                response, duration = session.worker({'op':'documents','texts':[document],'title':'Agreement'})
                singles.extend(vectors(response, 1))
                single_times.append(duration)
            response, batch_ms = session.worker({'op':'documents','texts':documents,'title':'Agreement'})
            batched = vectors(response, 8)
            cosines = [sum(a*b for a,b in zip(left,right)) for left,right in zip(singles,batched)]
            require(min(cosines) > .99999, 'Benchmark contract failed at benchmark.py:161')
            reports.append({'threads':threads,'worker_info':hello,'process_load_ms':load_ms,
                'first_query_ms':first_ms,'first_query_worker':first['performance'],
                'warm_queries':distribution([sample['round_trip_ms'] for sample in samples]),
                'query_samples':samples,'window_samples':window_samples,
                'documents':{'single_request_ms':single_times,'eight_single_total_ms':sum(single_times),
                    'eight_document_request_ms':batch_ms,'batch_cosine_min':min(cosines)}})
            print(f'CPU {threads} threads: warm p95 {reports[-1]["warm_queries"]["p95_ms"]:.2f} ms', flush=True)
        finally:
            session.close()
    return {'text_model_verification_ms':verification_ms,'thread_reports':reports,
            'scope':'Native actual Google Q8 worker; verified model; cached OS pages possible. '
                    'Includes NDJSON encode/decode; excludes Rust adapter, ANN and UI. '
                    'First query precedes warm samples. Documents measure IPC batch, not assumed parallelism.'}

def mcp(args):
    require(args.db and args.cli, 'Benchmark contract failed at benchmark.py:177')
    require(args.db.is_file(), 'Benchmark contract failed at benchmark.py:178')
    session = Session([str(args.cli.resolve()),'mcp','--db',str(args.db.resolve()),
        '--model-dir',str(args.models.resolve()),'--embedding-worker',str(args.worker.resolve()),
        '--cpu-threads','1'])
    try:
        session.mcp(1,'initialize',{'protocolVersion':'2025-11-25','capabilities':{},
            'clientInfo':{'name':'searchi-cpu-performance','version':'1'}})
        startup_ms = elapsed(session.started)
        session.send({'jsonrpc':'2.0','method':'notifications/initialized'})
        stats,_ = session.mcp(2,'tools/call',{'name':'stats','arguments':{}})
        require(stats['structuredContent']['mode'] == 'hybrid', 'Benchmark contract failed at benchmark.py:188')
        require(stats['structuredContent']['worker']['device'] == 'cpu', 'Benchmark contract failed at benchmark.py:189')
        samples=[]
        identifier=3
        cases = [('first_unseen',QUERIES[0],None)]
        cases += [('same_query',QUERIES[0],None)]*args.repeats
        cases += [('filter_change',QUERIES[0],ext) for ext in ['pdf','md','pdf','md']]
        cases += [('distinct_query',query,None) for query in QUERIES[1:]]
        for case,query,extension in cases:
            arguments={'query':query,'limit':3}
            if extension:
                arguments['extension']=extension
            result,duration=session.mcp(identifier,'tools/call',{'name':'search','arguments':arguments})
            identifier+=1
            require(not result.get('isError'), result)
            response=result['structuredContent']
            require(response['mode']=='hybrid', 'Benchmark contract failed at benchmark.py:204')
            if extension:
                require(all(item['filename'].lower().endswith('.'+extension) for item in response['results']), 'Benchmark contract failed at benchmark.py:206')
            samples.append({'case':case,'query':query,'extension':extension,'request_ms':duration,
                'engine_ms':response['metrics']['total_ms'],
                'matches':[{'filename':item['filename'],'source':item['source'],
                            'snippet':item['snippet'],'score':item['score']}
                           for item in response['results']]})
        return {'startup_to_initialized_ms':startup_ms,'stats':stats['structuredContent'],
            'samples':samples,'same_query':distribution([s['engine_ms'] for s in samples if s['case']=='same_query']),
            'filter_change':distribution([s['engine_ms'] for s in samples if s['case']=='filter_change']),
            'scope':'Actual shared service through real MCP. Startup includes verified models and background warmup. '
                    'Seven-file README corpus when using that explicit fixture DB; no UI scheduling or paint.'}
    finally:
        session.close()

def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--models',type=Path,required=True)
    parser.add_argument('--worker',type=Path,required=True)
    parser.add_argument('--output',type=Path,required=True)
    parser.add_argument('--mode',choices=['native','mcp'],default='native')
    parser.add_argument('--threads',type=int,nargs='+',default=[1,2,4])
    parser.add_argument('--repeats',type=int,default=3)
    parser.add_argument('--db',type=Path)
    parser.add_argument('--cli',type=Path)
    parser.add_argument('--rust-profile',choices=['debug','release','not-applicable'],default='not-applicable')
    parser.add_argument('--label',default='')
    args=parser.parse_args()
    require(not args.output.exists(), 'Output must be new; preserve previous evidence')
    require(1<=args.repeats<=20 and all(1<=n<=8 for n in args.threads), 'Benchmark contract failed at benchmark.py:233')
    def digest(path):
        with path.open('rb') as stream:
            return hashlib.file_digest(stream,'sha256').hexdigest()
    report={'captured_at_utc':time.strftime('%Y-%m-%dT%H:%M:%SZ',time.gmtime()),
        'os':platform.platform(),'machine':platform.machine(),'logical_cpus':os.cpu_count(),
        'cpu_only':True,'rust_profile':args.rust_profile,'label':args.label,
        'worker_sha256':digest(args.worker),
        'cli_sha256':digest(args.cli) if args.cli else None,
        'results':native(args) if args.mode=='native' else mcp(args)}
    args.output.parent.mkdir(parents=True,exist_ok=True)
    with args.output.open('x', encoding='utf-8') as output:
        output.write(json.dumps(report,indent=2,ensure_ascii=False)+'\n')
    print(f'Saved {args.output}',flush=True)

if __name__=='__main__':
    main()