wehub-resource-sync
424 lines
16 KiB
Python
424 lines
16 KiB
Python
import argparse
|
||
import heapq
|
||
import multiprocessing as mp
|
||
import os
|
||
import queue
|
||
import shutil
|
||
import threading
|
||
import traceback
|
||
from collections import defaultdict
|
||
from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait
|
||
from copy import deepcopy
|
||
from typing import Optional
|
||
|
||
from bfcl_eval.constants.eval_config import (
|
||
PROJECT_ROOT,
|
||
RESULT_FILE_PATTERN,
|
||
RESULT_PATH,
|
||
TEST_IDS_TO_GENERATE_PATH,
|
||
)
|
||
from bfcl_eval.constants.model_config import MODEL_CONFIG_MAPPING
|
||
from bfcl_eval.eval_checker.eval_runner_helper import load_file
|
||
from bfcl_eval.model_handler.base_handler import BaseHandler
|
||
from bfcl_eval.model_handler.local_inference.base_oss_handler import OSSHandler
|
||
from bfcl_eval.utils import *
|
||
from tqdm import tqdm
|
||
|
||
|
||
def get_args():
|
||
parser = argparse.ArgumentParser()
|
||
# Refer to model_choice for supported models.
|
||
parser.add_argument("--model", type=str, default="gorilla-openfunctions-v2", nargs="+")
|
||
# Refer to test_categories for supported categories.
|
||
parser.add_argument("--test-category", type=str, default="all", nargs="+")
|
||
|
||
# Parameters for the model that you want to test.
|
||
parser.add_argument("--temperature", type=float, default=0.001)
|
||
parser.add_argument("--include-input-log", action="store_true", default=False)
|
||
parser.add_argument("--exclude-state-log", action="store_true", default=False)
|
||
parser.add_argument("--num-threads", required=False, type=int)
|
||
parser.add_argument("--num-gpus", default=1, type=int)
|
||
parser.add_argument("--backend", default="vllm", type=str, choices=["vllm", "sglang"])
|
||
parser.add_argument("--gpu-memory-utilization", default=0.9, type=float)
|
||
parser.add_argument("--result-dir", default=None, type=str)
|
||
parser.add_argument("--run-ids", action="store_true", default=False)
|
||
parser.add_argument("--allow-overwrite", "-o", action="store_true", default=False)
|
||
parser.add_argument(
|
||
"--skip-server-setup",
|
||
action="store_true",
|
||
default=False,
|
||
help="Skip vLLM/SGLang server setup and use existing endpoint specified by the LOCAL_SERVER_ENDPOINT and LOCAL_SERVER_PORT environment variables.",
|
||
)
|
||
# Optional local model path
|
||
parser.add_argument(
|
||
"--local-model-path",
|
||
type=str,
|
||
default=None,
|
||
help="Specify the path to a local directory containing the model's config/tokenizer/weights for fully offline inference. Use this only if the model weights are stored in a location other than the default HF_HOME directory.",
|
||
)
|
||
parser.add_argument(
|
||
"--lora-modules",
|
||
type=str,
|
||
default=None,
|
||
nargs="*",
|
||
help="Specify the path to the LoRA modules for vLLM backend in name=\"path\" format. Can be specified multiple times.",
|
||
)
|
||
parser.add_argument(
|
||
"--enable-lora",
|
||
action="store_true",
|
||
default=False,
|
||
help="Enable LoRA for vLLM backend.",
|
||
)
|
||
parser.add_argument(
|
||
"--max-lora-rank",
|
||
type=int,
|
||
default=None,
|
||
help="Specify the maximum LoRA rank for vLLM backend.",
|
||
)
|
||
args = parser.parse_args()
|
||
print(f"Parsed arguments: {args}")
|
||
|
||
return args
|
||
|
||
|
||
def build_handler(model_name, temperature):
|
||
config = MODEL_CONFIG_MAPPING[model_name]
|
||
handler = config.model_handler(
|
||
model_name=config.model_name,
|
||
temperature=temperature,
|
||
registry_name=model_name,
|
||
is_fc_model=config.is_fc_model,
|
||
)
|
||
return handler
|
||
|
||
|
||
def get_involved_test_entries(test_category_args, run_ids):
|
||
all_test_categories, all_test_entries_involved = [], []
|
||
if run_ids:
|
||
all_test_categories, all_test_entries_involved = load_test_entries_from_id_file(
|
||
TEST_IDS_TO_GENERATE_PATH
|
||
)
|
||
|
||
else:
|
||
all_test_categories = parse_test_category_argument(test_category_args)
|
||
for test_category in all_test_categories:
|
||
all_test_entries_involved.extend(load_dataset_entry(test_category))
|
||
|
||
return (
|
||
all_test_categories,
|
||
all_test_entries_involved,
|
||
)
|
||
|
||
|
||
def collect_test_cases(args, model_name, all_test_categories, all_test_entries_involved):
|
||
model_name_dir = model_name.replace("/", "_")
|
||
model_result_dir = args.result_dir / model_name_dir
|
||
|
||
existing_result = []
|
||
for test_category in all_test_categories:
|
||
# TODO: Simplify the handling of memory prerequisite entries/categories
|
||
result_file_paths = [
|
||
model_result_dir
|
||
/ get_directory_structure_by_category(test_category)
|
||
/ get_file_name_by_category(test_category, is_result_file=True)
|
||
]
|
||
if is_memory(test_category):
|
||
# Memory test cases have the pre-requisite entries in a separate file
|
||
result_file_paths.append(
|
||
model_result_dir
|
||
/ get_directory_structure_by_category(test_category)
|
||
/ get_file_name_by_category(f"{test_category}_prereq", is_result_file=True)
|
||
)
|
||
|
||
for file_path in result_file_paths:
|
||
if file_path.exists():
|
||
# Not allowing overwrite, we will load the existing results
|
||
if not args.allow_overwrite:
|
||
existing_result.extend(load_file(file_path))
|
||
# Allow overwrite and not running specific test ids, we will delete the existing result file before generating new results
|
||
elif not args.run_ids:
|
||
file_path.unlink()
|
||
# Allow overwrite and running specific test ids, we will do nothing here
|
||
else:
|
||
pass
|
||
|
||
if is_memory(test_category):
|
||
# We also need to special handle the pre-requisite entries and the snapshot result for memory test cases
|
||
snapshot_folder = model_result_dir / "memory_snapshot" / test_category
|
||
if snapshot_folder.exists():
|
||
if not args.allow_overwrite:
|
||
pass
|
||
elif not args.run_ids:
|
||
shutil.rmtree(snapshot_folder)
|
||
else:
|
||
# TODO: If run_ids and id involes prereq entries, we should just delete those snapshot files
|
||
# It's not implemented yet, but it won't affect the accuracy, as those files will be overwritten anyway (assume generation success)
|
||
pass
|
||
|
||
existing_ids = [entry["id"] for entry in existing_result]
|
||
|
||
test_cases_to_generate = [
|
||
test_case
|
||
for test_case in all_test_entries_involved
|
||
if test_case["id"] not in existing_ids
|
||
]
|
||
|
||
# Skip format sensitivity test cases for FC models
|
||
if (
|
||
any(is_format_sensitivity(test_category) for test_category in all_test_categories)
|
||
and MODEL_CONFIG_MAPPING[model_name].is_fc_model
|
||
):
|
||
test_cases_to_generate = [
|
||
test_case
|
||
for test_case in test_cases_to_generate
|
||
if not is_format_sensitivity(test_case["id"])
|
||
]
|
||
|
||
test_cases_to_generate = clean_up_memory_prereq_entries(test_cases_to_generate)
|
||
# TODO: Should we move these to the load_dataset_entry function?
|
||
test_cases_to_generate = populate_initial_settings_for_memory_test_cases(
|
||
test_cases_to_generate, model_result_dir
|
||
)
|
||
test_cases_to_generate = populate_initial_settings_for_web_search_test_cases(
|
||
test_cases_to_generate
|
||
)
|
||
|
||
return sorted(test_cases_to_generate, key=sort_key)
|
||
|
||
|
||
def multi_threaded_inference(handler, test_case, include_input_log, exclude_state_log):
|
||
|
||
assert type(test_case["function"]) is list
|
||
|
||
try:
|
||
result, metadata = handler.inference(
|
||
test_case, include_input_log, exclude_state_log
|
||
)
|
||
except Exception as e:
|
||
# This is usually the case when the model getting stuck on one particular test case.
|
||
# For example, timeout error or FC model returning invalid JSON response.
|
||
# Since temperature is already set to 0.001, retrying the same test case will not help.
|
||
# So we continue the generation process and record the error message as the model response
|
||
error_block = (
|
||
"-" * 100
|
||
+ "\n❗️❗️ Error occurred during inference. Continuing to next test case.\n"
|
||
+ f"❗️❗️ Test case ID: {test_case['id']}, Error: {str(e)}\n"
|
||
+ traceback.format_exc(limit=10)
|
||
+ "-" * 100
|
||
)
|
||
tqdm.write(error_block)
|
||
|
||
result = f"Error during inference: {str(e)}"
|
||
metadata = {"traceback": traceback.format_exc()}
|
||
|
||
result_to_write = {
|
||
"id": test_case["id"],
|
||
"result": result,
|
||
**metadata,
|
||
}
|
||
|
||
return result_to_write
|
||
|
||
|
||
def generate_results(args, model_name, test_cases_total):
|
||
handler = build_handler(model_name, args.temperature)
|
||
|
||
if isinstance(handler, OSSHandler):
|
||
handler: OSSHandler
|
||
is_oss_model = True
|
||
# For OSS models, if the user didn't explicitly set the number of threads,
|
||
# we default to 100 threads to speed up the inference.
|
||
num_threads = (
|
||
args.num_threads
|
||
if args.num_threads is not None
|
||
else LOCAL_SERVER_MAX_CONCURRENT_REQUEST
|
||
)
|
||
else:
|
||
handler: BaseHandler
|
||
is_oss_model = False
|
||
num_threads = args.num_threads if args.num_threads is not None else 1
|
||
|
||
# Use a separate thread to write the results to the file to avoid concurrent IO issues
|
||
def _writer():
|
||
"""Consume result dicts from the queue and write them with exclusive access."""
|
||
while True:
|
||
item = write_queue.get()
|
||
if item is None:
|
||
break
|
||
handler.write(item, result_dir=args.result_dir, update_mode=args.run_ids)
|
||
write_queue.task_done()
|
||
|
||
write_queue: queue.Queue = queue.Queue()
|
||
|
||
writer_thread = threading.Thread(target=_writer, daemon=True)
|
||
writer_thread.start()
|
||
|
||
try:
|
||
if is_oss_model:
|
||
handler.spin_up_local_server(
|
||
num_gpus=args.num_gpus,
|
||
gpu_memory_utilization=args.gpu_memory_utilization,
|
||
backend=args.backend,
|
||
skip_server_setup=args.skip_server_setup,
|
||
local_model_path=args.local_model_path,
|
||
lora_modules=args.lora_modules,
|
||
enable_lora=args.enable_lora,
|
||
max_lora_rank=args.max_lora_rank,
|
||
)
|
||
|
||
# ───── dependency bookkeeping ──────────────────────────────
|
||
dependencies = {
|
||
test_case["id"]: set(test_case.get("depends_on", []))
|
||
for test_case in test_cases_total
|
||
}
|
||
children_of = defaultdict(list)
|
||
for test_case in test_cases_total:
|
||
for dependency_id in test_case.get("depends_on", []):
|
||
children_of[dependency_id].append(test_case["id"])
|
||
|
||
id_to_test_case = {test_case["id"]: test_case for test_case in test_cases_total}
|
||
|
||
ready_queue = [
|
||
(sort_key(id_to_test_case[test_case_id]), test_case_id)
|
||
for test_case_id, dependency_ids in dependencies.items()
|
||
if not dependency_ids
|
||
]
|
||
heapq.heapify(ready_queue)
|
||
in_flight: dict[Future, str] = {} # future -> test_case_id
|
||
completed = set()
|
||
|
||
with ThreadPoolExecutor(max_workers=num_threads) as pool, tqdm(
|
||
total=len(test_cases_total),
|
||
desc=f"Generating results for {model_name}",
|
||
position=0,
|
||
leave=True,
|
||
dynamic_ncols=True,
|
||
mininterval=0.2,
|
||
smoothing=0.1,
|
||
bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]",
|
||
) as pbar:
|
||
|
||
# seed initial ready tasks
|
||
while ready_queue and len(in_flight) < num_threads:
|
||
_, test_case_id = heapq.heappop(ready_queue)
|
||
test_case = id_to_test_case[test_case_id]
|
||
future = pool.submit(
|
||
multi_threaded_inference,
|
||
handler,
|
||
test_case,
|
||
args.include_input_log,
|
||
args.exclude_state_log,
|
||
)
|
||
in_flight[future] = test_case_id
|
||
|
||
# main scheduler loop
|
||
while in_flight:
|
||
done, _ = wait(in_flight, return_when=FIRST_COMPLETED)
|
||
for future in done:
|
||
test_case_id = in_flight.pop(future)
|
||
result_dict = future.result()
|
||
|
||
# Enqueue the result for the writer thread to handle file IO
|
||
write_queue.put(result_dict)
|
||
|
||
# Update progress bar right after inference completes
|
||
pbar.update()
|
||
completed.add(test_case_id)
|
||
|
||
# unlock children
|
||
for child_id in children_of[test_case_id]:
|
||
dependencies[child_id].discard(test_case_id)
|
||
if not dependencies[child_id]:
|
||
heapq.heappush(
|
||
ready_queue,
|
||
(sort_key(id_to_test_case[child_id]), child_id),
|
||
)
|
||
|
||
# refill the pool up to max_workers
|
||
while ready_queue and len(in_flight) < num_threads:
|
||
_, test_case_id = heapq.heappop(ready_queue)
|
||
test_case = id_to_test_case[test_case_id]
|
||
future = pool.submit(
|
||
multi_threaded_inference,
|
||
handler,
|
||
test_case,
|
||
args.include_input_log,
|
||
args.exclude_state_log,
|
||
)
|
||
in_flight[future] = test_case_id
|
||
|
||
finally:
|
||
# Signal writer thread to finish and wait for it
|
||
write_queue.put(None)
|
||
writer_thread.join()
|
||
|
||
if is_oss_model:
|
||
handler.shutdown_local_server()
|
||
|
||
|
||
def main(args):
|
||
|
||
# Note: The following environment variables are needed for the memory vector store implementation
|
||
# Otherwise you get segfault or huggingface tokenizer warnings
|
||
# disable HuggingFace tokenizers’ thread pool
|
||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||
# limit all OpenMP/MKL threads to 1
|
||
os.environ["OMP_NUM_THREADS"] = "1"
|
||
os.environ["MKL_NUM_THREADS"] = "1"
|
||
# use spawn method for multiprocessing
|
||
mp.set_start_method("spawn", force=True)
|
||
|
||
if type(args.model) is not list:
|
||
args.model = [args.model]
|
||
if type(args.test_category) is not list:
|
||
args.test_category = [args.test_category]
|
||
|
||
(
|
||
all_test_categories,
|
||
all_test_entries_involved,
|
||
) = get_involved_test_entries(args.test_category, args.run_ids)
|
||
|
||
for model_name in args.model:
|
||
if model_name not in MODEL_CONFIG_MAPPING:
|
||
raise ValueError(
|
||
f"Unknown model_name '{model_name}'.\n"
|
||
"• For officially supported models, please refer to `SUPPORTED_MODELS.md`.\n"
|
||
"• For running new models, please refer to `README.md` and `CONTRIBUTING.md`."
|
||
)
|
||
tqdm.write(f"Generating results for {args.model}")
|
||
if args.run_ids:
|
||
tqdm.write("Running specific test cases. Ignoring `--test-category` argument.")
|
||
else:
|
||
tqdm.write(f"Running full test cases for categories: {all_test_categories}.")
|
||
|
||
if any(is_format_sensitivity(test_category) for test_category in all_test_categories):
|
||
for model_name in args.model:
|
||
if MODEL_CONFIG_MAPPING[model_name].is_fc_model:
|
||
tqdm.write(
|
||
"⚠️ Warning: Format sensitivity test cases are only supported for prompting (non-FC) models. "
|
||
f"Since {model_name} is a FC model based on its config, the format sensitivity test cases will be skipped."
|
||
)
|
||
|
||
if args.result_dir is not None:
|
||
args.result_dir = PROJECT_ROOT / args.result_dir
|
||
else:
|
||
args.result_dir = RESULT_PATH
|
||
|
||
for model_name in args.model:
|
||
test_cases_total = collect_test_cases(
|
||
args,
|
||
model_name,
|
||
all_test_categories,
|
||
deepcopy(all_test_entries_involved),
|
||
)
|
||
|
||
if len(test_cases_total) == 0:
|
||
tqdm.write(
|
||
f"✅ All selected test cases have been previously generated for {model_name}. No new test cases to generate."
|
||
)
|
||
else:
|
||
generate_results(args, model_name, test_cases_total)
|
||
# Sort the result files by id at the end
|
||
for model_result_json in args.result_dir.rglob(RESULT_FILE_PATTERN):
|
||
sort_file_content_by_id(model_result_json)
|