chore: vendor sglang v0.5.10 snapshot
This commit is contained in:
372
third_party/sglang/sgl-model-gateway/bindings/golang/src/preprocessor.rs
vendored
Normal file
372
third_party/sglang/sgl-model-gateway/bindings/golang/src/preprocessor.rs
vendored
Normal file
@@ -0,0 +1,372 @@
|
||||
//! Preprocessing FFI functions for chat requests
|
||||
//!
|
||||
//! This module provides C-compatible functions for preprocessing chat completion requests:
|
||||
//! - Apply chat_template to messages
|
||||
//! - Tokenize the processed text
|
||||
//! - Generate tool constraints
|
||||
//!
|
||||
//! These functions are designed to be called once per request, reducing FFI overhead.
|
||||
|
||||
use std::ffi::{CStr, CString};
|
||||
use std::os::raw::{c_char, c_int};
|
||||
use std::ptr;
|
||||
use std::os::raw::c_uint;
|
||||
|
||||
use smg::tokenizer::create_tokenizer_from_file;
|
||||
use smg::protocols::chat::ChatCompletionRequest;
|
||||
use smg::routers::grpc::utils::{process_chat_messages, generate_tool_constraints};
|
||||
|
||||
use super::error::{SglErrorCode, set_error_message};
|
||||
use super::memory::{sgl_free_string, sgl_free_token_ids};
|
||||
use super::tokenizer::TokenizerHandle;
|
||||
|
||||
/// Handle for preprocessed request
|
||||
#[repr(C)]
|
||||
pub struct PreprocessedRequestHandle {
|
||||
pub(crate) prompt_text: CString,
|
||||
pub(crate) token_ids: Vec<i32>,
|
||||
pub(crate) tool_constraints_json: Option<CString>,
|
||||
pub(crate) prompt_tokens: i32,
|
||||
}
|
||||
|
||||
/// Preprocess a chat completion request
|
||||
///
|
||||
/// This function:
|
||||
/// 1. Applies chat_template to messages
|
||||
/// 2. Tokenizes the processed text
|
||||
/// 3. Generates tool constraints (if tools are present)
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `request_json` - OpenAI ChatCompletionRequest as JSON string
|
||||
/// * `tokenizer_path` - Path to tokenizer directory
|
||||
/// * `prompt_text_out` - Pointer to receive prompt text (C string, must be freed with sgl_free_string)
|
||||
/// * `token_ids_out` - Pointer to receive token IDs array (must be freed with sgl_free_token_ids)
|
||||
/// * `token_ids_len_out` - Pointer to receive token IDs array length
|
||||
/// * `tool_constraints_json_out` - Optional pointer to receive tool constraints JSON (must be freed with sgl_free_string)
|
||||
/// * `prompt_tokens_out` - Pointer to receive prompt token count
|
||||
/// * `error_out` - Optional pointer to receive error message
|
||||
///
|
||||
/// # Returns
|
||||
/// * SglErrorCode::Success on success, error code on failure
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_preprocess_chat_request(
|
||||
request_json: *const c_char,
|
||||
tokenizer_path: *const c_char,
|
||||
prompt_text_out: *mut *mut c_char,
|
||||
token_ids_out: *mut *mut c_uint,
|
||||
token_ids_len_out: *mut usize,
|
||||
tool_constraints_json_out: *mut *mut c_char,
|
||||
prompt_tokens_out: *mut c_int,
|
||||
error_out: *mut *mut c_char,
|
||||
) -> SglErrorCode {
|
||||
if request_json.is_null()
|
||||
|| tokenizer_path.is_null()
|
||||
|| prompt_text_out.is_null()
|
||||
|| token_ids_out.is_null()
|
||||
|| token_ids_len_out.is_null()
|
||||
|| prompt_tokens_out.is_null()
|
||||
{
|
||||
set_error_message(error_out, "Invalid arguments: null pointer");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
|
||||
// Parse input strings
|
||||
let request_str = match CStr::from_ptr(request_json).to_str() {
|
||||
Ok(s) => s,
|
||||
Err(_) => {
|
||||
set_error_message(error_out, "Invalid UTF-8 in request_json");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
};
|
||||
|
||||
let tokenizer_path_str = match CStr::from_ptr(tokenizer_path).to_str() {
|
||||
Ok(s) => s,
|
||||
Err(_) => {
|
||||
set_error_message(error_out, "Invalid UTF-8 in tokenizer_path");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
};
|
||||
|
||||
// Parse ChatCompletionRequest
|
||||
let chat_request: ChatCompletionRequest = match serde_json::from_str(request_str) {
|
||||
Ok(req) => req,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to parse request JSON: {}", e));
|
||||
return SglErrorCode::ParsingError;
|
||||
}
|
||||
};
|
||||
|
||||
// Create tokenizer
|
||||
let tokenizer = match create_tokenizer_from_file(tokenizer_path_str) {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to create tokenizer: {}", e));
|
||||
return SglErrorCode::TokenizationError;
|
||||
}
|
||||
};
|
||||
|
||||
// Process chat messages (apply chat_template)
|
||||
let processed_messages = match process_chat_messages(&chat_request, tokenizer.as_ref()) {
|
||||
Ok(msgs) => msgs,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to process chat messages: {}", e));
|
||||
return SglErrorCode::ParsingError;
|
||||
}
|
||||
};
|
||||
|
||||
// Tokenize the processed text
|
||||
let encoding = match tokenizer.encode(&processed_messages.text, false) {
|
||||
Ok(enc) => enc,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Tokenization failed: {}", e));
|
||||
return SglErrorCode::TokenizationError;
|
||||
}
|
||||
};
|
||||
|
||||
let token_ids_vec: Vec<i32> = encoding
|
||||
.token_ids()
|
||||
.iter()
|
||||
.map(|&id| id as i32)
|
||||
.collect();
|
||||
|
||||
let prompt_tokens = token_ids_vec.len() as i32;
|
||||
|
||||
// Generate tool constraints if tools are present
|
||||
let tool_constraints_json = if let Some(tools) = chat_request.tools.as_ref() {
|
||||
match generate_tool_constraints(tools, &chat_request.tool_choice, &chat_request.model) {
|
||||
Ok(Some(constraints)) => {
|
||||
match serde_json::to_string(&constraints) {
|
||||
Ok(json_str) => Some(CString::new(json_str).unwrap()),
|
||||
Err(e) => {
|
||||
set_error_message(
|
||||
error_out,
|
||||
&format!("Failed to serialize tool constraints: {}", e),
|
||||
);
|
||||
return SglErrorCode::ParsingError;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(None) => None,
|
||||
Err(e) => {
|
||||
set_error_message(
|
||||
error_out,
|
||||
&format!("Failed to generate tool constraints: {}", e),
|
||||
);
|
||||
return SglErrorCode::ParsingError;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Allocate memory for outputs
|
||||
let prompt_text_cstr = match CString::new(processed_messages.text) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to create C string: {}", e));
|
||||
return SglErrorCode::MemoryError;
|
||||
}
|
||||
};
|
||||
|
||||
let token_ids_len = token_ids_vec.len();
|
||||
// Convert i32 to u32 for token IDs (as expected by the memory management functions)
|
||||
let token_ids_u32: Vec<u32> = token_ids_vec.iter().map(|&id| id as u32).collect();
|
||||
let token_ids_ptr = if token_ids_u32.is_empty() {
|
||||
ptr::null_mut()
|
||||
} else {
|
||||
let boxed = token_ids_u32.into_boxed_slice();
|
||||
Box::into_raw(boxed) as *mut c_uint
|
||||
};
|
||||
|
||||
// Set output values
|
||||
*prompt_text_out = prompt_text_cstr.into_raw();
|
||||
*token_ids_out = token_ids_ptr;
|
||||
*token_ids_len_out = token_ids_len;
|
||||
*prompt_tokens_out = prompt_tokens;
|
||||
|
||||
if !tool_constraints_json_out.is_null() {
|
||||
if let Some(constraints) = tool_constraints_json {
|
||||
*tool_constraints_json_out = constraints.into_raw();
|
||||
} else {
|
||||
*tool_constraints_json_out = ptr::null_mut();
|
||||
}
|
||||
}
|
||||
|
||||
SglErrorCode::Success
|
||||
}
|
||||
|
||||
/// Preprocess a chat completion request using an existing tokenizer handle
|
||||
///
|
||||
/// This function is similar to sgl_preprocess_chat_request, but accepts a TokenizerHandle
|
||||
/// instead of creating a new tokenizer. This allows reusing a cached tokenizer instance,
|
||||
/// significantly reducing initialization overhead in concurrent scenarios.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `request_json` - OpenAI ChatCompletionRequest as JSON string
|
||||
/// * `tokenizer_handle` - Existing tokenizer handle (must be valid)
|
||||
/// * `prompt_text_out` - Pointer to receive prompt text (C string, must be freed with sgl_free_string)
|
||||
/// * `token_ids_out` - Pointer to receive token IDs array (must be freed with sgl_free_token_ids)
|
||||
/// * `token_ids_len_out` - Pointer to receive token IDs array length
|
||||
/// * `tool_constraints_json_out` - Optional pointer to receive tool constraints JSON (must be freed with sgl_free_string)
|
||||
/// * `prompt_tokens_out` - Pointer to receive prompt token count
|
||||
/// * `error_out` - Optional pointer to receive error message
|
||||
///
|
||||
/// # Returns
|
||||
/// * SglErrorCode::Success on success, error code on failure
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_preprocess_chat_request_with_tokenizer(
|
||||
request_json: *const c_char,
|
||||
tokenizer_handle: *mut TokenizerHandle,
|
||||
prompt_text_out: *mut *mut c_char,
|
||||
token_ids_out: *mut *mut c_uint,
|
||||
token_ids_len_out: *mut usize,
|
||||
tool_constraints_json_out: *mut *mut c_char,
|
||||
prompt_tokens_out: *mut c_int,
|
||||
error_out: *mut *mut c_char,
|
||||
) -> SglErrorCode {
|
||||
if request_json.is_null()
|
||||
|| tokenizer_handle.is_null()
|
||||
|| prompt_text_out.is_null()
|
||||
|| token_ids_out.is_null()
|
||||
|| token_ids_len_out.is_null()
|
||||
|| prompt_tokens_out.is_null()
|
||||
{
|
||||
set_error_message(error_out, "Invalid arguments: null pointer");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
|
||||
// Parse input string
|
||||
let request_str = match CStr::from_ptr(request_json).to_str() {
|
||||
Ok(s) => s,
|
||||
Err(_) => {
|
||||
set_error_message(error_out, "Invalid UTF-8 in request_json");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
};
|
||||
|
||||
// Parse ChatCompletionRequest
|
||||
let chat_request: ChatCompletionRequest = match serde_json::from_str(request_str) {
|
||||
Ok(req) => req,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to parse request JSON: {}", e));
|
||||
return SglErrorCode::ParsingError;
|
||||
}
|
||||
};
|
||||
|
||||
// Use existing tokenizer from handle (no need to create new one!)
|
||||
let handle_ref = &*tokenizer_handle;
|
||||
let tokenizer = &handle_ref.tokenizer;
|
||||
|
||||
// Process chat messages (apply chat_template)
|
||||
let processed_messages = match process_chat_messages(&chat_request, tokenizer.as_ref()) {
|
||||
Ok(msgs) => msgs,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to process chat messages: {}", e));
|
||||
return SglErrorCode::ParsingError;
|
||||
}
|
||||
};
|
||||
|
||||
// Tokenize the processed text
|
||||
let encoding = match tokenizer.encode(&processed_messages.text, false) {
|
||||
Ok(enc) => enc,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Tokenization failed: {}", e));
|
||||
return SglErrorCode::TokenizationError;
|
||||
}
|
||||
};
|
||||
|
||||
let token_ids_vec: Vec<i32> = encoding
|
||||
.token_ids()
|
||||
.iter()
|
||||
.map(|&id| id as i32)
|
||||
.collect();
|
||||
|
||||
let prompt_tokens = token_ids_vec.len() as i32;
|
||||
|
||||
// Generate tool constraints if tools are present
|
||||
let tool_constraints_json = if let Some(tools) = chat_request.tools.as_ref() {
|
||||
match generate_tool_constraints(tools, &chat_request.tool_choice, &chat_request.model) {
|
||||
Ok(Some(constraints)) => {
|
||||
match serde_json::to_string(&constraints) {
|
||||
Ok(json_str) => Some(CString::new(json_str).unwrap()),
|
||||
Err(e) => {
|
||||
set_error_message(
|
||||
error_out,
|
||||
&format!("Failed to serialize tool constraints: {}", e),
|
||||
);
|
||||
return SglErrorCode::ParsingError;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(None) => None,
|
||||
Err(e) => {
|
||||
set_error_message(
|
||||
error_out,
|
||||
&format!("Failed to generate tool constraints: {}", e),
|
||||
);
|
||||
return SglErrorCode::ParsingError;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Allocate memory for outputs
|
||||
let prompt_text_cstr = match CString::new(processed_messages.text) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to create C string: {}", e));
|
||||
return SglErrorCode::MemoryError;
|
||||
}
|
||||
};
|
||||
|
||||
let token_ids_len = token_ids_vec.len();
|
||||
// Convert i32 to u32 for token IDs (as expected by the memory management functions)
|
||||
let token_ids_u32: Vec<u32> = token_ids_vec.iter().map(|&id| id as u32).collect();
|
||||
let token_ids_ptr = if token_ids_u32.is_empty() {
|
||||
ptr::null_mut()
|
||||
} else {
|
||||
let boxed = token_ids_u32.into_boxed_slice();
|
||||
Box::into_raw(boxed) as *mut c_uint
|
||||
};
|
||||
|
||||
// Set output values
|
||||
*prompt_text_out = prompt_text_cstr.into_raw();
|
||||
*token_ids_out = token_ids_ptr;
|
||||
*token_ids_len_out = token_ids_len;
|
||||
*prompt_tokens_out = prompt_tokens;
|
||||
|
||||
if !tool_constraints_json_out.is_null() {
|
||||
if let Some(constraints) = tool_constraints_json {
|
||||
*tool_constraints_json_out = constraints.into_raw();
|
||||
} else {
|
||||
*tool_constraints_json_out = ptr::null_mut();
|
||||
}
|
||||
}
|
||||
|
||||
SglErrorCode::Success
|
||||
}
|
||||
|
||||
/// Free a preprocessed request handle (cleanup function)
|
||||
///
|
||||
/// This function frees the memory allocated by sgl_preprocess_chat_request.
|
||||
/// It should be called after the preprocessed data is no longer needed.
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_preprocessed_request_free(
|
||||
prompt_text: *mut c_char,
|
||||
token_ids: *mut c_uint,
|
||||
token_ids_len: usize,
|
||||
tool_constraints_json: *mut c_char,
|
||||
) {
|
||||
if !prompt_text.is_null() {
|
||||
sgl_free_string(prompt_text);
|
||||
}
|
||||
|
||||
if !token_ids.is_null() && token_ids_len > 0 {
|
||||
sgl_free_token_ids(token_ids, token_ids_len);
|
||||
}
|
||||
|
||||
if !tool_constraints_json.is_null() {
|
||||
sgl_free_string(tool_constraints_json);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user