bench_run.py

from itertools import chain
import socket
import subprocess
import traceback
from functools import cache
import os
import json
from pathlib import Path
import asyncio

from openai import AsyncClient
from tqdm.auto import tqdm
import psutil

from arxiv import filter_license_redistributable
from bench_tasks import BenchTask, CitationListTask, FillNumberTask, GlobalConfig, ParagraphOrderingTask, SectionCountTask, XRefTask
from bench_utils import get_questions
from daily_papers import get_arxiv_month_top10, get_daily_papers
from model_registry import MODELS


@cache
def get_one_shot_question() -> dict:
    one_shot_question = {
        'type': 'bench',
        'bench': json.load(open(f'generated_bench_v3/2502.20811.json')),
    }
    return one_shot_question


def get_primary_ip():
    if os.getenv('LOCAL') == '1':
        return '127.1.2.3'
    
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.settimeout(0)
    try:
        s.connect(('8.8.8.8', 1))
        IP = s.getsockname()[0]
    except Exception:
        IP = '127.0.0.1'
    finally:
        s.close()
    return IP


def create_server(model: str) -> subprocess.Popen:
    cmd = f'bash {str(Path(os.path.abspath(__file__)).parent)}/launch_sglang.sh {model}'.split()
    proc = subprocess.Popen(cmd)
    return proc


def kill_process_tree(pid):
    try:
        parent = psutil.Process(pid)
        children = parent.children(recursive=True)
        
        for child in children:
            try:
                child.terminate()
            except psutil.NoSuchProcess:
                pass
        
        try:
            parent.terminate()
        except psutil.NoSuchProcess:
            pass
        
        gone, alive = psutil.wait_procs([parent] + children, timeout=3)

        for p in alive:
            p.kill()
            
    except psutil.NoSuchProcess:
        pass


async def gather(*args: list) -> list[list]:
    argv = list(chain(*args))
    if os.getenv('EXT_MODEL'):
        v = []
        for arg in tqdm(argv):
            v.append(asyncio.create_task(arg))
            # await asyncio.sleep(0.25)
    else:
        v = argv
    outputs = await tqdm.gather(*v)
    v = []
    for arg in args:
        assert len(arg) == 0 or len(outputs) > 0
        v.append(outputs[:len(arg)])
        outputs = outputs[len(arg):]
    assert len(outputs) == 0
    return v


async def run(client: AsyncClient, urls: list[str], output_dir: Path, bench_dir: str):
    output_dir: Path = Path(output_dir)
    if len(urls) > 0 and not GlobalConfig.dry_run:
        output_dir.mkdir(parents=True, exist_ok=True)
    
    max_tasks = 1
    running = []
    for url in tqdm(urls):
        while len(running) >= max_tasks:
            done, pending = await asyncio.wait(running, return_when=asyncio.FIRST_COMPLETED)
            running = list(pending)
        running.append(asyncio.create_task(run_one(client, url, output_dir, bench_dir)))
    
    if len(running) > 0:
        await asyncio.wait(running)


async def run_one(client: AsyncClient, url: str, output_dir: Path, bench_dir: str, variant: str = 'main'):
    abs = url.replace('src', 'abs')
    arxiv_id = url.split('/')[-1]
    output_path = output_dir / f'{arxiv_id}.json'
    
    if os.environ.get('SKIP_IF_EXISTS', None) is not None:
        if output_path.exists():
            print('Result exists, continuing')
            return
    
    cache = {}
    if output_path.exists():
        cache = json.load(open(output_path, 'r', encoding='utf-8'))

    try:
        bench = json.load(open(f'{bench_dir}/{arxiv_id}.json', 'r', encoding='utf-8'))
        
        title = bench['title']
        print(title)
        print(abs)
        
        full_text = bench['text']
        
        lprob_sum, avg_nll = float('nan'), float('nan')
        
        if variant == 'main':
            bench_tasks: list[BenchTask] = [
                XRefTask(             arxiv_id=arxiv_id),
                CitationListTask(     arxiv_id=arxiv_id),
                FillNumberTask(       arxiv_id=arxiv_id),
                SectionCountTask(     arxiv_id=arxiv_id),
                ParagraphOrderingTask(arxiv_id=arxiv_id),
            ]
        elif variant == 'no_context':
            bench_tasks: list[BenchTask] = [
                XRefTask(             arxiv_id=arxiv_id, add_context=False),
                CitationListTask(     arxiv_id=arxiv_id, add_context=False),
                FillNumberTask(       arxiv_id=arxiv_id, add_context=False),
                SectionCountTask(     arxiv_id=arxiv_id, add_context=False),
                ParagraphOrderingTask(arxiv_id=arxiv_id, add_context=False),
            ]
        elif variant == 'one_shot':
             bench_tasks: list[BenchTask] = [
                XRefTask(             arxiv_id=arxiv_id, add_one_shot_question=get_one_shot_question()),
                CitationListTask(     arxiv_id=arxiv_id, add_one_shot_question=get_one_shot_question()),
                FillNumberTask(       arxiv_id=arxiv_id, add_one_shot_question=get_one_shot_question()),
                SectionCountTask(     arxiv_id=arxiv_id, add_one_shot_question=get_one_shot_question()),
                ParagraphOrderingTask(arxiv_id=arxiv_id, add_one_shot_question=get_one_shot_question()),
            ]
        else:
            assert 0, f'Unknown variant: {variant}'

        num_filters = GlobalConfig.num_filters
        
        
        eval_tasks = [
            task.run_task_cached(cache, client, full_text, get_questions(bench[task.json_name], num_filters=num_filters))
            for task in bench_tasks
        ]
        
        eval_tasks_outputs = await gather(*eval_tasks,)
        
        bench_tasks_outputs = {
            task.json_output_name: {
                **task.dump_task(full_text, get_questions(bench[task.json_name], num_filters=num_filters), task_outputs),
                'one_shot_messages': task.one_shot_messages,
            }
            for task, task_outputs in zip(bench_tasks, eval_tasks_outputs)
        }

        for task_name, results in bench_tasks_outputs.items():
            if "correct" in results:
                print(f'{task_name}: {results["correct"]}/{results["total"]} = {results["accuracy"]:.2%}')

        
        result = {
            **cache,
            'title': title,
            'url': abs,
            'lprob_sum': lprob_sum,
            'avg_nll': avg_nll,
            **bench_tasks_outputs,
            'text': full_text,
        }
        
        if not GlobalConfig.dry_run:
            json.dump(
                result,
                open(output_dir / f'{arxiv_id}.json', 'w', encoding='utf-8'),
                ensure_ascii=False,
                indent=4,
            )

    except Exception as e:
        print(f"Error: {str(e)}")
        traceback.print_exc()
        pass
    finally:
        pass
    print('-----------------------------------------')


async def main(short_model_name = None, variant='main', exp_profile='main'):
    if short_model_name is None:
        short_model_name = os.getenv('MODEL', None)

    if short_model_name is not None:
        ip = get_primary_ip()
        proc = create_server(short_model_name)
    else:
        proc = None
    
    bench_dir = 'generated_bench_v3/dummy'
    ref_name = 'ref-dummy'
    variant = 'main'
    assert variant in ['main', 'no_context', 'one_shot']
    ref_name += f'-{variant}'
    
    if (model_name := os.getenv('EXT_MODEL')) is None:
        port, model_name, tokenizer_name = MODELS[short_model_name]
        
        if (lora_name := os.environ.get('LORA', None)) is not None:
            model_name = f'{model_name}-{lora_name}'
        
        if model_name == 'qwen-3-8b-128k':
            GlobalConfig.invoke_extra_body_kwargs = {"chat_template_kwargs": {"enable_thinking": False}}

        base_url = f'http://{ip}:{port}/v1'
        
        # t = os.environ['all_proxy'], os.environ['http_proxy'], os.environ['https_proxy']
        # del os.environ['all_proxy'], os.environ['http_proxy'], os.environ['https_proxy']
        async_client = AsyncClient(api_key='a', base_url=base_url, timeout=86400)
        # os.environ['all_proxy'], os.environ['http_proxy'], os.environ['https_proxy'] = t
        
        while True:
            try:
                await async_client.models.list(timeout=1.0)
                print("server is up!")
                break
            except Exception:
                pass
    
    else:
        model_name = os.getenv('EXT_MODEL')
        async_client = AsyncClient(
            api_key=os.getenv('EXT_API_KEY'),
            base_url=os.getenv('EXT_BASE_URL'),
            timeout=300,
        )
        GlobalConfig.is_thinking_model = os.getenv('EXT_THINKING_MODEL', '0') == '1'

    if exp_profile == 'history':
        # History Data Experiment
        assert variant == 'main'
        for year in range(2015, 2025):
            for month in range(1, 13):
                urls = get_arxiv_month_top10(year, month)
                output_dir = f'outputs-v3bench/arxiv-top10-{year}-{month}/{model_name}/{ref_name}'
                await run(async_client, urls, output_dir, bench_dir, 'main')
    elif exp_profile == 'main':
        # Main experiment
        dates = [f'2025-03-{i:02}' for i in range(1, 32)]
        # dates = [f'2026-03-{i:02}' for i in range(1, 32)]
        for date in dates:
            urls = filter_license_redistributable(get_daily_papers([date]))
            print(f'{date}: {len(urls)}')
        for date in dates:
            urls = filter_license_redistributable(get_daily_papers([date]))
            output_dir = f'outputs-v3bench/{date}/{model_name}/{ref_name}'
            await run(async_client, urls, output_dir, bench_dir, variant)
    else:
        assert 0, f'Unknown exp_profile: {exp_profile}'
    
    if proc is not None:
        await asyncio.sleep(5)
        kill_process_tree(proc.pid)


async def main_all():
    if os.getenv('MODEL', None) != 'all':
        await main(exp_profile='main')
        # await main(exp_profile='history')
        return
        
    short_model_names = [
        'llama',
        'qwen',
        'mistral',
        'phi',
        'qwen-7b-1m',
        'qwen-14b-1m',
        'prolong',
        'megabeam',
        'qwen3',
    ]
    for short_model_name in short_model_names:
        await main(short_model_name)

if __name__ == '__main__':
    asyncio.run(main_all())