lm-eval-harness vllm_norm.py 源码精读

@register_model("vllm_norm")
class VLLM(LM):
    _DEFAULT_MAX_LENGTH = 2048

    def __init__(self, url, max_tokens_to_sample: int = 2048, **kwargs):
        super().__init__()

pretrained = os.environ.get(
    "MODEL_DIR", "/data/minimax-dialogue/users/shennai/model/DeepSeek-V3-Base"
)
try:
    self.tokenizer = AutoTokenizer.from_pretrained(pretrained, trust_remote_code=True)
except:
    self.tokenizer = AutoTokenizer.from_pretrained(
        pretrained, use_fast=False, trust_remote_code=True
    )

self.urls = []
server_ip_dir = os.environ.get(
    "SERVER_IP_DIR", "/minimax-dialogue/users/yize/data/server_ip_address"
)
# Wait for SERVER_IP_DIR to be populated (race with server startup).
# Retry every 5 seconds up to 5 minutes, then fail loudly.
_wait_interval_s = 5
_max_wait_s = 5 * 60
_waited_s = 0
while True:
    urls: list[str] = []
    for file in glob(os.path.join(server_ip_dir, "*.txt")):
        try:
            with open(file, "r") as f:
                v = f.read().strip()
            if v:
                urls.append(v)
        except Exception:
            continue

    if urls:
        self.urls = urls
        break

    if _waited_s >= _max_wait_s:
        raise RuntimeError(
            "SERVER_IP_DIR has no non-empty '*.txt' ip files after waiting "
            f"{_max_wait_s}s. SERVER_IP_DIR={server_ip_dir!r}. "
            f"dir_listing={dir_listing!r}"
        )

    time.sleep(_wait_interval_s)
    _waited_s += _wait_interval_s

wait_after_server_start = os.environ.get("WAIT_AFTER_SERVER_START", "0")
try:
    wait_after_server_start_s = float(wait_after_server_start)
except Exception:
    wait_after_server_start_s = 0.0
if wait_after_server_start_s > 0:
    print(
        f"[vllm_norm] WAIT_AFTER_SERVER_START={wait_after_server_start_s}s; sleeping...",
        flush=True,
    )
    time.sleep(wait_after_server_start_s)
    # Re-scan SERVER_IP_DIR after sleeping to pick up newly written ip files.
    refreshed_urls: list[str] = []
    for file in glob(os.path.join(server_ip_dir, "*.txt")):
        try:
            with open(file, "r") as f:
                v = f.read().strip()
            if v:
                refreshed_urls.append(v)
        except Exception:
            continue
    if refreshed_urls:
        self.urls = refreshed_urls

self.port = int(os.getenv("SERVING_PORT", 8000))
self.temperature = 0.01
self.max_tokens = 2000
self.top_p = 0.9
self._max_gen_toks = max_tokens_to_sample
self.thread_num = int(os.environ.get("NUM_PARRALLEL", "8"))

self._include_path = os.environ.get(
    "INCLUDE_PATH",
    "/output/lm-evaluation-harness/lm_eval/tasks/pretrain/",
)
self._is_pretrain = "pretrain" in self._include_path
self._add_assistant_thinking_prefix = (
    int(os.environ.get("VLLM_ADD_THINKING_PREFIX", 0)) == 1
    or int(os.environ.get("MUST_THINK", 0)) == 1
)

sys_prompt_raw = os.environ.get("VLLM_SYSTEM_PROMPT", '')
self._sys_prompt = sys_prompt_raw.encode().decode('unicode_escape') if sys_prompt_raw else ''

self._thinking_message = (
    {"role": "assistant", "content": "<think>\n"}
    if self._add_assistant_thinking_prefix else None
)
self._system_message = (
    [{"role": "system", "content": self._sys_prompt}]
    if self._sys_prompt else None
)

# HTTP 连接池配置优化
self.connector_limit = int(os.environ.get(
    "HTTP_CONNECTOR_LIMIT", str(self.thread_num * 2)
))
self.connector_limit_per_host = int(os.environ.get(
    "HTTP_CONNECTOR_LIMIT_PER_HOST", str(self.thread_num)
))

def _preprocess_payload(self, payload: dict) -> dict:
    """批量预处理 payload，减少每个请求的处理时间"""
    prompt = payload.get("prompt", None)

    payload["model"] = self._model_dir
    if self._is_pretrain:
        payload['logprobs'] = 1
        return payload

    if payload.get("prompt"):
        payload.pop("prompt")
    payload["return_entropy"] = self._return_entropy

    if isinstance(prompt, list):
        payload["messages"] = prompt
    else:
        prefix = self._prefix
        suffix = self._suffix
        if prefix and suffix and prompt.startswith(prefix) and prompt.endswith(suffix):
            prompt = prompt[len(prefix) : -len(suffix)]

        # 构建 messages
        if BOS_TOKEN in prompt or EOS_TOKEN in prompt:
            payload["messages"] = build_serving_from_chatml(prompt)
        else:
            payload["messages"] = [{"role": "user", "content": prompt}]

        if self._add_assistant_thinking_prefix:
            payload['continue_final_message'] = True
            payload['add_generation_prompt'] = False
            payload["messages"].append(self._thinking_message)

        if self._system_message:
            payload["messages"] = self._system_message + payload["messages"]

    payload = specify_extra_kwargs(payload)
    return payload

BOS_TOKEN = ']~b]'
EOS_TOKEN = '[e~['

def build_serving_from_chatml(prompt):
    system_data = []
    settings = prompt.split(BOS_TOKEN)
    for setting in settings:
        if len(setting) == 0:
            continue
        if not setting.startswith('system'):
            break
        setting = setting[len('system '):]
        fields = setting.split('=')
        field_key = fields[0]
        field_value = '='.join(fields[1:]).split('\n')
        if field_key not in ('ai_setting', 'user_setting'):
            continue
        system_data.append({
            'role': 'system',
            field_key: field_value[0],
            'text': '\n'.join(field_value[1:]).replace(EOS_TOKEN, '').strip()
        })

    dialogue_data = []
    dialogues = prompt.split(BOS_TOKEN)[1:]
    for i, dialogue in enumerate(dialogues):
        if len(dialogue) == 0 or dialogue.startswith('system'):
            continue
        role = dialogue.split(' ')[0]
        fields = dialogue.split('\n')
        dialogue_data.append({
            'role': role,
            'name': fields[0].split('name=')[1],
            'text': '\n'.join(fields[1:]).replace(EOS_TOKEN, '').strip(),
            'trainable': False
        })

    if len(system_data) > 0:
        conversations = [{'role': 'system', 'content': system_data[0]['text']}]
    else:
        conversations = [{'role': 'system', 'content': 'You are a helpful assistant.'}]
    for dd in dialogue_data:
        conversations.append({
            'role': 'user' if dd['role'] == 'user' else 'assistant',
            'content': dd['text']
        })
    if conversations[-1]['role'] == 'assistant':
        conversations = conversations[:-1]
    return conversations

def specify_extra_kwargs(payload):
    timeout = os.environ.get("TIMEOUT", None)
    if timeout:
        payload["timeout"] = timeout

    reasoning_effort = os.environ.get("REASONING_EFFORT", None)
    if reasoning_effort:
        payload["reasoning_effort"] = reasoning_effort

    payload["chat_template_kwargs"] = {}

    chat_template_file = os.environ.get("TOKENIZER_CHAT_TEMPLATE_FILE", None)
    if chat_template_file:
        with open(chat_template_file, "r") as f:
            chat_template = f.read().strip()
            payload["chat_template_kwargs"]["chat_template"] = chat_template

    if os.environ.get("ENABLE_THINKING", False):
        payload["chat_template_kwargs"]["enable_thinking"] = True
    return payload

def _collate(x):
    # the negative sign on len(toks) sorts descending
    toks = self.tok_encode(str(x[0]["text"]))
    return -len(toks), x[0]["text"]

re_ords = Collator([reg.args for reg in requests], _collate, grouping=True)
chunks = re_ords.get_batched(n=len(requests), batch_fn=None)

try:
    # vlm任务
    contexts, all_gen_kwargs, doc_to_visual, docs = zip(*chunk)
    visuals = [
        doc_to_visual[i](doc) for i, doc in enumerate(docs)
    ]
    vision = True
except:
    contexts, all_gen_kwargs = zip(*chunk)
    vision = False

gen_kwargs = all_gen_kwargs[0]
kwargs = copy.deepcopy(gen_kwargs)
if "until" in kwargs.keys():
    until = kwargs.pop("until")
    if isinstance(until, str):
        until = [until]
    elif not isinstance(until, list):
        raise ValueError(
            f"Expected `kwargs['until']` to be of type Union[str,list] but got {until}"
        )

if not until:
    until = [self.tokenizer.decode(self.eot_token_id)]
else:
    until += [self.tokenizer.decode(self.eot_token_id)]

if "max_gen_toks" in kwargs.keys():
    max_gen_toks = kwargs.pop("max_gen_toks")
else:
    max_gen_toks = self.max_gen_toks

pload = {
    "prompt": "",
    "temperature": kwargs.get("temperature", self.temperature),
    "top_p": kwargs.get("top_p", self.top_p),
    "max_tokens": max_gen_toks,
    "stream": False,
    "stop": until,
}
for i, context in enumerate(contexts):
    if vision and DEFAULT_IMAGE_TOKEN not in context:
        image_tokens = [DEFAULT_IMAGE_TOKEN] * len(visuals[i])
        image_tokens = " ".join(image_tokens)
        context = f"{image_tokens}\n{context}"

    if vision:
        text = sft_prefix + context + sft_suffix
    else:
        text = context["text"]
    pload["prompt"] = text

    if vision:
        pload["multi_modal_data"] = dict(
            image=[
                self._image_to_base64(visuals[i][j])
                for j in range(len(visuals[i]))
            ]
        )

    requests.append(copy.deepcopy(pload))

# perform batched generation
if os.environ.get("USE_ASYNC"):
    cont, log_probs_list = self._model_generate_async(
        requests=requests,
        generate=True,
        max_tokens=max_gen_toks,
        stop=until,
        **kwargs,
    )
else:
    cont, log_probs_list = self._model_generate(
        requests=requests,
        generate=True,
        max_tokens=max_gen_toks,
        stop=until,
        **kwargs,
    )

def loglikelihood(self, requests: List[Instance]) -> List[Tuple[float, bool]]:
    new_reqs = []
    for i, req in enumerate(requests):
        context, continuation, gen_args = req.args
        context = context["text"]

        # 处理字符位置标记的格式
        if isinstance(context, dict) and context.get("__char_spans__"):
            full_text = context["full_text"]
            target_char_spans = context["target_char_spans"]
            tokens, mask = self._char_spans_to_token_mask(
                full_text, target_char_spans
            )
            new_reqs.append(((context, continuation), tokens, mask))
            continue

        elif isinstance(context, list):  # encoded texts with masks
            text = context
            mask = continuation

        else:
            if context == "":
                context_enc, continuation_enc = (
                    [self.eot_token_id],
                    self.tok_encode(continuation),
                )
            else:
                strip_context = gen_args.get("strip_context", True) if gen_args else True
                context_enc, continuation_enc = self._encode_pair(
                    context, continuation, strip_context
                )
            text = context_enc + continuation_enc
            mask = [0] * len(context_enc) + [1] * len(continuation_enc)

        new_reqs.append(((context, continuation), text, mask))
    return self._loglikelihood_tokens(new_reqs)

def _loglikelihood_tokens(
    self,
    requests: List[Tuple[Tuple[str, str], List[int], List[int]]],
    disable_tqdm: bool = False,
) -> List[Tuple[float, bool]]:
    res = []

    def _collate(x):
        toks = x[1] + x[2]
        return -len(toks), tuple(toks)

    re_ord = Collator(requests, sort_fn=_collate)
    chunks = re_ord.get_batched(n=len(requests), batch_fn=None)

    for index, chunk in enumerate(chunks):
        inputs = []
        masks = []
        ctxlens = []
        for cache_key, text, mask in chunk:
            ctxlen = mask.index(1)
            try:
                prompt = cache_key[0] + cache_key[1]
            except:
                prompt = cache_key[0]['full_text']
            pload = {
                "prompt": prompt,
                "max_tokens": 1,
                "stream": False,
                "logprobs": 1,
                "echo": True,
            }
            inputs.append(pload)
            masks.append(mask)
            ctxlens.append(ctxlen)

        outputs, _ = self._model_generate(requests=inputs, generate=False)

        for output, mask, ctxlen, (cache_key, text, _), pload in zip(
            outputs, masks, ctxlens, chunk, inputs
        ):
            answer = self._parse_logprobs(
                tokens=text, outputs=output, ctxlen=ctxlen, masks=mask
            )
            res.append(answer)

    return re_ord.get_original(res)

@staticmethod
def _parse_logprobs(
    tokens: List, outputs, ctxlen: int, masks: List
) -> Tuple[float, bool, str]:
    # 过滤 None 值（第一个 token 无条件概率）
    if len(outputs["logprobs"]["token_logprobs"]) != len(tokens):
        outputs["logprobs"]["token_logprobs"] = [
            logprob_dict for logprob_dict in outputs["logprobs"]["token_logprobs"]
            if logprob_dict is not None
        ]
    continuation_logprobs_dicts = convert_keys_to_int(outputs["logprobs"])

    # 计算 continuation_logprobs：只累加 mask=1 部分的 logprob
    continuation_logprobs = sum(
        logprob_dict[1] * mask
        for token, logprob_dict, mask in zip(
            tokens[ctxlen:],
            continuation_logprobs_dicts[ctxlen:],
            masks[ctxlen:]
        )
    )

    # 判断 is_greedy：top logprob token 是否匹配 continuation token
    toplogprobs = outputs["logprobs"]["top_logprobs"]
    is_greedy = True
    for toplogprob, logprob_dict, mask in zip(
        toplogprobs[ctxlen:], continuation_logprobs_dicts[ctxlen:], masks[ctxlen:]
    ):
        if toplogprob and mask:
            top_token = max(toplogprob, key=lambda x: toplogprob[x])
            if top_token != logprob_dict[0]:
                is_greedy = False
                break

    return (continuation_logprobs, is_greedy, json.dumps(info_dict_list))

def _model_generate_async(
    self,
    requests: List[List[int]] = None,
    generate: bool = False,
    max_tokens: int = None,
    stop: Optional[List[str]] = None,
    **kwargs,
):
    if generate:
        sampling_kwargs = deepcopy(kwargs)
        sampling_kwargs["max_tokens"] = max_tokens
        if "max_new_tokens" in sampling_kwargs:
            sampling_kwargs["max_tokens"] = sampling_kwargs["max_new_tokens"]
            sampling_kwargs.pop("max_new_tokens")
        sampling_kwargs["stop"] = stop
    else:
        sampling_kwargs = dict(temperature=0, prompt_logprobs=2, max_tokens=1)

    res = [None] * len(requests)
    log_probs_list = [None] * len(requests)
    pbar = tqdm(total=len(requests), disable=(self.rank != 0))
    assert len(self.urls) == 1, f"Only one URL is supported"
    url = self.urls[0]
    request_tuples = [(_idx, req) for _idx, req in enumerate(requests)]
    asyncio.run(
        self.__post_requests_to_server_async(url, request_tuples, pbar, res, log_probs_list)
    )
    return res, log_probs_list

timeout = aiohttp.ClientTimeout(total=None)
connector = aiohttp.TCPConnector(
    limit=self.connector_limit,            # 总连接池大小
    limit_per_host=self.connector_limit_per_host,  # 每个host的连接数
    ttl_dns_cache=1e9,                     # DNS缓存时间（秒）
    keepalive_timeout=300,                 # 连接保活时间（秒）
    force_close=False,                     # 保持连接复用
    enable_cleanup_closed=True,            # 清理已关闭的连接
    use_dns_cache=True,                    # 启用DNS缓存
)

semaphore = asyncio.Semaphore(self.thread_num)
async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session:
    tasks = [
        async_post_func(url_completions, url_chat, session, semaphore, req, ...)
        for req in preprocessed_tuples
    ]
    await asyncio.gather(*tasks)

def _model_generate(
    self,
    requests: List[List[int]] = None,
    generate: bool = False,
    max_tokens: int = None,
    stop: Optional[List[str]] = None,
    **kwargs,
):
    if generate:
        sampling_kwargs = deepcopy(kwargs)
        sampling_kwargs["max_tokens"] = max_tokens
        sampling_kwargs["stop"] = stop
    else:
        sampling_kwargs = dict(temperature=0, prompt_logprobs=2, max_tokens=1)

    request_queue = queue.Queue()
    res = [None] * len(requests)
    log_probs_list = [None] * len(requests)
    for _idx, req in enumerate(requests):
        request_queue.put((_idx, req))

    # 每个 URL 一个线程，随机打乱避免热点
    with ThreadPoolExecutor(max_workers=len(self.urls)) as executor:
        futures = []
        for url in random.sample(self.urls, len(self.urls)):
            f = executor.submit(
                self.__post_requests_to_server, url, request_queue, res, log_probs_list
            )
            futures.append(f)
        for f in futures:
            f.result()

    return res, log_probs_list

def __post_requests_to_server(self, url, request_queue, results, log_probs_list):
    # 每个线程内部再创建 thread_num 个子线程
    with ThreadPoolExecutor(max_workers=self.thread_num) as executor:
        futures = []
        for _ in range(self.thread_num):
            f = executor.submit(
                self.__post_one_request, url, request_queue, results, log_probs_list
            )
            futures.append(f)
        for f in futures:
            f.result()

error_message = response.json()["message"]
if (
    "Please reduce the length of the messages or completion" in error_message
    or "Please reduce the number of tokens in the input messages or the completion"
       in error_message
):
    numbers = re.findall(r'\d+', error_message)
    if int(numbers[0]) > int(numbers[2]):
        payload["max_tokens"] = int(numbers[0]) - int(numbers[2]) - 1
        print(f"payload['max_tokens'] is set to: {payload['max_tokens']}")
    else:
        response = "输入token长度超过模型长度限制！！！！"
        break

response_message = response_json["choices"][0]["message"]
if response_message.get("reasoning_content", None) is not None:
    response.update({
        "content": response_message["content"],
        "reasoning_content": response_message.get("reasoning_content", None)
    })
else:
    response.update({"content": response_message["content"]})

if self._return_entropy:
    entropy_list = response_json["choices"][0]['entropy']
    entropy_list = entropy_list[-num_completion_tokens:]
    response.update({
        'entropy_list': [np.mean([e for e in entropy_list if isinstance(e, float)])]
    })

assert int(os.environ.get("ADD_BOD", 0)) == 0, \
    "ADD_BOD must be 0 for vllm_norm, which should be supported in hf chat_template."

prefix = self._prefix
suffix = self._suffix
if prefix and suffix and prompt.startswith(prefix) and prompt.endswith(suffix):
    prompt = prompt[len(prefix) : -len(suffix)]

def _image_to_base64(self, img, format="jpeg"):
    try:
        buffer = BytesIO()
        img.save(buffer, format=format)
        img_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
        return img_str
    except Exception as e:
        print(f"转换过程出错: {str(e)}")
        return None

def _char_spans_to_token_mask(self, full_text, target_char_spans):
    enc = self.tokenizer(
        full_text, return_offsets_mapping=True, add_special_tokens=False
    )
    mask = [0] * len(enc["input_ids"])
    offsets = enc["offset_mapping"]

    for i, (start, end) in enumerate(offsets):
        for t_start, t_end in target_char_spans:
            # 区间重叠判断
            if max(start, t_start) < min(end, t_end):
                mask[i] = 1
                break

    return enc["input_ids"], mask