concurrency_runner.py

"""Runner integration for the concurrency sweep (benchmark_type='concurrency')."""
import asyncio
import json
import time
import uuid

import concurrency_benchmark as cb
import metrics_capture
import telemetry

ALLOWED = {'benchmark_type', 'name', 'model_ids', 'concurrency_levels', 'requests_per_level', 'context_sizes',
           'max_tokens', 'timeout_s', 'max_calls', 'temperature', 'seed', 'repeats', 'max_cost_usd', 'allow_paid',
           'fixed_output_length'}
DEFAULTS = {'name': '', 'model_ids': None, 'concurrency_levels': list(cb.DEFAULT_LEVELS), 'requests_per_level': 16,
            'context_sizes': [2048], 'max_tokens': 256, 'timeout_s': 120, 'max_calls': 400, 'temperature': 0.0,
            'seed': 42, 'repeats': 1, 'max_cost_usd': 0, 'allow_paid': False, 'fixed_output_length': True}
PROMPT_BUDGET_BYTES = 128 * 1024 * 1024
IDLE_WAIT_S = 15.0


def _int(v, lo, hi):
    return isinstance(v, int) and not isinstance(v, bool) and lo <= v <= hi


class ConcurrencyMixin:
    def _create_concurrency(self, d):
        if set(d) - ALLOWED: raise ValueError('Unknown run fields for a concurrency test: ' + ', '.join(sorted(set(d) - ALLOWED)))
        cfg = dict(DEFAULTS); cfg.update(d); cfg['benchmark_type'] = 'concurrency'
        cfg.update(cb.validate_concurrency(cfg))
        ids = cfg['model_ids']
        if not isinstance(ids, list) or not ids or len(ids) > 8 or any(not isinstance(x, str) or not x for x in ids) or len(set(ids)) != len(ids):
            raise ValueError('Invalid model_ids')
        if cfg['repeats'] != 1: raise ValueError('Concurrency tests use repeats=1; raise requests_per_level instead')
        if not isinstance(cfg['context_sizes'], list) or len(cfg['context_sizes']) != 1 or not _int(cfg['context_sizes'][0], 64, 1048576):
            raise ValueError('Concurrency tests take exactly one context size (64-1048576 approximate tokens)')
        if not _int(cfg['max_tokens'], 1, 32768) or not _int(cfg['timeout_s'], 1, 1800) or not _int(cfg['max_calls'], 1, 20000) or not _int(cfg['seed'], -2**31, 2**31 - 1):
            raise ValueError('Invalid numeric limits')
        if not isinstance(cfg['temperature'], (int, float)) or isinstance(cfg['temperature'], bool) or not 0 <= cfg['temperature'] <= 2:
            raise ValueError('Invalid temperature')
        if cfg['allow_paid'] is not False or cfg['max_cost_usd'] not in (0, 0.0):
            raise ValueError('Concurrency tests never use paid providers (allow_paid must be false, max_cost_usd 0)')
        if not isinstance(cfg['fixed_output_length'], bool): raise ValueError('Invalid fixed_output_length')
        if not isinstance(cfg['name'], str) or len(cfg['name']) > 120: raise ValueError('Invalid run name')
        with self.s.connect() as c:
            models = [dict(x) for x in c.execute('select * from models where id in (%s)' % ','.join('?' * len(ids)), ids)]
        if len(models) != len(ids): raise ValueError('Unknown model')
        models.sort(key=lambda m: ids.index(m['id']))
        prov = {p['id']: p for p in self.s.providers()}
        if any(prov.get(m['provider_id'], {}).get('kind') not in ('openai', 'openwebui') for m in models):
            raise ValueError('Concurrency tests run only against private OpenAI-compatible/OpenWebUI endpoints, never OpenRouter')
        ctx = cfg['context_sizes'][0]
        for m in models:
            if ctx + cfg['max_tokens'] + 256 > m['context_limit']:
                raise ValueError(f"Context {ctx} + max_tokens {cfg['max_tokens']} exceeds the registered limit {m['context_limit']} of {m['name']}")
        per_model = cb.planned_calls(cfg['concurrency_levels'], cfg['requests_per_level'])
        if per_model * len(models) > cfg['max_calls']:
            raise ValueError(f"Planned {per_model * len(models)} requests exceed max_calls {cfg['max_calls']}")
        if per_model * len(models) * ctx * 4 > PROMPT_BUDGET_BYTES:
            raise ValueError('Planned prompts exceed the 128 MiB budget; lower context, requests or levels')
        rid = uuid.uuid4().hex; created_at = time.time()
        if cfg['name'].strip().lower() in ('', 'run', 'broodbench run'):
            stamp = time.strftime('%Y%m%dT%H%M%SZ', time.gmtime(created_at))
            cfg['name'] = (f"{' + '.join(m['name'] for m in models)} | concurrency {','.join(map(str, cfg['concurrency_levels']))} "
                           f"| ctx {ctx} | out {cfg['max_tokens']} | {stamp} | {rid[:8]}")[:120]
        cfg.update(models=models, providers=[{k: p[k] for k in ('id', 'name', 'kind', 'base_url')} for p in prov.values() if p['id'] in {m['provider_id'] for m in models}],
                   suite='concurrency', context_mode='fixed', pricing={}, spend_usd=0.0, spend_source='none_private_endpoint',
                   concurrency_results={}, concurrency_method={
                       'probe': 'Per level: one wave of N simultaneous max_tokens=1 requests. aggregate_prefill_tps = computed prompt tokens / probe wall seconds (includes queueing, transport and one sampled token).',
                       'generation': 'Per level: requests_per_level requests with at most N in flight. output_tps = completion tokens / wave wall seconds. decode_tps = per-request stream speed (client estimate); server_decode_tps = vLLM per-request decode time when attributable.',
                       'prefill_tps': 'Per-request prefill rate under load from vLLM counter deltas; only when the endpoint is idle at both boundaries and request/token counts match exactly.',
                       'input_tps': 'Prompt tokens / generation wall seconds. Not prefill throughput.',
                       'prompts': 'Unique leading nonce per request so prefix caching cannot share KV; cached tokens are recorded.'})
        with self.s.connect() as c:
            c.execute('insert into runs(id,name,status,created_at,config,progress_total) values(?,?,?,?,?,?)', (rid, cfg['name'], 'queued', created_at, json.dumps(cfg), per_model * len(models)))
            c.execute('insert into queue(run_id,config) values(?,?)', (rid, json.dumps(cfg)))
        return self.get(rid)

    async def _wait_idle(self, p, model_id, baseline_success=None, expected=0):
        """Poll /metrics until no requests are running/waiting (and our completions are counted)."""
        deadline = time.monotonic() + IDLE_WAIT_S; snap = None
        while True:
            snap = await metrics_capture.snapshot(p, model_id)
            if snap is None: return None
            v = snap['values']
            idle = v['num_requests_running'] == 0 and v['num_requests_waiting'] == 0
            counted = baseline_success is None or v['request_success_total'] - baseline_success >= expected
            if (idle and counted) or time.monotonic() > deadline: return snap
            await asyncio.sleep(0.2)

    async def _concurrency_wave(self, rid, cfg, p, model, phase, level, total, run_tag,
                                ctx=None, timeout_s=None, case_prefix='conc', regime=None):
        ctx = ctx or cfg['context_sizes'][0]
        call_cfg = dict(cfg, _no_call_metrics=True, context_mode='fixed')
        if timeout_s:
            call_cfg['timeout_s'] = timeout_s
        if phase == 'probe':
            call_cfg['max_tokens'] = 1
        elif cfg.get('fixed_output_length'):
            call_cfg['_extra_payload'] = {'ignore_eos': True}
        before = await self._wait_idle(p, model['model_id'])
        wave_t0 = time.monotonic()

        def persist(rec):
            with self.s.connect() as c:
                c.execute('insert into results values(?,?,?,?,?)', (rec['id'], rid, model['id'], rec['case_id'], json.dumps(rec)))
                c.execute('update runs set progress_done=progress_done+1 where id=?', (rid,))

        async def one(i):
            msgs = cb.build_messages(ctx, cfg['seed'], run_tag, phase, level, i, context_tag=f'ctx{ctx}')
            case_id = (f'conc-{phase}-c{level}-r{i:04d}' if case_prefix == 'conc'
                       else f'{case_prefix}-{phase}-ctx{ctx}-c{level}-r{i:04d}')
            rec = {'id': uuid.uuid4().hex, 'model_id': model['id'], 'model_name': model['name'], 'benchmark_type': cfg['benchmark_type'],
                   'case_id': case_id, 'phase': 'probe' if phase == 'probe' else 'generation',
                   'concurrency': level, 'request_index': i, 'track': 'concurrency', 'family': 'concurrency',
                   'context_size': ctx, 'prompt_chars': len(msgs[0]['content']), 'prompt_head': msgs[0]['content'][:150],
                   'max_tokens': call_cfg['max_tokens'], 'effective_timeout_s': call_cfg['timeout_s'],
                   'status': 'error', 'error': None, 'error_detail': None,
                   'calls': [], 'response_chars': 0, 'response_head': ''}
            if regime:
                rec['regime'] = regime
            started = time.monotonic(); rec['started_offset_s'] = started - wave_t0
            try:
                data = await self._call(p, model, msgs, call_cfg)
                record = data.get('_telemetry') or {}
                msg = data['choices'][0]['message'] if data.get('choices') else {}
                text = (msg.get('content') or '') + (msg.get('reasoning_content') or '')
                rec.update(response_chars=len(text), response_head=text[:200])
                rec['calls'] = [record]
                if not record.get('stream_complete', True) or record.get('finish_reason') in (None, 'error', 'content_filter'):
                    rec.update(status='error', error='IncompleteStream', error_detail=f"finish_reason={record.get('finish_reason')} stream_complete={record.get('stream_complete')}")
                else:
                    rec['status'] = 'completed'
            except asyncio.CancelledError:
                rec.update(status='cancelled', error='Cancelled', latency_s=time.monotonic() - started, finished_offset_s=time.monotonic() - wave_t0)
                persist(rec); raise
            except Exception as e:
                rec.update(error=type(e).__name__, error_detail=str(e)[:500])
                if getattr(e, 'response', None) is not None and hasattr(e.response, 'status_code'): rec['error'] += f' HTTP {e.response.status_code}'
            finished = time.monotonic()
            call = rec['calls'][0] if rec['calls'] else {}
            rec.update(latency_s=finished - started, finished_offset_s=finished - wave_t0,
                       prompt_tokens=call.get('prompt_tokens'), completion_tokens=call.get('completion_tokens'),
                       cached_tokens=call.get('cached_tokens'), finish_reason=call.get('finish_reason'),
                       performance=telemetry.aggregate(rec['calls']))
            persist(rec)
            return rec

        wave = await cb.run_wave(one, total, level)
        records = [r['value'] if r['value'] is not None else {'status': 'error', 'error': r['error'], 'calls': []} for r in wave['results']]
        done = [r for r in records if r.get('status') == 'completed']
        prompt = sum(r['prompt_tokens'] for r in done) if done and all(isinstance(r.get('prompt_tokens'), int) for r in done) else None
        out = sum(r['completion_tokens'] for r in done) if done and all(isinstance(r.get('completion_tokens'), int) for r in done) else None
        after = await self._wait_idle(p, model['model_id'], before['values']['request_success_total'] if before else None, len(done))
        if len(done) != total:
            metrics, reason = None, f'{total - len(done)} request(s) failed; server counters cannot be attributed to completed requests'
        else:
            metrics, reason = cb.batch_measurement(before, after, total, prompt, out)
        return records, wave['wall_s'], wave['max_in_flight'], metrics, reason

    async def _run_concurrency(self, rid, cfg):
        status = 'done'
        try:
            for model in cfg['models']:
                p = next((x for x in self.s.providers() if x['id'] == model['provider_id']), None)
                if not p: raise RuntimeError('Configured provider missing')
                rows = cfg['concurrency_results'].setdefault(model['id'], [])
                for level in cfg['concurrency_levels']:
                    if rid in self.cancel: raise asyncio.CancelledError()
                    probe = await self._concurrency_wave(rid, cfg, p, model, 'probe', level, level, rid[:8])
                    gen = await self._concurrency_wave(rid, cfg, p, model, 'gen', level, cfg['requests_per_level'], rid[:8])
                    row = cb.summarize_level(level, gen[0], gen[1], gen[2], gen[3], gen[4], probe[0], probe[1], probe[2], probe[3], probe[4])
                    row['measured_at'] = time.time(); rows.append(row)
                    with self.s.connect() as c: c.execute('update runs set config=? where id=?', (json.dumps(cfg), rid))
        except asyncio.CancelledError:
            status = 'cancelled'
        except Exception as e:
            status = 'error'; cfg['run_error'] = f'{type(e).__name__}: {str(e)[:500]}'
        if rid in self.cancel: status = 'cancelled'
        with self.s.connect() as c: c.execute('update runs set status=?,config=? where id=?', (status, json.dumps(cfg), rid))
        self.cancel.discard(rid)

    def _concurrency_summary(self, cfg, results):
        out = []
        for m in cfg.get('models', []):
            items = [x for x in results if x.get('model_id') == m['id']]
            gen_calls = [c for x in items if x.get('phase') == 'generation' for c in x.get('calls', [])]
            rows = (cfg.get('concurrency_results') or {}).get(m['id'], [])
            best = max((r for r in rows if r.get('output_tps') is not None), key=lambda r: r['output_tps'], default=None)
            out.append({'model_id': m['id'], 'name': m['name'], 'quant': m.get('quant', ''), 'benchmark_type': cfg.get('benchmark_type', 'concurrency'),
                        'context_sizes': sorted({r['context_size'] for r in rows if r.get('context_size')}),
                        'concurrency_levels': rows, 'total': len(items),
                        'completed': sum(x.get('status') == 'completed' for x in items),
                        'errors': sum(x.get('status') in ('error', 'cancelled') for x in items),
                        'peak_output_tps': best['output_tps'] if best else None,
                        'peak_output_concurrency': best['concurrency'] if best else None,
                        'performance': telemetry.aggregate(gen_calls)})
        return out