Function main

experiments/run_humaneval.py:94–212

()

async def main():
    args = parse_args()
    output_dir = Path(args.output_dir).expanduser()
    output_dir.mkdir(parents=True, exist_ok=True)
    configure_logging(log_path=output_dir / "logs/log.txt")
    dataset = JSONLReader.parse_file(args.dataset_json)

    current_time = Time.instance().value or time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
    Time.instance().value = current_time
    result_file = output_dir / f"{args.domain}_{args.llm_name}_{current_time}.json"
    latency_target = str(output_dir)

    agent_names = [name for name, num in zip(args.agent_names, args.agent_nums) for _ in range(num)]
    kwargs = get_kwargs(args.mode, len(agent_names))

    kv_config: Optional[KVCommConfig] = None
    if args.execution_mode == "allow_kv_reuse":
        kv_config = KVCommConfig.from_env().apply_overrides(
            threshold=args.kv_threshold,
            max_anchor_num=args.kv_max_anchor_num,
            window_size=args.kv_window_size,
            thread_pool_workers=args.kv_thread_workers,
            worker_timeout=args.kv_worker_timeout,
        )
    else:
        kv_config = KVCommConfig.from_env()

    graph = Graph(
        domain=args.domain,
        llm_name=args.llm_name,
        agent_names=agent_names,
        decision_method=args.decision_method,
        kv_config=kv_config,
        **kwargs,
    )

    num_batches = int(len(dataset) / args.batch_size)
    total_solved, total_executed = 0, 0

    for i_batch in range(num_batches):
        logger.opt(colors=True).info(f"<blue>[BATCH]</blue> {i_batch} {'-' * 40}")
        start_ts = time.time()
        current_batch = dataloader(dataset, args.batch_size, i_batch)
        if not current_batch:
            logger.warning("No more data available.")
            break

        tasks = []
        meta_info = []
        for record in current_batch:
            realized_graph = copy.deepcopy(graph)
            realized_graph.spatial_logits = graph.spatial_logits
            realized_graph.temporal_logits = graph.temporal_logits
            task = record["prompt"]
            tests = record["test"]
            input_dict = {"task": task, "_batch_index": i_batch}

            mode_kwargs = {}