mirror of
https://github.com/mudler/LocalAI.git
synced 2026-06-25 00:59:28 -04:00
5e3774dfe3
Audit of the Score path against vllm-metal (MLX on macOS): the engine accepts SamplingParams(prompt_logprobs=1) but returns an all-None prompt_logprobs list rather than computing it, so scoring is not supported there. The old guard treated the truthy [None] list as valid and silently scored every candidate as 0. Detect the all-None case and return UNIMPLEMENTED instead. No-op on Linux/CUDA, which populate real entries. Signed-off-by: Ettore Di Giacinto <mudler@localai.io> Assisted-by: Claude:opus-4.8 [Claude Code]
931 lines
39 KiB
Python
931 lines
39 KiB
Python
#!/usr/bin/env python3
|
|
import asyncio
|
|
import dataclasses
|
|
import difflib
|
|
from concurrent import futures
|
|
import argparse
|
|
import signal
|
|
import sys
|
|
import os
|
|
import json
|
|
import time
|
|
import gc
|
|
import tempfile
|
|
from typing import List
|
|
from PIL import Image
|
|
|
|
import backend_pb2
|
|
import backend_pb2_grpc
|
|
|
|
import grpc
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common'))
|
|
from grpc_auth import get_auth_interceptors
|
|
|
|
from vllm.engine.arg_utils import AsyncEngineArgs
|
|
from vllm.engine.async_llm_engine import AsyncLLMEngine
|
|
from vllm.sampling_params import SamplingParams
|
|
from vllm.utils import random_uuid
|
|
try:
|
|
from vllm.tokenizers import get_tokenizer # vLLM >= 0.22
|
|
except ImportError:
|
|
from vllm.transformers_utils.tokenizer import get_tokenizer # vLLM < 0.22
|
|
from vllm.multimodal.utils import fetch_image
|
|
from vllm.assets.video import VideoAsset
|
|
import base64
|
|
import io
|
|
|
|
# Version-compat imports — wrap in try/except for older vLLM versions
|
|
try:
|
|
from vllm.tool_parsers import ToolParserManager
|
|
HAS_TOOL_PARSERS = True
|
|
except ImportError:
|
|
HAS_TOOL_PARSERS = False
|
|
|
|
try:
|
|
from vllm.reasoning import ReasoningParserManager
|
|
HAS_REASONING_PARSERS = True
|
|
except ImportError:
|
|
HAS_REASONING_PARSERS = False
|
|
|
|
# vLLM >= 0.23 renamed GuidedDecodingParams -> StructuredOutputsParams and the
|
|
# SamplingParams field guided_decoding -> structured_outputs.
|
|
try:
|
|
from vllm.sampling_params import StructuredOutputsParams
|
|
HAS_GUIDED_DECODING = True
|
|
except ImportError:
|
|
HAS_GUIDED_DECODING = False
|
|
|
|
_ONE_DAY_IN_SECONDS = 60 * 60 * 24
|
|
|
|
# If MAX_WORKERS are specified in the environment use it, otherwise default to 1
|
|
MAX_WORKERS = int(os.environ.get('PYTHON_GRPC_MAX_WORKERS', '1'))
|
|
|
|
# Implement the BackendServicer class with the service methods
|
|
class BackendServicer(backend_pb2_grpc.BackendServicer):
|
|
"""
|
|
A gRPC servicer that implements the Backend service defined in backend.proto.
|
|
"""
|
|
def generate(self,prompt, max_new_tokens):
|
|
"""
|
|
Generates text based on the given prompt and maximum number of new tokens.
|
|
|
|
Args:
|
|
prompt (str): The prompt to generate text from.
|
|
max_new_tokens (int): The maximum number of new tokens to generate.
|
|
|
|
Returns:
|
|
str: The generated text.
|
|
"""
|
|
self.generator.end_beam_search()
|
|
|
|
# Tokenizing the input
|
|
ids = self.generator.tokenizer.encode(prompt)
|
|
|
|
self.generator.gen_begin_reuse(ids)
|
|
initial_len = self.generator.sequence[0].shape[0]
|
|
has_leading_space = False
|
|
decoded_text = ''
|
|
for i in range(max_new_tokens):
|
|
token = self.generator.gen_single_token()
|
|
if i == 0 and self.generator.tokenizer.tokenizer.IdToPiece(int(token)).startswith('▁'):
|
|
has_leading_space = True
|
|
|
|
decoded_text = self.generator.tokenizer.decode(self.generator.sequence[0][initial_len:])
|
|
if has_leading_space:
|
|
decoded_text = ' ' + decoded_text
|
|
|
|
if token.item() == self.generator.tokenizer.eos_token_id:
|
|
break
|
|
return decoded_text
|
|
|
|
def _parse_options(self, options_list):
|
|
"""Parse Options[] key:value string list into a dict."""
|
|
opts = {}
|
|
for opt in options_list:
|
|
if ":" not in opt:
|
|
continue
|
|
key, value = opt.split(":", 1)
|
|
opts[key.strip()] = value.strip()
|
|
return opts
|
|
|
|
def _apply_engine_args(self, engine_args, engine_args_json):
|
|
"""Apply user-supplied engine_args (JSON object) onto an AsyncEngineArgs.
|
|
|
|
Returns a new AsyncEngineArgs with the typed fields preserved and the
|
|
user's overrides layered on top. Uses ``dataclasses.replace`` so vLLM's
|
|
``__post_init__`` re-runs and auto-converts dict-valued fields like
|
|
``compilation_config`` / ``attention_config`` into their dataclass form.
|
|
``speculative_config`` and ``kv_transfer_config`` are accepted as dicts
|
|
directly (vLLM converts them at engine init).
|
|
|
|
Unknown keys raise ValueError with the closest valid field as a hint.
|
|
"""
|
|
if not engine_args_json:
|
|
return engine_args
|
|
try:
|
|
extra = json.loads(engine_args_json)
|
|
except json.JSONDecodeError as e:
|
|
raise ValueError(f"engine_args is not valid JSON: {e}") from e
|
|
if not isinstance(extra, dict):
|
|
raise ValueError(
|
|
f"engine_args must be a JSON object, got {type(extra).__name__}"
|
|
)
|
|
valid = {f.name for f in dataclasses.fields(type(engine_args))}
|
|
for key in extra:
|
|
if key not in valid:
|
|
suggestion = difflib.get_close_matches(key, valid, n=1)
|
|
hint = f" did you mean {suggestion[0]!r}?" if suggestion else ""
|
|
raise ValueError(f"unknown engine_args key {key!r}.{hint}")
|
|
return dataclasses.replace(engine_args, **extra)
|
|
|
|
def _messages_to_dicts(self, messages):
|
|
"""Convert proto Messages to list of dicts suitable for apply_chat_template()."""
|
|
result = []
|
|
for msg in messages:
|
|
d = {"role": msg.role, "content": msg.content or ""}
|
|
if msg.name:
|
|
d["name"] = msg.name
|
|
if msg.tool_call_id:
|
|
d["tool_call_id"] = msg.tool_call_id
|
|
if msg.reasoning_content:
|
|
d["reasoning_content"] = msg.reasoning_content
|
|
if msg.tool_calls:
|
|
try:
|
|
tool_calls = json.loads(msg.tool_calls)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
else:
|
|
# OpenAI wire format carries function.arguments as a
|
|
# JSON-encoded string, but chat templates (e.g. Qwen3)
|
|
# iterate over it as a mapping. vLLM's own OpenAI server
|
|
# parses arguments before applying the template, so do
|
|
# the same here.
|
|
if isinstance(tool_calls, list):
|
|
for tc in tool_calls:
|
|
func = tc.get("function") if isinstance(tc, dict) else None
|
|
if isinstance(func, dict) and isinstance(func.get("arguments"), str):
|
|
try:
|
|
func["arguments"] = json.loads(func["arguments"])
|
|
except json.JSONDecodeError:
|
|
pass
|
|
d["tool_calls"] = tool_calls
|
|
result.append(d)
|
|
return result
|
|
|
|
def Health(self, request, context):
|
|
"""
|
|
Returns a health check message.
|
|
|
|
Args:
|
|
request: The health check request.
|
|
context: The gRPC context.
|
|
|
|
Returns:
|
|
backend_pb2.Reply: The health check reply.
|
|
"""
|
|
return backend_pb2.Reply(message=bytes("OK", 'utf-8'))
|
|
|
|
async def LoadModel(self, request, context):
|
|
"""
|
|
Loads a language model.
|
|
|
|
Args:
|
|
request: The load model request.
|
|
context: The gRPC context.
|
|
|
|
Returns:
|
|
backend_pb2.Result: The load model result.
|
|
"""
|
|
engine_args = AsyncEngineArgs(
|
|
model=request.Model,
|
|
)
|
|
|
|
if request.Quantization != "":
|
|
engine_args.quantization = request.Quantization
|
|
if request.LoadFormat != "":
|
|
engine_args.load_format = request.LoadFormat
|
|
if request.GPUMemoryUtilization != 0:
|
|
engine_args.gpu_memory_utilization = request.GPUMemoryUtilization
|
|
if request.TrustRemoteCode:
|
|
engine_args.trust_remote_code = request.TrustRemoteCode
|
|
if request.EnforceEager:
|
|
engine_args.enforce_eager = request.EnforceEager
|
|
if request.TensorParallelSize:
|
|
engine_args.tensor_parallel_size = request.TensorParallelSize
|
|
if request.SwapSpace != 0:
|
|
engine_args.swap_space = request.SwapSpace
|
|
if request.MaxModelLen != 0:
|
|
engine_args.max_model_len = request.MaxModelLen
|
|
if request.DisableLogStatus:
|
|
engine_args.disable_log_status = request.DisableLogStatus
|
|
if request.DType != "":
|
|
engine_args.dtype = request.DType
|
|
if request.LimitImagePerPrompt != 0 or request.LimitVideoPerPrompt != 0 or request.LimitAudioPerPrompt != 0:
|
|
# limit-mm-per-prompt defaults to 1 per modality, based on vLLM docs
|
|
engine_args.limit_mm_per_prompt = {
|
|
"image": max(request.LimitImagePerPrompt, 1),
|
|
"video": max(request.LimitVideoPerPrompt, 1),
|
|
"audio": max(request.LimitAudioPerPrompt, 1)
|
|
}
|
|
|
|
# engine_args from YAML overrides typed fields above so operators can
|
|
# tune anything the AsyncEngineArgs dataclass exposes without waiting
|
|
# on protobuf changes.
|
|
try:
|
|
engine_args = self._apply_engine_args(engine_args, request.EngineArgs)
|
|
except ValueError as err:
|
|
print(f"engine_args error: {err}", file=sys.stderr)
|
|
return backend_pb2.Result(success=False, message=str(err))
|
|
|
|
try:
|
|
self.llm = AsyncLLMEngine.from_engine_args(engine_args)
|
|
except Exception as err:
|
|
print(f"Unexpected {err=}, {type(err)=}", file=sys.stderr)
|
|
return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")
|
|
|
|
try:
|
|
# vLLM >= 0.14 removed get_model_config() on AsyncLLM; the tokenizer
|
|
# is either already loaded on the engine or can be built from the
|
|
# Model name directly.
|
|
tokenizer = None
|
|
if hasattr(self.llm, "get_tokenizer"):
|
|
try:
|
|
tokenizer = await self.llm.get_tokenizer()
|
|
except TypeError:
|
|
tokenizer = self.llm.get_tokenizer()
|
|
except Exception:
|
|
tokenizer = None
|
|
if tokenizer is None and hasattr(self.llm, "tokenizer"):
|
|
tokenizer = self.llm.tokenizer
|
|
if tokenizer is None:
|
|
tokenizer = get_tokenizer(
|
|
request.Model,
|
|
trust_remote_code=bool(request.TrustRemoteCode),
|
|
truncation_side="left",
|
|
)
|
|
self.tokenizer = tokenizer
|
|
except Exception as err:
|
|
return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")
|
|
|
|
# Parse options for parser selection
|
|
opts = self._parse_options(request.Options)
|
|
|
|
# Instantiate tool/reasoning parser classes (they'll be instantiated per-request with tokenizer)
|
|
self.tool_parser_cls = None
|
|
self.reasoning_parser_cls = None
|
|
if HAS_TOOL_PARSERS and opts.get("tool_parser"):
|
|
try:
|
|
self.tool_parser_cls = ToolParserManager.get_tool_parser(opts["tool_parser"])
|
|
print(f"Loaded tool_parser: {opts['tool_parser']}", file=sys.stderr)
|
|
except Exception as e:
|
|
print(f"Failed to load tool_parser {opts.get('tool_parser')}: {e}", file=sys.stderr)
|
|
|
|
if HAS_REASONING_PARSERS and opts.get("reasoning_parser"):
|
|
try:
|
|
self.reasoning_parser_cls = ReasoningParserManager.get_reasoning_parser(opts["reasoning_parser"])
|
|
print(f"Loaded reasoning_parser: {opts['reasoning_parser']}", file=sys.stderr)
|
|
except Exception as e:
|
|
print(f"Failed to load reasoning_parser {opts.get('reasoning_parser')}: {e}", file=sys.stderr)
|
|
|
|
print("Model loaded successfully", file=sys.stderr)
|
|
return backend_pb2.Result(message="Model loaded successfully", success=True)
|
|
|
|
async def Predict(self, request, context):
|
|
"""
|
|
Generates text based on the given prompt and sampling parameters.
|
|
|
|
Args:
|
|
request: The predict request.
|
|
context: The gRPC context.
|
|
|
|
Returns:
|
|
backend_pb2.Reply: The predict result.
|
|
"""
|
|
gen = self._predict(request, context, streaming=False)
|
|
res = await gen.__anext__()
|
|
return res
|
|
|
|
def Embedding(self, request, context):
|
|
"""
|
|
A gRPC method that calculates embeddings for a given sentence.
|
|
|
|
Args:
|
|
request: An EmbeddingRequest object that contains the request parameters.
|
|
context: A grpc.ServicerContext object that provides information about the RPC.
|
|
|
|
Returns:
|
|
An EmbeddingResult object that contains the calculated embeddings.
|
|
"""
|
|
print("Calculated embeddings for: " + request.Embeddings, file=sys.stderr)
|
|
outputs = self.model.encode(request.Embeddings)
|
|
# Check if we have one result at least
|
|
if len(outputs) == 0:
|
|
context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
|
|
context.set_details("No embeddings were calculated.")
|
|
return backend_pb2.EmbeddingResult()
|
|
return backend_pb2.EmbeddingResult(embeddings=outputs[0].outputs.embedding)
|
|
|
|
async def PredictStream(self, request, context):
|
|
"""
|
|
Generates text based on the given prompt and sampling parameters, and streams the results.
|
|
|
|
Args:
|
|
request: The predict stream request.
|
|
context: The gRPC context.
|
|
|
|
Returns:
|
|
backend_pb2.Result: The predict stream result.
|
|
"""
|
|
iterations = self._predict(request, context, streaming=True)
|
|
try:
|
|
async for iteration in iterations:
|
|
yield iteration
|
|
finally:
|
|
await iterations.aclose()
|
|
|
|
async def TokenizeString(self, request, context):
|
|
if not hasattr(self, 'tokenizer') or self.tokenizer is None:
|
|
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
|
|
context.set_details("Model/tokenizer not loaded")
|
|
return backend_pb2.TokenizationResponse()
|
|
try:
|
|
tokens = self.tokenizer.encode(request.Prompt)
|
|
return backend_pb2.TokenizationResponse(length=len(tokens), tokens=tokens)
|
|
except Exception as e:
|
|
context.set_code(grpc.StatusCode.INTERNAL)
|
|
context.set_details(str(e))
|
|
return backend_pb2.TokenizationResponse()
|
|
|
|
async def Free(self, request, context):
|
|
try:
|
|
if hasattr(self, 'llm'):
|
|
del self.llm
|
|
if hasattr(self, 'tokenizer'):
|
|
del self.tokenizer
|
|
self.tool_parser_cls = None
|
|
self.reasoning_parser_cls = None
|
|
gc.collect()
|
|
try:
|
|
import torch
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
except ImportError:
|
|
pass
|
|
return backend_pb2.Result(success=True, message="Model freed")
|
|
except Exception as e:
|
|
return backend_pb2.Result(success=False, message=str(e))
|
|
|
|
async def Score(self, request, context):
|
|
"""
|
|
Joint log-probability of each candidate continuation given the
|
|
shared prompt. Used by routing-policy multi-label classification
|
|
(read the distribution rather than asking the model to emit a
|
|
single argmax label), reranking, and reward-model scoring.
|
|
|
|
Implementation uses vLLM's `prompt_logprobs` to recover the
|
|
per-token log P(token_i | tokens_<i) for the full concatenated
|
|
sequence; the candidate's tokens are the suffix whose logprobs
|
|
get summed. max_tokens=1 because vLLM requires at least one
|
|
generated token; the generated token is discarded.
|
|
"""
|
|
if not hasattr(self, 'llm') or self.llm is None:
|
|
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
|
|
context.set_details("Model not loaded")
|
|
return backend_pb2.ScoreResponse()
|
|
if not hasattr(self, 'tokenizer') or self.tokenizer is None:
|
|
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
|
|
context.set_details("Tokenizer not available")
|
|
return backend_pb2.ScoreResponse()
|
|
if len(request.candidates) == 0:
|
|
context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
|
|
context.set_details("candidates must be non-empty")
|
|
return backend_pb2.ScoreResponse()
|
|
|
|
try:
|
|
prompt = request.prompt or ""
|
|
prompt_token_ids = self.tokenizer.encode(prompt)
|
|
prompt_len = len(prompt_token_ids)
|
|
results = []
|
|
|
|
for candidate in request.candidates:
|
|
# Tokenise the concatenated sequence. We can't naively
|
|
# use len(prompt_tokens) + len(tokenizer.encode(candidate))
|
|
# because BPE merges at the boundary may produce a
|
|
# different tokenisation. Encoding the joined text and
|
|
# walking the divergence point is the correct primitive.
|
|
full_text = prompt + candidate
|
|
full_token_ids = self.tokenizer.encode(full_text)
|
|
|
|
divergence = prompt_len
|
|
min_len = min(prompt_len, len(full_token_ids))
|
|
for i in range(min_len):
|
|
if prompt_token_ids[i] != full_token_ids[i]:
|
|
divergence = i
|
|
break
|
|
|
|
candidate_token_ids = full_token_ids[divergence:]
|
|
num_candidate_tokens = len(candidate_token_ids)
|
|
if num_candidate_tokens == 0:
|
|
results.append(backend_pb2.CandidateScore(
|
|
log_prob=0.0,
|
|
length_normalized_log_prob=0.0,
|
|
num_tokens=0,
|
|
))
|
|
continue
|
|
|
|
sampling = SamplingParams(
|
|
max_tokens=1,
|
|
temperature=0.0,
|
|
prompt_logprobs=1,
|
|
detokenize=False,
|
|
)
|
|
|
|
request_id = random_uuid()
|
|
last_output = None
|
|
outputs_iter = self.llm.generate(
|
|
{"prompt": full_text},
|
|
sampling_params=sampling,
|
|
request_id=request_id,
|
|
)
|
|
try:
|
|
async for out in outputs_iter:
|
|
last_output = out
|
|
finally:
|
|
try:
|
|
await outputs_iter.aclose()
|
|
except Exception:
|
|
pass
|
|
|
|
_pl = getattr(last_output, "prompt_logprobs", None) if last_output is not None else None
|
|
# Some engines accept the prompt_logprobs request but return a
|
|
# list of all-None entries instead of computing them (observed
|
|
# with vllm-metal's MLX backend on macOS). Treat that as
|
|
# unsupported rather than silently scoring every candidate as 0.
|
|
if not _pl or all(e is None for e in _pl):
|
|
context.set_code(grpc.StatusCode.UNIMPLEMENTED)
|
|
context.set_details("This backend did not return prompt_logprobs; scoring is unsupported on this engine (e.g. vllm-metal / MLX on macOS).")
|
|
return backend_pb2.ScoreResponse()
|
|
|
|
prompt_logprobs = last_output.prompt_logprobs
|
|
total = 0.0
|
|
tokens_proto = []
|
|
for offset, tok_id in enumerate(candidate_token_ids):
|
|
position = divergence + offset
|
|
if position >= len(prompt_logprobs) or prompt_logprobs[position] is None:
|
|
continue
|
|
entry = prompt_logprobs[position]
|
|
lp_obj = entry.get(tok_id)
|
|
if lp_obj is not None:
|
|
lp = lp_obj.logprob
|
|
else:
|
|
# Token not in top-K; vLLM's top-1 may miss it.
|
|
# Fall back to the lowest available logprob in the
|
|
# entry — a conservative lower-bound on the true
|
|
# log P, biased against this candidate.
|
|
lp = min(v.logprob for v in entry.values())
|
|
total += lp
|
|
if request.include_token_logprobs:
|
|
tokens_proto.append(backend_pb2.TokenLogProb(
|
|
token=self.tokenizer.decode([tok_id]),
|
|
log_prob=lp,
|
|
))
|
|
|
|
cs = backend_pb2.CandidateScore(
|
|
log_prob=total,
|
|
num_tokens=num_candidate_tokens,
|
|
)
|
|
if request.length_normalize and num_candidate_tokens > 0:
|
|
cs.length_normalized_log_prob = total / num_candidate_tokens
|
|
if tokens_proto:
|
|
cs.tokens.extend(tokens_proto)
|
|
results.append(cs)
|
|
|
|
return backend_pb2.ScoreResponse(candidates=results)
|
|
except Exception as e:
|
|
print(f"Score error: {e}", file=sys.stderr)
|
|
context.set_code(grpc.StatusCode.INTERNAL)
|
|
context.set_details(str(e))
|
|
return backend_pb2.ScoreResponse()
|
|
|
|
async def _predict(self, request, context, streaming=False):
|
|
# Build the sampling parameters
|
|
# NOTE: this must stay in sync with the vllm backend
|
|
request_to_sampling_params = {
|
|
"N": "n",
|
|
"PresencePenalty": "presence_penalty",
|
|
"FrequencyPenalty": "frequency_penalty",
|
|
"RepetitionPenalty": "repetition_penalty",
|
|
"Temperature": "temperature",
|
|
"TopP": "top_p",
|
|
"TopK": "top_k",
|
|
"MinP": "min_p",
|
|
"Seed": "seed",
|
|
"StopPrompts": "stop",
|
|
"StopTokenIds": "stop_token_ids",
|
|
"BadWords": "bad_words",
|
|
"IncludeStopStrInOutput": "include_stop_str_in_output",
|
|
"IgnoreEOS": "ignore_eos",
|
|
"Tokens": "max_tokens",
|
|
"MinTokens": "min_tokens",
|
|
"Logprobs": "logprobs",
|
|
"PromptLogprobs": "prompt_logprobs",
|
|
"SkipSpecialTokens": "skip_special_tokens",
|
|
"SpacesBetweenSpecialTokens": "spaces_between_special_tokens",
|
|
"TruncatePromptTokens": "truncate_prompt_tokens",
|
|
}
|
|
|
|
sampling_params = SamplingParams(top_p=0.9, max_tokens=200)
|
|
|
|
for request_field, param_field in request_to_sampling_params.items():
|
|
if hasattr(request, request_field):
|
|
value = getattr(request, request_field)
|
|
if value not in (None, 0, [], False):
|
|
setattr(sampling_params, param_field, value)
|
|
|
|
# Structured-output decoding: use Grammar field to pass JSON schema or BNF
|
|
if HAS_GUIDED_DECODING and request.Grammar:
|
|
try:
|
|
json.loads(request.Grammar) # valid JSON = JSON schema
|
|
sampling_params.structured_outputs = StructuredOutputsParams(json=request.Grammar)
|
|
except json.JSONDecodeError:
|
|
sampling_params.structured_outputs = StructuredOutputsParams(grammar=request.Grammar)
|
|
|
|
# Extract image paths and process images
|
|
prompt = request.Prompt
|
|
|
|
image_paths = request.Images
|
|
image_data = [self.load_image(img_path) for img_path in image_paths]
|
|
|
|
videos_path = request.Videos
|
|
video_data = [self.load_video(video_path) for video_path in videos_path]
|
|
|
|
# If tokenizer template is enabled and messages are provided instead of prompt, apply the tokenizer template
|
|
if not request.Prompt and request.UseTokenizerTemplate and request.Messages:
|
|
messages_dicts = self._messages_to_dicts(request.Messages)
|
|
template_kwargs = {"tokenize": False, "add_generation_prompt": True}
|
|
|
|
# Pass tools for tool calling
|
|
if request.Tools:
|
|
try:
|
|
template_kwargs["tools"] = json.loads(request.Tools)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
# Enable thinking mode if requested
|
|
if request.Metadata.get("enable_thinking", "").lower() == "true":
|
|
template_kwargs["enable_thinking"] = True
|
|
|
|
try:
|
|
prompt = self.tokenizer.apply_chat_template(messages_dicts, **template_kwargs)
|
|
except TypeError:
|
|
# Some tokenizers don't support tools/enable_thinking kwargs — retry without them
|
|
prompt = self.tokenizer.apply_chat_template(
|
|
messages_dicts, tokenize=False, add_generation_prompt=True
|
|
)
|
|
|
|
# Generate text using the LLM engine
|
|
request_id = random_uuid()
|
|
print(f"Generating text with request_id: {request_id}", file=sys.stderr)
|
|
multi_modal_data = {}
|
|
if image_data:
|
|
multi_modal_data["image"] = image_data
|
|
if video_data:
|
|
multi_modal_data["video"] = video_data
|
|
outputs = self.llm.generate(
|
|
{
|
|
"prompt": prompt,
|
|
"multi_modal_data": multi_modal_data if multi_modal_data else None,
|
|
},
|
|
sampling_params=sampling_params,
|
|
request_id=request_id,
|
|
)
|
|
|
|
# Stream the results
|
|
generated_text = ""
|
|
generated_token_ids: list[int] = []
|
|
last_output = None
|
|
|
|
# Tool-parsing strategy decision (made once, before the loop):
|
|
#
|
|
# When a tool parser is active, the model's raw tool-call markup
|
|
# (e.g. <tool_call>...) must not be streamed verbatim as delta.content
|
|
# — clients would see the unparsed syntax. Two paths:
|
|
#
|
|
# (A) native streaming via parser.extract_tool_calls_streaming. All
|
|
# concrete tool parsers shipped with vLLM 0.23+ implement this
|
|
# (Granite4, Qwen3Coder, DeepSeekV31, Jamba, Ernie45, Hermes,
|
|
# llama3_json, mistral, …). The parser decides per-delta whether
|
|
# to emit content or suppress tool-call markup, and emits a
|
|
# structured DeltaMessage(tool_calls=[...]) when a call is ready.
|
|
# (B) buffer fallback — used only when the parser surprisingly lacks
|
|
# the streaming method or it raises mid-stream. The post-loop
|
|
# extract_tool_calls assembles the final chat_delta. Same correctness
|
|
# guarantee as a non-streaming response, at the cost of a delayed
|
|
# final chunk.
|
|
has_tool_parser = bool(self.tool_parser_cls and request.Tools)
|
|
tp_instance = None
|
|
tp_request = None
|
|
native_streaming = False
|
|
native_streaming_error = False
|
|
if has_tool_parser:
|
|
try:
|
|
tools_for_parser = json.loads(request.Tools)
|
|
except json.JSONDecodeError:
|
|
tools_for_parser = []
|
|
try:
|
|
tp_instance = self.tool_parser_cls(self.tokenizer, tools=tools_for_parser)
|
|
except TypeError:
|
|
tp_instance = self.tool_parser_cls(self.tokenizer)
|
|
# Build a minimal ChatCompletionRequest so the streaming method
|
|
# sees the tools list. We do not need any other request fields —
|
|
# parsers only read .tools (and sometimes .tool_choice, which we
|
|
# leave at default).
|
|
try:
|
|
from vllm.entrypoints.openai.chat_completion.protocol import (
|
|
ChatCompletionRequest as _CCR,
|
|
)
|
|
tp_request = _CCR(
|
|
model="local",
|
|
messages=[{"role": "user", "content": ""}],
|
|
tools=tools_for_parser or None,
|
|
)
|
|
except Exception as e:
|
|
print(f"Could not build ChatCompletionRequest for streaming parser: {e}",
|
|
file=sys.stderr)
|
|
tp_request = None
|
|
native_streaming = (
|
|
tp_request is not None
|
|
and hasattr(tp_instance, "extract_tool_calls_streaming")
|
|
)
|
|
|
|
try:
|
|
async for request_output in outputs:
|
|
iteration_text = request_output.outputs[0].text
|
|
last_output = request_output
|
|
|
|
if streaming:
|
|
delta_iteration_text = iteration_text.removeprefix(generated_text)
|
|
new_token_ids = list(request_output.outputs[0].token_ids)
|
|
delta_token_ids = new_token_ids[len(generated_token_ids):]
|
|
|
|
if not has_tool_parser:
|
|
# Plain streaming — unchanged from pre-tool-parser path.
|
|
yield backend_pb2.Reply(
|
|
message=bytes(delta_iteration_text, encoding='utf-8'),
|
|
chat_deltas=[backend_pb2.ChatDelta(content=delta_iteration_text)],
|
|
)
|
|
elif native_streaming and not native_streaming_error:
|
|
# (A) Native vLLM extract_tool_calls_streaming.
|
|
try:
|
|
msg = tp_instance.extract_tool_calls_streaming(
|
|
previous_text=generated_text,
|
|
current_text=iteration_text,
|
|
delta_text=delta_iteration_text,
|
|
previous_token_ids=generated_token_ids,
|
|
current_token_ids=new_token_ids,
|
|
delta_token_ids=delta_token_ids,
|
|
request=tp_request,
|
|
)
|
|
except Exception as e:
|
|
print(f"Streaming tool parser error (falling back to "
|
|
f"buffer for the rest of the stream): {e}",
|
|
file=sys.stderr)
|
|
native_streaming_error = True
|
|
msg = None
|
|
if msg is not None:
|
|
tc_protos = []
|
|
for tc in (msg.tool_calls or []):
|
|
fn = tc.function or None
|
|
tc_protos.append(backend_pb2.ToolCallDelta(
|
|
index=tc.index,
|
|
id=tc.id or "",
|
|
name=(fn.name if fn and fn.name else "") or "",
|
|
arguments=(fn.arguments if fn and fn.arguments else "") or "",
|
|
))
|
|
cd_kwargs = {}
|
|
if msg.content:
|
|
cd_kwargs["content"] = msg.content
|
|
if msg.reasoning:
|
|
cd_kwargs["reasoning_content"] = msg.reasoning
|
|
if tc_protos:
|
|
cd_kwargs["tool_calls"] = tc_protos
|
|
if cd_kwargs:
|
|
yield backend_pb2.Reply(
|
|
message=bytes(msg.content or "", encoding='utf-8'),
|
|
chat_deltas=[backend_pb2.ChatDelta(**cd_kwargs)],
|
|
)
|
|
# (B) buffer fallback — emit nothing during the stream.
|
|
# The post-loop extract_tool_calls block builds the final chunk.
|
|
|
|
# Keep track of text + token_ids generated
|
|
generated_text = iteration_text
|
|
generated_token_ids = list(request_output.outputs[0].token_ids)
|
|
finally:
|
|
await outputs.aclose()
|
|
|
|
# Remove the image files from /tmp folder
|
|
for img_path in image_paths:
|
|
try:
|
|
os.remove(img_path)
|
|
except Exception as e:
|
|
print(f"Error removing image file: {img_path}, {e}", file=sys.stderr)
|
|
|
|
# Parse reasoning and tool calls from final text using vLLM's native parsers
|
|
content = generated_text
|
|
reasoning_content = ""
|
|
tool_calls_proto = []
|
|
|
|
if self.reasoning_parser_cls:
|
|
try:
|
|
rp = self.reasoning_parser_cls(self.tokenizer)
|
|
r, c = rp.extract_reasoning(generated_text, request=None)
|
|
reasoning_content = r or ""
|
|
content = c if c is not None else generated_text
|
|
except Exception as e:
|
|
print(f"Reasoning parser error: {e}", file=sys.stderr)
|
|
|
|
# When (A) native streaming ran cleanly, per-delta yields above already
|
|
# delivered everything — do NOT extract again on the full text or we'd
|
|
# duplicate content/tool_calls into the final chunk.
|
|
if has_tool_parser and not (native_streaming and not native_streaming_error):
|
|
try:
|
|
tp = tp_instance
|
|
if tp is None:
|
|
# Defensive: tp_instance build failed earlier; reconstruct.
|
|
tools = json.loads(request.Tools)
|
|
try:
|
|
tp = self.tool_parser_cls(self.tokenizer, tools=tools)
|
|
except TypeError:
|
|
tp = self.tool_parser_cls(self.tokenizer)
|
|
info = tp.extract_tool_calls(content, request=None)
|
|
if info.tools_called:
|
|
content = info.content or ""
|
|
for i, tc in enumerate(info.tool_calls):
|
|
tool_calls_proto.append(backend_pb2.ToolCallDelta(
|
|
index=i,
|
|
id=tc.id,
|
|
name=tc.function.name,
|
|
arguments=tc.function.arguments,
|
|
))
|
|
except Exception as e:
|
|
print(f"Tool parser error: {e}", file=sys.stderr)
|
|
elif native_streaming and not native_streaming_error:
|
|
# Per-delta path already emitted content + tool_calls; the final
|
|
# chat_delta should carry only metadata (token counts, logprobs).
|
|
content = ""
|
|
|
|
# Extract token counts
|
|
prompt_tokens = 0
|
|
completion_tokens = 0
|
|
if last_output is not None:
|
|
try:
|
|
prompt_tokens = len(last_output.prompt_token_ids or [])
|
|
except Exception:
|
|
pass
|
|
try:
|
|
completion_tokens = len(last_output.outputs[0].token_ids or [])
|
|
except Exception:
|
|
pass
|
|
|
|
# Extract logprobs if requested
|
|
logprobs_bytes = b""
|
|
if last_output is not None and request.Logprobs > 0:
|
|
try:
|
|
lp = last_output.outputs[0].logprobs
|
|
if lp:
|
|
logprobs_data = {"content": []}
|
|
for token_lp_dict in lp:
|
|
if token_lp_dict:
|
|
first_tok_id, first_lp = next(iter(token_lp_dict.items()))
|
|
logprobs_data["content"].append({
|
|
"token": getattr(first_lp, "decoded_token", str(first_tok_id)),
|
|
"logprob": first_lp.logprob,
|
|
})
|
|
logprobs_bytes = json.dumps(logprobs_data).encode("utf-8")
|
|
except Exception as e:
|
|
print(f"Logprobs extraction error: {e}", file=sys.stderr)
|
|
|
|
chat_delta = backend_pb2.ChatDelta(
|
|
content=content,
|
|
reasoning_content=reasoning_content,
|
|
tool_calls=tool_calls_proto,
|
|
)
|
|
|
|
if streaming:
|
|
# Final chunk with structured data.
|
|
#
|
|
# If we used the buffer fallback (has_tool_parser=True AND native
|
|
# streaming did NOT run cleanly) and the parser found no tool call,
|
|
# flush the buffered content as ONE content delta — and clear the
|
|
# final chat_delta's content so the metadata chunk does not repeat
|
|
# what we just sent. This is the plain-text-with-tool-parser path.
|
|
buffered_fallback = (
|
|
has_tool_parser
|
|
and not (native_streaming and not native_streaming_error)
|
|
)
|
|
if buffered_fallback and not tool_calls_proto and content:
|
|
yield backend_pb2.Reply(
|
|
message=bytes(content, encoding='utf-8'),
|
|
chat_deltas=[backend_pb2.ChatDelta(content=content)],
|
|
)
|
|
chat_delta = backend_pb2.ChatDelta(
|
|
reasoning_content=reasoning_content,
|
|
tool_calls=tool_calls_proto,
|
|
)
|
|
yield backend_pb2.Reply(
|
|
message=b"",
|
|
prompt_tokens=prompt_tokens,
|
|
tokens=completion_tokens,
|
|
chat_deltas=[chat_delta],
|
|
logprobs=logprobs_bytes,
|
|
)
|
|
return
|
|
|
|
# Non-streaming: single Reply with everything
|
|
yield backend_pb2.Reply(
|
|
message=bytes(content, encoding='utf-8'),
|
|
prompt_tokens=prompt_tokens,
|
|
tokens=completion_tokens,
|
|
chat_deltas=[chat_delta],
|
|
logprobs=logprobs_bytes,
|
|
)
|
|
|
|
def load_image(self, image_path: str):
|
|
"""
|
|
Load an image from the given file path or base64 encoded data.
|
|
|
|
Args:
|
|
image_path (str): The path to the image file or base64 encoded data.
|
|
|
|
Returns:
|
|
Image: The loaded image.
|
|
"""
|
|
try:
|
|
|
|
image_data = base64.b64decode(image_path)
|
|
image = Image.open(io.BytesIO(image_data))
|
|
return image
|
|
except Exception as e:
|
|
print(f"Error loading image {image_path}: {e}", file=sys.stderr)
|
|
return None
|
|
|
|
def load_video(self, video_path: str):
|
|
"""
|
|
Load a video from the given file path.
|
|
|
|
Args:
|
|
video_path (str): The path to the image file.
|
|
|
|
Returns:
|
|
Video: The loaded video.
|
|
"""
|
|
try:
|
|
timestamp = str(int(time.time() * 1000)) # Generate timestamp
|
|
p = os.path.join(tempfile.gettempdir(), f"vl-{timestamp}.data")
|
|
with open(p, "wb") as f:
|
|
f.write(base64.b64decode(video_path))
|
|
video = VideoAsset(name=p).np_ndarrays
|
|
os.remove(p)
|
|
return video
|
|
except Exception as e:
|
|
print(f"Error loading video {video_path}: {e}", file=sys.stderr)
|
|
return None
|
|
|
|
async def serve(address):
|
|
# Start asyncio gRPC server
|
|
server = grpc.aio.server(migration_thread_pool=futures.ThreadPoolExecutor(max_workers=MAX_WORKERS),
|
|
options=[
|
|
('grpc.max_message_length', 50 * 1024 * 1024), # 50MB
|
|
('grpc.max_send_message_length', 50 * 1024 * 1024), # 50MB
|
|
('grpc.max_receive_message_length', 50 * 1024 * 1024), # 50MB
|
|
],
|
|
interceptors=get_auth_interceptors(aio=True),
|
|
)
|
|
# Add the servicer to the server
|
|
backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server)
|
|
# Bind the server to the address
|
|
server.add_insecure_port(address)
|
|
|
|
# Gracefully shutdown the server on SIGTERM or SIGINT
|
|
loop = asyncio.get_event_loop()
|
|
for sig in (signal.SIGINT, signal.SIGTERM):
|
|
loop.add_signal_handler(
|
|
sig, lambda: asyncio.ensure_future(server.stop(5))
|
|
)
|
|
|
|
# Start the server
|
|
await server.start()
|
|
print("Server started. Listening on: " + address, file=sys.stderr)
|
|
# Wait for the server to be terminated
|
|
await server.wait_for_termination()
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Run the gRPC server.")
|
|
parser.add_argument(
|
|
"--addr", default="localhost:50051", help="The address to bind the server to."
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
asyncio.run(serve(args.addr))
|