gahow/agentic-pd-hybrid
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: vendor sglang v0.5.10 snapshot

Browse Source
This commit is contained in:
Gahow Wang
2026-04-24 12:29:36 +00:00
parent 78f0d15221
commit bded08301f
4308 changed files with 1200894 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
third_party/sglang/sgl-model-gateway/bindings/golang/.gitignore vendored Normal file
View File

@@ -0,0 +1,24 @@
# Build artifacts
target/
lib/
# Compiled binaries
examples/simple/simple
examples/streaming/streaming
# Go build artifacts
*.o
*.a
*.so
*.dylib
# IDE and editor files
.vscode/
.idea/
*.swp
*.swo
*~
# Environment files
.env
.env.local

48
third_party/sglang/sgl-model-gateway/bindings/golang/Cargo.toml vendored Normal file
View File

@@ -0,0 +1,48 @@
[package]
name = "sgl-model-gateway-golang"
version = "0.3.2"
edition = "2021"
[lib]
name = "sgl_model_gateway_go"
crate-type = ["cdylib"]
[dependencies]
tokio = { version = "1.42.0", features = ["full"] }
serde_json = { version = "1.0", default-features = false, features = [
"std",
"preserve_order",
] }
uuid = { version = "1.10", features = ["v4", "serde"] }
once_cell = "1.21.3"
futures-util = "0.3"
tracing = "0.1"
libc = "0.2.179"
[dependencies.sgl-model-gateway]
path = "../.."
default-features = true
[features]
default = []
vendored-openssl = ["sgl-model-gateway/vendored-openssl"]
[profile.release]
opt-level = "z" # Optimize for size
lto = "fat" # Full LTO for smaller binaries
codegen-units = 1 # Better optimization, slower compile
strip = true # Strip debug symbols
[profile.ci]
inherits = "release"
opt-level = 2 # Lighter optimization (still fast runtime, much faster compile)
lto = "thin" # Thin LTO - good balance
codegen-units = 16 # More parallelization for faster builds
strip = true
[profile.dev]
opt-level = 0
debug = 1
split-debuginfo = "unpacked"
incremental = true
codegen-units = 256

103
third_party/sglang/sgl-model-gateway/bindings/golang/Makefile vendored Normal file
View File

@@ -0,0 +1,103 @@
# Makefile for sgl-model-gateway golang bindings
# This builds the Rust FFI library and provides convenience targets for Go development
# Configuration
CARGO_BUILD_DIR ?= $(shell pwd)/target
BUILD_MODE ?= release
LIB_NAME = libsgl_model_gateway_go
# Detect OS
UNAME_S := $(shell uname -s)
ifeq ($(UNAME_S),Linux)
LIB_EXT = .so
LD_LIBRARY_PATH_VAR = LD_LIBRARY_PATH
endif
ifeq ($(UNAME_S),Darwin)
LIB_EXT = .dylib
LD_LIBRARY_PATH_VAR = DYLD_LIBRARY_PATH
endif
# Paths
ROOT_DIR := $(shell pwd)
RUST_SRC_DIR := $(ROOT_DIR)/src
LIB_BUILD_DIR := $(CARGO_BUILD_DIR)/$(BUILD_MODE)
LIB_BUILD_PATH := $(LIB_BUILD_DIR)/$(LIB_NAME)$(LIB_EXT)
LIB_EXPORT_DIR := $(ROOT_DIR)/lib
LIB_EXPORT_PATH := $(LIB_EXPORT_DIR)/$(LIB_NAME)$(LIB_EXT)
# Python LDFLAGS (needed for Rust FFI that depends on Python)
PYTHON_LDFLAGS := $(shell python3-config --ldflags --embed 2>/dev/null || python3-config --ldflags 2>/dev/null || echo "")
# CGO flags - use exported lib directory if available, otherwise build directory
LIB_DIR := $(if $(wildcard $(LIB_EXPORT_PATH)),$(LIB_EXPORT_DIR),$(LIB_BUILD_DIR))
export CGO_LDFLAGS = -L$(LIB_DIR) -lsgl_model_gateway_go $(PYTHON_LDFLAGS) -ldl
export $(LD_LIBRARY_PATH_VAR) := $(LIB_DIR):$($(LD_LIBRARY_PATH_VAR))
.PHONY: all build build-dev lib lib-clean clean test examples help run-simple run-streaming check-lib
help:
@echo "Available targets:"
@echo " build - Build release version of Rust FFI library"
@echo " build-dev - Build debug version of Rust FFI library"
@echo " lib - Copy built library to ./lib directory"
@echo " lib-clean - Clean ./lib directory"
@echo " clean - Clean build artifacts"
@echo " test - Run Go tests"
@echo " examples - Build example programs"
@echo " run-simple - Run simple example"
@echo " run-streaming - Run streaming example"
all: build
build:
@echo "Building Rust FFI library (release mode)..."
@CARGO_TARGET_DIR=$(CARGO_BUILD_DIR) cargo build --release --manifest-path Cargo.toml
@echo "Library built at: $(LIB_BUILD_PATH)"
build-dev:
@echo "Building Rust FFI library (debug mode)..."
@CARGO_TARGET_DIR=$(CARGO_BUILD_DIR) cargo build --manifest-path Cargo.toml
@echo "Library built at: $(LIB_BUILD_DIR)/debug/$(LIB_NAME)$(LIB_EXT)"
lib: build
@echo "Copying library to ./lib directory..."
@mkdir -p $(LIB_EXPORT_DIR)
@cp $(LIB_BUILD_PATH) $(LIB_EXPORT_PATH)
@echo "Library exported at: $(LIB_EXPORT_PATH)"
lib-clean:
@echo "Cleaning ./lib directory..."
@rm -rf $(LIB_EXPORT_DIR)
@echo "Lib directory cleaned"
clean: lib-clean
@echo "Cleaning build artifacts..."
@CARGO_TARGET_DIR=$(CARGO_BUILD_DIR) cargo clean --manifest-path Cargo.toml
@echo "Clean complete"
test: build
@echo "Running Go tests..."
@go test ./...
examples: build
@echo "Building example programs..."
@cd examples/simple && go build -o simple main.go
@cd examples/streaming && go build -o streaming main.go
@echo "Examples built"
run-simple: build
@echo "Running simple example..."
@cd examples/simple && bash run.sh
run-streaming: build
@echo "Running streaming example..."
@cd examples/streaming && bash run.sh
# Check if library exists (either in lib dir or build dir)
check-lib:
@if [ ! -f "$(LIB_EXPORT_PATH)" ] && [ ! -f "$(LIB_BUILD_PATH)" ]; then \
echo "Error: Library not found at $(LIB_EXPORT_PATH) or $(LIB_BUILD_PATH)"; \
echo "Run 'make build' or 'make lib' first"; \
exit 1; \
fi
@echo "Library found at: $(LIB_DIR)/$(LIB_NAME)$(LIB_EXT)"

482
third_party/sglang/sgl-model-gateway/bindings/golang/README.md vendored Normal file
View File

@@ -0,0 +1,482 @@
# SGLang Go gRPC SDK
A high-level Go SDK for interacting with SGLang gRPC API, designed with an OpenAI-style API for familiarity and ease of use.
**Location**: `sgl-model-gateway/bindings/golang/`
## Table of Contents
- [Features](#features)
- [Installation](#installation)
- [Quick Start](#quick-start)
- [Basic Usage](#basic-usage)
- [Streaming Usage](#streaming-usage)
- [Examples](#examples)
- [Configuration](#configuration)
- [API Reference](#api-reference)
- [Testing](#testing)
- [Unit Tests](#unit-tests)
- [Integration Tests](#integration-tests)
- [Benchmarks](#benchmarks)
- [Documentation](#documentation)
- [Development](#development)
- [Troubleshooting](#troubleshooting)
- [License](#license)
## Features
- **OpenAI-style API**: Familiar interface similar to OpenAI Go SDK
- **Streaming Support**: Real-time streaming chat completions
- **Non-streaming Support**: Simple request/response API
- **Tool Calling**: Support for function calling and tool use
- **Type-safe**: Full Go type definitions for requests and responses
- **Comprehensive Testing**: 18+ unit and integration tests
- **Thread-safe**: All public methods are safe for concurrent use
- **Well-documented**: Full API documentation with examples
## Installation
```bash
go get github.com/sglang/sglang-go-grpc-sdk
```
### Sync Dependencies
```bash
cd sgl-model-gateway/bindings/golang
go mod tidy
```
### Build Requirements
- Go 1.21+, Rust toolchain, Python 3.x
## Quick Start
### Benchmark
Run the OpenAI-compatible server and benchmark:
```bash
# Set environment variables
export SGL_TOKENIZER_PATH="/Users/yangyanbo/tokenizer"
export SGL_GRPC_ENDPOINT="grpc://10.109.185.20:8001"
# Run server
cd examples/oai_server
bash run.sh
# Run E2E benchmark
cd ../..
make e2e E2E_MODEL=/work/models/qwencoder-3b E2E_TOKENIZER=/Users/yangyanbo/tokenizer E2E_INPUT_LEN=1024 E2E_OUTPUT_LEN=512
```
## Examples
The SDK includes several examples in the `examples/` directory:
- **simple**: Basic non-streaming chat completion example
- **streaming**: Real-time streaming with performance metrics
### Running Examples
```bash
# Run simple example
cd bindings/golang/examples/simple
bash run.sh
# Run streaming example
cd bindings/golang/examples/streaming
bash run.sh
# Or use Makefile from bindings/golang directory
cd bindings/golang
make run-simple
make run-streaming
```
### Basic Usage (Non-streaming)
```go
package main
import (
"context"
"fmt"
"log"
"github.com/sglang/sglang-go-grpc-sdk"
)
func main() {
// Create client
client, err := sglang.NewClient(sglang.ClientConfig{
Endpoint: "grpc://localhost:20000",
TokenizerPath: "/path/to/tokenizer",
})
if err != nil {
log.Fatal(err)
}
defer client.Close()
// Create completion
resp, err := client.CreateChatCompletion(context.Background(), sglang.ChatCompletionRequest{
Model: "default",
Messages: []sglang.ChatMessage{
{Role: "user", Content: "Hello!"},
},
Stream: false,
})
if err != nil {
log.Fatal(err)
}
fmt.Println(resp.Choices[0].Message.Content)
fmt.Printf("Usage: Prompt=%d, Completion=%d, Total=%d\n",
resp.Usage.PromptTokens,
resp.Usage.CompletionTokens,
resp.Usage.TotalTokens)
}
```
### Streaming Usage
```go
package main
import (
"context"
"fmt"
"io"
"log"
"github.com/sglang/sglang-go-grpc-sdk"
)
func main() {
// Create client
client, err := sglang.NewClient(sglang.ClientConfig{
Endpoint: "grpc://localhost:20000",
TokenizerPath: "/path/to/tokenizer",
})
if err != nil {
log.Fatal(err)
}
defer client.Close()
// Create streaming completion
ctx := context.Background()
stream, err := client.CreateChatCompletionStream(ctx, sglang.ChatCompletionRequest{
Model: "default",
Messages: []sglang.ChatMessage{
{Role: "user", Content: "Tell me a story"},
},
Stream: true,
MaxCompletionTokens: intPtr(500),
})
if err != nil {
log.Fatal(err)
}
defer stream.Close()
// Read streaming response
for {
chunk, err := stream.Recv()
if err == io.EOF {
break
}
if err != nil {
log.Fatal(err)
}
for _, choice := range chunk.Choices {
if choice.Delta.Content != "" {
fmt.Print(choice.Delta.Content)
}
}
}
fmt.Println() // newline
}
// Helper functions for optional pointer fields
func intPtr(i int) *int {
return &i
}
func float32Ptr(f float32) *float32 {
return &f
}
```
Examples automatically detect the server endpoint and tokenizer path via environment variables or defaults.
## Configuration
### Environment Variables
- `SGL_GRPC_ENDPOINT`: gRPC server endpoint (default: `grpc://localhost:20000`)
- `SGL_TOKENIZER_PATH`: Path to tokenizer directory (required)
- `CARGO_BUILD_DIR`: Rust build output directory (auto-detected if not set)
### ClientConfig
```go
type ClientConfig struct {
// Endpoint is the gRPC endpoint URL (e.g., "grpc://localhost:20000")
// Required field. Must include the scheme (grpc://) and port number.
Endpoint string
// TokenizerPath is the path to the tokenizer directory containing
// tokenizer configuration files (e.g., tokenizer.json, vocab.json)
// Required field.
TokenizerPath string
}
```
## API Reference
### Client Methods
```go
type Client struct {
// Thread-safe client for SGLang gRPC API
}
// Creates a new client with the given configuration
func NewClient(config ClientConfig) (*Client, error)
// Closes the client and releases all resources
func (c *Client) Close() error
// Creates a non-streaming chat completion
func (c *Client) CreateChatCompletion(ctx context.Context, req ChatCompletionRequest) (*ChatCompletionResponse, error)
// Creates a streaming chat completion
func (c *Client) CreateChatCompletionStream(ctx context.Context, req ChatCompletionRequest) (*ChatCompletionStream, error)
```
### Request Types
- `ChatCompletionRequest`: Main request type for chat completions
- Model, Messages, Stream, Temperature, TopP, MaxCompletionTokens, Tools, etc.
- `ChatMessage`: Individual message in a conversation
- Role, Content
- `Tool`: Tool/function definition for function calling
- Type, Function (name, description, parameters)
### Response Types
- `ChatCompletionResponse`: Non-streaming response
- ID, Model, Created, Choices, Usage
- `ChatCompletionStreamResponse`: Streaming response chunk
- Same structure as above but for incremental updates
- `Message`: Complete message with content and tool calls
- `ToolCall`: Tool call information with function and arguments
- `Usage`: Token usage statistics
- PromptTokens, CompletionTokens, TotalTokens
## Testing
The SDK includes comprehensive testing infrastructure with both unit and integration tests.
### Unit Tests
Unit tests are located in `client_test.go` and test individual components without requiring a server.
#### Running Unit Tests
```bash
# Run all unit tests
go test ./...
# Run with verbose output
go test -v ./...
# Run specific test
go test -run TestClientConfig
# Run tests with race detector (detects concurrency issues)
go test -race ./...
# Run with coverage analysis
go test -cover ./...
# Generate detailed coverage report
go test -coverprofile=coverage.out ./...
go tool cover -html=coverage.out -o coverage.html
```
#### Unit Test Coverage
- Configuration validation, type structures, response handling, concurrent operations, and benchmarks
- `client_test.go` - 10 unit tests covering core functionality
### Integration Tests
Integration tests require a running SGLang server and test the full client-server interaction.
#### Prerequisites
1. Start SGLang server: `python -m sglang.launch_server --model-path <model_path>`
2. Set environment variables:
```bash
export SGL_GRPC_ENDPOINT=grpc://localhost:20000
export SGL_TOKENIZER_PATH=/path/to/tokenizer
```
#### Running Integration Tests
```bash
# Run all integration tests
go test -tags=integration ./...
# Run specific integration test
go test -tags=integration -run TestIntegrationNonStreamingCompletion
# Run with verbose output
go test -tags=integration -v ./...
# Run with race detector
go test -tags=integration -race ./...
```
#### Integration Test Coverage
**Test File**: `integration_test.go` - 4 integration tests
- `TestIntegrationNonStreamingCompletion` - Basic non-streaming request/response
- `TestIntegrationStreamingCompletion` - Streaming response handling
- `TestIntegrationConcurrentRequests` - Multiple simultaneous requests
- `TestIntegrationContextCancellation` - Context timeout and cancellation
### Benchmarks
```bash
go test -bench=. -benchmem ./...
```
## Documentation
All public types and functions include comprehensive documentation with usage examples.
### Key Documented Components
- `Client` - Main client with thread-safety notes
- `ClientConfig` - Configuration requirements and validation rules
- `ChatCompletionRequest` - Request structure with field descriptions
- `ChatCompletionResponse` - Response structure and usage
- `ChatCompletionStreamResponse` - Streaming response format
- `Usage` - Token usage information structure
- `Tool`, `Function`, `ToolCall` - Tool call structures
### Viewing Documentation
```bash
godoc -http=:6060
# Visit: http://localhost:6060/pkg/github.com/sglang/sglang-go-grpc-sdk/
```
## Development
```bash
cd bindings/golang
make build # Build Go bindings
go vet ./... # Check code quality
go fmt ./... # Format code
go test -race ./... # Run tests
```
### Project Structure
```
bindings/golang/
├── client.go # Main client implementation
├── client_test.go # Unit tests
├── integration_test.go # Integration tests
├── README.md # This file
├── Makefile # Build automation
├── Cargo.toml # Rust FFI dependencies
├── examples/ # Example programs
│ ├── simple/ # Non-streaming example
│ └── streaming/ # Streaming example
├── src/ # Rust FFI source
│ ├── client.rs # Client FFI
│ ├── stream.rs # Stream handling
│ ├── grpc_converter.rs # Response conversion
│ └── ...
└── internal/ # Internal packages
└── ffi/ # FFI bindings
```
## Troubleshooting
### Missing Dependencies
Run `go mod tidy` to sync dependencies.
### Connection Errors
Ensure SGLang server is running and check `SGL_GRPC_ENDPOINT`.
### Tokenizer Not Found
Set `SGL_TOKENIZER_PATH` environment variable.
2. Verify path contains required files: `ls $SGL_TOKENIZER_PATH`
3. Files should include: `tokenizer.json`, `vocab.json`, `config.json`
### Build Failures
**Error**: `library 'sgl_model_gateway_go' not found`
**Solution**:
1. Rebuild Rust library: `cd sgl-model-gateway/bindings/golang && make build`
2. Or manually with cargo: `cd sgl-model-gateway/bindings/golang && cargo build --release`
3. Set `CARGO_BUILD_DIR` if using non-standard build location
4. Ensure Rust toolchain is installed: `rustup toolchain list`
### Tests Hanging
**Error**: Tests seem to hang indefinitely
**Solution**:
1. Use timeout for hanging tests: `timeout 30s go test ./...`
2. Run with verbose output to see which test hangs: `go test -v ./...`
3. Ensure server is responsive: `grpcurl -plaintext localhost:20000 list`
### Memory Issues
**Error**: Out of memory during tests
**Solution**:
```bash
# Run with memory limit for long-running tests
GODEBUG=madvdontneed=1 go test -timeout 5m ./...
# Monitor memory during tests
watch -n1 'ps aux | grep test'
```
## Contributing
When adding new features:
1. Add comprehensive documentation to public types/functions
2. Include usage examples for complex APIs
3. Add unit tests covering happy path and error cases
4. Add integration tests if server interaction required
5. Ensure code passes `go vet` and `go test -race`
6. Update this README if adding new features
## License
See LICENSE file for details.
---
**Need Help?**
- Check examples in `examples/` directory
- Run tests to see working code: `go test -v ./...`
- Review function documentation: `godoc` or inline comments
- Check troubleshooting section above

483
third_party/sglang/sgl-model-gateway/bindings/golang/client.go vendored Normal file
View File

@@ -0,0 +1,483 @@
// Package sglang provides a Go SDK for SGLang gRPC API.
//
// SGLang is a fast language model serving framework. This package provides a Go client
// library for interacting with SGLang's gRPC API, following the style of OpenAI's Go SDK.
//
// Basic usage:
//
// client, err := sglang.NewClient(sglang.ClientConfig{
// Endpoint: "grpc://localhost:20000",
// TokenizerPath: "/path/to/tokenizer",
// })
// if err != nil {
// log.Fatal(err)
// }
// defer client.Close()
//
// resp, err := client.CreateChatCompletion(ctx, sglang.ChatCompletionRequest{
// Model: "default",
// Messages: []sglang.ChatMessage{
// {Role: "user", Content: "Hello"},
// },
// })
//
// For streaming responses, use CreateChatCompletionStream instead.
package sglang
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"strings"
"sync"
"time"
grpcclient "github.com/sglang/sglang-go-grpc-sdk/internal/grpc"
)
// Client is the main client for interacting with SGLang gRPC API.
// It manages the connection to the SGLang server and handles both streaming
// and non-streaming chat completions.
//
// Thread-safe: All public methods are safe for concurrent use.
type Client struct {
endpoint string
tokenizerPath string
grpcClient *grpcclient.GrpcClient // gRPC-based client
mu sync.RWMutex
}
// ClientConfig holds configuration for creating a new client.
type ClientConfig struct {
// Endpoint is the gRPC endpoint URL (e.g., "grpc://localhost:20000").
// Required field. Must include the scheme (grpc://) and port number.
Endpoint string
// TokenizerPath is the path to the tokenizer directory containing
// tokenizer configuration files (e.g., tokenizer.json, vocab.json).
// Required field.
TokenizerPath string
// ChannelBufferSizes configures buffer sizes for internal channels.
// If nil, default values will be used (optimized for high concurrency).
ChannelBufferSizes *ChannelBufferSizes
// Timeouts configures timeout values for various operations.
// If nil, default values will be used.
Timeouts *Timeouts
}
// ChannelBufferSizes configures buffer sizes for internal channels.
// These affect concurrency and memory usage. Larger buffers allow more
// concurrent operations but use more memory.
type ChannelBufferSizes = grpcclient.ChannelBufferSizes
// Timeouts configures timeout values for various operations.
type Timeouts = grpcclient.Timeouts
// defaultChannelBufferSizes returns default channel buffer sizes optimized for high concurrency (10k+).
// These values are designed to handle thousands of concurrent requests without blocking.
func defaultChannelBufferSizes() ChannelBufferSizes {
return ChannelBufferSizes{
ResultJSONChan: 10000, // Increased for high concurrency: each request may produce 200-500 chunks
ErrChan: 100, // Errors are rare, 100 is sufficient
RecvChan: 2000, // Increased for high concurrency: more gRPC responses to buffer
}
}
// defaultTimeouts returns default timeout values.
func defaultTimeouts() Timeouts {
return Timeouts{
KeepaliveTime: 300 * time.Second, // Increased to reduce ping frequency and avoid "too many pings" errors
KeepaliveTimeout: 20 * time.Second,
CloseTimeout: 5 * time.Second,
}
}
// NewClient creates a new SGLang client with the given configuration.
//
// The client maintains a long-lived connection to the SGLang server and should
// be reused for multiple requests. Call Close() to release resources.
//
// Returns an error if:
// - Endpoint is empty
// - TokenizerPath is empty
// - Connection to the server fails
func NewClient(config ClientConfig) (*Client, error) {
if config.Endpoint == "" {
return nil, errors.New("endpoint is required")
}
if config.TokenizerPath == "" {
return nil, errors.New("tokenizer path is required")
}
bufferSizes := defaultChannelBufferSizes()
if config.ChannelBufferSizes != nil {
if config.ChannelBufferSizes.ResultJSONChan > 0 {
bufferSizes.ResultJSONChan = config.ChannelBufferSizes.ResultJSONChan
}
if config.ChannelBufferSizes.ErrChan > 0 {
bufferSizes.ErrChan = config.ChannelBufferSizes.ErrChan
}
if config.ChannelBufferSizes.RecvChan > 0 {
bufferSizes.RecvChan = config.ChannelBufferSizes.RecvChan
}
}
timeouts := defaultTimeouts()
if config.Timeouts != nil {
if config.Timeouts.KeepaliveTime > 0 {
timeouts.KeepaliveTime = config.Timeouts.KeepaliveTime
}
if config.Timeouts.KeepaliveTimeout > 0 {
timeouts.KeepaliveTimeout = config.Timeouts.KeepaliveTimeout
}
if config.Timeouts.CloseTimeout > 0 {
timeouts.CloseTimeout = config.Timeouts.CloseTimeout
}
}
grpcClient, err := grpcclient.NewGrpcClient(config.Endpoint, config.TokenizerPath, bufferSizes, timeouts)
if err != nil {
return nil, fmt.Errorf("failed to create gRPC client: %w", err)
}
return &Client{
endpoint: config.Endpoint,
tokenizerPath: config.TokenizerPath,
grpcClient: grpcClient,
}, nil
}
// Close closes the client and releases all resources.
//
// After Close() is called, the client cannot be used for further requests.
// Calling Close() multiple times is safe and idempotent.
func (c *Client) Close() error {
c.mu.Lock()
defer c.mu.Unlock()
if c.grpcClient != nil {
if err := c.grpcClient.Close(); err != nil {
return err
}
c.grpcClient = nil
}
return nil
}
// ChatCompletionRequest represents a request for chat completion.
// It follows the OpenAI API style for familiar usage.
type ChatCompletionRequest struct {
// Model specifies the model to use for completion (e.g., "default")
Model string `json:"model"`
// Messages is the list of messages in the conversation
Messages []ChatMessage `json:"messages"`
Temperature *float32 `json:"temperature,omitempty"`
TopP *float32 `json:"top_p,omitempty"`
TopK *int `json:"top_k,omitempty"`
MaxCompletionTokens *int `json:"max_completion_tokens,omitempty"`
Stream bool `json:"stream"`
Tools []Tool `json:"tools,omitempty"`
ToolChoice interface{} `json:"tool_choice,omitempty"`
Stop interface{} `json:"stop,omitempty"`
StopTokenIDs []int `json:"stop_token_ids,omitempty"`
SkipSpecialTokens bool `json:"skip_special_tokens,omitempty"`
FrequencyPenalty *float32 `json:"frequency_penalty,omitempty"`
PresencePenalty *float32 `json:"presence_penalty,omitempty"`
ResponseFormat *ResponseFormat `json:"response_format,omitempty"`
Seed *int `json:"seed,omitempty"`
Logprobs bool `json:"logprobs,omitempty"`
TopLogprobs *int `json:"top_logprobs,omitempty"`
User string `json:"user,omitempty"`
}
// ChatMessage represents a single message in a chat conversation
type ChatMessage struct {
Role string `json:"role"`
Content interface{} `json:"content"`
Name string `json:"name,omitempty"`
}
// Tool represents a tool/function that can be called
type Tool struct {
Type string `json:"type"`
Function Function `json:"function"`
}
// Function represents a function definition
type Function struct {
Name string `json:"name"`
Description string `json:"description,omitempty"`
Parameters map[string]interface{} `json:"parameters"`
}
// ResponseFormat represents the response format
type ResponseFormat struct {
Type string `json:"type"`
}
// ChatCompletionResponse represents a non-streaming chat completion response
type ChatCompletionResponse struct {
ID string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
Model string `json:"model"`
SystemFingerprint string `json:"system_fingerprint,omitempty"`
Choices []Choice `json:"choices"`
Usage Usage `json:"usage"`
}
// Choice represents a choice in the completion response
type Choice struct {
Index int `json:"index"`
Message Message `json:"message"`
FinishReason string `json:"finish_reason"`
}
// Message represents a message in the response
type Message struct {
Role string `json:"role"`
Content string `json:"content"`
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
}
// ToolCall represents a tool call in the response
type ToolCall struct {
ID string `json:"id"`
Type string `json:"type"`
Function FunctionCall `json:"function"`
}
// FunctionCall represents a function call
type FunctionCall struct {
Name string `json:"name"`
Arguments string `json:"arguments"`
}
// Usage represents token usage information
type Usage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
}
// ChatCompletionStreamResponse represents a streaming chat completion response
type ChatCompletionStreamResponse struct {
ID string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
Model string `json:"model"`
SystemFingerprint string `json:"system_fingerprint,omitempty"`
Choices []StreamChoice `json:"choices"`
Usage *Usage `json:"usage,omitempty"`
}
// StreamChoice represents a choice in a streaming response
type StreamChoice struct {
Index int `json:"index"`
Delta MessageDelta `json:"delta"`
FinishReason string `json:"finish_reason,omitempty"`
}
// MessageDelta represents incremental message updates
type MessageDelta struct {
Role string `json:"role,omitempty"`
Content string `json:"content,omitempty"`
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
}
// CreateChatCompletion creates a non-streaming chat completion with context support.
//
// Context Support:
// The ctx parameter is fully supported for cancellation and timeouts:
// - If ctx is cancelled, the request will be interrupted on the next stream.RecvJSON() call
// - If ctx times out, the request will return context.DeadlineExceeded
//
// Example with timeout:
//
// ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
// defer cancel()
// resp, err := client.CreateChatCompletion(ctx, req)
//
// Note: Internally, this creates a stream and collects all chunks,
// so context monitoring happens at the chunk level.
func (c *Client) CreateChatCompletion(ctx context.Context, req ChatCompletionRequest) (*ChatCompletionResponse, error) {
// For non-streaming, we'll collect all chunks and return the final response
req.Stream = true // We still use streaming internally, but collect all chunks
if len(req.Tools) == 0 {
req.Tools = nil
}
stream, err := c.CreateChatCompletionStream(ctx, req)
if err != nil {
return nil, err
}
defer stream.Close()
var fullContent strings.Builder
var fullToolCalls []ToolCall
var finishReason string
var usage Usage
var responseID string
var created int64
var model string
var systemFingerprint string
for {
chunkJSON, err := stream.RecvJSON()
if err == io.EOF {
break
}
if err != nil {
return nil, err
}
var chunk ChatCompletionStreamResponse
if err := json.Unmarshal([]byte(chunkJSON), &chunk); err != nil {
return nil, fmt.Errorf("failed to parse chunk: %w", err)
}
if chunk.ID != "" {
responseID = chunk.ID
}
if chunk.Created > 0 {
created = chunk.Created
}
if chunk.Model != "" {
model = chunk.Model
}
if chunk.SystemFingerprint != "" {
systemFingerprint = chunk.SystemFingerprint
}
for _, choice := range chunk.Choices {
if choice.Delta.Content != "" {
fullContent.WriteString(choice.Delta.Content)
}
if len(choice.Delta.ToolCalls) > 0 {
fullToolCalls = append(fullToolCalls, choice.Delta.ToolCalls...)
}
if choice.FinishReason != "" {
finishReason = choice.FinishReason
}
}
if chunk.Usage != nil {
usage = *chunk.Usage
}
}
message := Message{
Role: "assistant",
Content: fullContent.String(),
}
if len(fullToolCalls) > 0 {
message.ToolCalls = fullToolCalls
}
if finishReason == "" {
finishReason = "stop"
}
return &ChatCompletionResponse{
ID: responseID,
Object: "chat.completion",
Created: created,
Model: model,
SystemFingerprint: systemFingerprint,
Choices: []Choice{
{
Index: 0,
Message: message,
FinishReason: finishReason,
},
},
Usage: usage,
}, nil
}
// ChatCompletionStream represents a streaming chat completion
type ChatCompletionStream struct {
grpcStream *grpcclient.GrpcChatCompletionStream
ctx context.Context
cancel context.CancelFunc
}
func (s *ChatCompletionStream) RecvJSON() (string, error) {
return s.grpcStream.RecvJSON()
}
// Close closes the stream and cancels any pending operations.
func (s *ChatCompletionStream) Close() error {
if s.cancel != nil {
s.cancel()
}
if s.grpcStream != nil {
return s.grpcStream.Close()
}
return nil
}
// CreateChatCompletionStream creates a streaming chat completion with context cancellation support.
//
// Context Support:
// The ctx parameter is now fully supported for cancellation and timeouts:
// - If ctx is cancelled, stream.RecvJSON() will return context.Canceled on the next call
// - If ctx times out (WithTimeout), stream.RecvJSON() will return context.DeadlineExceeded
// - Calling stream.Close() also cancels the context
//
// Example with timeout:
//
// ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
// defer cancel()
// stream, err := client.CreateChatCompletionStream(ctx, req)
// // Stream will auto-close if 30 seconds elapse
//
// Example with cancellation:
//
// ctx, cancel := context.WithCancel(context.Background())
// stream, err := client.CreateChatCompletionStream(ctx, req)
// go func() {
// time.Sleep(5*time.Second)
// cancel() // Cancel after 5 seconds
// }()
func (c *Client) CreateChatCompletionStream(ctx context.Context, req ChatCompletionRequest) (*ChatCompletionStream, error) {
reqJSON, err := json.Marshal(req)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
}
var reqMap map[string]interface{}
if err := json.Unmarshal(reqJSON, &reqMap); err != nil {
return nil, fmt.Errorf("failed to unmarshal request to map: %w", err)
}
if _, exists := reqMap["tools"]; !exists {
reqMap["tools"] = []interface{}{}
}
reqJSON, err = json.Marshal(reqMap)
if err != nil {
return nil, fmt.Errorf("failed to marshal request map to JSON: %w", err)
}
if c.grpcClient == nil {
return nil, errors.New("gRPC client is closed")
}
grpcStream, err := c.grpcClient.CreateChatCompletionStream(ctx, string(reqJSON))
if err != nil {
return nil, fmt.Errorf("failed to create gRPC stream: %w", err)
}
streamCtx, cancel := context.WithCancel(ctx)
return &ChatCompletionStream{
grpcStream: grpcStream,
ctx: streamCtx,
cancel: cancel,
}, nil
}

325
third_party/sglang/sgl-model-gateway/bindings/golang/client_test.go vendored Normal file
View File

@@ -0,0 +1,325 @@
package sglang
import (
"context"
"testing"
)
// TestClientConfig tests ClientConfig validation
func TestClientConfig(t *testing.T) {
tests := []struct {
name string
config ClientConfig
wantErr bool
}{
{
name: "valid config",
config: ClientConfig{
Endpoint: "grpc://localhost:20000",
TokenizerPath: "/path/to/tokenizer",
},
wantErr: false,
},
{
name: "missing endpoint",
config: ClientConfig{
Endpoint: "",
TokenizerPath: "/path/to/tokenizer",
},
wantErr: true,
},
{
name: "missing tokenizer path",
config: ClientConfig{
Endpoint: "grpc://localhost:20000",
TokenizerPath: "",
},
wantErr: true,
},
{
name: "both missing",
config: ClientConfig{
Endpoint: "",
TokenizerPath: "",
},
wantErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
_, err := NewClient(tt.config)
if (err != nil) != tt.wantErr {
t.Errorf("NewClient() error = %v, wantErr %v", err, tt.wantErr)
}
})
}
}
// TestChatMessageTypes tests ChatMessage struct and its variants
func TestChatMessageTypes(t *testing.T) {
msg := ChatMessage{
Role: "user",
Content: "Hello",
}
if msg.Role != "user" {
t.Errorf("Expected role 'user', got '%s'", msg.Role)
}
if msg.Content != "Hello" {
t.Errorf("Expected content 'Hello', got '%s'", msg.Content)
}
}
// TestChatCompletionRequestValidation tests ChatCompletionRequest validation
func TestChatCompletionRequestValidation(t *testing.T) {
// Test valid request
req := ChatCompletionRequest{
Model: "default",
Messages: []ChatMessage{
{Role: "user", Content: "test"},
},
Stream: false,
}
if req.Model == "" {
t.Error("Expected model to be set")
}
if len(req.Messages) == 0 {
t.Error("Expected messages to be non-empty")
}
if req.Messages[0].Role != "user" {
t.Errorf("Expected first message role 'user', got '%s'", req.Messages[0].Role)
}
}
// TestClientClose tests that Close can be called multiple times safely
func TestClientClose(t *testing.T) {
// Create a mock client (note: in real tests, you might want to skip this
// if it requires actual server connection)
config := ClientConfig{
Endpoint: "grpc://localhost:20000",
TokenizerPath: "/path/to/tokenizer",
}
// Skip if connection fails (expected in unit test environment)
client, err := NewClient(config)
if err != nil {
t.Skip("Skipping client close test: server not available")
}
// First close should succeed
if err := client.Close(); err != nil {
t.Errorf("First Close() failed: %v", err)
}
// Second close should also succeed (idempotent)
if err := client.Close(); err != nil {
t.Errorf("Second Close() failed: %v", err)
}
}
// TestChatCompletionResponseTypes tests response type structures
func TestChatCompletionResponseTypes(t *testing.T) {
resp := ChatCompletionResponse{
ID: "test-id",
Model: "default",
Created: 1234567890,
Choices: []Choice{
{
Message: Message{
Role: "assistant",
Content: "Hello",
},
FinishReason: "stop",
},
},
Usage: Usage{
PromptTokens: 10,
CompletionTokens: 20,
TotalTokens: 30,
},
}
if resp.ID != "test-id" {
t.Errorf("Expected ID 'test-id', got '%s'", resp.ID)
}
if len(resp.Choices) != 1 {
t.Errorf("Expected 1 choice, got %d", len(resp.Choices))
}
if resp.Choices[0].Message.Content != "Hello" {
t.Errorf("Expected content 'Hello', got '%s'", resp.Choices[0].Message.Content)
}
if resp.Usage.TotalTokens != 30 {
t.Errorf("Expected total tokens 30, got %d", resp.Usage.TotalTokens)
}
}
// TestStreamingResponseTypes tests streaming response structures
func TestStreamingResponseTypes(t *testing.T) {
chunk := ChatCompletionStreamResponse{
ID: "stream-id",
Created: 1234567890,
Choices: []StreamChoice{
{
Index: 0,
Delta: MessageDelta{
Content: "Hello",
},
FinishReason: "",
},
},
}
if chunk.ID != "stream-id" {
t.Errorf("Expected ID 'stream-id', got '%s'", chunk.ID)
}
if len(chunk.Choices) == 0 {
t.Error("Expected at least one choice")
}
if chunk.Choices[0].Delta.Content != "Hello" {
t.Errorf("Expected delta content 'Hello', got '%s'", chunk.Choices[0].Delta.Content)
}
}
// TestToolCallStructure tests Tool and ToolCall structures
func TestToolCallStructure(t *testing.T) {
tool := Tool{
Type: "function",
Function: Function{
Name: "get_weather",
Description: "Get the weather",
Parameters: map[string]interface{}{
"location": "string",
},
},
}
if tool.Type != "function" {
t.Errorf("Expected tool type 'function', got '%s'", tool.Type)
}
if tool.Function.Name != "get_weather" {
t.Errorf("Expected function name 'get_weather', got '%s'", tool.Function.Name)
}
toolCall := ToolCall{
ID: "call-123",
Type: "function",
Function: FunctionCall{
Name: "get_weather",
Arguments: `{"location": "San Francisco"}`,
},
}
if toolCall.ID != "call-123" {
t.Errorf("Expected tool call ID 'call-123', got '%s'", toolCall.ID)
}
}
// TestConcurrentClientOperations tests thread safety
// This is a basic test that just verifies concurrent calls don't panic
func TestConcurrentClientOperations(t *testing.T) {
config := ClientConfig{
Endpoint: "grpc://localhost:20000",
TokenizerPath: "/path/to/tokenizer",
}
client, err := NewClient(config)
if err != nil {
t.Skip("Skipping concurrent operations test: server not available")
}
defer client.Close()
// Try concurrent Close calls (should not panic or race)
done := make(chan bool, 2)
go func() {
client.Close()
done <- true
}()
go func() {
client.Close()
done <- true
}()
<-done
<-done
}
// BenchmarkChatCompletionRequest benchmarks request creation
func BenchmarkChatCompletionRequest(b *testing.B) {
for i := 0; i < b.N; i++ {
_ = ChatCompletionRequest{
Model: "default",
Messages: []ChatMessage{
{Role: "user", Content: "test message"},
},
Stream: false,
Temperature: floatPtr(0.7),
MaxCompletionTokens: intPtr(100),
}
}
}
// Helper functions for benchmarks
func floatPtr(f float32) *float32 {
return &f
}
func intPtr(i int) *int {
return &i
}
// TestContextCancellation tests that cancelled context is handled gracefully.
//
// NOTE: Currently, the FFI layer is blocking and doesn't actively monitor context cancellation.
// This test verifies that the client at least returns an error rather than panicking or
// hanging indefinitely when a pre-cancelled context is passed.
//
// Future: When FFI supports context cancellation (via signals or async operations),
// this test should be updated to assert that the error is context.Canceled or wrapped
// context cancellation error.
func TestContextCancellation(t *testing.T) {
config := ClientConfig{
Endpoint: "grpc://localhost:20000",
TokenizerPath: "/path/to/tokenizer",
}
client, err := NewClient(config)
if err != nil {
t.Skip("Skipping context cancellation test: server not available")
}
defer client.Close()
// Create a pre-cancelled context
ctx, cancel := context.WithCancel(context.Background())
cancel()
req := ChatCompletionRequest{
Model: "default",
Messages: []ChatMessage{
{Role: "user", Content: "test"},
},
}
// Attempt request with cancelled context
// Since FFI is blocking, we expect either:
// 1. An error from the server/network
// 2. The call to complete normally (FFI doesn't check context)
// What we DON'T expect is a panic or indefinite hang
_, err = client.CreateChatCompletion(ctx, req)
if err != nil {
t.Logf("Request with cancelled context returned error: %v", err)
} else {
t.Logf("Request with cancelled context completed (FFI may not support context cancellation)")
}
}

239
third_party/sglang/sgl-model-gateway/bindings/golang/examples/oai_server/Makefile vendored Normal file
View File

@@ -0,0 +1,239 @@
# Makefile for OAI Server
# Builds binary, runs tests, and provides basic targets
# Configuration
APP_NAME = oai_server
VERSION ?= $(shell git describe --tags --always --dirty 2>/dev/null || echo "dev")
BUILD_TIME := $(shell date -u '+%Y-%m-%d_%H:%M:%S')
GIT_COMMIT := $(shell git rev-parse --short HEAD 2>/dev/null || echo "unknown")
# Paths
ROOT_DIR := $(shell pwd)
BINDINGS_DIR := $(shell cd $(ROOT_DIR)/../.. && pwd)
BUILD_DIR := $(ROOT_DIR)/build
BINARY := $(BUILD_DIR)/$(APP_NAME)
# Rust FFI library paths
LIB_DIR := $(BINDINGS_DIR)/lib
LIB_NAME = libsgl_model_gateway_go
# Detect OS
UNAME_S := $(shell uname -s)
ifeq ($(UNAME_S),Linux)
LIB_EXT = .so
LD_LIBRARY_PATH_VAR = LD_LIBRARY_PATH
ARCH := $(shell uname -m)
ifeq ($(ARCH),x86_64)
GOARCH = amd64
else ifeq ($(ARCH),aarch64)
GOARCH = arm64
endif
endif
ifeq ($(UNAME_S),Darwin)
LIB_EXT = .dylib
LD_LIBRARY_PATH_VAR = DYLD_LIBRARY_PATH
ARCH := $(shell uname -m)
ifeq ($(ARCH),x86_64)
GOARCH = amd64
else ifeq ($(ARCH),arm64)
GOARCH = arm64
endif
endif
# Build flags
LDFLAGS = -X main.Version=$(VERSION) -X main.BuildTime=$(BUILD_TIME) -X main.GitCommit=$(GIT_COMMIT)
GO_BUILD_FLAGS = -ldflags "$(LDFLAGS)"
# Python LDFLAGS (needed for Rust FFI that depends on Python)
PYTHON_LDFLAGS := $(shell python3-config --ldflags --embed 2>/dev/null || python3-config --ldflags 2>/dev/null || python-config --ldflags --embed 2>/dev/null || python-config --ldflags 2>/dev/null || echo "")
# CGO flags
CGO_LDFLAGS = -L$(LIB_DIR) $(PYTHON_LDFLAGS)
.PHONY: all build build-dev test e2e clean help lib run stream check-rust-lib check-server
# E2E test configuration
E2E_HOST ?= localhost
E2E_PORT ?= 8080
E2E_MODEL ?= default
E2E_TOKENIZER ?= $(shell echo $$SGL_TOKENIZER_PATH || echo "./examples/tokenizer")
E2E_NUM_PROMPTS ?= 100
E2E_INPUT_LEN ?= 1024
E2E_OUTPUT_LEN ?= 512
E2E_REQUEST_RATE ?= 20
E2E_MAX_CONCURRENCY ?= 20
E2E_BASE_URL ?= http://$(E2E_HOST):$(E2E_PORT)
help:
@echo "OAI Server Makefile"
@echo ""
@echo "Available targets:"
@echo " lib - Build Rust FFI library"
@echo " build - Build binary (release mode)"
@echo " build-dev - Build binary (debug mode)"
@echo " test - Run tests"
@echo " e2e - Run end-to-end test with bench_serving.py"
@echo " run - Run the server (development)"
@echo " stream - Run streaming example"
@echo " clean - Clean build artifacts"
@echo ""
@echo "E2E test variables:"
@echo " E2E_HOST - OAI Server host (default: localhost)"
@echo " E2E_PORT - OAI Server port (default: 8080)"
@echo " E2E_MODEL - Model name (default: default)"
@echo " E2E_TOKENIZER - Tokenizer path"
@echo " E2E_NUM_PROMPTS - Number of prompts (default: 100)"
@echo " E2E_INPUT_LEN - Input token length (default: 1024)"
@echo " E2E_OUTPUT_LEN - Output token length (default: 512)"
@echo " E2E_REQUEST_RATE - Request rate per second (default: 20)"
@echo " E2E_MAX_CONCURRENCY - Max concurrent requests (default: 20)"
all: build
# Build Rust FFI library
lib:
@echo "Building Rust FFI library..."
@cd $(BINDINGS_DIR) && $(MAKE) lib
@echo "✓ Rust FFI library built"
# Check if Rust FFI library exists
check-rust-lib:
@if [ ! -f "$(LIB_DIR)/$(LIB_NAME)$(LIB_EXT)" ]; then \
echo "Error: Rust FFI library not found at $(LIB_DIR)/$(LIB_NAME)$(LIB_EXT)"; \
echo "Building Rust library..."; \
cd $(BINDINGS_DIR) && $(MAKE) lib; \
fi
@echo "✓ Rust FFI library found"
# Build binary (release)
build: check-rust-lib
@echo "Building $(APP_NAME) (release mode)..."
@mkdir -p $(BUILD_DIR)
@CGO_ENABLED=1 \
CGO_LDFLAGS="$(CGO_LDFLAGS)" \
GOOS=$(shell go env GOOS) \
GOARCH=$(GOARCH) \
go build $(GO_BUILD_FLAGS) -o $(BINARY) .
@echo "✓ Binary built: $(BINARY)"
# Build binary (debug)
build-dev: check-rust-lib
@echo "Building $(APP_NAME) (debug mode)..."
@mkdir -p $(BUILD_DIR)
@CGO_ENABLED=1 \
CGO_LDFLAGS="$(CGO_LDFLAGS)" \
go build -o $(BINARY) .
@echo "✓ Binary built (debug): $(BINARY)"
# Run tests
test: check-rust-lib
@echo "Running tests..."
@CGO_ENABLED=1 \
CGO_LDFLAGS="$(CGO_LDFLAGS)" \
export $(LD_LIBRARY_PATH_VAR)="$(LIB_DIR):$$$(LD_LIBRARY_PATH_VAR)" && \
go test -v ./...
@echo "✓ Tests completed"
# Check if OAI Server is running
check-server:
@echo "Checking if OAI Server is running at $(E2E_BASE_URL)..."
@if curl -s -f $(E2E_BASE_URL)/health > /dev/null 2>&1; then \
echo "✓ OAI Server is running"; \
exit 0; \
else \
echo "✗ OAI Server is not running at $(E2E_BASE_URL)"; \
echo " Start it with: make run"; \
exit 1; \
fi
# Find sglang project root (4 levels up from oai_server)
SGLANG_ROOT := $(shell cd $(ROOT_DIR)/../../../../.. && pwd)
# Run end-to-end test with bench_serving.py
e2e: check-server
@echo "Checking if bench_serving.py is available..."
@if python -m sglang.bench_serving --help > /dev/null 2>&1; then \
echo "✓ Using installed bench_serving.py module"; \
USE_SGLANG_ROOT=false; \
elif [ -f "$(SGLANG_ROOT)/python/sglang/bench_serving.py" ]; then \
echo "✓ Using bench_serving.py from $(SGLANG_ROOT)"; \
USE_SGLANG_ROOT=true; \
else \
echo "✗ bench_serving.py is not available"; \
echo " Install dependencies: pip install aiohttp numpy datasets transformers tqdm pillow pybase64"; \
exit 1; \
fi
@echo "Running end-to-end test with bench_serving.py..."
@echo "Configuration:"
@echo " Server: $(E2E_BASE_URL)"
@if [ "$(E2E_MODEL)" != "default" ]; then \
echo " Model: $(E2E_MODEL)"; \
fi
@if [ -n "$(E2E_TOKENIZER)" ]; then \
echo " Tokenizer: $(E2E_TOKENIZER)"; \
fi
@echo " Prompts: $(E2E_NUM_PROMPTS)"
@echo " Input/Output: $(E2E_INPUT_LEN)/$(E2E_OUTPUT_LEN) tokens"
@echo " Request rate: $(E2E_REQUEST_RATE) req/s"
@echo " Max concurrency: $(E2E_MAX_CONCURRENCY)"
@echo ""
@TOKENIZER_ABS=$$(cd $(ROOT_DIR) && python3 -c "import os; path='$(E2E_TOKENIZER)'; print(os.path.abspath(path) if not os.path.isabs(path) else path)" 2>/dev/null || echo "$(E2E_TOKENIZER)"); \
if [ -n "$(E2E_TOKENIZER)" ]; then \
if [ -n "$$TOKENIZER_ABS" ] && ([ -d "$$TOKENIZER_ABS" ] || [ -f "$$TOKENIZER_ABS" ]); then \
TOKENIZER_ARG="--tokenizer $$TOKENIZER_ABS"; \
else \
TOKENIZER_ARG="--tokenizer $(E2E_TOKENIZER)"; \
fi; \
else \
TOKENIZER_ARG=""; \
fi; \
if [ "$$USE_SGLANG_ROOT" = "true" ]; then \
cd $(SGLANG_ROOT) && PYTHONPATH=$(SGLANG_ROOT)/python:$$PYTHONPATH python python/sglang/bench_serving.py \
--backend sglang-oai-chat \
--base-url $(E2E_BASE_URL) \
$$([ "$(E2E_MODEL)" != "default" ] && echo "--model $(E2E_MODEL)") \
$$TOKENIZER_ARG \
--dataset-name random \
--num-prompts $(E2E_NUM_PROMPTS) \
--random-input-len $(E2E_INPUT_LEN) \
--random-output-len $(E2E_OUTPUT_LEN) \
--request-rate $(E2E_REQUEST_RATE) \
--max-concurrency $(E2E_MAX_CONCURRENCY) \
--warmup-requests 5 \
--disable-tqdm || (echo "✗ E2E test failed"; exit 1); \
else \
python -m sglang.bench_serving \
--backend sglang-oai-chat \
--base-url $(E2E_BASE_URL) \
$$([ "$(E2E_MODEL)" != "default" ] && echo "--model $(E2E_MODEL)") \
$$TOKENIZER_ARG \
--dataset-name random \
--num-prompts $(E2E_NUM_PROMPTS) \
--random-input-len $(E2E_INPUT_LEN) \
--random-output-len $(E2E_OUTPUT_LEN) \
--request-rate $(E2E_REQUEST_RATE) \
--max-concurrency $(E2E_MAX_CONCURRENCY) \
--warmup-requests 5 \
--disable-tqdm || (echo "✗ E2E test failed"; exit 1); \
fi
@echo ""
@echo "✓ E2E test completed"
# Run the server (development)
run: build-dev
@echo "Running server..."
@export $(LD_LIBRARY_PATH_VAR)="$(LIB_DIR):$$$(LD_LIBRARY_PATH_VAR)" && \
$(BINARY)
# Run streaming example
stream: check-rust-lib
@echo "Running streaming example..."
@cd $(BINDINGS_DIR)/examples/streaming && \
export $(LD_LIBRARY_PATH_VAR)="$(LIB_DIR):$$$(LD_LIBRARY_PATH_VAR)" && \
bash run.sh
# Clean build artifacts
clean:
@echo "Cleaning build artifacts..."
@rm -rf $(BUILD_DIR)
@echo "✓ Clean complete"

305
third_party/sglang/sgl-model-gateway/bindings/golang/examples/oai_server/README.md vendored Normal file
View File

@@ -0,0 +1,305 @@
# Go SGLang Router - OpenAI Compatible API Server
Go SGLang Router is a high-performance OpenAI-compatible API server that communicates with the SGLang backend via gRPC and performs efficient preprocessing and postprocessing through Rust FFI.
## Features
- ✅ **OpenAI API Compatible**: Fully compatible with OpenAI Chat Completions API
- ✅ **High Performance**: Low latency and high throughput using gRPC and Rust FFI
- ✅ **Streaming Support**: Server-Sent Events (SSE) streaming responses
- ✅ **Thread-Safe**: Pre-created tokenizer handle, lock-free concurrency
- ✅ **Graceful Shutdown**: Context cancellation mechanism to avoid resource leaks and panics
- ✅ **Configurable**: Supports configuring channel buffer sizes and timeout durations
## Architecture Overview
**Important Note**: gRPC mode **still calls FFI**, which is used for:
- **Preprocessing**: chat_template and tokenization (request phase)
- **Postprocessing**: token decoding and tool parsing (response phase)
gRPC is only used for communication with the SGLang backend, while input/output processing completely relies on Rust FFI.
```
┌─────────────────────────────────────────────────────────────────┐
│ HTTP Client │
│ (OpenAI API Format) │
└────────────────────────────┬────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ FastHTTP Server │
│ handlers/chat.go:HandleChatCompletion │
│ - Parse request JSON │
│ - SetBodyStreamWriter (SSE) │
└────────────────────────────┬────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ SGLang Client (client.go) │
│ CreateChatCompletionStream(ctx, req) │
│ - Wraps gRPC client │
└────────────────────────────┬────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ gRPC Client (internal/grpc/client_grpc.go) │
│ CreateChatCompletionStream(ctx, reqJSON) │
│ │
│ ┌──────────────────────────────────────────────────────────┐ │
│ │ Step 1: FFI Preprocess (Rust FFI) │ │
│ │ - ffi.PreprocessChatRequestWithTokenizer() │ │
│ │ - chat_template application │ │
│ │ - tokenization │ │
│ │ - tool constraints generation │ │
│ │ Returns: PromptText, TokenIDs, ToolConstraintsJSON, │ │
│ │ PromptTokens │ │
│ └────────────────────┬─────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌──────────────────────────────────────────────────────────┐ │
│ │ Step 2: Build gRPC Request │ │
│ │ - Parse request JSON (model, temperature, etc.) │ │
│ │ - Create proto.GenerateRequest │ │
│ │ - Set TokenizedInput (PromptText, TokenIDs) │ │
│ │ - Set SamplingParams (temperature, top_p, top_k, etc.) │ │
│ │ - Set Constraints (from ToolConstraintsJSON) │ │
│ └────────────────────┬─────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌──────────────────────────────────────────────────────────┐ │
│ │ Step 3: Create gRPC Stream │ │
│ │ - client.Generate(generateReq) → gRPC stream │ │
│ │ - Connects to SGLang Backend (Rust) │ │
│ └────────────────────┬─────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌──────────────────────────────────────────────────────────┐ │
│ │ Step 4: Create Converter & BatchPostprocessor │ │
│ │ - ffi.CreateGrpcResponseConverterWithTokenizer() │ │
│ │ - Uses preprocessed.PromptTokens for initial count │ │
│ │ - ffi.NewBatchPostprocessor(batchSize=1, immediate) │ │
│ └────────────────────┬─────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌──────────────────────────────────────────────────────────┐ │
│ │ Step 5: Start readLoop (Background Goroutine) │ │
│ │ - go grpcStream.readLoop() │ │
│ │ - Returns GrpcChatCompletionStream immediately │ │
│ └────────────────────┬─────────────────────────────────────┘ │
└───────────────────────┼────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ GrpcChatCompletionStream.readLoop() │
│ (Background Goroutine) │
│ │
│ ┌──────────────────────────────────────────────────────────┐ │
│ │ Recv() Goroutine (Dedicated) │ │
│ │ - Continuously calls stream.Recv() │ │
│ │ - Sends results to recvChan (buffered, 2000) │ │
│ │ - Exits on ctx.Done() or error │ │
│ │ - Calls stream.CloseSend() on ctx.Done() │ │
│ └────────────────────┬─────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌──────────────────────────────────────────────────────────┐ │
│ │ Main Loop │ │
│ │ - Reads from recvChan │ │
│ │ - For each proto.GenerateResponse: │ │
│ │ → go processAndSendResponse() (async) │ │
│ │ - protoToJSON() converts proto to JSON string │ │
│ │ - batchPostprocessor.AddChunk(protoJSON) │ │
│ │ → FFI postprocessing (token decoding, tool parsing)│ │
│ │ → Returns OpenAI-format JSON strings │ │
│ │ - Sends JSON to resultJSONChan (buffered, 10000) │ │
│ │ - All operations check ctx.Done() for cancellation │ │
│ │ - On EOF: flush batch, send remaining results, return │ │
│ │ - On error: send to errChan (buffered, 100) │ │
│ │ - defer: cancel ctx, wait goroutines, close channels │ │
│ └────────────────────┬─────────────────────────────────────┘ │
└───────────────────────┼────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ resultJSONChan (Buffered Channel, 10000) │
│ - Contains OpenAI-format JSON strings │
│ - Ready for consumption │
└────────────────────────────┬────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ ChatCompletionStream.RecvJSON() │
│ (client.go:410) │
│ - Direct wrapper: return grpcStream.RecvJSON() │
│ - No intermediate processing │
└────────────────────────────┬────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ FastHTTP SetBodyStreamWriter │
│ (handlers/chat.go:159) │
│ - Loop: stream.RecvJSON() → format SSE → flush │
│ - Format: "data: {json}\n\n" │
│ - Final: "data: [DONE]\n\n" │
│ - Immediate flush after each chunk │
└────────────────────────────┬────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ HTTP Client │
│ (SSE Stream) │
│ Receives: data: {...}\n\n │
└─────────────────────────────────────────────────────────────────┘
```
## Quick Start
### Start Server
```bash
./run.sh
```
The server will start on port `:8080`.
### Usage Example
```bash
curl http://localhost:8080/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "/path/to/model",
"messages": [{"role": "user", "content": "Hello!"}],
"stream": true
}'
```
## Key Design
### 1. Thread-Safe Tokenizer
- Pre-create `TokenizerHandle` at startup
- Rust side uses `Arc<dyn TokenizerTrait>`, thread-safe
- Lock-free concurrency, eliminating lock contention
### 2. Context Cancellation Mechanism (Graceful Shutdown)
- Use `context.Context` cancellation mechanism
- In `readLoop`'s `defer`: cancel context first, then wait for all goroutines to complete, finally close channels
- `processAndSendResponse` checks `ctx.Done()` at function start, all `select` statements include `case <-s.ctx.Done()`
- Avoids "send on closed channel" panic
### 3. Cancellable Recv()
- Use dedicated goroutine to execute `Recv()`
- Pass results through `recvChan`
- Call `CloseSend()` when context is cancelled to make `Recv()` return error
### 4. Simplified Channel Design
- `resultJSONChan`: Main data channel (gRPC layer)
- `errChan`: Error channel (gRPC layer)
- `recvChan`: Internal communication channel (gRPC layer)
- Removed redundant channels and duplicate reads
## Configuration
### Channel Buffer Sizes
```go
type ChannelBufferSizes struct {
ResultJSONChan int // Default: 10000
ErrChan int // Default: 100
RecvChan int // Default: 2000
}
```
### Timeout Configuration
```go
type Timeouts struct {
KeepaliveTime time.Duration // Default: 300s
KeepaliveTimeout time.Duration // Default: 20s
CloseTimeout time.Duration // Default: 5s
}
```
## Performance Optimizations
1. **Pre-create Tokenizer**: Created at startup to avoid first request latency
2. **Lock-Free Concurrency**: Tokenizer is thread-safe, no locks needed
3. **Lazy Parsing**: JSON parsing deferred until needed
4. **Direct JSON Passing**: `RecvJSON()` avoids parse/serialize overhead
5. **Immediate Batching**: batchSize=1, no delay
6. **Async Processing**: `readLoop` processes in background, doesn't block request handling
7. **Configurable Buffers**: Adjust channel sizes based on concurrency needs
## File Structure
```
sgl-model-gateway/bindings/golang/
├── client.go # High-level client API
├── internal/
│ ├── grpc/
│ │ └── client_grpc.go # gRPC client implementation
│ ├── ffi/ # FFI bindings (Rust)
│ └── proto/ # Protobuf definitions
└── examples/
└── oai_server/
├── handlers/
│ └── chat.go # HTTP request handling
├── models/
│ └── chat.go # Request/response models
└── service/
└── sglang_service.go # Service layer
```
## Error Handling
### Context Cancellation Mechanism
1. **Client disconnects** → `SetBodyStreamWriter` detects flush error
2. **Cancel streamCtx** → `readLoop` detects `ctx.Done()`
3. **Call stream.CloseSend()** → `Recv()` goroutine returns error
4. **readLoop defer executes**:
- Set `closed` flag
- Cancel context (if not already cancelled)
- Wait for all `processAndSendResponse` goroutines to complete (`processWg.Wait()`)
- Close all channels (`resultJSONChan`, `errChan`, `readLoopDone`)
5. **Clean up resources and exit**
### Channel Blocking and Race Condition Prevention
- **Context cancellation mechanism**: All channel sends use `select` statements with `case <-s.ctx.Done()`
- **Graceful exit**: When context is cancelled, all blocking send operations can return immediately
- **WaitGroup synchronization**: `readLoop`'s `defer` uses `processWg.Wait()` to ensure all goroutines complete before closing channels
- **Avoid panic**: Through context cancellation and WaitGroup synchronization, avoids "send on closed channel" panic
## Key Functions
### CreateChatCompletionStream
**Location**: `internal/grpc/client_grpc.go:108`
- Preprocess request (FFI)
- Build gRPC request
- Create converter and batch processor
- Start `readLoop`
### readLoop
**Location**: `internal/grpc/client_grpc.go:290`
- Start Recv() goroutine (continuously calls `stream.Recv()`)
- Process proto responses
- Asynchronously call `processAndSendResponse` (tracked with `processWg`)
- **Graceful shutdown in defer**:
- Set `closed` flag
- Cancel context (if not already cancelled)
- Wait for all `processAndSendResponse` goroutines to complete (`processWg.Wait()`)
- Close all channels (`resultJSONChan`, `errChan`, `readLoopDone`)
### processAndSendResponse
**Location**: `internal/grpc/client_grpc.go:379`
- Check `ctx.Done()` at function start, return immediately if cancelled
- Convert proto to JSON
- Call FFI batch processor
- All `select` statements include `case <-s.ctx.Done()` for graceful shutdown handling
- Send JSON to channel
### RecvJSON
**Location**:
- `internal/grpc/client_grpc.go:412`: gRPC layer implementation
- `client.go:410`: Client wrapper layer
- Read from `resultJSONChan`
- Directly return JSON string, no parsing needed

55
third_party/sglang/sgl-model-gateway/bindings/golang/examples/oai_server/config/config.go vendored Normal file
View File

@@ -0,0 +1,55 @@
package config
import (
"os"
)
// Config holds the application configuration
type Config struct {
Endpoint string
TokenizerPath string
Port string
LogDir string
LogLevel string
}
// Load loads configuration from environment variables with defaults
func Load() *Config {
// Get tokenizer path from environment or use default
tokenizerPath := os.Getenv("SGL_TOKENIZER_PATH")
if tokenizerPath == "" {
tokenizerPath = "../tokenizer"
}
// Get endpoint from environment or use default
endpoint := os.Getenv("SGL_GRPC_ENDPOINT")
if endpoint == "" {
endpoint = "grpc://localhost:20000"
}
// Get port from environment or use default
port := os.Getenv("PORT")
if port == "" {
port = "8080"
}
// Get log directory from environment or use default
logDir := os.Getenv("LOG_DIR")
if logDir == "" {
logDir = "./logs"
}
// Get log level from environment or use default
logLevel := os.Getenv("LOG_LEVEL")
if logLevel == "" {
logLevel = "info"
}
return &Config{
Endpoint: endpoint,
TokenizerPath: tokenizerPath,
Port: port,
LogDir: logDir,
LogLevel: logLevel,
}
}

121
third_party/sglang/sgl-model-gateway/bindings/golang/examples/oai_server/docs/benchmark_result.md vendored Normal file
View File

@@ -0,0 +1,121 @@
/tmp/ShareGPT_V3_unfiltered_cleaned_split.json: 100%|████████████████████| 642M/642M [10:02<00:00, 1.12MB/s]
#Input tokens: 50561
#Output tokens: 25883
Starting warmup with 5 sequences...
Warmup completed with 5 sequences. Starting main benchmark run...
============ Serving Benchmark Result ============
Backend: sglang-oai-chat
Traffic request rate: 20.0
Max request concurrency: 20
Successful requests: 100
Benchmark duration (s): 107.24
Total input tokens: 50561
Total input text tokens: 50561
Total input vision tokens: 0
Total generated tokens: 25883
Total generated tokens (retokenized): 129591
Request throughput (req/s): 0.93
Input token throughput (tok/s): 471.48
Output token throughput (tok/s): 241.36
Total token throughput (tok/s): 712.84
Concurrency: 16.42
----------------End-to-End Latency----------------
Mean E2E Latency (ms): 17609.46
Median E2E Latency (ms): 12343.82
---------------Time to First Token----------------
Mean TTFT (ms): 190.71
Median TTFT (ms): 164.86
P99 TTFT (ms): 397.72
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms): 162.55
Median TPOT (ms): 63.51
P99 TPOT (ms): 1337.20
---------------Inter-Token Latency----------------
Mean ITL (ms): 25.85
Median ITL (ms): 24.26
P95 ITL (ms): 48.26
P99 ITL (ms): 119.04
Max ITL (ms): 194.58
==================================================
✓ E2E test completed
## Rust
============ Serving Benchmark Result ============
Backend: sglang-oai-chat
Traffic request rate: 20.0
Max request concurrency: 20
Successful requests: 100
Benchmark duration (s): 37.71
Total input tokens: 50561
Total input text tokens: 50561
Total input vision tokens: 0
Total generated tokens: 25883
Total generated tokens (retokenized): 25599
Request throughput (req/s): 2.65
Input token throughput (tok/s): 1340.75
Output token throughput (tok/s): 686.35
Total token throughput (tok/s): 2027.10
Concurrency: 18.58
----------------End-to-End Latency----------------
Mean E2E Latency (ms): 7008.05
Median E2E Latency (ms): 7061.24
---------------Time to First Token----------------
Mean TTFT (ms): 156.09
Median TTFT (ms): 133.81
P99 TTFT (ms): 318.53
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms): 26.59
Median TPOT (ms): 26.75
P99 TPOT (ms): 29.18
---------------Inter-Token Latency----------------
Mean ITL (ms): 26.71
Median ITL (ms): 23.61
P95 ITL (ms): 66.11
P99 ITL (ms): 115.30
Max ITL (ms): 201.08
==================================================
## golang
#Input tokens: 50561
#Output tokens: 25883
Starting warmup with 5 sequences...
Warmup completed with 5 sequences. Starting main benchmark run...
============ Serving Benchmark Result ============
Backend: sglang-oai-chat
Traffic request rate: 20.0
Max request concurrency: 20
Successful requests: 100
Benchmark duration (s): 34.22
Total input tokens: 50561
Total input text tokens: 50561
Total input vision tokens: 0
Total generated tokens: 22970
Total generated tokens (retokenized): 31740
Request throughput (req/s): 2.92
Input token throughput (tok/s): 1477.70
Output token throughput (tok/s): 671.32
Total token throughput (tok/s): 2149.03
Concurrency: 18.42
----------------End-to-End Latency----------------
Mean E2E Latency (ms): 6303.33
Median E2E Latency (ms): 6294.46
---------------Time to First Token----------------
Mean TTFT (ms): 157.10
Median TTFT (ms): 149.16
P99 TTFT (ms): 251.98
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms): 26.49
Median TPOT (ms): 27.15
P99 TPOT (ms): 28.73
---------------Inter-Token Latency----------------
Mean ITL (ms): 26.97
Median ITL (ms): 24.61
P95 ITL (ms): 52.39
P99 ITL (ms): 86.52
Max ITL (ms): 194.55
==================================================

60
third_party/sglang/sgl-model-gateway/bindings/golang/examples/oai_server/go.sum vendored Normal file
View File

@@ -0,0 +1,60 @@
github.com/andybalholm/brotli v1.1.0 h1:eLKJA0d02Lf0mVpIDgYnqXcUn0GqVmEFny3VuID1U3M=
github.com/andybalholm/brotli v1.1.0/go.mod h1:sms7XGricyQI9K10gOSf56VKKWS4oLer58Q+mhRPtnY=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI=
github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/klauspost/compress v1.17.9 h1:6KIumPrER1LHsvBVuDa0r5xaG0Es51mhhB9BQB2qeMA=
github.com/klauspost/compress v1.17.9/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
github.com/valyala/fasthttp v1.52.0 h1:wqBQpxH71XW0e2g+Og4dzQM8pk34aFYlA1Ga8db7gU0=
github.com/valyala/fasthttp v1.52.0/go.mod h1:hf5C4QnVMkNXMspnsUlfM3WitlgYflyhHYoKol/szxQ=
go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=
go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
go.opentelemetry.io/otel v1.38.0 h1:RkfdswUDRimDg0m2Az18RKOsnI8UDzppJAtj01/Ymk8=
go.opentelemetry.io/otel v1.38.0/go.mod h1:zcmtmQ1+YmQM9wrNsTGV/q/uyusom3P8RxwExxkZhjM=
go.opentelemetry.io/otel/metric v1.38.0 h1:Kl6lzIYGAh5M159u9NgiRkmoMKjvbsKtYRwgfrA6WpA=
go.opentelemetry.io/otel/metric v1.38.0/go.mod h1:kB5n/QoRM8YwmUahxvI3bO34eVtQf2i4utNVLr9gEmI=
go.opentelemetry.io/otel/sdk v1.38.0 h1:l48sr5YbNf2hpCUj/FoGhW9yDkl+Ma+LrVl8qaM5b+E=
go.opentelemetry.io/otel/sdk v1.38.0/go.mod h1:ghmNdGlVemJI3+ZB5iDEuk4bWA3GkTpW+DOoZMYBVVg=
go.opentelemetry.io/otel/sdk/metric v1.38.0 h1:aSH66iL0aZqo//xXzQLYozmWrXxyFkBJ6qT5wthqPoM=
go.opentelemetry.io/otel/sdk/metric v1.38.0/go.mod h1:dg9PBnW9XdQ1Hd6ZnRz689CbtrUp0wMMs9iPcgT9EZA=
go.opentelemetry.io/otel/trace v1.38.0 h1:Fxk5bKrDZJUH+AMyyIXGcFAPah0oRcT+LuNtJrmcNLE=
go.opentelemetry.io/otel/trace v1.38.0/go.mod h1:j1P9ivuFsTceSWe1oY+EeW3sc+Pp42sO++GHkg4wwhs=
go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
go.uber.org/multierr v1.10.0 h1:S0h4aNzvfcFsC3dRF1jLoaov7oRaKqRGC/pUEJ2yvPQ=
go.uber.org/multierr v1.10.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y=
go.uber.org/zap v1.27.0 h1:aJMhYGrd5QSmlpLMr2MftRKl7t8J8PTZPA732ud/XR8=
go.uber.org/zap v1.27.0/go.mod h1:GB2qFLM7cTU87MWRP2mPIjqfIDnGu+VIO4V/SdhGo2E=
golang.org/x/net v0.46.1-0.20251013234738-63d1a5100f82 h1:6/3JGEh1C88g7m+qzzTbl3A0FtsLguXieqofVLU/JAo=
golang.org/x/net v0.46.1-0.20251013234738-63d1a5100f82/go.mod h1:Q9BGdFy1y4nkUwiLvT5qtyhAnEHgnQ/zd8PfU6nc210=
golang.org/x/sys v0.37.0 h1:fdNQudmxPjkdUTPnLn5mdQv7Zwvbvpaxqs831goi9kQ=
golang.org/x/sys v0.37.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
golang.org/x/text v0.30.0 h1:yznKA/E9zq54KzlzBEAWn1NXSQ8DIp/NYMy88xJjl4k=
golang.org/x/text v0.30.0/go.mod h1:yDdHFIX9t+tORqspjENWgzaCVXgk0yYnYuSZ8UzzBVM=
gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk=
gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E=
google.golang.org/genproto/googleapis/rpc v0.0.0-20251022142026-3a174f9686a8 h1:M1rk8KBnUsBDg1oPGHNCxG4vc1f49epmTO7xscSajMk=
google.golang.org/genproto/googleapis/rpc v0.0.0-20251022142026-3a174f9686a8/go.mod h1:7i2o+ce6H/6BluujYR+kqX3GKH+dChPTQU19wjRPiGk=
google.golang.org/grpc v1.77.0 h1:wVVY6/8cGA6vvffn+wWK5ToddbgdU3d8MNENr4evgXM=
google.golang.org/grpc v1.77.0/go.mod h1:z0BY1iVj0q8E1uSQCjL9cppRj+gnZjzDnzV0dHhrNig=
google.golang.org/protobuf v1.36.10 h1:AYd7cD/uASjIL6Q9LiTjz8JLcrh/88q5UObnmY3aOOE=
google.golang.org/protobuf v1.36.10/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
gopkg.in/natefinch/lumberjack.v2 v2.2.1 h1:bBRl1b0OH9s/DuPhuXpNl+VtCaJXFZ5/uEFST95x9zc=
gopkg.in/natefinch/lumberjack.v2 v2.2.1/go.mod h1:YD8tP3GAjkrDg1eZH7EGmyESg/lsYskCTPBJVb9jqSc=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

556
third_party/sglang/sgl-model-gateway/bindings/golang/examples/oai_server/handlers/chat.go vendored Normal file
View File

@@ -0,0 +1,556 @@
package handlers
import (
"bufio"
"context"
"encoding/json"
"fmt"
"io"
"strings"
"time"
sglang "github.com/sglang/sglang-go-grpc-sdk"
"github.com/valyala/fasthttp"
"go.uber.org/zap"
"oai_server/models"
"oai_server/service"
"oai_server/utils"
)
// ChatHandler handles chat completion requests
type ChatHandler struct {
logger *zap.Logger
service *service.SGLangService
}
// NewChatHandler creates a new chat handler
func NewChatHandler(logger *zap.Logger, svc *service.SGLangService) *ChatHandler {
return &ChatHandler{
logger: logger,
service: svc,
}
}
// recvResult holds the result of a RecvJSON() call
type recvResult struct {
chunkJSON string
err error
}
// HandleChatCompletion handles POST /v1/chat/completions
func (h *ChatHandler) HandleChatCompletion(ctx *fasthttp.RequestCtx) {
var req models.ChatRequest
if err := json.Unmarshal(ctx.PostBody(), &req); err != nil {
h.logger.Warn("Invalid chat completion request", zap.Error(err))
utils.RespondError(ctx, 400, fmt.Sprintf("Invalid request: %v", err), "invalid_request_error")
return
}
path := string(ctx.Path())
defer func() {
statusCode := ctx.Response.StatusCode()
if statusCode == 0 {
statusCode = 200
}
h.logHTTPResponse(statusCode, path)
}()
// Convert to SGLang format
messages := make([]sglang.ChatMessage, len(req.Messages))
for i, msg := range req.Messages {
role, roleOk := msg["role"]
content, contentOk := msg["content"]
// Validate role
if !roleOk || role == "" {
h.logger.Warn("Missing or empty role in message", zap.Int("message_index", i))
utils.RespondError(ctx, 400, "Message role is required and cannot be empty", "invalid_request_error")
return
}
// Ensure content is always a string (not null)
// Chat template requires content field to be present, even if empty
// If content is missing or null, use empty string
contentStr := ""
if contentOk && content != "" {
contentStr = content
}
messages[i] = sglang.ChatMessage{
Role: role,
Content: contentStr,
}
}
sglReq := sglang.ChatCompletionRequest{
Model: req.Model,
Messages: messages,
Stream: req.Stream,
}
if req.Temperature != nil {
temp := float32(*req.Temperature)
sglReq.Temperature = &temp
}
if req.TopP != nil {
topP := float32(*req.TopP)
sglReq.TopP = &topP
}
if req.MaxCompletionTokens != nil {
sglReq.MaxCompletionTokens = req.MaxCompletionTokens
} else if req.MaxTokens != nil {
sglReq.MaxCompletionTokens = req.MaxTokens
}
requestCtx := context.Background()
if req.Stream {
h.handleStreamingCompletion(ctx, requestCtx, sglReq)
} else {
h.handleNonStreamingCompletion(ctx, requestCtx, sglReq)
}
}
// isBrokenPipeError checks if the error is a broken pipe error (client disconnected)
func isBrokenPipeError(err error) bool {
if err == nil {
return false
}
errStr := err.Error()
return strings.Contains(errStr, "broken pipe") ||
strings.Contains(errStr, "connection reset by peer") ||
strings.Contains(errStr, "connection closed") ||
strings.Contains(errStr, "write: connection closed")
}
// logHTTPResponse logs HTTP response with colored output
func (h *ChatHandler) logHTTPResponse(statusCode int, path string) {
var statusText string
var colorCode string
switch {
case statusCode >= 200 && statusCode < 300:
colorCode = "\033[32m" // Green
statusText = "OK"
case statusCode >= 300 && statusCode < 400:
colorCode = "\033[33m" // Yellow
statusText = "Redirect"
case statusCode >= 400 && statusCode < 500:
colorCode = "\033[33m" // Yellow
statusText = "Client Error"
case statusCode >= 500:
colorCode = "\033[31m" // Red
statusText = "Server Error"
default:
colorCode = "\033[37m" // White
statusText = "Unknown"
}
resetCode := "\033[0m"
msg := fmt.Sprintf("%s[%d %s]%s %s", colorCode, statusCode, statusText, resetCode, path)
h.logger.Info(msg)
}
func (h *ChatHandler) handleStreamingCompletion(ctx *fasthttp.RequestCtx, requestCtx context.Context, req sglang.ChatCompletionRequest) {
ctx.SetContentType("text/event-stream")
ctx.Response.Header.Set("Cache-Control", "no-cache")
ctx.Response.Header.Set("Connection", "keep-alive")
ctx.Response.Header.Set("X-Accel-Buffering", "no")
ctx.SetStatusCode(200)
var clientDisconnected bool
// Flush timeout: prevent deadlock if client is slow or disconnected
// This timeout should be longer than typical network latency but shorter than client timeout
const flushTimeout = 5 * time.Second
ctx.SetBodyStreamWriter(func(w *bufio.Writer) {
streamCtx, cancel := context.WithCancel(context.Background())
defer cancel()
stream, err := h.service.Client().CreateChatCompletionStream(streamCtx, req)
if err != nil {
h.logger.Error("Failed to create chat completion stream",
zap.Error(err),
zap.String("model", req.Model),
)
// Use sendSSEError to send error in consistent format
errInfo, sendErr := h.sendSSEError(w, err)
if sendErr != nil {
h.logger.Warn("Failed to send SSE error", zap.Error(sendErr))
} else if errInfo.IsTimeout {
h.logger.Error("Stream creation timeout", zap.Error(err))
}
return
}
defer func() {
if closeErr := stream.Close(); closeErr != nil {
h.logger.Warn("Failed to close stream", zap.Error(closeErr))
}
}()
// Use a single dedicated goroutine to continuously call RecvJSON() and send results via channel
recvChan := make(chan recvResult, 20)
recvGoroutineDone := make(chan struct{})
go func() {
defer func() {
close(recvChan)
close(recvGoroutineDone)
}()
for {
// Check context before calling RecvJSON() to avoid blocking if context is cancelled
select {
case <-streamCtx.Done():
return
default:
}
// Call RecvJSON() - this may block, but stream.Close() will unblock it
// when context is cancelled (called from main loop)
chunkJSON, err := stream.RecvJSON()
// Check context again after RecvJSON() returns
select {
case <-streamCtx.Done():
return
default:
}
// Send to channel (may block if channel is full)
// If channel is full, this will block until main loop reads from it
// This is acceptable because main loop should be actively reading
select {
case recvChan <- recvResult{chunkJSON: chunkJSON, err: err}:
if err != nil {
// EOF or other error, stop the goroutine
return
}
case <-streamCtx.Done():
// Context cancelled while sending, stop the goroutine
return
}
}
}()
for {
if clientDisconnected {
cancel()
// Close stream immediately to unblock RecvJSON() calls
stream.Close()
return
}
select {
case <-streamCtx.Done():
// Close stream to ensure RecvJSON() goroutine can exit
stream.Close()
return
case result, ok := <-recvChan:
if !ok {
// Channel closed, stream ended
return
}
if result.err == io.EOF {
if !clientDisconnected {
w.WriteString("data: [DONE]\n\n")
// Flush with timeout to prevent deadlock
flushDone := make(chan error, 1)
go func() {
flushDone <- w.Flush()
}()
flushCtx, flushCancel := context.WithTimeout(streamCtx, flushTimeout)
defer flushCancel()
select {
case flushErr := <-flushDone:
if flushErr != nil && !isBrokenPipeError(flushErr) {
h.logger.Warn("Final flush error", zap.Error(flushErr))
}
case <-flushCtx.Done():
if flushCtx.Err() == context.DeadlineExceeded {
h.logger.Warn("Final flush timeout", zap.Duration("timeout", flushTimeout))
}
case <-streamCtx.Done():
// Context cancelled, skip flush
}
}
return
}
if result.err != nil {
if result.err == context.Canceled || result.err == context.DeadlineExceeded {
return
}
// Send error to client before closing
errInfo, sendErr := h.sendSSEError(w, result.err)
if sendErr != nil {
h.logger.Warn("Failed to send SSE error", zap.Error(sendErr))
}
if errInfo.IsTimeout {
h.logger.Error("Stream timeout error", zap.Error(result.err))
} else {
h.logger.Error("Stream error", zap.Error(result.err))
}
return
}
if result.chunkJSON == "" {
continue
}
w.WriteString("data: ")
w.WriteString(result.chunkJSON)
w.WriteString("\n\n")
// Flush with timeout to prevent deadlock:
// If Flush blocks indefinitely (slow client), RecvJSON goroutine may fill recvChan
// and then block trying to send, causing deadlock
// Note: bufio.Writer.Flush() doesn't have a timeout parameter, so we use
// a goroutine + select pattern to implement timeout behavior
flushDone := make(chan error, 1)
go func() {
flushDone <- w.Flush()
}()
flushCtx, flushCancel := context.WithTimeout(streamCtx, flushTimeout)
defer flushCancel()
select {
case err := <-flushDone:
if err != nil {
if isBrokenPipeError(err) {
clientDisconnected = true
cancel()
// Close stream immediately to unblock RecvJSON() calls
stream.Close()
return
}
h.logger.Warn("Flush error", zap.Error(err))
}
case <-flushCtx.Done():
// Flush timeout: client may be slow or disconnected
// Continue processing to avoid deadlock, but mark as disconnected
if flushCtx.Err() == context.DeadlineExceeded {
h.logger.Warn("Flush timeout, client may be slow or disconnected", zap.Duration("timeout", flushTimeout))
}
clientDisconnected = true
cancel()
stream.Close()
return
case <-streamCtx.Done():
// Context cancelled, stop flushing
return
}
}
}
})
}
func (h *ChatHandler) handleNonStreamingCompletion(ctx *fasthttp.RequestCtx, requestCtx context.Context, req sglang.ChatCompletionRequest) {
resp, err := h.service.Client().CreateChatCompletion(requestCtx, req)
if err != nil {
h.logger.Error("Failed to create chat completion",
zap.Error(err),
zap.String("model", req.Model),
)
utils.RespondError(ctx, 500, fmt.Sprintf("Failed to create completion: %v", err), "server_error")
return
}
// Convert to OpenAI format
response := utils.BuildResponseBase(resp.ID, resp.Created, resp.Model)
response["object"] = "chat.completion"
choices := make([]map[string]interface{}, len(resp.Choices))
for i, choice := range resp.Choices {
choiceMap := map[string]interface{}{
"index": choice.Index,
"message": map[string]interface{}{
"role": choice.Message.Role,
"content": choice.Message.Content,
},
"finish_reason": choice.FinishReason,
}
if len(choice.Message.ToolCalls) > 0 {
toolCalls := make([]map[string]interface{}, len(choice.Message.ToolCalls))
for j, tc := range choice.Message.ToolCalls {
toolCalls[j] = map[string]interface{}{
"id": tc.ID,
"type": tc.Type,
"function": map[string]interface{}{"name": tc.Function.Name, "arguments": tc.Function.Arguments},
}
}
choiceMap["message"].(map[string]interface{})["tool_calls"] = toolCalls
}
choices[i] = choiceMap
}
response["choices"] = choices
// Usage is always present (not a pointer)
response["usage"] = map[string]interface{}{
"prompt_tokens": resp.Usage.PromptTokens,
"completion_tokens": resp.Usage.CompletionTokens,
"total_tokens": resp.Usage.TotalTokens,
}
ctx.SetStatusCode(200)
ctx.SetContentType("application/json")
jsonData, _ := json.Marshal(response)
ctx.Write(jsonData)
}
// StreamErrorInfo holds parsed error information
type StreamErrorInfo struct {
Message string
Type string
Code int
IsTimeout bool
}
// parseStreamError parses error type and code
func parseStreamError(err error) StreamErrorInfo {
if err == nil {
return StreamErrorInfo{}
}
errorMsg := err.Error()
// Check timeout error by message prefix
isTimeout := strings.HasPrefix(errorMsg, "stream.Recv() timeout") || strings.Contains(errorMsg, "timeout after")
errorType := "server_error"
errorCode := 500
if isTimeout {
errorType = "timeout_error"
errorCode = 504
}
return StreamErrorInfo{
Message: errorMsg,
Type: errorType,
Code: errorCode,
IsTimeout: isTimeout,
}
}
// formatErrorJSON formats error as OpenAI JSON
func formatErrorJSON(errInfo StreamErrorInfo) string {
errorObj := map[string]interface{}{
"error": map[string]interface{}{
"message": errInfo.Message,
"type": errInfo.Type,
"code": errInfo.Code,
},
}
jsonBytes, _ := json.Marshal(errorObj)
return string(jsonBytes)
}
// sendSSEError sends SSE error response. Callers should log errors.
func (h *ChatHandler) sendSSEError(w *bufio.Writer, err error) (StreamErrorInfo, error) {
errInfo := parseStreamError(err)
errorJSON := formatErrorJSON(errInfo)
w.WriteString("data: ")
w.WriteString(errorJSON)
w.WriteString("\n\n")
if flushErr := w.Flush(); flushErr != nil && !isBrokenPipeError(flushErr) {
h.logger.Warn("Failed to flush error response", zap.Error(flushErr))
return errInfo, flushErr
}
return errInfo, nil
}
// HandleGenerate handles POST /generate (SGLang native API)
func (h *ChatHandler) HandleGenerate(ctx *fasthttp.RequestCtx) {
path := string(ctx.Path())
defer func() {
statusCode := ctx.Response.StatusCode()
if statusCode == 0 {
statusCode = 200
}
h.logHTTPResponse(statusCode, path)
}()
// Parse request body
var req map[string]interface{}
if err := json.Unmarshal(ctx.PostBody(), &req); err != nil {
h.logger.Warn("Invalid generate request", zap.Error(err))
utils.RespondError(ctx, 400, fmt.Sprintf("Invalid request: %v", err), "invalid_request_error")
return
}
// Extract text and sampling_params
text, ok := req["text"].(string)
if !ok || text == "" {
utils.RespondError(ctx, 400, "Missing or invalid 'text' field", "invalid_request_error")
return
}
samplingParams, _ := req["sampling_params"].(map[string]interface{})
if samplingParams == nil {
samplingParams = make(map[string]interface{})
}
// Convert to chat completion format for processing
chatReq := sglang.ChatCompletionRequest{
Model: "default",
Messages: []sglang.ChatMessage{{Role: "user", Content: text}},
Stream: false,
}
// Copy sampling params
if maxNewTokens, ok := samplingParams["max_new_tokens"].(float64); ok {
tokens := int(maxNewTokens)
chatReq.MaxCompletionTokens = &tokens
}
if temp, ok := samplingParams["temperature"].(float64); ok {
temp32 := float32(temp)
chatReq.Temperature = &temp32
}
if topP, ok := samplingParams["top_p"].(float64); ok {
topP32 := float32(topP)
chatReq.TopP = &topP32
}
if topK, ok := samplingParams["top_k"].(float64); ok {
topKInt := int(topK)
chatReq.TopK = &topKInt
}
requestCtx := context.Background()
// Use non-streaming completion for /generate endpoint
resp, err := h.service.Client().CreateChatCompletion(requestCtx, chatReq)
if err != nil {
h.logger.Error("Failed to create completion",
zap.Error(err),
)
utils.RespondError(ctx, 500, fmt.Sprintf("Failed to create completion: %v", err), "server_error")
return
}
// Convert to SGLang /generate response format
// meta_info must match SGLang's expected format with completion_tokens at top level
finishReason := resp.Choices[0].FinishReason
if finishReason == "" {
finishReason = "stop"
}
response := map[string]interface{}{
"text": resp.Choices[0].Message.Content,
"meta_info": map[string]interface{}{
"id": resp.ID,
"finish_reason": finishReason,
"prompt_tokens": resp.Usage.PromptTokens,
"completion_tokens": resp.Usage.CompletionTokens,
"cached_tokens": 0, // Not available from chat completion API
"weight_version": "", // Not available from chat completion API
},
}
ctx.SetStatusCode(200)
ctx.SetContentType("application/json")
jsonData, _ := json.Marshal(response)
ctx.Write(jsonData)
}

33
third_party/sglang/sgl-model-gateway/bindings/golang/examples/oai_server/handlers/health.go vendored Normal file
View File

@@ -0,0 +1,33 @@
package handlers
import (
"encoding/json"
"github.com/valyala/fasthttp"
"go.uber.org/zap"
)
// HealthHandler handles health check requests
type HealthHandler struct {
logger *zap.Logger
}
// NewHealthHandler creates a new health handler
func NewHealthHandler(logger *zap.Logger) *HealthHandler {
return &HealthHandler{
logger: logger,
}
}
// Check handles GET /health
func (h *HealthHandler) Check(ctx *fasthttp.RequestCtx) {
ctx.SetStatusCode(200)
ctx.SetContentType("application/json")
response := map[string]string{
"status": "ok",
}
jsonData, _ := json.Marshal(response)
ctx.Write(jsonData)
}

67
third_party/sglang/sgl-model-gateway/bindings/golang/examples/oai_server/handlers/models.go vendored Normal file
View File

@@ -0,0 +1,67 @@
package handlers
import (
"encoding/json"
"github.com/valyala/fasthttp"
"go.uber.org/zap"
)
// ModelsHandler handles model list requests
type ModelsHandler struct {
logger *zap.Logger
tokenizerPath string
}
// NewModelsHandler creates a new models handler
func NewModelsHandler(logger *zap.Logger, tokenizerPath string) *ModelsHandler {
return &ModelsHandler{
logger: logger,
tokenizerPath: tokenizerPath,
}
}
// List handles GET /v1/models
func (h *ModelsHandler) List(ctx *fasthttp.RequestCtx) {
// Return a default model for OpenAI compatibility
ctx.SetStatusCode(200)
ctx.SetContentType("application/json")
response := map[string]interface{}{
"object": "list",
"data": []map[string]interface{}{
{
"id": "default",
"object": "model",
"created": 1677610602,
"owned_by": "sglang",
},
},
}
jsonData, _ := json.Marshal(response)
ctx.Write(jsonData)
}
// GetModelInfo handles GET /get_model_info
// Returns model information compatible with SGLang RuntimeEndpoint
func (h *ModelsHandler) GetModelInfo(ctx *fasthttp.RequestCtx) {
ctx.SetStatusCode(200)
ctx.SetContentType("application/json")
// Return model info compatible with SGLang RuntimeEndpoint expectations
response := map[string]interface{}{
"model_path": h.tokenizerPath, // Use tokenizer path as model path
"tokenizer_path": h.tokenizerPath,
"is_generation": true,
"preferred_sampling_params": "",
"weight_version": "",
"has_image_understanding": false,
"has_audio_understanding": false,
"model_type": "",
"architectures": nil,
}
jsonData, _ := json.Marshal(response)
ctx.Write(jsonData)
}

67
third_party/sglang/sgl-model-gateway/bindings/golang/examples/oai_server/logger/logger.go vendored Normal file
View File

@@ -0,0 +1,67 @@
package logger
import (
"os"
"path/filepath"
"time"
"go.uber.org/zap"
"go.uber.org/zap/zapcore"
"gopkg.in/natefinch/lumberjack.v2"
)
// Init initializes the logger with file and console output
func Init(logDir, logLevel string) (*zap.Logger, error) {
// Ensure log directory exists
if err := os.MkdirAll(logDir, 0755); err != nil {
return nil, err
}
// Parse log level
var level zapcore.Level
if err := level.UnmarshalText([]byte(logLevel)); err != nil {
level = zapcore.InfoLevel
}
// Create log file path with date
logFile := filepath.Join(logDir, "oai_server-"+time.Now().Format("2006-01-02")+".log")
// File writer with rotation
fileWriter := zapcore.AddSync(&lumberjack.Logger{
Filename: logFile,
MaxSize: 100, // megabytes
MaxBackups: 10,
MaxAge: 30, // days
Compress: true,
})
// Console writer
consoleWriter := zapcore.AddSync(os.Stdout)
// Encoder config
encoderConfig := zap.NewProductionEncoderConfig()
encoderConfig.TimeKey = "timestamp"
encoderConfig.EncodeTime = zapcore.ISO8601TimeEncoder
encoderConfig.EncodeLevel = zapcore.CapitalLevelEncoder
// Create cores
fileCore := zapcore.NewCore(
zapcore.NewJSONEncoder(encoderConfig),
fileWriter,
level,
)
consoleCore := zapcore.NewCore(
zapcore.NewConsoleEncoder(encoderConfig),
consoleWriter,
level,
)
// Combine cores
core := zapcore.NewTee(fileCore, consoleCore)
// Create logger
logger := zap.New(core, zap.AddCaller(), zap.AddStacktrace(zapcore.ErrorLevel))
return logger, nil
}

116
third_party/sglang/sgl-model-gateway/bindings/golang/examples/oai_server/main.go vendored Normal file
View File

@@ -0,0 +1,116 @@
// OpenAI-compatible chat server using SGLang Go SDK and fasthttp framework
package main
import (
"fmt"
"net/http"
"os"
_ "net/http/pprof" // Enable pprof endpoints
"github.com/valyala/fasthttp"
"go.uber.org/zap"
"oai_server/config"
"oai_server/handlers"
"oai_server/logger"
"oai_server/service"
)
// Version information (set at build time via ldflags)
var (
Version = "dev"
BuildTime = "unknown"
GitCommit = "unknown"
)
func main() {
// Load configuration
cfg := config.Load()
// Initialize logger
appLogger, err := logger.Init(cfg.LogDir, cfg.LogLevel)
if err != nil {
panic(fmt.Sprintf("Failed to initialize logger: %v", err))
}
defer appLogger.Sync()
appLogger.Info("Starting OpenAI-compatible server",
zap.String("endpoint", cfg.Endpoint),
zap.String("tokenizer", cfg.TokenizerPath),
zap.String("port", cfg.Port),
)
// Initialize SGLang service
sglangService, err := service.NewSGLangService(cfg.Endpoint, cfg.TokenizerPath)
if err != nil {
appLogger.Fatal("Failed to create SGLang client", zap.Error(err))
}
defer sglangService.Close()
appLogger.Info("SGLang client created successfully")
// Enable pprof if requested
if os.Getenv("PPROF_ENABLED") == "true" {
pprofPort := os.Getenv("PPROF_PORT")
if pprofPort == "" {
pprofPort = "6060"
}
go func() {
pprofAddr := ":" + pprofPort
appLogger.Info("Starting pprof server", zap.String("address", pprofAddr))
if err := http.ListenAndServe(pprofAddr, nil); err != nil {
appLogger.Error("pprof server failed", zap.Error(err))
}
}()
appLogger.Info("pprof enabled", zap.String("port", pprofPort), zap.String("endpoint", fmt.Sprintf("http://localhost:%s/debug/pprof/", pprofPort)))
}
// Initialize handlers
healthHandler := handlers.NewHealthHandler(appLogger)
modelsHandler := handlers.NewModelsHandler(appLogger, cfg.TokenizerPath)
chatHandler := handlers.NewChatHandler(appLogger, sglangService)
// Setup fasthttp router
router := func(ctx *fasthttp.RequestCtx) {
path := string(ctx.Path())
method := string(ctx.Method())
switch {
case method == "GET" && path == "/health":
healthHandler.Check(ctx)
case method == "GET" && path == "/v1/models":
modelsHandler.List(ctx)
case method == "GET" && path == "/get_model_info":
modelsHandler.GetModelInfo(ctx)
case method == "POST" && path == "/v1/chat/completions":
chatHandler.HandleChatCompletion(ctx)
case (method == "POST" || method == "PUT") && path == "/generate":
chatHandler.HandleGenerate(ctx)
default:
ctx.Error("Not Found", fasthttp.StatusNotFound)
}
}
// Start server
serverAddr := ":" + cfg.Port
baseURL := fmt.Sprintf("http://localhost:%s", cfg.Port)
appLogger.Info("Server starting",
zap.String("address", serverAddr),
zap.String("base_url", baseURL),
)
// Print available HTTP endpoints (similar to FastAPI startup)
appLogger.Info("Available HTTP endpoints:")
appLogger.Info(fmt.Sprintf(" GET %s/health", baseURL))
appLogger.Info(fmt.Sprintf(" GET %s/v1/models", baseURL))
appLogger.Info(fmt.Sprintf(" GET %s/get_model_info", baseURL))
appLogger.Info(fmt.Sprintf(" POST %s/v1/chat/completions", baseURL))
appLogger.Info(fmt.Sprintf(" POST %s/generate", baseURL))
appLogger.Info(fmt.Sprintf("Application startup complete. Listening on %s", baseURL))
if err := fasthttp.ListenAndServe(serverAddr, router); err != nil {
appLogger.Fatal("Server failed", zap.Error(err))
}
}

14
third_party/sglang/sgl-model-gateway/bindings/golang/examples/oai_server/models/chat.go vendored Normal file
View File

@@ -0,0 +1,14 @@
package models
// ChatRequest represents an OpenAI-compatible chat completion request
type ChatRequest struct {
Model string `json:"model" binding:"required"`
Messages []map[string]string `json:"messages" binding:"required"`
Stream bool `json:"stream,omitempty"`
Temperature *float64 `json:"temperature,omitempty"`
TopP *float64 `json:"top_p,omitempty"`
MaxTokens *int `json:"max_tokens,omitempty"` // OpenAI API standard field
MaxCompletionTokens *int `json:"max_completion_tokens,omitempty"` // SGLang-specific field (used by bench_serving.py)
Tools []map[string]interface{} `json:"tools,omitempty"`
ToolChoice interface{} `json:"tool_choice,omitempty"`
}

111
third_party/sglang/sgl-model-gateway/bindings/golang/examples/oai_server/run.sh vendored Executable file
View File

@@ -0,0 +1,111 @@
#!/bin/bash
# OpenAI-compatible server runner
# Usage: ./run.sh [tokenizer_path] [endpoint] [port] [--profile] [--pprof-port PORT]
#
# Options:
# --profile Enable pprof profiling (default port: 6060)
# --pprof-port PORT Set pprof port (default: 6060, requires --profile)
# Set library path for Rust FFI library
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
BINDINGS_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)"
LIB_DIR="${BINDINGS_DIR}/lib"
if [ ! -d "$LIB_DIR" ]; then
echo "Error: Library directory not found at $LIB_DIR"
echo "Please run 'make lib' first to build and export the library"
exit 1
fi
# Get Python LDFLAGS (needed for Rust FFI that depends on Python)
PYTHON_LDFLAGS=$(python3-config --ldflags --embed 2>/dev/null || python3-config --ldflags 2>/dev/null || echo "")
# Set CGO_LDFLAGS to link with the Rust library
# Note: -lsgl_model_gateway_go and -ldl are already in the #cgo directive in internal/ffi/client.go
# We only need to add the library path (-L) and Python flags
export CGO_LDFLAGS="-L${LIB_DIR} ${PYTHON_LDFLAGS}"
# macOS uses DYLD_LIBRARY_PATH, Linux uses LD_LIBRARY_PATH
if [[ "$OSTYPE" == "darwin"* ]]; then
export DYLD_LIBRARY_PATH="${LIB_DIR}:${DYLD_LIBRARY_PATH}"
else
export LD_LIBRARY_PATH="${LIB_DIR}:${LD_LIBRARY_PATH}"
fi
# Parse arguments
ENABLE_PROFILE=false
PPROF_PORT="6060"
TOKENIZER_PATH=""
ENDPOINT=""
PORT=""
while [[ $# -gt 0 ]]; do
case $1 in
--profile)
ENABLE_PROFILE=true
shift
;;
--pprof-port)
ENABLE_PROFILE=true
PPROF_PORT="$2"
shift 2
;;
*)
if [[ -z "$TOKENIZER_PATH" ]]; then
TOKENIZER_PATH="$1"
elif [[ -z "$ENDPOINT" ]]; then
ENDPOINT="$1"
elif [[ -z "$PORT" ]]; then
PORT="$1"
fi
shift
;;
esac
done
# Default configuration
DEFAULT_TOKENIZER_PATH="${SGL_TOKENIZER_PATH:-../tokenizer}"
DEFAULT_ENDPOINT="${SGL_GRPC_ENDPOINT:-grpc://localhost:20000}"
DEFAULT_PORT="${PORT:-8080}"
TOKENIZER_PATH="${TOKENIZER_PATH:-${DEFAULT_TOKENIZER_PATH}}"
ENDPOINT="${ENDPOINT:-${DEFAULT_ENDPOINT}}"
PORT="${PORT:-${DEFAULT_PORT}}"
echo "Running OpenAI-compatible server..."
echo "Library path: ${LIB_DIR}"
echo "Tokenizer: $TOKENIZER_PATH"
echo "Endpoint: $ENDPOINT"
echo "Port: $PORT"
echo "Client Mode: gRPC (default)"
echo "FFI Postprocessing: ENABLED (normal mode)"
echo "FFI Preprocessing: ENABLED (normal mode)"
if [[ "$ENABLE_PROFILE" == "true" ]]; then
echo "Profiling: enabled (port: $PPROF_PORT)"
echo " pprof endpoint: http://localhost:$PPROF_PORT/debug/pprof/"
export PPROF_ENABLED=true
export PPROF_PORT="$PPROF_PORT"
else
echo "Profiling: disabled"
fi
echo ""
# Change to script directory
cd "$(dirname "${BASH_SOURCE[0]}")"
# Ensure Go module is properly initialized
if [ ! -f "go.mod" ]; then
echo "Error: go.mod not found in $(pwd)"
exit 1
fi
# Ensure Go modules are enabled
export GO111MODULE=on
# Sync Go module dependencies
echo "Syncing Go module dependencies..."
go mod tidy
# Run the server (use ./main.go to ensure module context is correct)
SGL_TOKENIZER_PATH="$TOKENIZER_PATH" SGL_GRPC_ENDPOINT="$ENDPOINT" PORT="$PORT" go run ./main.go

554
third_party/sglang/sgl-model-gateway/bindings/golang/examples/oai_server/scripts/analyze_tpot.sh vendored Executable file
View File

@@ -0,0 +1,554 @@
#!/bin/bash
# TPOT performance bottleneck analysis script
# Specifically designed to analyze why Go Router is twice as slow as Rust Router
#
# Usage:
# ./scripts/analyze_tpot.sh [options]
#
# Options:
# --duration SECONDS CPU profile duration (default: 60)
# --requests NUM Number of requests (default: 100)
# --concurrency NUM Concurrency level (default: 20)
# --pprof-port PORT pprof port (default: 6060)
# --server-url URL Server URL (default: http://localhost:8080)
set -e
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
PROFILE_DIR="${PROJECT_ROOT}/profiles"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
OUTPUT_DIR="${PROFILE_DIR}/tpot_analysis_${TIMESTAMP}"
# Colors
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
BLUE='\033[0;34m'
NC='\033[0m'
# Default values
DURATION=${DURATION:-60}
NUM_REQUESTS=${NUM_REQUESTS:-100}
CONCURRENCY=${CONCURRENCY:-20}
PPROF_PORT=${PPROF_PORT:-6060}
SERVER_URL=${SERVER_URL:-http://localhost:8080}
# Parse arguments
while [[ $# -gt 0 ]]; do
case $1 in
--duration)
DURATION="$2"
shift 2
;;
--requests)
NUM_REQUESTS="$2"
shift 2
;;
--concurrency)
CONCURRENCY="$2"
shift 2
;;
--pprof-port)
PPROF_PORT="$2"
shift 2
;;
--server-url)
SERVER_URL="$2"
shift 2
;;
*)
echo "Unknown option: $1"
exit 1
;;
esac
done
mkdir -p "$OUTPUT_DIR"
# Check for graphviz (optional, needed for some pprof visualizations)
HAS_GRAPHVIZ=false
if command -v dot >/dev/null 2>&1; then
HAS_GRAPHVIZ=true
fi
echo -e "${BLUE}========================================${NC}"
echo -e "${BLUE}TPOT Performance Bottleneck Analysis${NC}"
echo -e "${BLUE}========================================${NC}"
echo ""
echo "Configuration:"
echo " Duration: ${DURATION}s"
echo " Requests: $NUM_REQUESTS"
echo " Concurrency: $CONCURRENCY"
echo " Server URL: $SERVER_URL"
echo " pprof Port: $PPROF_PORT"
echo " Output Dir: $OUTPUT_DIR"
if [ "$HAS_GRAPHVIZ" = "false" ]; then
echo ""
echo -e "${YELLOW}Note: graphviz not found. Some pprof visualizations may not work.${NC}"
echo -e "${YELLOW}To install graphviz:${NC}"
echo -e "${YELLOW} macOS: brew install graphviz${NC}"
echo -e "${YELLOW} Ubuntu: sudo apt-get install graphviz${NC}"
echo -e "${YELLOW} CentOS: sudo yum install graphviz${NC}"
echo -e "${YELLOW}Text reports will still be generated without graphviz.${NC}"
fi
echo ""
# Check if server is running
echo -e "${YELLOW}[Check] Verifying server is running...${NC}"
if ! curl -s "${SERVER_URL}/health" > /dev/null 2>&1; then
echo -e "${RED}Error: Server not responding at ${SERVER_URL}${NC}"
echo ""
echo "Please start the server first with profiling enabled:"
echo " ./run.sh --profile --pprof-port $PPROF_PORT"
echo " or"
echo " PPROF_ENABLED=true PPROF_PORT=$PPROF_PORT make run"
exit 1
fi
echo -e "${GREEN}✓ Server is running${NC}"
echo ""
# Check if pprof is enabled
echo -e "${YELLOW}[Check] Verifying pprof is enabled...${NC}"
if ! curl -s "http://localhost:${PPROF_PORT}/debug/pprof/" > /dev/null 2>&1; then
echo -e "${RED}Error: pprof not accessible at http://localhost:${PPROF_PORT}/debug/pprof/${NC}"
echo ""
echo "Please start the server with profiling enabled:"
echo " ./run.sh --profile --pprof-port $PPROF_PORT"
exit 1
fi
echo -e "${GREEN}✓ pprof is enabled${NC}"
echo ""
# ============================================
# Step 1: Collect baseline profiles
# ============================================
echo -e "${GREEN}[Step 1/8] Collecting baseline profiles...${NC}"
# Baseline memory
go tool pprof -proto -output="${OUTPUT_DIR}/heap_before.pb.gz" \
"http://localhost:${PPROF_PORT}/debug/pprof/heap" > /dev/null 2>&1 || true
# Baseline goroutine
go tool pprof -proto -output="${OUTPUT_DIR}/goroutine_before.pb.gz" \
"http://localhost:${PPROF_PORT}/debug/pprof/goroutine" > /dev/null 2>&1 || true
echo -e "${GREEN}✓ Baseline profiles collected${NC}"
echo ""
# ============================================
# Step 2: Start CPU profile collection
# ============================================
echo -e "${GREEN}[Step 2/8] Starting CPU profile collection (${DURATION}s)...${NC}"
go tool pprof -proto -output="${OUTPUT_DIR}/cpu_${DURATION}s.pb.gz" \
"http://localhost:${PPROF_PORT}/debug/pprof/profile?seconds=${DURATION}" &
CPU_PID=$!
sleep 2
echo -e "${GREEN}✓ CPU profile collection started${NC}"
echo ""
# ============================================
# Step 3: Run load test with streaming requests
# ============================================
echo -e "${GREEN}[Step 3/8] Running load test ($NUM_REQUESTS streaming requests, concurrency=$CONCURRENCY)...${NC}"
# Function to run a single streaming request
run_streaming_request() {
local request_id=$1
local start_time=$(date +%s)
local start_nanos=$(date +%N 2>/dev/null || echo "000000000")
curl -N -s -X POST "${SERVER_URL}/v1/chat/completions" \
-H "Content-Type: application/json" \
-d "{
\"model\": \"default\",
\"messages\": [{\"role\": \"user\", \"content\": \"Write a 500-word story with character dialogue and scene descriptions\"}],
\"stream\": true,
\"max_tokens\": 300,
\"temperature\": 0.7
}" > /dev/null
local end_time=$(date +%s)
local end_nanos=$(date +%N 2>/dev/null || echo "000000000")
local duration=$((end_time - start_time))
echo "$duration" >> "${OUTPUT_DIR}/request_times.txt"
}
# Run requests with controlled concurrency
# Use a temporary file to track job PIDs to avoid conflicts with CPU_PID
JOB_PIDS_FILE="${OUTPUT_DIR}/.job_pids_$$"
> "$JOB_PIDS_FILE"
for i in $(seq 1 $NUM_REQUESTS); do
# Wait if we've reached concurrency limit
while [ $(wc -l < "$JOB_PIDS_FILE" 2>/dev/null || echo 0) -ge $CONCURRENCY ]; do
# Check and remove completed jobs
while IFS= read -r pid; do
if [ -n "$pid" ] && ! kill -0 "$pid" 2>/dev/null; then
# Process completed, remove from file
grep -v "^${pid}$" "$JOB_PIDS_FILE" > "${JOB_PIDS_FILE}.tmp" && \
mv "${JOB_PIDS_FILE}.tmp" "$JOB_PIDS_FILE" || true
fi
done < "$JOB_PIDS_FILE"
sleep 0.1
done
# Start new request
run_streaming_request $i &
echo $! >> "$JOB_PIDS_FILE"
# Progress indicator
if [ $((i % 10)) -eq 0 ]; then
echo " Progress: $i/$NUM_REQUESTS requests sent..."
fi
done
# Wait for all remaining jobs (excluding CPU_PID)
while IFS= read -r pid; do
if [ -n "$pid" ] && [ "$pid" != "$CPU_PID" ]; then
wait "$pid" 2>/dev/null || true
fi
done < "$JOB_PIDS_FILE"
# Clean up
rm -f "$JOB_PIDS_FILE" "${JOB_PIDS_FILE}.tmp" 2>/dev/null || true
echo -e "${GREEN}✓ Load test completed${NC}"
echo ""
# ============================================
# Step 4: Wait for CPU profile to complete
# ============================================
echo -e "${GREEN}[Step 4/8] Waiting for CPU profile to complete...${NC}"
# Wait for the process, but handle the case where it might have already completed
if kill -0 $CPU_PID 2>/dev/null; then
wait $CPU_PID 2>/dev/null || true
else
# Process already completed, just wait a bit to ensure file is written
sleep 1
fi
echo -e "${GREEN}✓ CPU profile collection completed${NC}"
echo ""
# ============================================
# Step 5: Collect final profiles
# ============================================
echo -e "${GREEN}[Step 5/8] Collecting final profiles...${NC}"
# Final memory
go tool pprof -proto -output="${OUTPUT_DIR}/heap_after.pb.gz" \
"http://localhost:${PPROF_PORT}/debug/pprof/heap" > /dev/null 2>&1 || true
# Final goroutine
go tool pprof -proto -output="${OUTPUT_DIR}/goroutine_after.pb.gz" \
"http://localhost:${PPROF_PORT}/debug/pprof/goroutine" > /dev/null 2>&1 || true
# Mutex profile
go tool pprof -proto -output="${OUTPUT_DIR}/mutex.pb.gz" \
"http://localhost:${PPROF_PORT}/debug/pprof/mutex" > /dev/null 2>&1 || true
# Block profile
go tool pprof -proto -output="${OUTPUT_DIR}/block.pb.gz" \
"http://localhost:${PPROF_PORT}/debug/pprof/block" > /dev/null 2>&1 || true
echo -e "${GREEN}✓ Final profiles collected${NC}"
echo ""
# ============================================
# Step 6: Generate analysis reports
# ============================================
echo -e "${GREEN}[Step 6/8] Generating analysis reports...${NC}"
# CPU analysis
echo " Generating CPU reports..."
go tool pprof -top -cum "${OUTPUT_DIR}/cpu_${DURATION}s.pb.gz" > "${OUTPUT_DIR}/01_cpu_top_cum.txt" 2>&1 || true
go tool pprof -top "${OUTPUT_DIR}/cpu_${DURATION}s.pb.gz" > "${OUTPUT_DIR}/02_cpu_top_flat.txt" 2>&1 || true
# Memory analysis
echo " Generating memory reports..."
if [ -f "${OUTPUT_DIR}/heap_after.pb.gz" ]; then
go tool pprof -top -alloc_space "${OUTPUT_DIR}/heap_after.pb.gz" > "${OUTPUT_DIR}/03_memory_alloc_space.txt" 2>&1 || true
go tool pprof -top -alloc_objects "${OUTPUT_DIR}/heap_after.pb.gz" > "${OUTPUT_DIR}/04_memory_alloc_objects.txt" 2>&1 || true
go tool pprof -top -inuse_space "${OUTPUT_DIR}/heap_after.pb.gz" > "${OUTPUT_DIR}/05_memory_inuse_space.txt" 2>&1 || true
fi
# Memory growth
if [ -f "${OUTPUT_DIR}/heap_before.pb.gz" ] && [ -f "${OUTPUT_DIR}/heap_after.pb.gz" ]; then
go tool pprof -top -base="${OUTPUT_DIR}/heap_before.pb.gz" \
"${OUTPUT_DIR}/heap_after.pb.gz" > "${OUTPUT_DIR}/06_memory_growth.txt" 2>&1 || true
fi
# FFI/CGO analysis
echo " Analyzing FFI/CGO calls..."
go tool pprof -top "${OUTPUT_DIR}/cpu_${DURATION}s.pb.gz" 2>&1 | \
grep -iE "(block_on|CGO|FFI|ffi|runtime\.cgo|_Cfunc)" > "${OUTPUT_DIR}/07_ffi_cgo_analysis.txt" || \
echo "No FFI/CGO related functions found" > "${OUTPUT_DIR}/07_ffi_cgo_analysis.txt"
# JSON serialization analysis
echo " Analyzing JSON serialization..."
go tool pprof -top "${OUTPUT_DIR}/cpu_${DURATION}s.pb.gz" 2>&1 | \
grep -iE "(json|Marshal|Unmarshal|Encode|Decode|sonic|jsoniter)" > "${OUTPUT_DIR}/08_json_analysis.txt" || \
echo "No JSON related functions found" > "${OUTPUT_DIR}/08_json_analysis.txt"
# Goroutine analysis
if [ -f "${OUTPUT_DIR}/goroutine_after.pb.gz" ]; then
echo " Analyzing goroutines..."
go tool pprof -top "${OUTPUT_DIR}/goroutine_after.pb.gz" > "${OUTPUT_DIR}/09_goroutine_analysis.txt" 2>&1 || true
fi
# Mutex analysis
if [ -f "${OUTPUT_DIR}/mutex.pb.gz" ]; then
echo " Analyzing mutex contention..."
go tool pprof -top "${OUTPUT_DIR}/mutex.pb.gz" > "${OUTPUT_DIR}/10_mutex_analysis.txt" 2>&1 || true
fi
# Block analysis
if [ -f "${OUTPUT_DIR}/block.pb.gz" ]; then
echo " Analyzing blocking operations..."
go tool pprof -top "${OUTPUT_DIR}/block.pb.gz" > "${OUTPUT_DIR}/11_block_analysis.txt" 2>&1 || true
fi
# Request timing statistics
if [ -f "${OUTPUT_DIR}/request_times.txt" ] && [ -s "${OUTPUT_DIR}/request_times.txt" ]; then
echo " Calculating request timing statistics..."
{
echo "Request Timing Statistics"
echo "========================"
echo ""
echo "Total requests: $(wc -l < "${OUTPUT_DIR}/request_times.txt" | tr -d ' ')"
echo ""
awk '{
sum+=$1
sumsq+=$1*$1
if(NR==1 || $1<min) min=$1
if(NR==1 || $1>max) max=$1
} END {
if(NR > 0) {
mean=sum/NR
variance=(sumsq/NR - mean*mean)
stddev=sqrt(variance)
print "Min: " min "s"
print "Max: " max "s"
print "Mean: " mean "s"
print "StdDev: " stddev "s"
}
}' "${OUTPUT_DIR}/request_times.txt"
} > "${OUTPUT_DIR}/12_request_timing.txt"
fi
echo -e "${GREEN}✓ Analysis reports generated${NC}"
echo ""
# ============================================
# Step 7: Generate summary report
# ============================================
echo -e "${GREEN}[Step 7/8] Generating summary report...${NC}"
SUMMARY_FILE="${OUTPUT_DIR}/00_SUMMARY.md"
cat > "$SUMMARY_FILE" <<EOF
# TPOT Performance Analysis Summary
**Analysis Date:** $(date)
**Duration:** ${DURATION}s
**Requests:** $NUM_REQUESTS
**Concurrency:** $CONCURRENCY
## Key Findings
### 1. CPU Hotspots (Top 10 Cumulative Time)
\`\`\`
$(head -15 "${OUTPUT_DIR}/01_cpu_top_cum.txt" | tail -10)
\`\`\`
### 2. CPU Hotspots (Top 10 Flat Time)
\`\`\`
$(head -15 "${OUTPUT_DIR}/02_cpu_top_flat.txt" | tail -10)
\`\`\`
### 3. FFI/CGO Overhead
\`\`\`
$(cat "${OUTPUT_DIR}/07_ffi_cgo_analysis.txt")
\`\`\`
### 4. JSON Serialization Overhead
\`\`\`
$(cat "${OUTPUT_DIR}/08_json_analysis.txt")
\`\`\`
### 5. Memory Allocation (Top 10 by Space)
\`\`\`
$(head -15 "${OUTPUT_DIR}/03_memory_alloc_space.txt" | tail -10)
\`\`\`
### 6. Memory Allocation (Top 10 by Objects)
\`\`\`
$(head -15 "${OUTPUT_DIR}/04_memory_alloc_objects.txt" | tail -10)
\`\`\`
### 7. Mutex Contention
\`\`\`
$(head -15 "${OUTPUT_DIR}/10_mutex_analysis.txt" | tail -10 2>/dev/null || echo "No significant mutex contention detected")
\`\`\`
### 8. Blocking Operations
\`\`\`
$(head -15 "${OUTPUT_DIR}/11_block_analysis.txt" | tail -10 2>/dev/null || echo "No significant blocking detected")
\`\`\`
## Performance Bottlenecks Identified
### High Priority Issues
1. **FFI/CGO Overhead**
- Check: \`cat ${OUTPUT_DIR}/07_ffi_cgo_analysis.txt\`
- Impact: FFI calls add overhead compared to native Rust code
- Recommendation: Minimize FFI calls, batch operations
2. **JSON Serialization**
- Check: \`cat ${OUTPUT_DIR}/08_json_analysis.txt\`
- Impact: JSON marshaling/unmarshaling can be expensive
- Recommendation: Use faster JSON library (jsoniter), reduce serialization frequency
3. **Memory Allocations**
- Check: \`cat ${OUTPUT_DIR}/03_memory_alloc_space.txt\`
- Impact: Frequent allocations cause GC pressure
- Recommendation: Use object pools, pre-allocate buffers
### Medium Priority Issues
4. **Goroutine Overhead**
- Check: \`cat ${OUTPUT_DIR}/09_goroutine_analysis.txt\`
- Impact: Too many goroutines can cause scheduling overhead
- Recommendation: Limit goroutine count, use worker pools
5. **Lock Contention**
- Check: \`cat ${OUTPUT_DIR}/10_mutex_analysis.txt\`
- Impact: Lock contention reduces parallelism
- Recommendation: Reduce lock granularity, use lock-free structures
## Comparison with Rust Router
### Expected Differences
1. **FFI Overhead**: Go → Rust FFI calls add ~100-500ns per call
2. **GC Overhead**: Go's GC can cause pauses (usually <1ms)
3. **JSON Library**: Go's standard library is slower than Rust's serde
4. **Memory Layout**: Go's GC affects cache locality
### Optimization Opportunities
1. **Reduce FFI Calls**
- Batch token processing
- Use async FFI (if possible)
- Cache frequently used FFI results
2. **Optimize JSON**
- Use jsoniter (already implemented)
- Pre-allocate JSON buffers
- Reduce serialization frequency
3. **Memory Management**
- Use sync.Pool for frequently allocated objects
- Pre-allocate slices with known capacity
- Avoid unnecessary string copies
4. **Concurrency**
- Use worker pools instead of spawning goroutines per request
- Limit concurrent FFI calls
- Use channels efficiently
## Next Steps
1. Review detailed reports in this directory
2. Use interactive pprof: \`go tool pprof -http=:8081 ${OUTPUT_DIR}/cpu_${DURATION}s.pb.gz\`
3. Compare with Rust router profiles (if available)
4. Implement optimizations based on findings
5. Re-run analysis to measure improvements
## Files Generated
- \`00_SUMMARY.md\` - This summary
- \`01_cpu_top_cum.txt\` - CPU top functions (cumulative)
- \`02_cpu_top_flat.txt\` - CPU top functions (flat)
- \`03_memory_alloc_space.txt\` - Memory allocation by space
- \`04_memory_alloc_objects.txt\` - Memory allocation by objects
- \`05_memory_inuse_space.txt\` - Memory in use by space
- \`06_memory_growth.txt\` - Memory growth during test
- \`07_ffi_cgo_analysis.txt\` - FFI/CGO overhead analysis
- \`08_json_analysis.txt\` - JSON serialization analysis
- \`09_goroutine_analysis.txt\` - Goroutine analysis
- \`10_mutex_analysis.txt\` - Mutex contention analysis
- \`11_block_analysis.txt\` - Blocking operations analysis
- \`12_request_timing.txt\` - Request timing statistics
- \`*.pb.gz\` - Raw profile files for interactive analysis
EOF
echo -e "${GREEN}✓ Summary report generated${NC}"
echo ""
# ============================================
# Step 8: Display summary
# ============================================
echo -e "${GREEN}[Step 8/8] Analysis Complete!${NC}"
echo ""
echo -e "${BLUE}========================================${NC}"
echo -e "${BLUE}Summary${NC}"
echo -e "${BLUE}========================================${NC}"
echo ""
echo -e "${YELLOW}Top CPU Hotspots (Cumulative):${NC}"
head -12 "${OUTPUT_DIR}/01_cpu_top_cum.txt" | tail -10
echo ""
echo -e "${YELLOW}FFI/CGO Overhead:${NC}"
cat "${OUTPUT_DIR}/07_ffi_cgo_analysis.txt"
echo ""
echo -e "${YELLOW}JSON Serialization Overhead:${NC}"
cat "${OUTPUT_DIR}/08_json_analysis.txt"
echo ""
echo -e "${YELLOW}Top Memory Allocations:${NC}"
head -12 "${OUTPUT_DIR}/03_memory_alloc_space.txt" | tail -10
echo ""
if [ -f "${OUTPUT_DIR}/12_request_timing.txt" ]; then
echo -e "${YELLOW}Request Timing:${NC}"
cat "${OUTPUT_DIR}/12_request_timing.txt"
echo ""
fi
echo -e "${GREEN}========================================${NC}"
echo ""
echo -e "${BLUE}Detailed Reports:${NC}"
echo " Summary: cat ${OUTPUT_DIR}/00_SUMMARY.md"
echo " CPU (cum): cat ${OUTPUT_DIR}/01_cpu_top_cum.txt"
echo " CPU (flat): cat ${OUTPUT_DIR}/02_cpu_top_flat.txt"
echo " FFI/CGO: cat ${OUTPUT_DIR}/07_ffi_cgo_analysis.txt"
echo " JSON: cat ${OUTPUT_DIR}/08_json_analysis.txt"
echo " Memory: cat ${OUTPUT_DIR}/03_memory_alloc_space.txt"
echo ""
echo -e "${BLUE}Interactive Analysis:${NC}"
echo " Run: go tool pprof -http=:8081 ${OUTPUT_DIR}/cpu_${DURATION}s.pb.gz"
echo " Then visit:"
echo " - http://localhost:8081/ui/flamegraph (Flame Graph - no graphviz needed)"
echo " - http://localhost:8081/ui/top (Top Functions - no graphviz needed)"
if [ "$HAS_GRAPHVIZ" = "true" ]; then
echo " - http://localhost:8081/ui/graph (Call Graph - requires graphviz)"
else
echo " - http://localhost:8081/ui/graph (Call Graph - requires graphviz, not available)"
fi
echo ""
if [ "$HAS_GRAPHVIZ" = "false" ]; then
echo -e "${YELLOW}Note: Install graphviz to enable call graph visualization:${NC}"
echo -e "${YELLOW} macOS: brew install graphviz${NC}"
echo -e "${YELLOW} Ubuntu: sudo apt-get install graphviz${NC}"
echo -e "${YELLOW} CentOS: sudo yum install graphviz${NC}"
echo ""
fi
echo -e "${GREEN}All files saved to: ${OUTPUT_DIR}${NC}"
echo ""

215
third_party/sglang/sgl-model-gateway/bindings/golang/examples/oai_server/scripts/pprof_analysis.sh vendored Executable file
View File

@@ -0,0 +1,215 @@
#!/bin/bash
# pprof performance analysis script
# Used to analyze performance bottlenecks of Go OpenAI server
set -e
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"
# Configuration
PPROF_PORT=${PPROF_PORT:-6060}
SERVER_PORT=${SERVER_PORT:-8080}
DURATION=${DURATION:-60} # Performance test duration (seconds)
OUTPUT_DIR="./pprof_results"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
# Create output directory
mkdir -p "$OUTPUT_DIR"
echo "=========================================="
echo "pprof Performance Analysis Tool"
echo "=========================================="
echo "PPROF_PORT: $PPROF_PORT"
echo "SERVER_PORT: $SERVER_PORT"
echo "DURATION: ${DURATION}s"
echo "OUTPUT_DIR: $OUTPUT_DIR"
echo ""
# Check if go tool pprof is available
if ! command -v go &> /dev/null; then
echo "Error: go command not found"
exit 1
fi
# Check if server is running
check_server() {
if curl -s "http://localhost:${SERVER_PORT}/health" > /dev/null 2>&1; then
return 0
else
return 1
fi
}
# Check if pprof is available
check_pprof() {
if curl -s "http://localhost:${PPROF_PORT}/debug/pprof/" > /dev/null 2>&1; then
return 0
else
return 1
fi
}
# Start server (if not running)
if ! check_server; then
echo "Server not running, please start the server first:"
echo " export PPROF_ENABLED=true"
echo " export PPROF_PORT=$PPROF_PORT"
echo " ./oai_server"
echo ""
echo "Or use the following command to start:"
echo " PPROF_ENABLED=true PPROF_PORT=$PPROF_PORT ./oai_server"
echo ""
read -p "Start server now? (y/n) " -n 1 -r
echo
if [[ $REPLY =~ ^[Yy]$ ]]; then
echo "Starting server..."
PPROF_ENABLED=true PPROF_PORT=$PPROF_PORT ./oai_server &
SERVER_PID=$!
echo "Server PID: $SERVER_PID"
# Wait for server to start
echo "Waiting for server to start..."
for i in {1..30}; do
if check_server; then
echo "Server started"
break
fi
sleep 1
done
if ! check_server; then
echo "Error: Server failed to start"
kill $SERVER_PID 2>/dev/null || true
exit 1
fi
else
exit 1
fi
fi
# Check if pprof is available
if ! check_pprof; then
echo "Error: pprof not enabled. Please set environment variables:"
echo " export PPROF_ENABLED=true"
echo " export PPROF_PORT=$PPROF_PORT"
exit 1
fi
echo "Starting to collect performance data..."
echo ""
# 1. CPU Profile (30 seconds)
echo "[1/6] Collecting CPU Profile (30 seconds)..."
go tool pprof -proto -output="$OUTPUT_DIR/cpu_${TIMESTAMP}.pb.gz" \
"http://localhost:${PPROF_PORT}/debug/pprof/profile?seconds=30" &
CPU_PID=$!
# 2. Collect Heap Profile simultaneously
echo "[2/6] Collecting Heap Profile..."
go tool pprof -proto -output="$OUTPUT_DIR/heap_${TIMESTAMP}.pb.gz" \
"http://localhost:${PPROF_PORT}/debug/pprof/heap" &
HEAP_PID=$!
# 3. Collect Goroutine Profile
echo "[3/6] Collecting Goroutine Profile..."
go tool pprof -proto -output="$OUTPUT_DIR/goroutine_${TIMESTAMP}.pb.gz" \
"http://localhost:${PPROF_PORT}/debug/pprof/goroutine" &
GOROUTINE_PID=$!
# 4. Collect Mutex Profile
echo "[4/6] Collecting Mutex Profile..."
go tool pprof -proto -output="$OUTPUT_DIR/mutex_${TIMESTAMP}.pb.gz" \
"http://localhost:${PPROF_PORT}/debug/pprof/mutex" &
MUTEX_PID=$!
# 5. Collect Block Profile
echo "[5/6] Collecting Block Profile..."
go tool pprof -proto -output="$OUTPUT_DIR/block_${TIMESTAMP}.pb.gz" \
"http://localhost:${PPROF_PORT}/debug/pprof/block" &
BLOCK_PID=$!
# 6. Run performance test (during CPU profile collection)
echo "[6/6] Running performance test..."
echo "Tip: Please use your performance testing tool (curl, ab, wrk, etc.) to send requests to the server"
echo " CPU profile will collect 30 seconds of performance data"
echo ""
# Wait for CPU profile to complete
wait $CPU_PID
echo "CPU Profile collection completed"
# Wait for other profiles
wait $HEAP_PID
wait $GOROUTINE_PID
wait $MUTEX_PID
wait $BLOCK_PID
echo ""
echo "=========================================="
echo "Performance data collection completed!"
echo "=========================================="
echo ""
echo "Generated analysis files:"
ls -lh "$OUTPUT_DIR"/*_${TIMESTAMP}.* 2>/dev/null || true
echo ""
# Generate analysis report
echo "Generating analysis report..."
echo ""
# CPU Top 20
echo "=== CPU Top 20 (sorted by flat time) ===" > "$OUTPUT_DIR/analysis_${TIMESTAMP}.txt"
go tool pprof -top -cum "$OUTPUT_DIR/cpu_${TIMESTAMP}.pb.gz" >> "$OUTPUT_DIR/analysis_${TIMESTAMP}.txt" 2>&1 || true
echo "" >> "$OUTPUT_DIR/analysis_${TIMESTAMP}.txt"
# Heap Top 20
echo "=== Heap Top 20 (sorted by allocation size) ===" >> "$OUTPUT_DIR/analysis_${TIMESTAMP}.txt"
go tool pprof -top "$OUTPUT_DIR/heap_${TIMESTAMP}.pb.gz" >> "$OUTPUT_DIR/analysis_${TIMESTAMP}.txt" 2>&1 || true
echo "" >> "$OUTPUT_DIR/analysis_${TIMESTAMP}.txt"
# Goroutine statistics
echo "=== Goroutine Statistics ===" >> "$OUTPUT_DIR/analysis_${TIMESTAMP}.txt"
go tool pprof -top "$OUTPUT_DIR/goroutine_${TIMESTAMP}.pb.gz" >> "$OUTPUT_DIR/analysis_${TIMESTAMP}.txt" 2>&1 || true
echo "" >> "$OUTPUT_DIR/analysis_${TIMESTAMP}.txt"
# Mutex statistics
echo "=== Mutex Wait Time ===" >> "$OUTPUT_DIR/analysis_${TIMESTAMP}.txt"
go tool pprof -top "$OUTPUT_DIR/mutex_${TIMESTAMP}.pb.gz" >> "$OUTPUT_DIR/analysis_${TIMESTAMP}.txt" 2>&1 || true
echo "" >> "$OUTPUT_DIR/analysis_${TIMESTAMP}.txt"
# Block statistics
echo "=== Block Wait Time ===" >> "$OUTPUT_DIR/analysis_${TIMESTAMP}.txt"
go tool pprof -top "$OUTPUT_DIR/block_${TIMESTAMP}.pb.gz" >> "$OUTPUT_DIR/analysis_${TIMESTAMP}.txt" 2>&1 || true
echo "Analysis report saved to: $OUTPUT_DIR/analysis_${TIMESTAMP}.txt"
echo ""
# Display key information
echo "=========================================="
echo "Key Performance Metrics Summary"
echo "=========================================="
echo ""
echo "View detailed report:"
echo " cat $OUTPUT_DIR/analysis_${TIMESTAMP}.txt"
echo ""
echo "Interactive CPU Profile view:"
echo " go tool pprof $OUTPUT_DIR/cpu_${TIMESTAMP}.pb.gz"
echo ""
echo "Interactive Heap Profile view:"
echo " go tool pprof $OUTPUT_DIR/heap_${TIMESTAMP}.pb.gz"
echo ""
echo "Generate flame graph (requires go-torch or pprof):"
echo " go tool pprof -http=:8080 $OUTPUT_DIR/cpu_${TIMESTAMP}.pb.gz"
echo ""
# If server was started, ask if it should be closed
if [ -n "$SERVER_PID" ]; then
read -p "Close server? (y/n) " -n 1 -r
echo
if [[ $REPLY =~ ^[Yy]$ ]]; then
kill $SERVER_PID 2>/dev/null || true
echo "Server closed"
fi
fi

52
third_party/sglang/sgl-model-gateway/bindings/golang/examples/oai_server/scripts/pprof_quick.sh vendored Executable file
View File

@@ -0,0 +1,52 @@
#!/bin/bash
# Quick pprof analysis script
# Collects 30-second CPU profile and immediately displays top results
set -e
PPROF_PORT=${PPROF_PORT:-6060}
DURATION=${DURATION:-30}
echo "=========================================="
echo "Quick pprof Analysis"
echo "=========================================="
echo "PPROF_PORT: $PPROF_PORT"
echo "DURATION: ${DURATION}s"
echo ""
echo "Tip: During data collection, please send requests to the server"
echo " You can use: ./pprof_test.sh"
echo ""
# Check if pprof is available
if ! curl -s "http://localhost:${PPROF_PORT}/debug/pprof/" > /dev/null 2>&1; then
echo "Error: pprof not enabled. Please set environment variables:"
echo " export PPROF_ENABLED=true"
echo " export PPROF_PORT=$PPROF_PORT"
exit 1
fi
echo "Starting to collect CPU Profile (${DURATION} seconds)..."
echo ""
# Collect CPU profile and directly display top results
go tool pprof -top -cum "http://localhost:${PPROF_PORT}/debug/pprof/profile?seconds=${DURATION}"
echo ""
echo "=========================================="
echo "Analysis Complete"
echo "=========================================="
echo ""
echo "More analysis options:"
echo " # Interactive view"
echo " go tool pprof http://localhost:${PPROF_PORT}/debug/pprof/profile?seconds=30"
echo ""
echo " # View heap memory"
echo " go tool pprof http://localhost:${PPROF_PORT}/debug/pprof/heap"
echo ""
echo " # View goroutines"
echo " go tool pprof http://localhost:${PPROF_PORT}/debug/pprof/goroutine"
echo ""
echo " # Generate Web UI"
echo " go tool pprof -http=:8080 http://localhost:${PPROF_PORT}/debug/pprof/profile?seconds=30"
echo ""

87
third_party/sglang/sgl-model-gateway/bindings/golang/examples/oai_server/scripts/pprof_test.sh vendored Executable file
View File

@@ -0,0 +1,87 @@
#!/bin/bash
# Simple performance test script for sending requests while collecting pprof data
set -e
SERVER_URL=${SERVER_URL:-"http://localhost:8080"}
DURATION=${DURATION:-30} # Test duration (seconds)
CONCURRENT=${CONCURRENT:-1} # Number of concurrent requests
echo "=========================================="
echo "Performance Test Script"
echo "=========================================="
echo "SERVER_URL: $SERVER_URL"
echo "DURATION: ${DURATION}s"
echo "CONCURRENT: $CONCURRENT"
echo ""
# Test request JSON
TEST_REQUEST='{
"model": "default",
"messages": [
{"role": "user", "content": "Hello, how are you?"}
],
"stream": true,
"max_tokens": 100
}'
# Check if server is available
if ! curl -s "${SERVER_URL}/health" > /dev/null 2>&1; then
echo "Error: Server not available (${SERVER_URL}/health)"
exit 1
fi
echo "Starting to send test requests..."
echo ""
# Function to send streaming request
send_stream_request() {
local request_num=$1
local start_time=$(date +%s.%N)
curl -s -N -X POST "${SERVER_URL}/v1/chat/completions" \
-H "Content-Type: application/json" \
-d "$TEST_REQUEST" \
> /dev/null 2>&1
local end_time=$(date +%s.%N)
local duration=$(echo "$end_time - $start_time" | bc)
echo "Request $request_num completed, duration: ${duration}s"
}
# Send requests concurrently
if [ "$CONCURRENT" -eq 1 ]; then
# Single-threaded mode: continuously send requests
end_time=$(($(date +%s) + DURATION))
request_count=0
while [ $(date +%s) -lt $end_time ]; do
request_count=$((request_count + 1))
send_stream_request $request_count
done
echo ""
echo "Test completed, sent $request_count requests"
else
# Multi-threaded mode: send requests concurrently
end_time=$(($(date +%s) + DURATION))
request_count=0
while [ $(date +%s) -lt $end_time ]; do
# Start concurrent requests
for i in $(seq 1 $CONCURRENT); do
request_count=$((request_count + 1))
send_stream_request $request_count &
done
# Wait for all requests to complete
wait
# Brief rest to avoid overload
sleep 0.1
done
echo ""
echo "Test completed, sent $request_count requests"
fi

140
third_party/sglang/sgl-model-gateway/bindings/golang/examples/oai_server/scripts/profile_tpot.sh vendored Executable file
View File

@@ -0,0 +1,140 @@
#!/bin/bash
# TPOT performance analysis script
# Quickly collect and analyze TPOT-related performance data
set -e
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
PROFILE_DIR="${PROJECT_ROOT}/profiles"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
OUTPUT_DIR="${PROFILE_DIR}/${TIMESTAMP}"
# Colors
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'
# Default values
PPROF_PORT=${PPROF_PORT:-6060}
SERVER_URL=${SERVER_URL:-http://localhost:8080}
DURATION=${DURATION:-30}
NUM_REQUESTS=${NUM_REQUESTS:-20}
mkdir -p "$OUTPUT_DIR"
echo -e "${GREEN}TPOT Performance Analysis${NC}"
echo "=========================="
echo "Profile directory: $OUTPUT_DIR"
echo "Duration: ${DURATION}s"
echo "Requests: $NUM_REQUESTS"
echo ""
# Check if server is running
if ! curl -s "${SERVER_URL}/health" > /dev/null 2>&1; then
echo -e "${YELLOW}Warning: Server not responding at ${SERVER_URL}${NC}"
echo "Please start the server first with profiling enabled:"
echo " PPROF_ENABLED=true PPROF_PORT=$PPROF_PORT make run"
exit 1
fi
# Collect baseline memory
echo -e "${GREEN}[1/5] Collecting baseline memory profile...${NC}"
go tool pprof -proto -output="${OUTPUT_DIR}/heap_before.pb.gz" \
"http://localhost:${PPROF_PORT}/debug/pprof/heap" > /dev/null 2>&1 || true
# Start CPU profile collection in background
echo -e "${GREEN}[2/5] Starting CPU profile collection (${DURATION}s)...${NC}"
go tool pprof -proto -output="${OUTPUT_DIR}/cpu_${DURATION}s.pb.gz" \
"http://localhost:${PPROF_PORT}/debug/pprof/profile?seconds=${DURATION}" &
CPU_PID=$!
# Wait a bit for profile to start
sleep 2
# Run load test
echo -e "${GREEN}[3/5] Running load test ($NUM_REQUESTS requests)...${NC}"
for i in $(seq 1 $NUM_REQUESTS); do
curl -N -s -X POST "${SERVER_URL}/v1/chat/completions" \
-H "Content-Type: application/json" \
-d "{
\"model\": \"default\",
\"messages\": [{\"role\": \"user\", \"content\": \"Write a story\"}],
\"stream\": true,
\"max_tokens\": 200
}" > /dev/null &
# Limit concurrency
if [ $((i % 5)) -eq 0 ]; then
wait
fi
done
wait
# Wait for CPU profile to complete
echo -e "${GREEN}[4/5] Waiting for CPU profile to complete...${NC}"
# Wait for the CPU profile process, but handle the case where it's not a child process
if kill -0 $CPU_PID 2>/dev/null; then
# Process is still running, wait for it
while kill -0 $CPU_PID 2>/dev/null; do
sleep 1
done
else
# Process already completed or not found, just wait a bit
sleep 2
fi
# Collect final memory
echo -e "${GREEN}[5/5] Collecting final memory profile...${NC}"
go tool pprof -proto -output="${OUTPUT_DIR}/heap_after.pb.gz" \
"http://localhost:${PPROF_PORT}/debug/pprof/heap" > /dev/null 2>&1 || true
# Generate reports
echo ""
echo -e "${GREEN}Generating reports...${NC}"
# CPU top (cumulative)
go tool pprof -top -cum "${OUTPUT_DIR}/cpu_${DURATION}s.pb.gz" > "${OUTPUT_DIR}/cpu_top_cum.txt" 2>&1 || true
# CPU top (flat)
go tool pprof -top "${OUTPUT_DIR}/cpu_${DURATION}s.pb.gz" > "${OUTPUT_DIR}/cpu_top_flat.txt" 2>&1 || true
# Memory growth
if [ -f "${OUTPUT_DIR}/heap_before.pb.gz" ] && [ -f "${OUTPUT_DIR}/heap_after.pb.gz" ]; then
go tool pprof -top -base="${OUTPUT_DIR}/heap_before.pb.gz" \
"${OUTPUT_DIR}/heap_after.pb.gz" > "${OUTPUT_DIR}/heap_growth.txt" 2>&1 || true
fi
# FFI/CGO related
go tool pprof -top "${OUTPUT_DIR}/cpu_${DURATION}s.pb.gz" 2>&1 | \
grep -E "(block_on|CGO|FFI|json|Marshal|Unmarshal)" > "${OUTPUT_DIR}/ffi_related.txt" || \
echo "No FFI/CGO related functions found" > "${OUTPUT_DIR}/ffi_related.txt"
# Summary
echo ""
echo -e "${GREEN}=== Analysis Summary ===${NC}"
echo ""
echo -e "${YELLOW}CPU Top (Cumulative) - Top 10:${NC}"
head -12 "${OUTPUT_DIR}/cpu_top_cum.txt" | tail -10 || true
echo ""
echo -e "${YELLOW}CPU Top (Flat) - Top 10:${NC}"
head -12 "${OUTPUT_DIR}/cpu_top_flat.txt" | tail -10 || true
echo ""
echo -e "${YELLOW}FFI/CGO Related Functions:${NC}"
cat "${OUTPUT_DIR}/ffi_related.txt" || true
echo ""
echo -e "${GREEN}=== Detailed Reports ===${NC}"
echo "CPU (cumulative): cat ${OUTPUT_DIR}/cpu_top_cum.txt"
echo "CPU (flat): cat ${OUTPUT_DIR}/cpu_top_flat.txt"
echo "Memory growth: cat ${OUTPUT_DIR}/heap_growth.txt"
echo "FFI related: cat ${OUTPUT_DIR}/ffi_related.txt"
echo ""
echo -e "${GREEN}=== Interactive Analysis ===${NC}"
echo "Run: go tool pprof -http=:8081 ${OUTPUT_DIR}/cpu_${DURATION}s.pb.gz"
echo "Then visit: http://localhost:8081/ui/flamegraph"
echo ""
echo "Profile files saved to: ${OUTPUT_DIR}"

37
third_party/sglang/sgl-model-gateway/bindings/golang/examples/oai_server/service/sglang.go vendored Normal file
View File

@@ -0,0 +1,37 @@
package service
import (
sglang "github.com/sglang/sglang-go-grpc-sdk"
)
// SGLangService wraps SGLang client
type SGLangService struct {
client *sglang.Client
}
func NewSGLangService(endpoint, tokenizerPath string) (*SGLangService, error) {
client, err := sglang.NewClient(sglang.ClientConfig{
Endpoint: endpoint,
TokenizerPath: tokenizerPath,
})
if err != nil {
return nil, err
}
return &SGLangService{
client: client,
}, nil
}
// Client returns the underlying SGLang client
func (s *SGLangService) Client() *sglang.Client {
return s.client
}
// Close closes the SGLang client
func (s *SGLangService) Close() error {
if s.client != nil {
return s.client.Close()
}
return nil
}

34
third_party/sglang/sgl-model-gateway/bindings/golang/examples/oai_server/utils/utils.go vendored Normal file
View File

@@ -0,0 +1,34 @@
package utils
import (
"encoding/json"
"github.com/valyala/fasthttp"
)
// RespondError sends an error response in OpenAI format
func RespondError(ctx *fasthttp.RequestCtx, statusCode int, message, errorType string) {
ctx.SetStatusCode(statusCode)
ctx.SetContentType("application/json")
response := map[string]interface{}{
"error": map[string]interface{}{
"message": message,
"type": errorType,
"code": statusCode,
},
}
jsonData, _ := json.Marshal(response)
ctx.Write(jsonData)
}
// BuildResponseBase builds the base response structure for OpenAI-compatible responses
func BuildResponseBase(id string, created int64, model string) map[string]interface{} {
return map[string]interface{}{
"id": id,
"object": "chat.completion",
"created": created,
"model": model,
}
}

85
third_party/sglang/sgl-model-gateway/bindings/golang/examples/simple/main.go vendored Normal file
View File

@@ -0,0 +1,85 @@
// Simple example demonstrating basic usage of SGLang Go SDK
package main
import (
"context"
"fmt"
"log"
"os"
"github.com/sglang/sglang-go-grpc-sdk"
)
func main() {
// Get configuration from environment or command line
endpoint := os.Getenv("SGL_GRPC_ENDPOINT")
if endpoint == "" {
endpoint = "grpc://localhost:20000"
}
tokenizerPath := os.Getenv("SGL_TOKENIZER_PATH")
if tokenizerPath == "" {
tokenizerPath = "./examples/tokenizer"
}
// Create client
client, err := sglang.NewClient(sglang.ClientConfig{
Endpoint: endpoint,
TokenizerPath: tokenizerPath,
})
if err != nil {
log.Fatalf("Failed to create client: %v", err)
}
defer client.Close()
// Create chat completion request
req := sglang.ChatCompletionRequest{
Model: "default",
Messages: []sglang.ChatMessage{
{
Role: "system",
Content: "You are a helpful assistant.",
},
{
Role: "user",
Content: "写一首歌关于夏天",
},
},
Stream: false,
Temperature: float32Ptr(0.7),
MaxCompletionTokens: intPtr(200),
SkipSpecialTokens: true,
Tools: nil, // Use nil instead of empty slice to avoid template errors
}
// Create completion
ctx := context.Background()
resp, err := client.CreateChatCompletion(ctx, req)
if err != nil {
log.Fatalf("Failed to create completion: %v", err)
}
// Print response
fmt.Println("=== Response ===")
fmt.Printf("ID: %s\n", resp.ID)
fmt.Printf("Model: %s\n", resp.Model)
fmt.Printf("Created: %d\n", resp.Created)
fmt.Println("\nContent:")
for _, choice := range resp.Choices {
fmt.Println(choice.Message.Content)
}
fmt.Printf("\nFinish Reason: %s\n", resp.Choices[0].FinishReason)
fmt.Printf("\nUsage: Prompt=%d, Completion=%d, Total=%d\n",
resp.Usage.PromptTokens,
resp.Usage.CompletionTokens,
resp.Usage.TotalTokens,
)
}
func float32Ptr(f float32) *float32 {
return &f
}
func intPtr(i int) *int {
return &i
}

46
third_party/sglang/sgl-model-gateway/bindings/golang/examples/simple/run.sh vendored Executable file
View File

@@ -0,0 +1,46 @@
#!/bin/bash
# Simple example runner
# Usage: ./run.sh [tokenizer_path] [endpoint]
# Set library path for Rust FFI library
# The library should be in ./lib directory (created by 'make lib')
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
LIB_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)/lib"
# Check if lib directory exists
if [ ! -d "$LIB_DIR" ]; then
echo "Error: Library directory not found at $LIB_DIR"
echo "Please run 'make lib' first to build and export the library"
exit 1
fi
# Get Python LDFLAGS (needed for Rust FFI that depends on Python)
PYTHON_LDFLAGS=$(python3-config --ldflags --embed 2>/dev/null || python3-config --ldflags 2>/dev/null || echo "")
# Set CGO_LDFLAGS to link with the Rust library
export CGO_LDFLAGS="-L${LIB_DIR} -lsgl_model_gateway_go ${PYTHON_LDFLAGS} -ldl"
# macOS uses DYLD_LIBRARY_PATH, Linux uses LD_LIBRARY_PATH
if [[ "$OSTYPE" == "darwin"* ]]; then
export DYLD_LIBRARY_PATH="${LIB_DIR}:${DYLD_LIBRARY_PATH}"
else
export LD_LIBRARY_PATH="${LIB_DIR}:${LD_LIBRARY_PATH}"
fi
# Default configuration (can be overridden by environment variables or command line arguments)
# Tokenizer path: ../tokenizer (relative to this script)
DEFAULT_TOKENIZER_PATH="${SGL_TOKENIZER_PATH:-../tokenizer}"
DEFAULT_ENDPOINT="${SGL_GRPC_ENDPOINT:-grpc://localhost:20000}"
TOKENIZER_PATH="${1:-${DEFAULT_TOKENIZER_PATH}}"
ENDPOINT="${2:-${DEFAULT_ENDPOINT}}"
echo "Running simple example..."
echo "Library path: ${LIB_DIR}"
echo "Tokenizer: $TOKENIZER_PATH"
echo "Endpoint: $ENDPOINT"
echo ""
cd "$(dirname "${BASH_SOURCE[0]}")"
SGL_TOKENIZER_PATH="$TOKENIZER_PATH" SGL_GRPC_ENDPOINT="$ENDPOINT" go run main.go

125
third_party/sglang/sgl-model-gateway/bindings/golang/examples/streaming/main.go vendored Normal file
View File

@@ -0,0 +1,125 @@
// Streaming example demonstrating real-time streaming with SGLang Go SDK
package main
import (
"context"
"fmt"
"io"
"log"
"os"
"strings"
"time"
"github.com/sglang/sglang-go-grpc-sdk"
)
func main() {
// Get configuration from environment or command line
endpoint := os.Getenv("SGL_GRPC_ENDPOINT")
if endpoint == "" {
endpoint = "grpc://localhost:20000"
}
tokenizerPath := os.Getenv("SGL_TOKENIZER_PATH")
if tokenizerPath == "" {
tokenizerPath = "./examples/tokenizer"
}
// Create client
client, err := sglang.NewClient(sglang.ClientConfig{
Endpoint: endpoint,
TokenizerPath: tokenizerPath,
})
if err != nil {
log.Fatalf("Failed to create client: %v", err)
}
defer client.Close()
// Create streaming chat completion request
req := sglang.ChatCompletionRequest{
Model: "default",
Messages: []sglang.ChatMessage{
{
Role: "system",
Content: "You are a helpful assistant.",
},
{
Role: "user",
Content: "写一首春天的诗歌",
},
},
Stream: true,
Temperature: float32Ptr(0.7),
MaxCompletionTokens: intPtr(500),
SkipSpecialTokens: true,
Tools: nil, // Use nil instead of empty slice to avoid template errors
}
// Create streaming completion
ctx := context.Background()
stream, err := client.CreateChatCompletionStream(ctx, req)
if err != nil {
log.Fatalf("Failed to create stream: %v", err)
}
defer stream.Close()
fmt.Println("=== Streaming Response ===")
fmt.Println()
var fullContent strings.Builder
chunkCount := 0
startTime := time.Now()
var firstTokenTime time.Time
firstTokenReceived := false
for {
chunk, err := stream.Recv()
if err == io.EOF {
break
}
if err != nil {
log.Fatalf("Stream error: %v", err)
}
chunkCount++
// Extract content from delta
for _, choice := range chunk.Choices {
if choice.Delta.Content != "" {
fmt.Print(choice.Delta.Content)
fullContent.WriteString(choice.Delta.Content)
// Track first token time (TTFT)
if !firstTokenReceived {
firstTokenTime = time.Now()
firstTokenReceived = true
ttft := firstTokenTime.Sub(startTime)
fmt.Printf("\n[TTFT: %v]\n", ttft)
}
}
if choice.FinishReason != "" {
fmt.Printf("\n\n[Finished: %s]\n", choice.FinishReason)
}
}
}
// Calculate metrics
if firstTokenReceived {
elapsed := time.Since(startTime)
tokensPerSecond := float64(fullContent.Len()) / elapsed.Seconds()
fmt.Printf("\n=== Metrics ===\n")
fmt.Printf("Total chunks: %d\n", chunkCount)
fmt.Printf("Total content length: %d characters\n", fullContent.Len())
fmt.Printf("Time elapsed: %v\n", elapsed)
fmt.Printf("Tokens per second: %.2f\n", tokensPerSecond)
}
}
func float32Ptr(f float32) *float32 {
return &f
}
func intPtr(i int) *int {
return &i
}

46
third_party/sglang/sgl-model-gateway/bindings/golang/examples/streaming/run.sh vendored Executable file
View File

@@ -0,0 +1,46 @@
#!/bin/bash
# Streaming example runner
# Usage: ./run.sh [tokenizer_path] [endpoint]
# Set library path for Rust FFI library
# The library should be in ./lib directory (created by 'make lib')
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
LIB_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)/lib"
# Check if lib directory exists
if [ ! -d "$LIB_DIR" ]; then
echo "Error: Library directory not found at $LIB_DIR"
echo "Please run 'make lib' first to build and export the library"
exit 1
fi
# Get Python LDFLAGS (needed for Rust FFI that depends on Python)
PYTHON_LDFLAGS=$(python3-config --ldflags --embed 2>/dev/null || python3-config --ldflags 2>/dev/null || echo "")
# Set CGO_LDFLAGS to link with the Rust library
export CGO_LDFLAGS="-L${LIB_DIR} -lsgl_model_gateway_go ${PYTHON_LDFLAGS} -ldl"
# macOS uses DYLD_LIBRARY_PATH, Linux uses LD_LIBRARY_PATH
if [[ "$OSTYPE" == "darwin"* ]]; then
export DYLD_LIBRARY_PATH="${LIB_DIR}:${DYLD_LIBRARY_PATH}"
else
export LD_LIBRARY_PATH="${LIB_DIR}:${LD_LIBRARY_PATH}"
fi
# Default configuration (can be overridden by environment variables or command line arguments)
# Tokenizer path: ../tokenizer (relative to this script)
DEFAULT_TOKENIZER_PATH="${SGL_TOKENIZER_PATH:-../tokenizer}"
DEFAULT_ENDPOINT="${SGL_GRPC_ENDPOINT:-grpc://localhost:20000}"
TOKENIZER_PATH="${1:-${DEFAULT_TOKENIZER_PATH}}"
ENDPOINT="${2:-${DEFAULT_ENDPOINT}}"
echo "Running streaming example..."
echo "Library path: ${LIB_DIR}"
echo "Tokenizer: $TOKENIZER_PATH"
echo "Endpoint: $ENDPOINT"
echo ""
cd "$(dirname "${BASH_SOURCE[0]}")"
SGL_TOKENIZER_PATH="$TOKENIZER_PATH" SGL_GRPC_ENDPOINT="$ENDPOINT" go run main.go

36
third_party/sglang/sgl-model-gateway/bindings/golang/go.sum vendored Normal file
View File

@@ -0,0 +1,36 @@
github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI=
github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=
go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
go.opentelemetry.io/otel v1.38.0 h1:RkfdswUDRimDg0m2Az18RKOsnI8UDzppJAtj01/Ymk8=
go.opentelemetry.io/otel v1.38.0/go.mod h1:zcmtmQ1+YmQM9wrNsTGV/q/uyusom3P8RxwExxkZhjM=
go.opentelemetry.io/otel/metric v1.38.0 h1:Kl6lzIYGAh5M159u9NgiRkmoMKjvbsKtYRwgfrA6WpA=
go.opentelemetry.io/otel/metric v1.38.0/go.mod h1:kB5n/QoRM8YwmUahxvI3bO34eVtQf2i4utNVLr9gEmI=
go.opentelemetry.io/otel/sdk v1.38.0 h1:l48sr5YbNf2hpCUj/FoGhW9yDkl+Ma+LrVl8qaM5b+E=
go.opentelemetry.io/otel/sdk v1.38.0/go.mod h1:ghmNdGlVemJI3+ZB5iDEuk4bWA3GkTpW+DOoZMYBVVg=
go.opentelemetry.io/otel/sdk/metric v1.38.0 h1:aSH66iL0aZqo//xXzQLYozmWrXxyFkBJ6qT5wthqPoM=
go.opentelemetry.io/otel/sdk/metric v1.38.0/go.mod h1:dg9PBnW9XdQ1Hd6ZnRz689CbtrUp0wMMs9iPcgT9EZA=
go.opentelemetry.io/otel/trace v1.38.0 h1:Fxk5bKrDZJUH+AMyyIXGcFAPah0oRcT+LuNtJrmcNLE=
go.opentelemetry.io/otel/trace v1.38.0/go.mod h1:j1P9ivuFsTceSWe1oY+EeW3sc+Pp42sO++GHkg4wwhs=
golang.org/x/net v0.46.1-0.20251013234738-63d1a5100f82 h1:6/3JGEh1C88g7m+qzzTbl3A0FtsLguXieqofVLU/JAo=
golang.org/x/net v0.46.1-0.20251013234738-63d1a5100f82/go.mod h1:Q9BGdFy1y4nkUwiLvT5qtyhAnEHgnQ/zd8PfU6nc210=
golang.org/x/sys v0.37.0 h1:fdNQudmxPjkdUTPnLn5mdQv7Zwvbvpaxqs831goi9kQ=
golang.org/x/sys v0.37.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
golang.org/x/text v0.30.0 h1:yznKA/E9zq54KzlzBEAWn1NXSQ8DIp/NYMy88xJjl4k=
golang.org/x/text v0.30.0/go.mod h1:yDdHFIX9t+tORqspjENWgzaCVXgk0yYnYuSZ8UzzBVM=
gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk=
gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E=
google.golang.org/genproto/googleapis/rpc v0.0.0-20251022142026-3a174f9686a8 h1:M1rk8KBnUsBDg1oPGHNCxG4vc1f49epmTO7xscSajMk=
google.golang.org/genproto/googleapis/rpc v0.0.0-20251022142026-3a174f9686a8/go.mod h1:7i2o+ce6H/6BluujYR+kqX3GKH+dChPTQU19wjRPiGk=
google.golang.org/grpc v1.77.0 h1:wVVY6/8cGA6vvffn+wWK5ToddbgdU3d8MNENr4evgXM=
google.golang.org/grpc v1.77.0/go.mod h1:z0BY1iVj0q8E1uSQCjL9cppRj+gnZjzDnzV0dHhrNig=
google.golang.org/protobuf v1.36.10 h1:AYd7cD/uASjIL6Q9LiTjz8JLcrh/88q5UObnmY3aOOE=
google.golang.org/protobuf v1.36.10/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=

228
third_party/sglang/sgl-model-gateway/bindings/golang/integration_test.go vendored Normal file
View File

@@ -0,0 +1,228 @@
//go:build integration
// +build integration
// integration_test.go contains integration tests that require a running SGLang server
//
// To run these tests:
// 1. Start an SGLang server: python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-hf
// 2. Run: go test -tags=integration -run TestIntegration
package sglang
import (
"context"
"io"
"os"
"testing"
"time"
)
// getTestConfig returns test configuration from environment or defaults
func getTestConfig(t *testing.T) ClientConfig {
endpoint := os.Getenv("SGL_GRPC_ENDPOINT")
if endpoint == "" {
endpoint = "grpc://localhost:20000"
}
tokenizerPath := os.Getenv("SGL_TOKENIZER_PATH")
if tokenizerPath == "" {
t.Skip("SGL_TOKENIZER_PATH not set")
}
return ClientConfig{
Endpoint: endpoint,
TokenizerPath: tokenizerPath,
}
}
// TestIntegrationNonStreamingCompletion tests non-streaming chat completion
func TestIntegrationNonStreamingCompletion(t *testing.T) {
config := getTestConfig(t)
client, err := NewClient(config)
if err != nil {
t.Fatalf("Failed to create client: %v", err)
}
defer client.Close()
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
req := ChatCompletionRequest{
Model: "default",
Messages: []ChatMessage{
{Role: "user", Content: "Say 'Hello, World!' only"},
},
Stream: false,
Temperature: float32Ptr(0.0),
MaxCompletionTokens: intPtr(50),
}
resp, err := client.CreateChatCompletion(ctx, req)
if err != nil {
t.Fatalf("CreateChatCompletion failed: %v", err)
}
if resp.ID == "" {
t.Error("Response ID is empty")
}
if len(resp.Choices) == 0 {
t.Error("Response has no choices")
}
if resp.Choices[0].Message.Content == "" {
t.Error("Response content is empty")
}
if resp.Usage == nil || resp.Usage.TotalTokens == 0 {
t.Error("Usage information is missing or invalid")
}
t.Logf("Response: %s", resp.Choices[0].Message.Content)
t.Logf("Usage: %+v", resp.Usage)
}
// TestIntegrationStreamingCompletion tests streaming chat completion
func TestIntegrationStreamingCompletion(t *testing.T) {
config := getTestConfig(t)
client, err := NewClient(config)
if err != nil {
t.Fatalf("Failed to create client: %v", err)
}
defer client.Close()
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
req := ChatCompletionRequest{
Model: "default",
Messages: []ChatMessage{
{Role: "user", Content: "Count from 1 to 5"},
},
Stream: true,
Temperature: float32Ptr(0.0),
MaxCompletionTokens: intPtr(100),
}
stream, err := client.CreateChatCompletionStream(ctx, req)
if err != nil {
t.Fatalf("CreateChatCompletionStream failed: %v", err)
}
defer stream.Close()
chunkCount := 0
totalContent := ""
for {
chunk, err := stream.Recv()
if err == io.EOF {
// io.EOF is expected at end of stream
break
}
if err != nil {
t.Fatalf("Stream error: %v", err)
}
chunkCount++
for _, choice := range chunk.Choices {
if choice.Delta.Content != "" {
totalContent += choice.Delta.Content
}
}
}
if chunkCount == 0 {
t.Error("Received no chunks from stream")
}
if totalContent == "" {
t.Error("Received no content from stream")
}
t.Logf("Received %d chunks with content: %s", chunkCount, totalContent)
}
// TestIntegrationConcurrentRequests tests multiple concurrent requests
func TestIntegrationConcurrentRequests(t *testing.T) {
config := getTestConfig(t)
client, err := NewClient(config)
if err != nil {
t.Fatalf("Failed to create client: %v", err)
}
defer client.Close()
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()
numRequests := 3
done := make(chan error, numRequests)
for i := 0; i < numRequests; i++ {
go func(idx int) {
req := ChatCompletionRequest{
Model: "default",
Messages: []ChatMessage{
{Role: "user", Content: "Say 'test'"},
},
Stream: false,
MaxCompletionTokens: intPtr(50),
}
_, err := client.CreateChatCompletion(ctx, req)
done <- err
}(i)
}
// Collect results
for i := 0; i < numRequests; i++ {
if err := <-done; err != nil {
t.Errorf("Request %d failed: %v", i, err)
}
}
t.Logf("All %d concurrent requests completed successfully", numRequests)
}
// TestIntegrationContextCancellation tests that context cancellation is handled
func TestIntegrationContextCancellation(t *testing.T) {
config := getTestConfig(t)
client, err := NewClient(config)
if err != nil {
t.Fatalf("Failed to create client: %v", err)
}
defer client.Close()
// Create a context that cancels immediately
ctx, cancel := context.WithCancel(context.Background())
cancel()
req := ChatCompletionRequest{
Model: "default",
Messages: []ChatMessage{
{Role: "user", Content: "test"},
},
Stream: false,
}
// Should handle cancelled context gracefully
_, err = client.CreateChatCompletion(ctx, req)
if err == nil {
t.Error("Expected error from cancelled context")
}
t.Logf("Cancelled context handled: %v", err)
}
// Helper functions
func float32Ptr(f float32) *float32 {
return &f
}
func intPtr(i int) *int {
return &i
}

126
third_party/sglang/sgl-model-gateway/bindings/golang/internal/ffi/batch_postprocessor.go vendored Normal file
View File

@@ -0,0 +1,126 @@
// Package ffi provides Go bindings for SGLang's Rust FFI (Foreign Function Interface).
package ffi
import (
"encoding/json"
"fmt"
"strings"
"time"
)
// BatchPostprocessor handles batch postprocessing of stream chunks to reduce FFI overhead
type BatchPostprocessor struct {
converter *GrpcResponseConverterHandle
buffer []string
batchSize int
flushInterval time.Duration
lastFlush time.Time
timer *time.Timer
}
// NewBatchPostprocessor creates a new batch postprocessor
func NewBatchPostprocessor(converter *GrpcResponseConverterHandle, batchSize int, flushInterval time.Duration) *BatchPostprocessor {
if batchSize <= 0 {
batchSize = 1
}
if flushInterval < 0 {
flushInterval = 0
}
return &BatchPostprocessor{
converter: converter,
buffer: make([]string, 0, batchSize),
batchSize: batchSize,
flushInterval: flushInterval,
lastFlush: time.Now(),
}
}
// AddChunk adds a chunk to the buffer and processes if batch is full
func (b *BatchPostprocessor) AddChunk(chunkJSON string) (results []string, shouldFlush bool, err error) {
if b.batchSize == 1 {
openaiJSON, _, err := PostprocessStreamChunk(b.converter, chunkJSON)
if err != nil {
return nil, false, err
}
return []string{openaiJSON}, false, nil
}
b.buffer = append(b.buffer, chunkJSON)
shouldProcess := len(b.buffer) >= b.batchSize
shouldFlushTimeout := b.flushInterval > 0 && time.Since(b.lastFlush) >= b.flushInterval
if shouldProcess || shouldFlushTimeout {
return b.processBatch()
}
return nil, false, nil
}
// Flush processes any remaining chunks in the buffer
func (b *BatchPostprocessor) Flush() (results []string, err error) {
if len(b.buffer) == 0 {
return nil, nil
}
res, _, err := b.processBatch()
return res, err
}
// processBatch processes the current buffer and returns results
func (b *BatchPostprocessor) processBatch() (results []string, shouldFlush bool, err error) {
if len(b.buffer) == 0 {
return nil, false, nil
}
var sb strings.Builder
sb.Grow(len(b.buffer) * 200)
sb.WriteString(`[`)
for i, chunkJSONStr := range b.buffer {
if i > 0 {
sb.WriteString(`,`)
}
sb.WriteString(chunkJSONStr)
}
sb.WriteString(`]`)
bufferJSON := sb.String()
resultJSON, _, err := PostprocessStreamChunksBatch(
b.converter,
bufferJSON,
b.batchSize*2,
)
if err != nil {
return nil, false, fmt.Errorf("batch postprocessing failed: %w", err)
}
var resultArray []json.RawMessage
if err := json.Unmarshal([]byte(resultJSON), &resultArray); err != nil {
return nil, false, fmt.Errorf("failed to unmarshal results array: %w", err)
}
resultStrings := make([]string, 0, len(resultArray))
for _, rawMsg := range resultArray {
resultStrings = append(resultStrings, string(rawMsg))
}
b.buffer = b.buffer[:0]
b.lastFlush = time.Now()
if b.timer != nil {
b.timer.Stop()
b.timer = nil
}
return resultStrings, false, nil
}
// Reset clears the buffer and resets the postprocessor state
func (b *BatchPostprocessor) Reset() {
b.buffer = b.buffer[:0]
b.lastFlush = time.Now()
if b.timer != nil {
b.timer.Stop()
b.timer = nil
}
}

228
third_party/sglang/sgl-model-gateway/bindings/golang/internal/ffi/client.go vendored Normal file
View File

@@ -0,0 +1,228 @@
// Package ffi provides Go bindings for SGLang's Rust FFI (Foreign Function Interface).
//
// This package wraps the Rust FFI layer of SGLang, providing low-level access to:
// - Client creation and connection management
// - Chat completion streaming
// - Stream reading and response conversion
// - Memory management for C strings
//
// Internal use only: This package is intended for internal use by the sglang package.
// End users should use the public sglang package instead.
package ffi
/*
#cgo LDFLAGS: -lsgl_model_gateway_go -ldl
#include <stdlib.h>
#include <stdint.h>
// Error codes
typedef enum {
SGL_ERROR_SUCCESS = 0,
SGL_ERROR_INVALID_ARGUMENT = 1,
SGL_ERROR_TOKENIZATION_ERROR = 2,
SGL_ERROR_PARSING_ERROR = 3,
SGL_ERROR_MEMORY_ERROR = 4,
SGL_ERROR_UNKNOWN = 99
} SglErrorCode;
// Opaque handles
typedef void* SglangClientHandle;
typedef void* SglangStreamHandle;
// Client SDK functions
SglangClientHandle* sgl_client_create(const char* endpoint, const char* tokenizer_path, char** error_out);
void sgl_client_free(SglangClientHandle* handle);
SglErrorCode sgl_client_chat_completion_stream(SglangClientHandle* client_handle, const char* request_json, SglangStreamHandle** stream_handle_out, char** error_out);
SglErrorCode sgl_stream_read_next(SglangStreamHandle* stream_handle, char** response_json_out, int* is_done_out, char** error_out);
void sgl_stream_free(SglangStreamHandle* handle);
void sgl_free_string(char* s);
*/
import "C"
import (
"fmt"
"unsafe"
)
// ErrorCode represents FFI error codes returned by Rust functions.
//
// These codes indicate the result of FFI operations. Use Error() to get a human-readable
// error message.
type ErrorCode int
const (
// ErrorSuccess indicates the operation completed successfully
ErrorSuccess ErrorCode = 0
// ErrorInvalidArgument indicates invalid arguments were passed to the FFI function
ErrorInvalidArgument ErrorCode = 1
// ErrorTokenizationError indicates an error during tokenization
ErrorTokenizationError ErrorCode = 2
// ErrorParsingError indicates an error parsing the response or request
ErrorParsingError ErrorCode = 3
// ErrorMemoryError indicates a memory allocation error
ErrorMemoryError ErrorCode = 4
// ErrorUnknown indicates an unclassified error
ErrorUnknown ErrorCode = 99
)
// Error implements the error interface for ErrorCode.
func (e ErrorCode) Error() string {
switch e {
case ErrorSuccess:
return "success"
case ErrorInvalidArgument:
return "invalid argument"
case ErrorTokenizationError:
return "tokenization error"
case ErrorParsingError:
return "parsing error"
case ErrorMemoryError:
return "memory error"
case ErrorUnknown:
return "unknown error"
default:
return fmt.Sprintf("unknown error code: %d", e)
}
}
// SglangClientHandle wraps the Rust client SDK FFI handle.
//
// This struct maintains a connection to the SGLang gRPC server and is used
// to create streams and manage the underlying Rust client resources.
type SglangClientHandle struct {
handle *C.SglangClientHandle
}
// NewClient creates a new SGLang client handle via FFI.
//
// This function initializes the Rust client with the given endpoint and tokenizer path.
//
// Parameters:
// - endpoint: gRPC endpoint URL (e.g., "grpc://localhost:20000")
// - tokenizerPath: Path to tokenizer directory
//
// Returns:
// - *SglangClientHandle: A new client handle
// - error: An error if client creation failed
func NewClient(endpoint, tokenizerPath string) (*SglangClientHandle, error) {
cEndpoint := C.CString(endpoint)
defer C.free(unsafe.Pointer(cEndpoint))
cTokenizerPath := C.CString(tokenizerPath)
defer C.free(unsafe.Pointer(cTokenizerPath))
var errorPtr *C.char
handle := C.sgl_client_create(cEndpoint, cTokenizerPath, &errorPtr)
if handle == nil {
errorMsg := ""
if errorPtr != nil {
errorMsg = C.GoString(errorPtr)
C.sgl_free_string(errorPtr)
}
if errorMsg == "" {
errorMsg = "failed to create client"
}
return nil, fmt.Errorf("%s", errorMsg)
}
return &SglangClientHandle{handle: handle}, nil
}
// Free releases the client handle
func (h *SglangClientHandle) Free() {
if h.handle != nil {
C.sgl_client_free(h.handle)
h.handle = nil
}
}
// ChatCompletionStream creates a streaming chat completion request
func (h *SglangClientHandle) ChatCompletionStream(requestJSON string) (*SglangStreamHandle, error) {
if h.handle == nil {
return nil, fmt.Errorf("client handle is nil")
}
cRequestJSON := C.CString(requestJSON)
defer C.free(unsafe.Pointer(cRequestJSON))
var streamHandle *C.SglangStreamHandle
var errorPtr *C.char
result := C.sgl_client_chat_completion_stream(
h.handle,
cRequestJSON,
&streamHandle,
&errorPtr,
)
if ErrorCode(result) != ErrorSuccess {
errorMsg := ""
if errorPtr != nil {
errorMsg = C.GoString(errorPtr)
C.sgl_free_string(errorPtr)
}
if errorMsg == "" {
errorMsg = fmt.Sprintf("error code %d", result)
}
return nil, fmt.Errorf("%s", errorMsg)
}
if streamHandle == nil {
return nil, fmt.Errorf("stream handle is nil")
}
return &SglangStreamHandle{handle: streamHandle}, nil
}
// SglangStreamHandle wraps the Rust stream FFI handle
type SglangStreamHandle struct {
handle *C.SglangStreamHandle
}
// ReadNext reads the next chunk from the stream
// Returns: (responseJSON, isDone, error)
func (h *SglangStreamHandle) ReadNext() (string, bool, error) {
if h.handle == nil {
return "", true, fmt.Errorf("stream handle is nil")
}
var responseJSON *C.char
var isDone C.int
var errorPtr *C.char
result := C.sgl_stream_read_next(
h.handle,
&responseJSON,
&isDone,
&errorPtr,
)
if ErrorCode(result) != ErrorSuccess {
errorMsg := ""
if errorPtr != nil {
errorMsg = C.GoString(errorPtr)
C.sgl_free_string(errorPtr)
}
if errorMsg == "" {
errorMsg = fmt.Sprintf("error code %d", result)
}
return "", isDone == 1, fmt.Errorf("%s", errorMsg)
}
responseStr := ""
if responseJSON != nil {
responseStr = C.GoString(responseJSON)
C.sgl_free_string(responseJSON)
}
return responseStr, isDone == 1, nil
}
// Free releases the stream handle
func (h *SglangStreamHandle) Free() {
if h.handle != nil {
C.sgl_stream_free(h.handle)
h.handle = nil
}
}

275
third_party/sglang/sgl-model-gateway/bindings/golang/internal/ffi/grpc_converter.go vendored Normal file
View File

@@ -0,0 +1,275 @@
package ffi
/*
#cgo LDFLAGS: -lsgl_model_gateway_go -ldl
#include <stdlib.h>
#include <stdint.h>
// Error codes (must match client.go)
typedef enum {
SGL_ERROR_SUCCESS = 0,
SGL_ERROR_INVALID_ARGUMENT = 1,
SGL_ERROR_TOKENIZATION_ERROR = 2,
SGL_ERROR_PARSING_ERROR = 3,
SGL_ERROR_MEMORY_ERROR = 4,
SGL_ERROR_UNKNOWN = 99
} SglErrorCode;
// Opaque handles
typedef void* TokenizerHandle;
typedef void* GrpcResponseConverterHandle;
// Converter functions
GrpcResponseConverterHandle* sgl_grpc_response_converter_create(
TokenizerHandle* tokenizer_handle,
const char* model,
const char* request_id,
const char* tools_json,
const char* tool_choice_json,
const char* stop,
const char* stop_token_ids,
int skip_special_tokens,
int initial_prompt_tokens,
char** error_out
);
void sgl_grpc_response_converter_free(GrpcResponseConverterHandle* handle);
// Tokenizer functions
TokenizerHandle* sgl_tokenizer_create_from_file(const char* tokenizer_path, char** error_out);
void sgl_tokenizer_free(TokenizerHandle* handle);
// Memory management
void sgl_free_string(char* s);
*/
import "C"
import (
"fmt"
"unsafe"
)
// CreateGrpcResponseConverter creates a gRPC response converter handle
// This function creates a new tokenizer handle each time (for backward compatibility)
// For better performance, use CreateGrpcResponseConverterWithTokenizer with a cached tokenizer
func CreateGrpcResponseConverter(
tokenizerPath string,
model string,
requestID string,
toolsJSON string,
toolChoiceJSON string,
stopJSON string,
stopTokenIDs []uint32,
skipSpecialTokens bool,
initialPromptTokens int32,
) (*GrpcResponseConverterHandle, error) {
// Create tokenizer handle
tokenizerHandle, err := createTokenizerHandle(tokenizerPath)
if err != nil {
return nil, fmt.Errorf("failed to create tokenizer handle: %w", err)
}
defer C.sgl_tokenizer_free(tokenizerHandle)
return createGrpcResponseConverterWithTokenizerHandle(
tokenizerHandle,
model,
requestID,
toolsJSON,
toolChoiceJSON,
stopJSON,
stopTokenIDs,
skipSpecialTokens,
initialPromptTokens,
)
}
// CreateGrpcResponseConverterWithTokenizer creates a gRPC response converter handle using a cached tokenizer
// This is more efficient as it reuses the tokenizer instead of creating a new one each time
func CreateGrpcResponseConverterWithTokenizer(
tokenizerHandle *TokenizerHandle,
model string,
requestID string,
toolsJSON string,
toolChoiceJSON string,
stopJSON string,
stopTokenIDs []uint32,
skipSpecialTokens bool,
initialPromptTokens int32,
) (*GrpcResponseConverterHandle, error) {
if tokenizerHandle == nil || tokenizerHandle.handle == nil {
return nil, fmt.Errorf("invalid tokenizer handle")
}
return createGrpcResponseConverterWithTokenizerHandle(
tokenizerHandle.handle,
model,
requestID,
toolsJSON,
toolChoiceJSON,
stopJSON,
stopTokenIDs,
skipSpecialTokens,
initialPromptTokens,
)
}
// createGrpcResponseConverterWithTokenizerHandle is the internal implementation
func createGrpcResponseConverterWithTokenizerHandle(
tokenizerHandle *C.TokenizerHandle,
model string,
requestID string,
toolsJSON string,
toolChoiceJSON string,
stopJSON string,
stopTokenIDs []uint32,
skipSpecialTokens bool,
initialPromptTokens int32,
) (*GrpcResponseConverterHandle, error) {
// Convert strings to C strings
modelC := C.CString(model)
defer C.free(unsafe.Pointer(modelC))
requestIDC := C.CString(requestID)
defer C.free(unsafe.Pointer(requestIDC))
var toolsJSONC *C.char
if toolsJSON != "" {
toolsJSONC = C.CString(toolsJSON)
defer C.free(unsafe.Pointer(toolsJSONC))
}
var toolChoiceJSONC *C.char
if toolChoiceJSON != "" {
toolChoiceJSONC = C.CString(toolChoiceJSON)
defer C.free(unsafe.Pointer(toolChoiceJSONC))
}
var stopJSONC *C.char
if stopJSON != "" {
stopJSONC = C.CString(stopJSON)
defer C.free(unsafe.Pointer(stopJSONC))
}
// Convert stop_token_ids to JSON string
stopTokenIDsJSON := ""
if len(stopTokenIDs) > 0 {
stopTokenIDsJSON = fmt.Sprintf("[%d", stopTokenIDs[0])
for i := 1; i < len(stopTokenIDs); i++ {
stopTokenIDsJSON += fmt.Sprintf(",%d", stopTokenIDs[i])
}
stopTokenIDsJSON += "]"
}
var stopTokenIDsJSONC *C.char
if stopTokenIDsJSON != "" {
stopTokenIDsJSONC = C.CString(stopTokenIDsJSON)
defer C.free(unsafe.Pointer(stopTokenIDsJSONC))
}
var errorOut *C.char
skipSpecialTokensC := C.int(0)
if skipSpecialTokens {
skipSpecialTokensC = C.int(1)
}
initialPromptTokensC := C.int(initialPromptTokens)
converterHandle := C.sgl_grpc_response_converter_create(
tokenizerHandle,
modelC,
requestIDC,
toolsJSONC,
toolChoiceJSONC,
stopJSONC,
stopTokenIDsJSONC,
skipSpecialTokensC,
initialPromptTokensC,
&errorOut,
)
if converterHandle == nil {
errorMsg := ""
if errorOut != nil {
errorMsg = C.GoString(errorOut)
C.sgl_free_string(errorOut)
}
if errorMsg == "" {
errorMsg = "failed to create converter handle"
}
return nil, fmt.Errorf("%s", errorMsg)
}
return &GrpcResponseConverterHandle{
handle: converterHandle,
}, nil
}
// FreeGrpcResponseConverter frees a gRPC response converter handle
func FreeGrpcResponseConverter(handle *GrpcResponseConverterHandle) {
if handle != nil && handle.handle != nil {
C.sgl_grpc_response_converter_free(handle.handle)
handle.handle = nil
}
}
// TokenizerHandle wraps the Rust tokenizer FFI handle
type TokenizerHandle struct {
handle *C.TokenizerHandle
}
// CreateTokenizerHandle creates a tokenizer handle (exported for caching)
func CreateTokenizerHandle(tokenizerPath string) (*TokenizerHandle, error) {
tokenizerPathC := C.CString(tokenizerPath)
defer C.free(unsafe.Pointer(tokenizerPathC))
var errorOut *C.char
tokenizerHandle := C.sgl_tokenizer_create_from_file(tokenizerPathC, &errorOut)
if tokenizerHandle == nil {
errorMsg := ""
if errorOut != nil {
errorMsg = C.GoString(errorOut)
C.sgl_free_string(errorOut)
}
if errorMsg == "" {
errorMsg = "failed to create tokenizer handle"
}
return nil, fmt.Errorf("%s", errorMsg)
}
return &TokenizerHandle{
handle: tokenizerHandle,
}, nil
}
// FreeTokenizerHandle frees a tokenizer handle
func FreeTokenizerHandle(handle *TokenizerHandle) {
if handle != nil && handle.handle != nil {
C.sgl_tokenizer_free(handle.handle)
handle.handle = nil
}
}
// createTokenizerHandle creates a tokenizer handle (helper function, internal use)
func createTokenizerHandle(tokenizerPath string) (*C.TokenizerHandle, error) {
tokenizerPathC := C.CString(tokenizerPath)
defer C.free(unsafe.Pointer(tokenizerPathC))
var errorOut *C.char
tokenizerHandle := C.sgl_tokenizer_create_from_file(tokenizerPathC, &errorOut)
if tokenizerHandle == nil {
errorMsg := ""
if errorOut != nil {
errorMsg = C.GoString(errorOut)
C.sgl_free_string(errorOut)
}
if errorMsg == "" {
errorMsg = "failed to create tokenizer handle"
}
return nil, fmt.Errorf("%s", errorMsg)
}
return tokenizerHandle, nil
}

156
third_party/sglang/sgl-model-gateway/bindings/golang/internal/ffi/postprocessor.go vendored Normal file
View File

@@ -0,0 +1,156 @@
// Package ffi provides Go bindings for SGLang's Rust FFI (Foreign Function Interface).
package ffi
/*
#cgo LDFLAGS: -lsgl_model_gateway_go -ldl
#include <stdlib.h>
#include <stdint.h>
// Error codes (must match client.go)
typedef enum {
SGL_ERROR_SUCCESS = 0,
SGL_ERROR_INVALID_ARGUMENT = 1,
SGL_ERROR_TOKENIZATION_ERROR = 2,
SGL_ERROR_PARSING_ERROR = 3,
SGL_ERROR_MEMORY_ERROR = 4,
SGL_ERROR_UNKNOWN = 99
} SglErrorCode;
// Opaque handle (must match grpc_converter.go)
typedef void* GrpcResponseConverterHandle;
// Postprocessor functions
SglErrorCode sgl_postprocess_stream_chunk(
GrpcResponseConverterHandle* converter_handle,
const char* proto_chunk_json,
char** openai_json_out,
int* is_done_out,
char** error_out
);
SglErrorCode sgl_postprocess_stream_chunks_batch(
GrpcResponseConverterHandle* converter_handle,
const char* proto_chunks_json_array,
int max_chunks,
char** openai_chunks_json_array_out,
int* chunks_count_out,
char** error_out
);
// Memory management
void sgl_free_string(char* s);
*/
import "C"
import (
"fmt"
"unsafe"
)
// GrpcResponseConverterHandle wraps the Rust gRPC response converter FFI handle
type GrpcResponseConverterHandle struct {
handle *C.GrpcResponseConverterHandle
}
// PostprocessStreamChunk postprocesses a gRPC stream chunk to OpenAI format
//
// This function:
// 1. Parses the proto chunk from JSON
// 2. Converts it to OpenAI format using the converter handle
// 3. Returns the OpenAI format JSON
//
// Returns the OpenAI format JSON, is_done flag, and any error.
func PostprocessStreamChunk(converterHandle *GrpcResponseConverterHandle, protoChunkJSON string) (openaiJSON string, isDone bool, err error) {
if converterHandle == nil || converterHandle.handle == nil {
return "", false, fmt.Errorf("invalid converter handle")
}
protoChunkJSONC := C.CString(protoChunkJSON)
defer C.free(unsafe.Pointer(protoChunkJSONC))
var openaiJSONOut *C.char
var isDoneOut C.int
var errorOut *C.char
errorCode := C.sgl_postprocess_stream_chunk(
converterHandle.handle,
protoChunkJSONC,
&openaiJSONOut,
&isDoneOut,
&errorOut,
)
if errorCode != C.SGL_ERROR_SUCCESS {
errorMsg := ""
if errorOut != nil {
errorMsg = C.GoString(errorOut)
C.sgl_free_string(errorOut)
}
return "", false, fmt.Errorf("postprocessing failed: %s", errorMsg)
}
openaiJSON = C.GoString(openaiJSONOut)
isDone = isDoneOut != 0
// Free the C string allocated by Rust
if openaiJSONOut != nil {
C.sgl_free_string(openaiJSONOut)
}
return openaiJSON, isDone, nil
}
// PostprocessStreamChunksBatch postprocesses multiple gRPC stream chunks in batch
//
// This function processes multiple chunks in a single FFI call, significantly reducing
// FFI overhead in streaming scenarios.
//
// Arguments:
// - converterHandle: Converter handle
// - protoChunksJSONArray: JSON array string of proto chunks
// - maxChunks: Maximum number of chunks to process (for safety, typically 10-20)
//
// Returns:
// - openaiChunksJSONArray: JSON array of OpenAI format chunks
// - chunksCount: Number of processed chunks
// - error: Any error that occurred
func PostprocessStreamChunksBatch(converterHandle *GrpcResponseConverterHandle, protoChunksJSONArray string, maxChunks int) (openaiChunksJSONArray string, chunksCount int, err error) {
if converterHandle == nil || converterHandle.handle == nil {
return "", 0, fmt.Errorf("invalid converter handle")
}
protoChunksJSONArrayC := C.CString(protoChunksJSONArray)
defer C.free(unsafe.Pointer(protoChunksJSONArrayC))
var openaiChunksJSONArrayOut *C.char
var chunksCountOut C.int
var errorOut *C.char
errorCode := C.sgl_postprocess_stream_chunks_batch(
converterHandle.handle,
protoChunksJSONArrayC,
C.int(maxChunks),
&openaiChunksJSONArrayOut,
&chunksCountOut,
&errorOut,
)
if errorCode != C.SGL_ERROR_SUCCESS {
errorMsg := ""
if errorOut != nil {
errorMsg = C.GoString(errorOut)
C.sgl_free_string(errorOut)
}
return "", 0, fmt.Errorf("batch postprocessing failed: %s", errorMsg)
}
openaiChunksJSONArray = C.GoString(openaiChunksJSONArrayOut)
chunksCount = int(chunksCountOut)
// Free the C string allocated by Rust
if openaiChunksJSONArrayOut != nil {
C.sgl_free_string(openaiChunksJSONArrayOut)
}
return openaiChunksJSONArray, chunksCount, nil
}

246
third_party/sglang/sgl-model-gateway/bindings/golang/internal/ffi/preprocessor.go vendored Normal file
View File

@@ -0,0 +1,246 @@
// Package ffi provides Go bindings for SGLang's Rust FFI (Foreign Function Interface).
package ffi
/*
#cgo LDFLAGS: -lsgl_model_gateway_go -ldl
#include <stdlib.h>
#include <stdint.h>
// Error codes (must match client.go)
typedef enum {
SGL_ERROR_SUCCESS = 0,
SGL_ERROR_INVALID_ARGUMENT = 1,
SGL_ERROR_TOKENIZATION_ERROR = 2,
SGL_ERROR_PARSING_ERROR = 3,
SGL_ERROR_MEMORY_ERROR = 4,
SGL_ERROR_UNKNOWN = 99
} SglErrorCode;
// Preprocessor functions
SglErrorCode sgl_preprocess_chat_request(
const char* request_json,
const char* tokenizer_path,
char** prompt_text_out,
uint32_t** token_ids_out,
size_t* token_ids_len_out,
char** tool_constraints_json_out,
int32_t* prompt_tokens_out,
char** error_out
);
// Opaque handle (must match grpc_converter.go)
typedef void* TokenizerHandle;
SglErrorCode sgl_preprocess_chat_request_with_tokenizer(
const char* request_json,
void* tokenizer_handle,
char** prompt_text_out,
uint32_t** token_ids_out,
size_t* token_ids_len_out,
char** tool_constraints_json_out,
int32_t* prompt_tokens_out,
char** error_out
);
void sgl_preprocessed_request_free(
char* prompt_text,
uint32_t* token_ids,
size_t token_ids_len,
char* tool_constraints_json
);
// Memory management
void sgl_free_string(char* s);
void sgl_free_token_ids(uint32_t* ptr, size_t count);
*/
import "C"
import (
"fmt"
"unsafe"
)
// PreprocessedRequest represents a preprocessed chat request
type PreprocessedRequest struct {
PromptText string
TokenIDs []uint32
ToolConstraintsJSON string
PromptTokens int32
// Internal pointers for memory management
promptTextPtr *C.char
tokenIDsPtr *C.uint32_t
tokenIDsLen uintptr
toolConstraintsJSONPtr *C.char
}
// PreprocessChatRequest preprocesses 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)
//
// Returns the preprocessed request data and any error.
func PreprocessChatRequest(requestJSON, tokenizerPath string) (*PreprocessedRequest, error) {
requestJSONC := C.CString(requestJSON)
defer C.free(unsafe.Pointer(requestJSONC))
tokenizerPathC := C.CString(tokenizerPath)
defer C.free(unsafe.Pointer(tokenizerPathC))
var promptTextOut *C.char
var tokenIDsOut *C.uint32_t
var tokenIDsLenOut C.size_t
var toolConstraintsJSONOut *C.char
var promptTokensOut C.int32_t
var errorOut *C.char
errorCode := C.sgl_preprocess_chat_request(
requestJSONC,
tokenizerPathC,
&promptTextOut,
&tokenIDsOut,
&tokenIDsLenOut,
&toolConstraintsJSONOut,
&promptTokensOut,
&errorOut,
)
if errorCode != C.SGL_ERROR_SUCCESS {
errorMsg := ""
if errorOut != nil {
errorMsg = C.GoString(errorOut)
C.sgl_free_string(errorOut)
}
return nil, fmt.Errorf("preprocessing failed: %s", errorMsg)
}
result := &PreprocessedRequest{
PromptText: C.GoString(promptTextOut),
TokenIDs: make([]uint32, tokenIDsLenOut),
ToolConstraintsJSON: "",
PromptTokens: int32(promptTokensOut),
}
// Copy token IDs
if tokenIDsOut != nil && tokenIDsLenOut > 0 {
tokenIDsSlice := (*[1 << 30]C.uint32_t)(unsafe.Pointer(tokenIDsOut))[:tokenIDsLenOut:tokenIDsLenOut]
for i := range result.TokenIDs {
result.TokenIDs[i] = uint32(tokenIDsSlice[i])
}
}
// Copy tool constraints JSON if present
if toolConstraintsJSONOut != nil {
result.ToolConstraintsJSON = C.GoString(toolConstraintsJSONOut)
}
// Store pointers for later cleanup
result.promptTextPtr = promptTextOut
result.tokenIDsPtr = tokenIDsOut
result.tokenIDsLen = uintptr(tokenIDsLenOut)
result.toolConstraintsJSONPtr = toolConstraintsJSONOut
return result, nil
}
// PreprocessChatRequestWithTokenizer preprocesses a chat completion request using an existing tokenizer handle
//
// This function is similar to PreprocessChatRequest, but accepts a TokenizerHandle
// instead of creating a new tokenizer. This allows reusing a cached tokenizer instance,
// significantly reducing initialization overhead in concurrent scenarios.
//
// Returns the preprocessed request data and any error.
func PreprocessChatRequestWithTokenizer(requestJSON string, tokenizerHandle *TokenizerHandle) (*PreprocessedRequest, error) {
requestJSONC := C.CString(requestJSON)
defer C.free(unsafe.Pointer(requestJSONC))
if tokenizerHandle == nil || tokenizerHandle.handle == nil {
return nil, fmt.Errorf("invalid tokenizer handle")
}
var promptTextOut *C.char
var tokenIDsOut *C.uint32_t
var tokenIDsLenOut C.size_t
var toolConstraintsJSONOut *C.char
var promptTokensOut C.int32_t
var errorOut *C.char
errorCode := C.sgl_preprocess_chat_request_with_tokenizer(
requestJSONC,
unsafe.Pointer(tokenizerHandle.handle), // Convert *C.TokenizerHandle to void*
&promptTextOut,
&tokenIDsOut,
&tokenIDsLenOut,
&toolConstraintsJSONOut,
&promptTokensOut,
&errorOut,
)
if errorCode != C.SGL_ERROR_SUCCESS {
errorMsg := ""
if errorOut != nil {
errorMsg = C.GoString(errorOut)
C.sgl_free_string(errorOut)
}
return nil, fmt.Errorf("preprocessing failed: %s", errorMsg)
}
result := &PreprocessedRequest{
PromptText: C.GoString(promptTextOut),
TokenIDs: make([]uint32, tokenIDsLenOut),
ToolConstraintsJSON: "",
PromptTokens: int32(promptTokensOut),
}
// Copy token IDs
if tokenIDsOut != nil && tokenIDsLenOut > 0 {
tokenIDsSlice := (*[1 << 30]C.uint32_t)(unsafe.Pointer(tokenIDsOut))[:tokenIDsLenOut:tokenIDsLenOut]
for i := range result.TokenIDs {
result.TokenIDs[i] = uint32(tokenIDsSlice[i])
}
}
// Copy tool constraints JSON if present
if toolConstraintsJSONOut != nil {
result.ToolConstraintsJSON = C.GoString(toolConstraintsJSONOut)
}
// Store pointers for later cleanup
result.promptTextPtr = promptTextOut
result.tokenIDsPtr = tokenIDsOut
result.tokenIDsLen = uintptr(tokenIDsLenOut)
result.toolConstraintsJSONPtr = toolConstraintsJSONOut
return result, nil
}
// Free frees the memory allocated for a preprocessed request
func (p *PreprocessedRequest) Free() {
if p.promptTextPtr != nil || p.tokenIDsPtr != nil || p.toolConstraintsJSONPtr != nil {
C.sgl_preprocessed_request_free(
p.promptTextPtr,
p.tokenIDsPtr,
C.size_t(p.tokenIDsLen),
p.toolConstraintsJSONPtr,
)
// Clear pointers to prevent double-free
p.promptTextPtr = nil
p.tokenIDsPtr = nil
p.tokenIDsLen = 0
p.toolConstraintsJSONPtr = nil
}
}
// FreePreprocessedRequest frees the memory allocated for a preprocessed request
// This is a convenience function for direct pointer management
func FreePreprocessedRequest(promptTextPtr *C.char, tokenIDsPtr *C.uint32_t, tokenIDsLen uintptr, toolConstraintsJSONPtr *C.char) {
if promptTextPtr != nil || tokenIDsPtr != nil || toolConstraintsJSONPtr != nil {
C.sgl_preprocessed_request_free(
promptTextPtr,
tokenIDsPtr,
C.size_t(tokenIDsLen),
toolConstraintsJSONPtr,
)
}
}

684
third_party/sglang/sgl-model-gateway/bindings/golang/internal/grpc/client_grpc.go vendored Normal file
View File

@@ -0,0 +1,684 @@
// Package grpc provides gRPC client implementation for SGLang
package grpc
import (
"context"
"encoding/json"
"fmt"
"io"
"strconv"
"strings"
"sync"
"sync/atomic"
"time"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/keepalive"
"google.golang.org/protobuf/types/known/timestamppb"
"github.com/sglang/sglang-go-grpc-sdk/internal/ffi"
"github.com/sglang/sglang-go-grpc-sdk/internal/proto"
)
type grpcClientStream interface {
Recv() (*proto.GenerateResponse, error)
CloseSend() error
}
// recvResult holds the result of a Recv() call
type recvResult struct {
resp *proto.GenerateResponse
err error
}
type GrpcClient struct {
conn *grpc.ClientConn
client proto.SglangSchedulerClient
tokenizerPath string
tokenizerHandle *ffi.TokenizerHandle
bufferSizes ChannelBufferSizes
timeouts Timeouts
requestCounter uint64 // Atomic counter to ensure unique request IDs
}
type ChannelBufferSizes struct {
ResultJSONChan int
ErrChan int
RecvChan int
}
type Timeouts struct {
KeepaliveTime time.Duration
KeepaliveTimeout time.Duration
CloseTimeout time.Duration
}
func NewGrpcClient(endpoint, tokenizerPath string, bufferSizes ChannelBufferSizes, timeouts Timeouts) (*GrpcClient, error) {
endpoint = strings.TrimPrefix(endpoint, "grpc://")
if !strings.Contains(endpoint, ":") {
return nil, fmt.Errorf("invalid endpoint format: %s (expected grpc://host:port)", endpoint)
}
keepaliveParams := keepalive.ClientParameters{
Time: timeouts.KeepaliveTime,
Timeout: timeouts.KeepaliveTimeout,
PermitWithoutStream: false,
}
opts := []grpc.DialOption{
grpc.WithTransportCredentials(insecure.NewCredentials()),
grpc.WithKeepaliveParams(keepaliveParams),
}
conn, err := grpc.NewClient(endpoint, opts...)
if err != nil {
return nil, fmt.Errorf("failed to connect to gRPC server: %w", err)
}
client := proto.NewSglangSchedulerClient(conn)
tokenizerHandle, err := ffi.CreateTokenizerHandle(tokenizerPath)
if err != nil {
conn.Close()
return nil, fmt.Errorf("failed to create tokenizer handle: %w", err)
}
return &GrpcClient{
conn: conn,
client: client,
tokenizerPath: tokenizerPath,
tokenizerHandle: tokenizerHandle,
bufferSizes: bufferSizes,
timeouts: timeouts,
}, nil
}
func (c *GrpcClient) Close() error {
if c.tokenizerHandle != nil {
ffi.FreeTokenizerHandle(c.tokenizerHandle)
c.tokenizerHandle = nil
}
if c.conn != nil {
return c.conn.Close()
}
return nil
}
func (c *GrpcClient) CreateChatCompletionStream(ctx context.Context, reqJSON string) (*GrpcChatCompletionStream, error) {
if c.tokenizerHandle == nil {
return nil, fmt.Errorf("tokenizer handle is nil (should be created at startup)")
}
preprocessed, err := ffi.PreprocessChatRequestWithTokenizer(reqJSON, c.tokenizerHandle)
if err != nil {
return nil, fmt.Errorf("preprocessing failed: %w", err)
}
defer func() {
if preprocessed != nil {
preprocessed.Free()
}
}()
// Parse request JSON to get parameters
var reqMap map[string]interface{}
if err := json.Unmarshal([]byte(reqJSON), &reqMap); err != nil {
return nil, fmt.Errorf("failed to parse request JSON: %w", err)
}
model, _ := reqMap["model"].(string)
if model == "" {
model = "default"
}
// Build GenerateRequest
// Generate unique request ID using timestamp + atomic counter to avoid collisions
// This matches Rust version's UUID-based approach for uniqueness
counter := atomic.AddUint64(&c.requestCounter, 1)
requestID := fmt.Sprintf("chatcmpl-%d-%d", time.Now().UnixNano(), counter)
generateReq := &proto.GenerateRequest{
RequestId: requestID,
Tokenized: &proto.TokenizedInput{
OriginalText: preprocessed.PromptText,
InputIds: preprocessed.TokenIDs,
},
Stream: true,
}
// Set sampling parameters
samplingParams := &proto.SamplingParams{
Temperature: 1.0,
TopP: 1.0,
TopK: -1,
SkipSpecialTokens: true,
}
if temp, ok := reqMap["temperature"].(float64); ok {
samplingParams.Temperature = float32(temp)
}
if topP, ok := reqMap["top_p"].(float64); ok {
samplingParams.TopP = float32(topP)
}
if topK, ok := reqMap["top_k"].(float64); ok {
samplingParams.TopK = int32(topK)
}
var maxTokensInt *int32
if maxCompletionTokens, ok := reqMap["max_completion_tokens"].(float64); ok {
tokens := int32(maxCompletionTokens)
maxTokensInt = &tokens
} else if maxTokens, ok := reqMap["max_tokens"].(float64); ok {
tokens := int32(maxTokens)
maxTokensInt = &tokens
}
if maxTokensInt != nil {
samplingParams.MaxNewTokens = maxTokensInt
}
// Parse tool constraints if available
if preprocessed.ToolConstraintsJSON != "" {
var toolConstraints map[string]interface{}
if err := json.Unmarshal([]byte(preprocessed.ToolConstraintsJSON), &toolConstraints); err == nil {
if regex, ok := toolConstraints["regex"].(string); ok {
samplingParams.Constraint = &proto.SamplingParams_Regex{Regex: regex}
} else if jsonSchema, ok := toolConstraints["json_schema"].(string); ok {
samplingParams.Constraint = &proto.SamplingParams_JsonSchema{JsonSchema: jsonSchema}
}
}
}
generateReq.SamplingParams = samplingParams
generateReq.Timestamp = timestamppb.Now()
stream, err := c.client.Generate(ctx, generateReq)
if err != nil {
return nil, fmt.Errorf("failed to create gRPC stream: %w", err)
}
toolsJSON := ""
if tools, ok := reqMap["tools"].([]interface{}); ok && len(tools) > 0 {
toolsBytes, _ := json.Marshal(tools)
toolsJSON = string(toolsBytes)
}
toolChoiceJSON := ""
if toolChoice, ok := reqMap["tool_choice"]; ok {
toolChoiceBytes, _ := json.Marshal(toolChoice)
toolChoiceJSON = string(toolChoiceBytes)
}
stopJSON := ""
if stop, ok := reqMap["stop"]; ok {
stopBytes, _ := json.Marshal(stop)
stopJSON = string(stopBytes)
}
stopTokenIDs := []uint32{}
if stopTokenIDsVal, ok := reqMap["stop_token_ids"].([]interface{}); ok {
for _, id := range stopTokenIDsVal {
if idFloat, ok := id.(float64); ok {
stopTokenIDs = append(stopTokenIDs, uint32(idFloat))
}
}
}
skipSpecialTokens := true
if skipSpecialTokensVal, ok := reqMap["skip_special_tokens"].(bool); ok {
skipSpecialTokens = skipSpecialTokensVal
}
if c.tokenizerHandle == nil {
stream.CloseSend()
return nil, fmt.Errorf("tokenizer handle is nil (should be created at startup)")
}
converterHandle, err := ffi.CreateGrpcResponseConverterWithTokenizer(
c.tokenizerHandle,
model,
generateReq.RequestId,
toolsJSON,
toolChoiceJSON,
stopJSON,
stopTokenIDs,
skipSpecialTokens,
preprocessed.PromptTokens, // Pass initial prompt tokens from preprocessing
)
if err != nil {
stream.CloseSend()
return nil, fmt.Errorf("failed to create converter handle: %w", err)
}
batchSize := 1
batchPostprocessor := ffi.NewBatchPostprocessor(converterHandle, batchSize, 0)
streamCtx, cancel := context.WithCancel(ctx)
grpcStream := &GrpcChatCompletionStream{
stream: stream,
converterHandle: converterHandle,
batchPostprocessor: batchPostprocessor,
batchSize: batchSize,
ctx: streamCtx,
cancel: cancel,
resultJSONChan: make(chan string, c.bufferSizes.ResultJSONChan),
errChan: make(chan error, c.bufferSizes.ErrChan),
readLoopDone: make(chan struct{}),
requestID: generateReq.RequestId,
model: model,
processWg: sync.WaitGroup{},
closeTimeout: c.timeouts.CloseTimeout,
bufferSizes: c.bufferSizes,
}
go grpcStream.readLoop()
return grpcStream, nil
}
// GrpcChatCompletionStream represents a streaming chat completion via gRPC
type GrpcChatCompletionStream struct {
stream grpcClientStream
converterHandle *ffi.GrpcResponseConverterHandle
batchPostprocessor *ffi.BatchPostprocessor
batchSize int
ctx context.Context
cancel context.CancelFunc
closed int32
resultJSONChan chan string
errChan chan error
readLoopDone chan struct{}
requestID string
model string
processWg sync.WaitGroup
closeTimeout time.Duration
bufferSizes ChannelBufferSizes
clientDisconnected int32 // Atomic flag: 1 if client disconnected, 0 otherwise
}
func (s *GrpcChatCompletionStream) readLoop() {
defer func() {
atomic.StoreInt32(&s.closed, 1)
s.processWg.Wait()
close(s.resultJSONChan)
close(s.errChan)
close(s.readLoopDone)
// Cancel context after channels are closed to ensure errors are read first
if s.cancel != nil {
s.cancel()
}
}()
recvChan := make(chan recvResult, s.bufferSizes.RecvChan)
const firstRecvTimeout = 60 * time.Second
go func() {
defer close(recvChan)
recvCount := 0
for {
select {
case <-s.ctx.Done():
// Skip CloseSend() if client disconnected
if atomic.LoadInt32(&s.clientDisconnected) == 0 {
_ = s.stream.CloseSend()
}
return
default:
}
recvCount++
var protoResp *proto.GenerateResponse
var err error
// First Recv() with timeout
if recvCount == 1 {
recvDone := make(chan recvResult, 1)
go func() {
resp, recvErr := s.stream.Recv()
recvDone <- recvResult{resp: resp, err: recvErr}
}()
select {
case result := <-recvDone:
protoResp = result.resp
err = result.err
case <-time.After(firstRecvTimeout):
timeoutErr := fmt.Errorf("stream.Recv() timeout after %v: backend may not be responding (request_id=%s)", firstRecvTimeout, s.requestID)
select {
case recvChan <- recvResult{resp: nil, err: timeoutErr}:
case <-s.ctx.Done():
}
return
case <-s.ctx.Done():
return
}
} else {
// Normal Recv()
protoResp, err = s.stream.Recv()
}
if err != nil {
select {
case recvChan <- recvResult{resp: nil, err: err}:
case <-s.ctx.Done():
return
}
return
}
select {
case <-s.ctx.Done():
// Skip CloseSend() if client disconnected
if atomic.LoadInt32(&s.clientDisconnected) == 0 {
_ = s.stream.CloseSend()
}
return
case recvChan <- recvResult{resp: protoResp, err: nil}:
}
}
}()
for {
select {
case <-s.ctx.Done():
// Skip CloseSend() if client disconnected
if atomic.LoadInt32(&s.clientDisconnected) == 0 {
_ = s.stream.CloseSend()
}
return
case result, ok := <-recvChan:
if !ok {
return
}
if result.err != nil {
if result.err == io.EOF {
results, flushErr := s.flushBatch()
if flushErr != nil {
select {
case s.errChan <- fmt.Errorf("failed to flush batch: %w", flushErr):
case <-s.ctx.Done():
}
return
}
for _, resultJSON := range results {
select {
case s.resultJSONChan <- resultJSON:
case <-s.ctx.Done():
return
}
}
return
}
select {
case s.errChan <- result.err:
case <-s.ctx.Done():
}
return
}
if result.resp != nil {
s.processWg.Add(1)
go func(resp *proto.GenerateResponse) {
defer s.processWg.Done()
s.processAndSendResponse(resp)
}(result.resp)
}
}
}
}
func (s *GrpcChatCompletionStream) processAndSendResponse(protoResp *proto.GenerateResponse) {
select {
case <-s.ctx.Done():
return
default:
}
if protoResp == nil {
return
}
protoJSON, err := protoToJSON(protoResp)
if err != nil {
select {
case s.errChan <- fmt.Errorf("failed to convert proto to JSON: %w", err):
case <-s.ctx.Done():
}
return
}
if s.batchPostprocessor == nil {
select {
case s.errChan <- fmt.Errorf("batch postprocessor is nil"):
case <-s.ctx.Done():
}
return
}
results, _, err := s.batchPostprocessor.AddChunk(protoJSON)
if err != nil {
select {
case s.errChan <- fmt.Errorf("batch postprocessing failed: %w", err):
case <-s.ctx.Done():
}
return
}
for _, resultJSON := range results {
select {
case s.resultJSONChan <- resultJSON:
case <-s.ctx.Done():
return
}
}
}
func (s *GrpcChatCompletionStream) RecvJSON() (string, error) {
// Use a loop instead of recursion to avoid stack overflow if there are many empty strings
for {
// Check errChan first to prioritize actual errors over context cancellation
select {
case err, ok := <-s.errChan:
if !ok {
return "", io.EOF
}
return "", err
default:
}
select {
case resultJSON, ok := <-s.resultJSONChan:
if !ok {
return "", io.EOF
}
// Skip empty strings and continue loop instead of recursing
if resultJSON != "" {
return resultJSON, nil
}
// Empty string, continue loop to get next result
continue
case err, ok := <-s.errChan:
if !ok {
return "", io.EOF
}
return "", err
case <-s.ctx.Done():
return "", s.ctx.Err()
}
}
}
// SetClientDisconnected marks that the client has disconnected.
// When Close() is called, it will not call CloseSend() to avoid aborting the request on server side.
func (s *GrpcChatCompletionStream) SetClientDisconnected() {
atomic.StoreInt32(&s.clientDisconnected, 1)
}
func (s *GrpcChatCompletionStream) Close() error {
if !atomic.CompareAndSwapInt32(&s.closed, 0, 1) {
return nil
}
if s.cancel != nil {
s.cancel()
}
clientDisconnected := atomic.LoadInt32(&s.clientDisconnected) == 1
select {
case <-s.readLoopDone:
// readLoop completed
default:
if !clientDisconnected {
// Call CloseSend() if client didn't disconnect
_ = s.stream.CloseSend()
}
select {
case <-s.readLoopDone:
case <-time.After(s.closeTimeout):
}
}
_, _ = s.flushBatch()
if s.converterHandle != nil {
ffi.FreeGrpcResponseConverter(s.converterHandle)
}
return nil
}
func (s *GrpcChatCompletionStream) flushBatch() ([]string, error) {
if s.batchPostprocessor != nil {
results, err := s.batchPostprocessor.Flush()
if err != nil {
return nil, fmt.Errorf("batch flush failed: %w", err)
}
return results, nil
}
return nil, nil
}
func protoToJSON(resp *proto.GenerateResponse) (string, error) {
var sb strings.Builder
sb.Grow(500)
sb.WriteString(`{"request_id":`)
if resp.RequestId == "" {
sb.WriteString(`""`)
} else {
requestIDJSON, err := json.Marshal(resp.RequestId)
if err != nil {
return "", err
}
sb.Write(requestIDJSON)
}
switch r := resp.Response.(type) {
case *proto.GenerateResponse_Chunk:
sb.WriteString(`,"chunk":{`)
sb.WriteString(`"token_ids":`)
tokenIDsJSON, err := json.Marshal(r.Chunk.TokenIds)
if err != nil {
return "", err
}
sb.Write(tokenIDsJSON)
sb.WriteString(`,"prompt_tokens":`)
sb.WriteString(strconv.FormatInt(int64(r.Chunk.PromptTokens), 10))
sb.WriteString(`,"completion_tokens":`)
sb.WriteString(strconv.FormatInt(int64(r.Chunk.CompletionTokens), 10))
sb.WriteString(`,"cached_tokens":`)
sb.WriteString(strconv.FormatInt(int64(r.Chunk.CachedTokens), 10))
sb.WriteString(`,"index":`)
sb.WriteString(strconv.FormatInt(int64(r.Chunk.Index), 10))
sb.WriteString(`}`)
case *proto.GenerateResponse_Complete:
sb.WriteString(`,"complete":{`)
sb.WriteString(`"output_ids":`)
outputIDsJSON, err := json.Marshal(r.Complete.OutputIds)
if err != nil {
return "", err
}
sb.Write(outputIDsJSON)
sb.WriteString(`,"finish_reason":`)
finishReasonJSON, err := json.Marshal(r.Complete.FinishReason)
if err != nil {
return "", err
}
sb.Write(finishReasonJSON)
sb.WriteString(`,"prompt_tokens":`)
sb.WriteString(strconv.FormatInt(int64(r.Complete.PromptTokens), 10))
sb.WriteString(`,"completion_tokens":`)
sb.WriteString(strconv.FormatInt(int64(r.Complete.CompletionTokens), 10))
sb.WriteString(`,"cached_tokens":`)
sb.WriteString(strconv.FormatInt(int64(r.Complete.CachedTokens), 10))
sb.WriteString(`}`)
case *proto.GenerateResponse_Error:
sb.WriteString(`,"error":{`)
sb.WriteString(`"message":`)
messageJSON, err := json.Marshal(r.Error.Message)
if err != nil {
return "", err
}
sb.Write(messageJSON)
sb.WriteString(`,"http_status_code":`)
httpStatusCodeJSON, err := json.Marshal(r.Error.HttpStatusCode)
if err != nil {
return "", err
}
sb.Write(httpStatusCodeJSON)
if r.Error.Details != "" {
sb.WriteString(`,"details":`)
detailsJSON, err := json.Marshal(r.Error.Details)
if err != nil {
return "", err
}
sb.Write(detailsJSON)
}
sb.WriteString(`}`)
}
sb.WriteString(`}`)
return sb.String(), nil
}
type ChatCompletionStreamResponse struct {
ID string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
Model string `json:"model"`
SystemFingerprint string `json:"system_fingerprint,omitempty"`
Choices []StreamChoice `json:"choices"`
Usage *Usage `json:"usage,omitempty"`
}
// StreamChoice represents a choice in a streaming response
type StreamChoice struct {
Index int `json:"index"`
Delta MessageDelta `json:"delta"`
FinishReason string `json:"finish_reason,omitempty"`
}
// MessageDelta represents incremental message updates
type MessageDelta struct {
Role string `json:"role,omitempty"`
Content string `json:"content,omitempty"`
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
}
// ToolCall represents a tool call in the response
type ToolCall struct {
ID string `json:"id"`
Type string `json:"type"`
Function FunctionCall `json:"function"`
}
// FunctionCall represents a function call
type FunctionCall struct {
Name string `json:"name"`
Arguments string `json:"arguments"`
}
// Usage represents token usage information
type Usage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
}

3325
third_party/sglang/sgl-model-gateway/bindings/golang/internal/proto/sglang_scheduler.pb.go vendored Normal file
View File

File diff suppressed because it is too large Load Diff

333
third_party/sglang/sgl-model-gateway/bindings/golang/internal/proto/sglang_scheduler_grpc.pb.go vendored Normal file
View File

@@ -0,0 +1,333 @@
// Code generated by protoc-gen-go-grpc. DO NOT EDIT.
// versions:
// - protoc-gen-go-grpc v1.5.1
// - protoc v3.21.12
// source: sglang_scheduler.proto
package proto
import (
context "context"
grpc "google.golang.org/grpc"
codes "google.golang.org/grpc/codes"
status "google.golang.org/grpc/status"
)
// This is a compile-time assertion to ensure that this generated file
// is compatible with the grpc package it is being compiled against.
// Requires gRPC-Go v1.64.0 or later.
const _ = grpc.SupportPackageIsVersion9
const (
SglangScheduler_Generate_FullMethodName = "/sglang.grpc.scheduler.SglangScheduler/Generate"
SglangScheduler_Embed_FullMethodName = "/sglang.grpc.scheduler.SglangScheduler/Embed"
SglangScheduler_HealthCheck_FullMethodName = "/sglang.grpc.scheduler.SglangScheduler/HealthCheck"
SglangScheduler_Abort_FullMethodName = "/sglang.grpc.scheduler.SglangScheduler/Abort"
SglangScheduler_GetModelInfo_FullMethodName = "/sglang.grpc.scheduler.SglangScheduler/GetModelInfo"
SglangScheduler_GetServerInfo_FullMethodName = "/sglang.grpc.scheduler.SglangScheduler/GetServerInfo"
)
// SglangSchedulerClient is the client API for SglangScheduler service.
//
// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream.
//
// Service definition for SGLang scheduler communication
// This protocol bridges the Rust router and Python scheduler
type SglangSchedulerClient interface {
// Submit a generation request (supports streaming)
Generate(ctx context.Context, in *GenerateRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[GenerateResponse], error)
// Submit an embedding request
Embed(ctx context.Context, in *EmbedRequest, opts ...grpc.CallOption) (*EmbedResponse, error)
// Health check and metrics
HealthCheck(ctx context.Context, in *HealthCheckRequest, opts ...grpc.CallOption) (*HealthCheckResponse, error)
// Abort a running request
Abort(ctx context.Context, in *AbortRequest, opts ...grpc.CallOption) (*AbortResponse, error)
// Get model information
GetModelInfo(ctx context.Context, in *GetModelInfoRequest, opts ...grpc.CallOption) (*GetModelInfoResponse, error)
// Get server information
GetServerInfo(ctx context.Context, in *GetServerInfoRequest, opts ...grpc.CallOption) (*GetServerInfoResponse, error)
}
type sglangSchedulerClient struct {
cc grpc.ClientConnInterface
}
func NewSglangSchedulerClient(cc grpc.ClientConnInterface) SglangSchedulerClient {
return &sglangSchedulerClient{cc}
}
func (c *sglangSchedulerClient) Generate(ctx context.Context, in *GenerateRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[GenerateResponse], error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
stream, err := c.cc.NewStream(ctx, &SglangScheduler_ServiceDesc.Streams[0], SglangScheduler_Generate_FullMethodName, cOpts...)
if err != nil {
return nil, err
}
x := &grpc.GenericClientStream[GenerateRequest, GenerateResponse]{ClientStream: stream}
if err := x.ClientStream.SendMsg(in); err != nil {
return nil, err
}
if err := x.ClientStream.CloseSend(); err != nil {
return nil, err
}
return x, nil
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type SglangScheduler_GenerateClient = grpc.ServerStreamingClient[GenerateResponse]
func (c *sglangSchedulerClient) Embed(ctx context.Context, in *EmbedRequest, opts ...grpc.CallOption) (*EmbedResponse, error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
out := new(EmbedResponse)
err := c.cc.Invoke(ctx, SglangScheduler_Embed_FullMethodName, in, out, cOpts...)
if err != nil {
return nil, err
}
return out, nil
}
func (c *sglangSchedulerClient) HealthCheck(ctx context.Context, in *HealthCheckRequest, opts ...grpc.CallOption) (*HealthCheckResponse, error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
out := new(HealthCheckResponse)
err := c.cc.Invoke(ctx, SglangScheduler_HealthCheck_FullMethodName, in, out, cOpts...)
if err != nil {
return nil, err
}
return out, nil
}
func (c *sglangSchedulerClient) Abort(ctx context.Context, in *AbortRequest, opts ...grpc.CallOption) (*AbortResponse, error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
out := new(AbortResponse)
err := c.cc.Invoke(ctx, SglangScheduler_Abort_FullMethodName, in, out, cOpts...)
if err != nil {
return nil, err
}
return out, nil
}
func (c *sglangSchedulerClient) GetModelInfo(ctx context.Context, in *GetModelInfoRequest, opts ...grpc.CallOption) (*GetModelInfoResponse, error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
out := new(GetModelInfoResponse)
err := c.cc.Invoke(ctx, SglangScheduler_GetModelInfo_FullMethodName, in, out, cOpts...)
if err != nil {
return nil, err
}
return out, nil
}
func (c *sglangSchedulerClient) GetServerInfo(ctx context.Context, in *GetServerInfoRequest, opts ...grpc.CallOption) (*GetServerInfoResponse, error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
out := new(GetServerInfoResponse)
err := c.cc.Invoke(ctx, SglangScheduler_GetServerInfo_FullMethodName, in, out, cOpts...)
if err != nil {
return nil, err
}
return out, nil
}
// SglangSchedulerServer is the server API for SglangScheduler service.
// All implementations must embed UnimplementedSglangSchedulerServer
// for forward compatibility.
//
// Service definition for SGLang scheduler communication
// This protocol bridges the Rust router and Python scheduler
type SglangSchedulerServer interface {
// Submit a generation request (supports streaming)
Generate(*GenerateRequest, grpc.ServerStreamingServer[GenerateResponse]) error
// Submit an embedding request
Embed(context.Context, *EmbedRequest) (*EmbedResponse, error)
// Health check and metrics
HealthCheck(context.Context, *HealthCheckRequest) (*HealthCheckResponse, error)
// Abort a running request
Abort(context.Context, *AbortRequest) (*AbortResponse, error)
// Get model information
GetModelInfo(context.Context, *GetModelInfoRequest) (*GetModelInfoResponse, error)
// Get server information
GetServerInfo(context.Context, *GetServerInfoRequest) (*GetServerInfoResponse, error)
mustEmbedUnimplementedSglangSchedulerServer()
}
// UnimplementedSglangSchedulerServer must be embedded to have
// forward compatible implementations.
//
// NOTE: this should be embedded by value instead of pointer to avoid a nil
// pointer dereference when methods are called.
type UnimplementedSglangSchedulerServer struct{}
func (UnimplementedSglangSchedulerServer) Generate(*GenerateRequest, grpc.ServerStreamingServer[GenerateResponse]) error {
return status.Errorf(codes.Unimplemented, "method Generate not implemented")
}
func (UnimplementedSglangSchedulerServer) Embed(context.Context, *EmbedRequest) (*EmbedResponse, error) {
return nil, status.Errorf(codes.Unimplemented, "method Embed not implemented")
}
func (UnimplementedSglangSchedulerServer) HealthCheck(context.Context, *HealthCheckRequest) (*HealthCheckResponse, error) {
return nil, status.Errorf(codes.Unimplemented, "method HealthCheck not implemented")
}
func (UnimplementedSglangSchedulerServer) Abort(context.Context, *AbortRequest) (*AbortResponse, error) {
return nil, status.Errorf(codes.Unimplemented, "method Abort not implemented")
}
func (UnimplementedSglangSchedulerServer) GetModelInfo(context.Context, *GetModelInfoRequest) (*GetModelInfoResponse, error) {
return nil, status.Errorf(codes.Unimplemented, "method GetModelInfo not implemented")
}
func (UnimplementedSglangSchedulerServer) GetServerInfo(context.Context, *GetServerInfoRequest) (*GetServerInfoResponse, error) {
return nil, status.Errorf(codes.Unimplemented, "method GetServerInfo not implemented")
}
func (UnimplementedSglangSchedulerServer) mustEmbedUnimplementedSglangSchedulerServer() {}
func (UnimplementedSglangSchedulerServer) testEmbeddedByValue() {}
// UnsafeSglangSchedulerServer may be embedded to opt out of forward compatibility for this service.
// Use of this interface is not recommended, as added methods to SglangSchedulerServer will
// result in compilation errors.
type UnsafeSglangSchedulerServer interface {
mustEmbedUnimplementedSglangSchedulerServer()
}
func RegisterSglangSchedulerServer(s grpc.ServiceRegistrar, srv SglangSchedulerServer) {
// If the following call pancis, it indicates UnimplementedSglangSchedulerServer was
// embedded by pointer and is nil. This will cause panics if an
// unimplemented method is ever invoked, so we test this at initialization
// time to prevent it from happening at runtime later due to I/O.
if t, ok := srv.(interface{ testEmbeddedByValue() }); ok {
t.testEmbeddedByValue()
}
s.RegisterService(&SglangScheduler_ServiceDesc, srv)
}
func _SglangScheduler_Generate_Handler(srv interface{}, stream grpc.ServerStream) error {
m := new(GenerateRequest)
if err := stream.RecvMsg(m); err != nil {
return err
}
return srv.(SglangSchedulerServer).Generate(m, &grpc.GenericServerStream[GenerateRequest, GenerateResponse]{ServerStream: stream})
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type SglangScheduler_GenerateServer = grpc.ServerStreamingServer[GenerateResponse]
func _SglangScheduler_Embed_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(EmbedRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(SglangSchedulerServer).Embed(ctx, in)
}
info := &grpc.UnaryServerInfo{
Server: srv,
FullMethod: SglangScheduler_Embed_FullMethodName,
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(SglangSchedulerServer).Embed(ctx, req.(*EmbedRequest))
}
return interceptor(ctx, in, info, handler)
}
func _SglangScheduler_HealthCheck_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(HealthCheckRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(SglangSchedulerServer).HealthCheck(ctx, in)
}
info := &grpc.UnaryServerInfo{
Server: srv,
FullMethod: SglangScheduler_HealthCheck_FullMethodName,
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(SglangSchedulerServer).HealthCheck(ctx, req.(*HealthCheckRequest))
}
return interceptor(ctx, in, info, handler)
}
func _SglangScheduler_Abort_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(AbortRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(SglangSchedulerServer).Abort(ctx, in)
}
info := &grpc.UnaryServerInfo{
Server: srv,
FullMethod: SglangScheduler_Abort_FullMethodName,
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(SglangSchedulerServer).Abort(ctx, req.(*AbortRequest))
}
return interceptor(ctx, in, info, handler)
}
func _SglangScheduler_GetModelInfo_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(GetModelInfoRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(SglangSchedulerServer).GetModelInfo(ctx, in)
}
info := &grpc.UnaryServerInfo{
Server: srv,
FullMethod: SglangScheduler_GetModelInfo_FullMethodName,
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(SglangSchedulerServer).GetModelInfo(ctx, req.(*GetModelInfoRequest))
}
return interceptor(ctx, in, info, handler)
}
func _SglangScheduler_GetServerInfo_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(GetServerInfoRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(SglangSchedulerServer).GetServerInfo(ctx, in)
}
info := &grpc.UnaryServerInfo{
Server: srv,
FullMethod: SglangScheduler_GetServerInfo_FullMethodName,
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(SglangSchedulerServer).GetServerInfo(ctx, req.(*GetServerInfoRequest))
}
return interceptor(ctx, in, info, handler)
}
// SglangScheduler_ServiceDesc is the grpc.ServiceDesc for SglangScheduler service.
// It's only intended for direct use with grpc.RegisterService,
// and not to be introspected or modified (even as a copy)
var SglangScheduler_ServiceDesc = grpc.ServiceDesc{
ServiceName: "sglang.grpc.scheduler.SglangScheduler",
HandlerType: (*SglangSchedulerServer)(nil),
Methods: []grpc.MethodDesc{
{
MethodName: "Embed",
Handler: _SglangScheduler_Embed_Handler,
},
{
MethodName: "HealthCheck",
Handler: _SglangScheduler_HealthCheck_Handler,
},
{
MethodName: "Abort",
Handler: _SglangScheduler_Abort_Handler,
},
{
MethodName: "GetModelInfo",
Handler: _SglangScheduler_GetModelInfo_Handler,
},
{
MethodName: "GetServerInfo",
Handler: _SglangScheduler_GetServerInfo_Handler,
},
},
Streams: []grpc.StreamDesc{
{
StreamName: "Generate",
Handler: _SglangScheduler_Generate_Handler,
ServerStreams: true,
},
},
Metadata: "sglang_scheduler.proto",
}

279
third_party/sglang/sgl-model-gateway/bindings/golang/src/client.rs vendored Normal file
View File

@@ -0,0 +1,279 @@
//! Client SDK FFI functions
use std::ffi::{CStr, CString};
use std::os::raw::{c_char};
use std::ptr;
use std::sync::Arc;
use tokio::runtime::Runtime;
use once_cell::sync::Lazy;
use uuid::Uuid;
use smg::tokenizer::create_tokenizer_from_file;
use smg::tokenizer::traits::Tokenizer;
use smg_grpc_client::sglang_scheduler::SglangSchedulerClient;
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::grpc_converter::sgl_grpc_response_converter_create;
use super::tokenizer::TokenizerHandle;
use super::stream::SglangStreamHandle;
/// Global tokio runtime for async operations
static RUNTIME: Lazy<Runtime> = Lazy::new(|| {
Runtime::new().expect("Failed to create tokio runtime for client FFI")
});
/// Handle for complete client SDK (gRPC client + tokenizer)
/// This handle manages the connection to sglang and provides a complete SDK interface
pub struct SglangClientHandle {
pub(crate) client: Arc<SglangSchedulerClient>,
pub(crate) tokenizer: Arc<dyn Tokenizer>,
}
/// Handle for streaming request (includes prompt token count)
#[allow(dead_code)]
pub struct StreamRequestState {
pub(crate) prompt_tokens: i32, // Number of prompt tokens for this request
}
/// Create a new SGLang client handle
///
/// # Arguments
/// * `endpoint` - gRPC endpoint (e.g., "grpc://localhost:20000")
/// * `tokenizer_path` - Path to tokenizer directory
/// * `error_out` - Optional pointer to receive error message
///
/// # Returns
/// * Pointer to SglangClientHandle on success, null on failure
#[no_mangle]
pub unsafe extern "C" fn sgl_client_create(
endpoint: *const c_char,
tokenizer_path: *const c_char,
error_out: *mut *mut c_char,
) -> *mut SglangClientHandle {
if endpoint.is_null() || tokenizer_path.is_null() {
set_error_message(error_out, "Invalid arguments: null pointer");
return ptr::null_mut();
}
let endpoint_str = match CStr::from_ptr(endpoint).to_str() {
Ok(s) => s,
Err(_) => {
set_error_message(error_out, "Invalid UTF-8 in endpoint");
return ptr::null_mut();
}
};
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 ptr::null_mut();
}
};
// 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 ptr::null_mut();
}
};
// Create gRPC client
let client = match RUNTIME.block_on(async {
SglangSchedulerClient::connect(endpoint_str).await
}) {
Ok(c) => Arc::new(c),
Err(e) => {
set_error_message(error_out, &format!("Failed to connect to endpoint: {}", e));
return ptr::null_mut();
}
};
Box::into_raw(Box::new(SglangClientHandle {
client,
tokenizer,
}))
}
/// Free a client handle
#[no_mangle]
pub unsafe extern "C" fn sgl_client_free(handle: *mut SglangClientHandle) {
if !handle.is_null() {
let _ = Box::from_raw(handle);
}
}
/// Send a chat completion request and start streaming
///
/// # Arguments
/// * `client_handle` - Client handle
/// * `request_json` - OpenAI ChatCompletionRequest as JSON string
/// * `stream_handle_out` - Pointer to receive stream handle
/// * `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_client_chat_completion_stream(
client_handle: *mut SglangClientHandle,
request_json: *const c_char,
stream_handle_out: *mut *mut SglangStreamHandle,
error_out: *mut *mut c_char,
) -> SglErrorCode {
if client_handle.is_null() || request_json.is_null() || stream_handle_out.is_null() {
set_error_message(error_out, "Invalid arguments: null pointer");
return SglErrorCode::InvalidArgument;
}
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 client_ref = &*client_handle;
let client = Arc::clone(&client_ref.client);
let tokenizer = Arc::clone(&client_ref.tokenizer);
// Parse OpenAI 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;
}
};
// Process messages and 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 messages: {}", e));
return SglErrorCode::TokenizationError;
}
};
// Tokenize
let token_ids = match tokenizer.encode(&processed_messages.text, false) {
Ok(encoding) => encoding.token_ids().to_vec(),
Err(e) => {
set_error_message(error_out, &format!("Failed to tokenize: {}", e));
return SglErrorCode::TokenizationError;
}
};
let prompt_tokens = token_ids.len() as i32; // Save prompt token count
// Generate tool constraints if needed
let tool_constraint = if let Some(tools) = chat_request.tools.as_ref() {
match generate_tool_constraints(tools, &chat_request.tool_choice, &chat_request.model) {
Ok(Some((constraint_type, constraint_value))) => Some((constraint_type, constraint_value)),
Ok(None) => None,
Err(e) => {
set_error_message(error_out, &format!("Failed to generate tool constraints: {}", e));
return SglErrorCode::ParsingError;
}
}
} else {
None
};
// Build GenerateRequest
let request_id = format!("chatcmpl-{}", Uuid::new_v4());
let proto_request = match client.build_generate_request_from_chat(
request_id.clone(),
&chat_request,
processed_messages.text,
token_ids,
processed_messages.multimodal_inputs,
tool_constraint,
) {
Ok(req) => req,
Err(e) => {
set_error_message(error_out, &format!("Failed to build generate request: {}", e));
return SglErrorCode::ParsingError;
}
};
// Send request and get stream
let stream = match RUNTIME.block_on(async {
client.generate(proto_request).await
}) {
Ok(s) => s,
Err(e) => {
set_error_message(error_out, &format!("Failed to send request: {}", e));
return SglErrorCode::UnknownError;
}
};
// Create response converter
let tools_json = chat_request.tools.as_ref()
.and_then(|t| serde_json::to_string(t).ok())
.map(|s| CString::new(s).unwrap().into_raw());
let tool_choice_json = chat_request.tool_choice.as_ref()
.and_then(|tc| serde_json::to_string(tc).ok())
.map(|s| CString::new(s).unwrap().into_raw());
let stop_json = chat_request.stop.as_ref()
.and_then(|s| serde_json::to_string(s).ok())
.map(|s| CString::new(s).unwrap().into_raw());
let stop_token_ids_json = chat_request.stop_token_ids.as_ref()
.and_then(|ids| serde_json::to_string(ids).ok())
.map(|s| CString::new(s).unwrap().into_raw());
// Create tokenizer handle for converter (we'll create a temporary one)
let tokenizer_handle = Box::into_raw(Box::new(TokenizerHandle {
tokenizer: Arc::clone(&tokenizer),
}));
let converter = sgl_grpc_response_converter_create(
tokenizer_handle,
CString::new(chat_request.model.clone()).unwrap().as_ptr(),
CString::new(request_id.clone()).unwrap().as_ptr(),
tools_json.unwrap_or(ptr::null_mut()),
tool_choice_json.unwrap_or(ptr::null_mut()),
stop_json.unwrap_or(ptr::null_mut()),
stop_token_ids_json.unwrap_or(ptr::null_mut()),
if chat_request.skip_special_tokens { 1 } else { 0 },
error_out,
);
// Free temporary tokenizer handle (converter now owns the tokenizer)
let _ = Box::from_raw(tokenizer_handle);
if converter.is_null() {
return SglErrorCode::MemoryError;
}
// Clean up temporary CStrings
if let Some(ptr) = tools_json {
let _ = CString::from_raw(ptr);
}
if let Some(ptr) = tool_choice_json {
let _ = CString::from_raw(ptr);
}
if let Some(ptr) = stop_json {
let _ = CString::from_raw(ptr);
}
if let Some(ptr) = stop_token_ids_json {
let _ = CString::from_raw(ptr);
}
// Create converter handle and set initial_prompt_tokens immediately
let mut converter_handle = *Box::from_raw(converter);
converter_handle.initial_prompt_tokens = Some(prompt_tokens);
// Create stream handle with prompt_tokens
*stream_handle_out = Box::into_raw(Box::new(SglangStreamHandle {
stream: Arc::new(tokio::sync::Mutex::new(stream)),
converter: Arc::new(tokio::sync::Mutex::new(converter_handle)),
client: Arc::clone(&client),
prompt_tokens,
}));
SglErrorCode::Success
}

50
third_party/sglang/sgl-model-gateway/bindings/golang/src/error.rs vendored Normal file
View File

@@ -0,0 +1,50 @@
//! Error handling for FFI functions
use std::ffi::CString;
use std::os::raw::c_char;
use std::ptr;
/// Error codes returned by FFI functions
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SglErrorCode {
Success = 0,
InvalidArgument = 1,
TokenizationError = 2,
ParsingError = 3,
MemoryError = 4,
UnknownError = 99,
}
/// Helper to set error message in FFI output parameter
pub fn set_error_message(error_out: *mut *mut c_char, message: &str) {
unsafe {
if !error_out.is_null() {
if let Ok(cstr) = CString::new(message) {
*error_out = cstr.into_raw();
} else {
*error_out = ptr::null_mut();
}
}
}
}
/// Helper to set error message from format string
pub fn set_error_message_fmt(error_out: *mut *mut c_char, fmt: std::fmt::Arguments) {
if !error_out.is_null() {
let msg = format!("{}", fmt);
set_error_message(error_out, &msg);
}
}
/// Helper to clear error message
pub fn clear_error_message(error_out: *mut *mut c_char) {
unsafe {
if !error_out.is_null() {
*error_out = ptr::null_mut();
}
}
}
// Helper functions for error handling
// Note: Some helper functions are kept for potential future use

758
third_party/sglang/sgl-model-gateway/bindings/golang/src/grpc_converter.rs vendored Normal file
View File

@@ -0,0 +1,758 @@
//! gRPC response converter FFI functions
use std::ffi::{CStr, CString};
use std::os::raw::{c_char, c_int};
use std::ptr;
use std::sync::Arc;
use std::collections::HashMap;
use serde_json::Value;
use tokio::runtime::Runtime;
use once_cell::sync::Lazy;
use smg::tokenizer::traits::Tokenizer;
use smg::tokenizer::stream::DecodeStream;
use smg::tool_parser::ToolParser;
use smg::protocols::common::{Tool, ToolChoice, ToolChoiceValue, ToolCallDelta, FunctionCallDelta, Usage, StringOrArray};
use smg::tokenizer::stop::StopSequenceDecoder;
use smg_grpc_client::sglang_proto as proto;
use super::error::{SglErrorCode, set_error_message, clear_error_message};
use super::tokenizer::TokenizerHandle;
use super::utils::generate_tool_call_id;
/// Global parser factory (initialized once)
// Use the re-exported ParserFactory from tool_parser module
static PARSER_FACTORY: Lazy<smg::tool_parser::ParserFactory> = Lazy::new(|| {
// ParserFactory is re-exported from tool_parser::factory, so we can use it directly
smg::tool_parser::ParserFactory::default()
});
/// Global tokio runtime for async operations
static RUNTIME: Lazy<Runtime> = Lazy::new(|| {
Runtime::new().expect("Failed to create tokio runtime for gRPC converter FFI")
});
/// Handle for gRPC response converter (maintains state for streaming)
#[repr(C)]
pub struct GrpcResponseConverterHandle {
pub(crate) tokenizer: Arc<dyn Tokenizer>,
pub(crate) tool_parser: Option<Arc<tokio::sync::Mutex<Box<dyn ToolParser>>>>,
pub(crate) stop_decoder: Option<Arc<tokio::sync::Mutex<StopSequenceDecoder>>>,
pub(crate) model: String,
pub(crate) request_id: String,
pub(crate) created: u64,
pub(crate) system_fingerprint: Option<String>,
pub(crate) tools: Option<Vec<Tool>>,
pub(crate) tool_choice: Option<ToolChoice>,
pub(crate) history_tool_calls_count: usize,
pub(crate) stream_buffers: HashMap<u32, String>, // Per-index text buffers
pub(crate) decode_streams: HashMap<u32, DecodeStream>, // Per-index incremental decoders
pub(crate) has_tool_calls: HashMap<u32, bool>, // Track if tool calls were emitted
pub(crate) is_first_chunk: HashMap<u32, bool>, // Track first chunk per index
pub(crate) prompt_tokens: HashMap<u32, i32>, // Track prompt tokens per index (from chunks)
pub(crate) completion_tokens: HashMap<u32, i32>, // Track completion tokens per index (cumulative)
pub(crate) initial_prompt_tokens: Option<i32>, // Initial prompt tokens from request (if available)
pub(crate) skip_special_tokens: bool, // Whether to skip special tokens when decoding
}
/// Create a gRPC response converter handle
///
/// # Arguments
/// * `tokenizer_handle` - Tokenizer handle (must be valid)
/// * `model` - Model name
/// * `request_id` - Request ID
/// * `tools_json` - Optional JSON array of tools
/// * `tool_choice_json` - Optional JSON object for tool_choice
/// * `stop` - Optional stop sequences (JSON array)
/// * `stop_token_ids` - Optional stop token IDs (JSON array)
/// * `skip_special_tokens` - Whether to skip special tokens
/// * `error_out` - Optional pointer to receive error message
///
/// # Returns
/// * Pointer to GrpcResponseConverterHandle on success, null on failure
#[no_mangle]
pub unsafe extern "C" fn sgl_grpc_response_converter_create(
tokenizer_handle: *mut TokenizerHandle,
model: *const c_char,
request_id: *const c_char,
tools_json: *const c_char,
tool_choice_json: *const c_char,
stop: *const c_char,
stop_token_ids: *const c_char,
skip_special_tokens: c_int,
error_out: *mut *mut c_char,
) -> *mut GrpcResponseConverterHandle {
if tokenizer_handle.is_null() || model.is_null() || request_id.is_null() {
set_error_message(error_out, "Invalid arguments: null pointer");
return ptr::null_mut();
}
let model_str = match CStr::from_ptr(model).to_str() {
Ok(s) => s,
Err(_) => {
set_error_message(error_out, "Invalid UTF-8 in model");
return ptr::null_mut();
}
};
let request_id_str = match CStr::from_ptr(request_id).to_str() {
Ok(s) => s,
Err(_) => {
set_error_message(error_out, "Invalid UTF-8 in request_id");
return ptr::null_mut();
}
};
let handle_ref = &*tokenizer_handle;
let tokenizer = Arc::clone(&handle_ref.tokenizer);
// Parse tools if provided
let tools: Option<Vec<Tool>> = if !tools_json.is_null() {
match CStr::from_ptr(tools_json).to_str() {
Ok(s) => serde_json::from_str::<Vec<Tool>>(s).ok(),
Err(_) => None,
}
} else {
None
};
// Parse tool_choice if provided
let tool_choice: Option<ToolChoice> = if !tool_choice_json.is_null() {
match CStr::from_ptr(tool_choice_json).to_str() {
Ok(s) => serde_json::from_str::<ToolChoice>(s).ok(),
Err(_) => None,
}
} else {
None
};
// Parse stop sequences
let stop: Option<StringOrArray> = if !stop.is_null() {
let stop_str = match CStr::from_ptr(stop).to_str() {
Ok(s) => s,
Err(_) => return ptr::null_mut(),
};
serde_json::from_str::<StringOrArray>(stop_str).ok()
} else {
None
};
// Parse stop token IDs
let stop_token_ids: Option<Vec<u32>> = if !stop_token_ids.is_null() {
let ids_str = match CStr::from_ptr(stop_token_ids).to_str() {
Ok(s) => s,
Err(_) => return ptr::null_mut(),
};
serde_json::from_str::<Vec<u32>>(ids_str).ok()
} else {
None
};
// Create stop decoder if needed
let stop_decoder = if stop.is_some() || stop_token_ids.is_some() {
Some(Arc::new(tokio::sync::Mutex::new(
smg::routers::grpc::utils::create_stop_decoder(
&tokenizer,
stop.as_ref(),
stop_token_ids.as_ref(),
skip_special_tokens != 0,
false, // no_stop_trim
),
)))
} else {
None
};
// Create tool parser if tools are provided
let tool_parser = if tools.is_some() {
PARSER_FACTORY.registry().create_for_model(model_str)
.map(|p| Arc::new(tokio::sync::Mutex::new(p)))
} else {
None
};
// Get system fingerprint from model (simplified)
let system_fingerprint = Some("fp_placeholder".to_string()); // TODO: Get actual fingerprint
Box::into_raw(Box::new(GrpcResponseConverterHandle {
tokenizer,
tool_parser,
stop_decoder,
model: model_str.to_string(),
request_id: request_id_str.to_string(),
created: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs(),
system_fingerprint,
tools,
tool_choice,
history_tool_calls_count: 0,
stream_buffers: HashMap::new(),
decode_streams: HashMap::new(),
has_tool_calls: HashMap::new(),
is_first_chunk: HashMap::new(),
prompt_tokens: HashMap::new(),
completion_tokens: HashMap::new(),
initial_prompt_tokens: None, // Will be set from stream handle
skip_special_tokens: skip_special_tokens != 0,
}))
}
/// Convert a gRPC GenerateResponse chunk to OpenAI format
///
/// # Arguments
/// * `handle` - Converter handle
/// * `response_json` - JSON string of proto.GenerateResponse
/// * `result_json_out` - Pointer to receive OpenAI format JSON (must be freed with sgl_free_string)
/// * `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_grpc_response_converter_convert_chunk(
handle: *mut GrpcResponseConverterHandle,
response_json: *const c_char,
result_json_out: *mut *mut c_char,
error_out: *mut *mut c_char,
) -> SglErrorCode {
if handle.is_null() || response_json.is_null() || result_json_out.is_null() {
set_error_message(error_out, "Invalid arguments: null pointer");
return SglErrorCode::InvalidArgument;
}
let response_str = match CStr::from_ptr(response_json).to_str() {
Ok(s) => s,
Err(_) => {
set_error_message(error_out, "Invalid UTF-8 in response_json");
return SglErrorCode::InvalidArgument;
}
};
// Parse proto.GenerateResponse from JSON
let json_value: Value = match serde_json::from_str(response_str) {
Ok(v) => v,
Err(e) => {
set_error_message(error_out, &format!("Failed to parse response JSON: {}", e));
return SglErrorCode::ParsingError;
}
};
// Build proto::GenerateResponse from JSON value
let mut proto_response = proto::GenerateResponse {
request_id: json_value.get("request_id")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
response: None,
};
// Parse the response oneof field
if let Some(chunk_json) = json_value.get("chunk") {
let chunk = proto::GenerateStreamChunk {
token_ids: chunk_json.get("token_ids")
.and_then(|v| v.as_array())
.map(|arr| arr.iter().filter_map(|v| v.as_u64().map(|n| n as u32)).collect())
.unwrap_or_default(),
prompt_tokens: chunk_json.get("prompt_tokens")
.and_then(|v| v.as_i64())
.map(|n| n as i32)
.unwrap_or(0),
completion_tokens: chunk_json.get("completion_tokens")
.and_then(|v| v.as_i64())
.map(|n| n as i32)
.unwrap_or(0),
cached_tokens: chunk_json.get("cached_tokens")
.and_then(|v| v.as_i64())
.map(|n| n as i32)
.unwrap_or(0),
output_logprobs: None,
hidden_states: vec![],
input_logprobs: None,
index: 0,
};
proto_response.response = Some(proto::generate_response::Response::Chunk(chunk));
} else if let Some(complete_json) = json_value.get("complete") {
let complete = proto::GenerateComplete {
output_ids: complete_json.get("output_ids")
.and_then(|v| v.as_array())
.map(|arr| arr.iter().filter_map(|v| v.as_u64().map(|n| n as u32)).collect())
.unwrap_or_default(),
finish_reason: complete_json.get("finish_reason")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
prompt_tokens: complete_json.get("prompt_tokens")
.and_then(|v| v.as_i64())
.map(|n| n as i32)
.unwrap_or(0),
completion_tokens: complete_json.get("completion_tokens")
.and_then(|v| v.as_i64())
.map(|n| n as i32)
.unwrap_or(0),
cached_tokens: complete_json.get("cached_tokens")
.and_then(|v| v.as_i64())
.map(|n| n as i32)
.unwrap_or(0),
output_logprobs: None,
all_hidden_states: vec![],
input_logprobs: None,
matched_stop: None,
index: 0,
};
proto_response.response = Some(proto::generate_response::Response::Complete(complete));
} else if let Some(error_json) = json_value.get("error") {
let error = proto::GenerateError {
message: error_json.get("message")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
http_status_code: error_json.get("http_status_code")
.and_then(|v| v.as_str())
.unwrap_or("500")
.to_string(),
details: error_json.get("details")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
};
proto_response.response = Some(proto::generate_response::Response::Error(error));
} else {
set_error_message(error_out, "Response JSON must contain 'chunk', 'complete', or 'error' field");
return SglErrorCode::ParsingError;
}
let handle_ref = &mut *handle;
let tokenizer = Arc::clone(&handle_ref.tokenizer);
let model = handle_ref.model.clone();
let request_id = handle_ref.request_id.clone();
let created = handle_ref.created;
let system_fingerprint = handle_ref.system_fingerprint.clone();
// Use tokio runtime to run async code
let result = RUNTIME.block_on(async {
convert_proto_chunk_to_openai(
proto_response,
handle_ref,
&tokenizer,
&model,
&request_id,
created,
system_fingerprint.as_deref(),
)
.await
});
match result {
Ok(Some(openai_response)) => {
// Serialize to JSON
let result_str = match serde_json::to_string(&openai_response) {
Ok(s) => s,
Err(e) => {
set_error_message(error_out, &format!("Failed to serialize response: {}", e));
return SglErrorCode::ParsingError;
}
};
let result_cstr = match CString::new(result_str) {
Ok(s) => s,
Err(e) => {
set_error_message(error_out, &format!("Failed to create result string: {}", e));
return SglErrorCode::MemoryError;
}
};
*result_json_out = result_cstr.into_raw();
clear_error_message(error_out);
SglErrorCode::Success
}
Ok(None) => {
// No response to send (e.g., empty chunk)
let empty = CString::new("").unwrap();
*result_json_out = empty.into_raw();
clear_error_message(error_out);
SglErrorCode::Success
}
Err(e) => {
set_error_message(error_out, &format!("Conversion error: {}", e));
SglErrorCode::ParsingError
}
}
}
/// Helper function to convert proto chunk to OpenAI format
pub(crate) async fn convert_proto_chunk_to_openai(
proto_response: proto::GenerateResponse,
handle: &mut GrpcResponseConverterHandle,
tokenizer: &Arc<dyn Tokenizer>,
model: &str,
request_id: &str,
created: u64,
system_fingerprint: Option<&str>,
) -> Result<Option<smg::protocols::chat::ChatCompletionStreamResponse>, String> {
use smg_grpc_client::sglang_proto::generate_response::Response::*;
use smg::protocols::chat::{ChatCompletionStreamResponse, ChatMessageDelta, ChatStreamChoice};
match proto_response.response {
Some(Chunk(chunk)) => {
let index = chunk.index;
// Mark as not first chunk if we've seen this index before
let is_first = handle.is_first_chunk.entry(index).or_insert(true);
let first_chunk = *is_first;
*is_first = false;
// Track token counts from chunks (cumulative values from proto)
// These are cumulative values, so we always use the latest value
// For prompt_tokens, if chunk value is 0, preserve existing value or use initial_prompt_tokens
// This prevents overwriting valid prompt_tokens with 0
if chunk.prompt_tokens > 0 {
handle.prompt_tokens.insert(index, chunk.prompt_tokens);
} else {
// If chunk.prompt_tokens is 0, try to preserve existing value or use initial_prompt_tokens
if !handle.prompt_tokens.contains_key(&index) {
// No existing value, try to use initial_prompt_tokens
if let Some(initial_prompt) = handle.initial_prompt_tokens {
handle.prompt_tokens.insert(index, initial_prompt);
}
}
// If existing value exists, keep it (don't overwrite with 0)
}
// For completion_tokens, always update (even if 0) as it's cumulative
handle.completion_tokens.insert(index, chunk.completion_tokens);
// Process tokens through stop decoder if available, otherwise use incremental decoder
let chunk_text = if let Some(ref stop_decoder) = handle.stop_decoder {
let mut decoder_guard = stop_decoder.lock().await;
let mut text = String::new();
for &token_id in &chunk.token_ids {
match decoder_guard.process_token(token_id).unwrap_or_else(|_| {
smg::tokenizer::stop::SequenceDecoderOutput::Held
}) {
smg::tokenizer::stop::SequenceDecoderOutput::Text(t) => {
text.push_str(&t);
}
smg::tokenizer::stop::SequenceDecoderOutput::StoppedWithText(t) => {
text.push_str(&t);
break;
}
smg::tokenizer::stop::SequenceDecoderOutput::Stopped => {
break;
}
smg::tokenizer::stop::SequenceDecoderOutput::Held => {}
}
}
text
} else {
// Use incremental decoder to handle multi-byte character boundaries
let decode_stream = handle.decode_streams.entry(index).or_insert_with(|| {
DecodeStream::new(
Arc::clone(&tokenizer),
&[], // No prompt tokens for completion
handle.skip_special_tokens,
)
});
// Process tokens incrementally
let mut text_parts = Vec::new();
for &token_id in &chunk.token_ids {
if let Ok(Some(text)) = decode_stream.step(token_id) {
text_parts.push(text);
}
}
text_parts.join("")
};
if chunk_text.is_empty() {
return Ok(None);
}
// Send first chunk with role
if first_chunk {
let first_response = ChatCompletionStreamResponse {
id: request_id.to_string(),
object: "chat.completion.chunk".to_string(),
created,
model: model.to_string(),
system_fingerprint: system_fingerprint.map(|s| s.to_string()),
choices: vec![ChatStreamChoice {
index,
delta: ChatMessageDelta {
role: Some("assistant".to_string()),
content: None,
tool_calls: None,
reasoning_content: None,
},
logprobs: None,
finish_reason: None,
matched_stop: None,
}],
usage: None,
};
return Ok(Some(first_response));
}
// Update stream buffer
let stream_buffer = handle.stream_buffers.entry(index).or_default();
stream_buffer.push_str(&chunk_text);
// Handle tool calls if tools are provided
if let (Some(ref tools), Some(ref tool_parser)) = (handle.tools.as_ref(), handle.tool_parser.as_ref()) {
let tool_choice_enabled = !matches!(
handle.tool_choice,
Some(ToolChoice::Value(ToolChoiceValue::None))
);
if tool_choice_enabled {
let mut parser_guard = tool_parser.lock().await;
match parser_guard.parse_incremental(&chunk_text, tools).await {
Ok(streaming_result) => {
if !streaming_result.calls.is_empty() {
handle.has_tool_calls.insert(index, true);
// Convert tool call items to OpenAI format
let tool_call_deltas: Vec<_> = streaming_result
.calls
.into_iter()
.map(|item| {
let id = if let Some(ref name) = item.name {
generate_tool_call_id(
model,
name,
item.tool_index,
handle.history_tool_calls_count,
)
} else {
format!("call_{}", item.tool_index)
};
ToolCallDelta {
index: item.tool_index as u32,
id: Some(id),
tool_type: if item.name.is_some() {
Some("function".to_string())
} else {
None
},
function: Some(FunctionCallDelta {
name: item.name,
arguments: if !item.parameters.is_empty() {
Some(item.parameters)
} else {
None
},
}),
}
})
.collect();
let tool_response = ChatCompletionStreamResponse {
id: request_id.to_string(),
object: "chat.completion.chunk".to_string(),
created,
model: model.to_string(),
system_fingerprint: system_fingerprint.map(|s| s.to_string()),
choices: vec![ChatStreamChoice {
index,
delta: ChatMessageDelta {
role: Some("assistant".to_string()),
content: None,
tool_calls: Some(tool_call_deltas),
reasoning_content: None,
},
logprobs: None,
finish_reason: None,
matched_stop: None,
}],
usage: None,
};
return Ok(Some(tool_response));
}
}
Err(e) => {
// Log error but continue with regular content
tracing::warn!("Tool parser error: {}", e);
}
}
}
}
// Regular content emission
let content_response = ChatCompletionStreamResponse {
id: request_id.to_string(),
object: "chat.completion.chunk".to_string(),
created,
model: model.to_string(),
system_fingerprint: system_fingerprint.map(|s| s.to_string()),
choices: vec![ChatStreamChoice {
index,
delta: ChatMessageDelta {
role: Some("assistant".to_string()),
content: Some(chunk_text),
tool_calls: None,
reasoning_content: None,
},
logprobs: None,
finish_reason: None,
matched_stop: None,
}],
usage: None,
};
Ok(Some(content_response))
}
Some(Complete(complete)) => {
let index = complete.index;
// Flush any remaining text
// Flush any remaining text from decode stream
let mut final_text = handle.stream_buffers.remove(&index).unwrap_or_default();
if let Some(ref mut decode_stream) = handle.decode_streams.get_mut(&index) {
if let Ok(Some(remaining)) = decode_stream.flush() {
final_text.push_str(&remaining);
}
}
handle.decode_streams.remove(&index);
// Determine finish reason - ensure it's never empty
// If finish_reason is empty, try to infer from other fields or use default
let finish_reason = if handle.has_tool_calls.get(&index).copied().unwrap_or(false)
&& (complete.finish_reason == "stop" || complete.finish_reason.is_empty())
{
"tool_calls".to_string()
} else if complete.finish_reason.is_empty() || complete.finish_reason.trim().is_empty() {
// If finish_reason is empty, try to infer from completion_tokens or use default
if complete.completion_tokens > 0 {
// If we have completion tokens, likely stopped normally
"stop".to_string()
} else if !complete.output_ids.is_empty() {
// If we have output_ids, likely stopped normally
"stop".to_string()
} else {
// Default fallback - always ensure we have a value
"stop".to_string()
}
} else {
complete.finish_reason.clone()
};
// Ensure finish_reason is never empty (defensive check)
let finish_reason = if finish_reason.is_empty() || finish_reason.trim().is_empty() {
"stop".to_string()
} else {
finish_reason
};
// Extract matched_stop
let matched_stop = match &complete.matched_stop {
Some(proto::generate_complete::MatchedStop::MatchedTokenId(token_id)) => {
Some(Value::Number(serde_json::Number::from(*token_id)))
}
Some(proto::generate_complete::MatchedStop::MatchedStopStr(stop_str)) => {
Some(Value::String(stop_str.clone()))
}
None => None,
};
// Build usage - prefer values from complete message, but fallback to accumulated values from chunks
// Complete message should have the final values, but sometimes they might be 0 or missing
// Always use the latest cumulative value from chunks if available, otherwise use complete message value
let mut prompt_tokens = handle.prompt_tokens.get(&index)
.copied()
.filter(|&v| v > 0)
.unwrap_or(complete.prompt_tokens);
let mut completion_tokens = handle.completion_tokens.get(&index)
.copied()
.filter(|&v| v > 0)
.unwrap_or(complete.completion_tokens);
// Always try to use initial_prompt_tokens if prompt_tokens is 0 or missing
// This is the most reliable source for prompt tokens since we calculate it from the request
if prompt_tokens == 0 {
if let Some(initial_prompt) = handle.initial_prompt_tokens {
prompt_tokens = initial_prompt;
}
}
// If completion_tokens is 0, try to infer from output_ids or accumulated chunks
if completion_tokens == 0 {
// Try to use completion_tokens from complete message even if 0
// Or calculate from output_ids
if complete.completion_tokens > 0 {
completion_tokens = complete.completion_tokens;
} else if !complete.output_ids.is_empty() {
completion_tokens = complete.output_ids.len() as i32;
} else if let Some(&last_completion) = handle.completion_tokens.get(&index) {
completion_tokens = last_completion;
}
}
// Final fallback: if both are still 0, try to use initial_prompt_tokens for prompt
// and calculate completion from output_ids
if prompt_tokens == 0 && completion_tokens == 0 {
// Try to infer from output_ids if available
let output_ids_len = complete.output_ids.len() as i32;
if output_ids_len > 0 {
completion_tokens = output_ids_len;
// Always try to use initial_prompt_tokens for prompt
if let Some(initial_prompt) = handle.initial_prompt_tokens {
prompt_tokens = initial_prompt;
}
}
}
// Final defensive check: ensure prompt_tokens is set if we have initial_prompt_tokens
if prompt_tokens == 0 {
if let Some(initial_prompt) = handle.initial_prompt_tokens {
prompt_tokens = initial_prompt;
}
}
// Always create usage, even if values are 0 (defensive)
let usage = Some(Usage {
prompt_tokens: prompt_tokens.max(0) as u32,
completion_tokens: completion_tokens.max(0) as u32,
total_tokens: (prompt_tokens.max(0) + completion_tokens.max(0)) as u32,
completion_tokens_details: None,
});
let finish_response = ChatCompletionStreamResponse {
id: request_id.to_string(),
object: "chat.completion.chunk".to_string(),
created,
model: model.to_string(),
system_fingerprint: system_fingerprint.map(|s| s.to_string()),
choices: vec![ChatStreamChoice {
index,
delta: ChatMessageDelta {
role: Some("assistant".to_string()),
content: if !final_text.is_empty() {
Some(final_text)
} else {
None
},
tool_calls: None,
reasoning_content: None,
},
logprobs: None,
finish_reason: Some(finish_reason),
matched_stop,
}],
usage,
};
Ok(Some(finish_response))
}
Some(Error(error)) => {
Err(format!("Server error: {} (status: {})", error.message, error.http_status_code))
}
None => Ok(None),
}
}
/// Free a gRPC response converter handle
#[no_mangle]
pub unsafe extern "C" fn sgl_grpc_response_converter_free(handle: *mut GrpcResponseConverterHandle) {
if !handle.is_null() {
let _ = Box::from_raw(handle);
}
}

103
third_party/sglang/sgl-model-gateway/bindings/golang/src/lib.rs vendored Normal file
View File

@@ -0,0 +1,103 @@
//! FFI module for exposing sgl-model-gateway preprocessing and postprocessing functions
//! to C-compatible languages (e.g., Golang via cgo)
//!
//! This module provides C-compatible function signatures for:
//! - Tokenizer operations (encode, decode, chat template)
//! - Tool parser operations (parse tool calls)
//! - Tool constraint generation
//! - gRPC client SDK (complete request-response flow)
//!
//! # Safety
//! All functions marked with `#[no_mangle]` and `extern "C"` must be called
//! with valid pointers and follow the documented memory management rules.
// Re-export error types
pub use error::{SglErrorCode, set_error_message, set_error_message_fmt, clear_error_message};
// Re-export memory management functions
pub use memory::{sgl_free_string, sgl_free_token_ids};
// Re-export tokenizer functions
pub use tokenizer::{
TokenizerHandle,
sgl_tokenizer_create_from_file,
sgl_tokenizer_encode,
sgl_tokenizer_apply_chat_template,
sgl_tokenizer_apply_chat_template_with_tools,
sgl_tokenizer_decode,
sgl_tokenizer_free,
};
// Re-export tool parser functions
pub use tool_parser::{
ToolParserHandle,
sgl_tool_parser_create,
sgl_tool_parser_parse_complete,
sgl_tool_parser_parse_incremental,
sgl_tool_parser_reset,
sgl_tool_parser_free,
};
// Re-export gRPC converter functions
pub use grpc_converter::{
GrpcResponseConverterHandle,
sgl_grpc_response_converter_create,
sgl_grpc_response_converter_convert_chunk,
sgl_grpc_response_converter_free,
};
// Re-export client SDK functions
pub use client::{
SglangClientHandle,
sgl_client_create,
sgl_client_free,
};
// Re-export stream functions
pub use stream::{
SglangStreamHandle,
sgl_stream_read_next,
sgl_stream_free,
};
// Re-export client stream function (defined in client.rs but used by stream)
pub use client::sgl_client_chat_completion_stream;
// Re-export preprocessor functions
pub use preprocessor::{
sgl_preprocess_chat_request,
sgl_preprocess_chat_request_with_tokenizer,
sgl_preprocessed_request_free,
};
// Re-export postprocessor functions
pub use postprocessor::{
sgl_postprocess_stream_chunk,
sgl_postprocess_stream_chunks_batch,
};
// Re-export utility functions
pub use utils::sgl_generate_tool_constraints;
// Sub-modules
mod error;
mod memory;
mod tokenizer;
mod tool_parser;
mod grpc_converter;
mod client;
mod stream;
mod utils;
mod preprocessor;
mod postprocessor;
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_error_codes() {
assert_eq!(SglErrorCode::Success as i32, 0);
assert_eq!(SglErrorCode::InvalidArgument as i32, 1);
}
}

28
third_party/sglang/sgl-model-gateway/bindings/golang/src/memory.rs vendored Normal file
View File

@@ -0,0 +1,28 @@
//! Memory management for FFI functions
use std::ffi::CString;
use std::os::raw::c_char;
/// Free a C string allocated by Rust
///
/// # Safety
/// This function must only be called with pointers returned by other FFI functions.
/// Calling with arbitrary pointers or multiple times on the same pointer is undefined behavior.
#[no_mangle]
pub unsafe extern "C" fn sgl_free_string(s: *mut c_char) {
if !s.is_null() {
let _ = CString::from_raw(s);
}
}
/// Free token IDs array allocated by Rust
///
/// # Safety
/// This function must only be called with pointers returned by `sgl_tokenizer_encode`.
/// The `count` parameter must match the length of the array.
#[no_mangle]
pub unsafe extern "C" fn sgl_free_token_ids(ptr: *mut u32, count: usize) {
if !ptr.is_null() && count > 0 {
let _ = Vec::from_raw_parts(ptr, count, count);
}
}

465
third_party/sglang/sgl-model-gateway/bindings/golang/src/postprocessor.rs vendored Normal file
View File

@@ -0,0 +1,465 @@
//! Postprocessing FFI functions for gRPC stream chunks
//!
//! This module provides C-compatible functions for postprocessing gRPC stream chunks:
//! - Parse tool calls from model output
//! - Convert proto format to OpenAI format
//! - Handle reasoning content parsing
//!
//! These functions are designed to be called for each stream chunk, but can be optimized
//! with batching in the future.
use std::ffi::{CStr, CString};
use std::os::raw::{c_char, c_int};
use std::ptr;
use std::sync::Arc;
use serde_json::Value;
use smg_grpc_client::sglang_proto as proto;
use super::error::{SglErrorCode, set_error_message};
use super::grpc_converter::GrpcResponseConverterHandle;
use tokio::runtime::Runtime;
use once_cell::sync::Lazy;
/// Global tokio runtime for async operations
static RUNTIME: Lazy<Runtime> = Lazy::new(|| {
Runtime::new().expect("Failed to create tokio runtime for postprocessor FFI")
});
/// Postprocess a gRPC stream chunk to OpenAI format
///
/// This function:
/// 1. Parses the proto chunk from JSON
/// 2. Converts it to OpenAI format using the converter handle
/// 3. Returns the OpenAI format JSON
///
/// # Arguments
/// * `converter_handle` - Converter handle (created with sgl_grpc_response_converter_create)
/// * `proto_chunk_json` - JSON string of proto.GenerateResponse
/// * `openai_json_out` - Pointer to receive OpenAI format JSON (must be freed with sgl_free_string)
/// * `is_done_out` - Pointer to receive is_done flag (1 if stream is complete, 0 otherwise)
/// * `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_postprocess_stream_chunk(
converter_handle: *mut GrpcResponseConverterHandle,
proto_chunk_json: *const c_char,
openai_json_out: *mut *mut c_char,
is_done_out: *mut c_int,
error_out: *mut *mut c_char,
) -> SglErrorCode {
if converter_handle.is_null()
|| proto_chunk_json.is_null()
|| openai_json_out.is_null()
|| is_done_out.is_null()
{
set_error_message(error_out, "Invalid arguments: null pointer");
return SglErrorCode::InvalidArgument;
}
let proto_chunk_str = match CStr::from_ptr(proto_chunk_json).to_str() {
Ok(s) => s,
Err(_) => {
set_error_message(error_out, "Invalid UTF-8 in proto_chunk_json");
return SglErrorCode::InvalidArgument;
}
};
// Parse proto.GenerateResponse from JSON
let json_value: Value = match serde_json::from_str(proto_chunk_str) {
Ok(v) => v,
Err(e) => {
set_error_message(error_out, &format!("Failed to parse proto chunk JSON: {}", e));
return SglErrorCode::ParsingError;
}
};
// Build proto::GenerateResponse from JSON value
let mut proto_response = proto::GenerateResponse {
request_id: json_value
.get("request_id")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
response: None,
};
// Parse the response oneof field
let is_done = if let Some(chunk_json) = json_value.get("chunk") {
let chunk = proto::GenerateStreamChunk {
token_ids: chunk_json
.get("token_ids")
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_u64().map(|n| n as u32))
.collect()
})
.unwrap_or_default(),
prompt_tokens: chunk_json
.get("prompt_tokens")
.and_then(|v| v.as_i64())
.map(|n| n as i32)
.unwrap_or(0),
completion_tokens: chunk_json
.get("completion_tokens")
.and_then(|v| v.as_i64())
.map(|n| n as i32)
.unwrap_or(0),
cached_tokens: chunk_json
.get("cached_tokens")
.and_then(|v| v.as_i64())
.map(|n| n as i32)
.unwrap_or(0),
output_logprobs: None,
hidden_states: vec![],
input_logprobs: None,
index: 0,
};
proto_response.response = Some(proto::generate_response::Response::Chunk(chunk));
false
} else if let Some(complete_json) = json_value.get("complete") {
let complete = proto::GenerateComplete {
output_ids: complete_json
.get("output_ids")
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_u64().map(|n| n as u32))
.collect()
})
.unwrap_or_default(),
finish_reason: complete_json
.get("finish_reason")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
prompt_tokens: complete_json
.get("prompt_tokens")
.and_then(|v| v.as_i64())
.map(|n| n as i32)
.unwrap_or(0),
completion_tokens: complete_json
.get("completion_tokens")
.and_then(|v| v.as_i64())
.map(|n| n as i32)
.unwrap_or(0),
cached_tokens: complete_json
.get("cached_tokens")
.and_then(|v| v.as_i64())
.map(|n| n as i32)
.unwrap_or(0),
output_logprobs: None,
all_hidden_states: vec![],
input_logprobs: None,
matched_stop: None,
index: 0,
};
proto_response.response = Some(proto::generate_response::Response::Complete(complete));
true
} else if let Some(error_json) = json_value.get("error") {
let error = proto::GenerateError {
message: error_json
.get("message")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
http_status_code: error_json
.get("http_status_code")
.and_then(|v| v.as_str())
.unwrap_or("500")
.to_string(),
details: error_json
.get("details")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
};
proto_response.response = Some(proto::generate_response::Response::Error(error));
true
} else {
set_error_message(
error_out,
"Proto chunk JSON must contain 'chunk', 'complete', or 'error' field",
);
return SglErrorCode::ParsingError;
};
// Convert proto chunk to OpenAI format using the converter's convert_chunk function
// We'll use the existing converter API instead of calling the internal function directly
let proto_chunk_json_cstr = match CString::new(proto_chunk_str) {
Ok(s) => s,
Err(e) => {
set_error_message(error_out, &format!("Failed to create C string: {}", e));
return SglErrorCode::MemoryError;
}
};
// Use the existing converter API
let mut openai_json_ptr: *mut c_char = ptr::null_mut();
let result = super::grpc_converter::sgl_grpc_response_converter_convert_chunk(
converter_handle,
proto_chunk_json_cstr.as_ptr(),
&mut openai_json_ptr,
error_out,
);
if result == SglErrorCode::Success {
*openai_json_out = openai_json_ptr;
*is_done_out = if is_done { 1 } else { 0 };
SglErrorCode::Success
} else {
*openai_json_out = ptr::null_mut();
*is_done_out = if is_done { 1 } else { 0 };
result
}
}
/// Postprocess multiple gRPC stream chunks in batch (reduces FFI overhead)
///
/// This function processes multiple chunks in a single FFI call, significantly reducing
/// FFI overhead in streaming scenarios.
///
/// # Arguments
/// * `converter_handle` - Converter handle (created with sgl_grpc_response_converter_create)
/// * `proto_chunks_json_array` - JSON array string of proto.GenerateResponse chunks
/// * `max_chunks` - Maximum number of chunks to process (for safety)
/// * `openai_chunks_json_array_out` - Pointer to receive JSON array of OpenAI format chunks (must be freed with sgl_free_string)
/// * `chunks_count_out` - Pointer to receive number of processed chunks
/// * `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_postprocess_stream_chunks_batch(
converter_handle: *mut GrpcResponseConverterHandle,
proto_chunks_json_array: *const c_char,
max_chunks: c_int,
openai_chunks_json_array_out: *mut *mut c_char,
chunks_count_out: *mut c_int,
error_out: *mut *mut c_char,
) -> SglErrorCode {
if converter_handle.is_null()
|| proto_chunks_json_array.is_null()
|| openai_chunks_json_array_out.is_null()
|| chunks_count_out.is_null()
{
set_error_message(error_out, "Invalid arguments: null pointer");
return SglErrorCode::InvalidArgument;
}
let chunks_array_str = match CStr::from_ptr(proto_chunks_json_array).to_str() {
Ok(s) => s,
Err(_) => {
set_error_message(error_out, "Invalid UTF-8 in proto_chunks_json_array");
return SglErrorCode::InvalidArgument;
}
};
// Parse JSON array of chunks
let chunks_array: Vec<Value> = match serde_json::from_str(chunks_array_str) {
Ok(arr) => arr,
Err(e) => {
set_error_message(
error_out,
&format!("Failed to parse chunks JSON array: {}", e),
);
return SglErrorCode::ParsingError;
}
};
// Limit batch size for safety
let max_chunks_usize = max_chunks as usize;
let chunks_to_process = if chunks_array.len() > max_chunks_usize {
&chunks_array[..max_chunks_usize]
} else {
&chunks_array
};
let handle_ref = &mut *converter_handle;
let tokenizer = Arc::clone(&handle_ref.tokenizer);
let model = handle_ref.model.clone();
let request_id = handle_ref.request_id.clone();
let created = handle_ref.created;
let system_fingerprint = handle_ref.system_fingerprint.clone();
// Process chunks in batch
let mut results = Vec::new();
let mut has_error = false;
let mut error_msg = String::new();
for chunk_json in chunks_to_process {
// Parse proto.GenerateResponse from JSON
let mut proto_response = proto::GenerateResponse {
request_id: chunk_json
.get("request_id")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
response: None,
};
// Parse the response oneof field (same logic as single chunk processing)
let _is_done = if let Some(chunk_json) = chunk_json.get("chunk") {
let chunk = proto::GenerateStreamChunk {
token_ids: chunk_json
.get("token_ids")
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_u64().map(|n| n as u32))
.collect()
})
.unwrap_or_default(),
prompt_tokens: chunk_json
.get("prompt_tokens")
.and_then(|v| v.as_i64())
.map(|n| n as i32)
.unwrap_or(0),
completion_tokens: chunk_json
.get("completion_tokens")
.and_then(|v| v.as_i64())
.map(|n| n as i32)
.unwrap_or(0),
cached_tokens: chunk_json
.get("cached_tokens")
.and_then(|v| v.as_i64())
.map(|n| n as i32)
.unwrap_or(0),
output_logprobs: None,
hidden_states: vec![],
input_logprobs: None,
index: 0,
};
proto_response.response = Some(proto::generate_response::Response::Chunk(chunk));
false
} else if let Some(complete_json) = chunk_json.get("complete") {
let complete = proto::GenerateComplete {
output_ids: complete_json
.get("output_ids")
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_u64().map(|n| n as u32))
.collect()
})
.unwrap_or_default(),
finish_reason: complete_json
.get("finish_reason")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
prompt_tokens: complete_json
.get("prompt_tokens")
.and_then(|v| v.as_i64())
.map(|n| n as i32)
.unwrap_or(0),
completion_tokens: complete_json
.get("completion_tokens")
.and_then(|v| v.as_i64())
.map(|n| n as i32)
.unwrap_or(0),
cached_tokens: complete_json
.get("cached_tokens")
.and_then(|v| v.as_i64())
.map(|n| n as i32)
.unwrap_or(0),
output_logprobs: None,
all_hidden_states: vec![],
input_logprobs: None,
matched_stop: None,
index: 0,
};
proto_response.response = Some(proto::generate_response::Response::Complete(complete));
true
} else if let Some(error_json) = chunk_json.get("error") {
let error = proto::GenerateError {
message: error_json
.get("message")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
http_status_code: error_json
.get("http_status_code")
.and_then(|v| v.as_str())
.unwrap_or("500")
.to_string(),
details: error_json
.get("details")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
};
proto_response.response = Some(proto::generate_response::Response::Error(error));
true
} else {
error_msg = format!(
"Chunk JSON must contain 'chunk', 'complete', or 'error' field: {}",
chunk_json
);
has_error = true;
break;
};
// Convert proto chunk to OpenAI format
let result = RUNTIME.block_on(async {
super::grpc_converter::convert_proto_chunk_to_openai(
proto_response,
handle_ref,
&tokenizer,
&model,
&request_id,
created,
system_fingerprint.as_deref(),
)
.await
});
match result {
Ok(Some(openai_response)) => {
results.push(openai_response);
}
Ok(None) => {
// Empty response, skip
}
Err(e) => {
error_msg = format!("Postprocessing failed for chunk: {}", e);
has_error = true;
break;
}
}
}
if has_error {
set_error_message(error_out, &error_msg);
return SglErrorCode::ParsingError;
}
// Serialize results to JSON array
let results_json = match serde_json::to_string(&results) {
Ok(s) => s,
Err(e) => {
set_error_message(
error_out,
&format!("Failed to serialize results JSON array: {}", e),
);
return SglErrorCode::ParsingError;
}
};
let results_cstr = match CString::new(results_json) {
Ok(s) => s,
Err(e) => {
set_error_message(error_out, &format!("Failed to create C string: {}", e));
return SglErrorCode::MemoryError;
}
};
*openai_chunks_json_array_out = results_cstr.into_raw();
*chunks_count_out = results.len() as c_int;
SglErrorCode::Success
}

372
third_party/sglang/sgl-model-gateway/bindings/golang/src/preprocessor.rs vendored Normal file
View 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);
}
}

288
third_party/sglang/sgl-model-gateway/bindings/golang/src/stream.rs vendored Normal file
View File

@@ -0,0 +1,288 @@
//! Stream handling FFI functions
//!
//! This module provides FFI (Foreign Function Interface) functions for managing
//! streaming responses from the SGLang gRPC API. It handles:
//!
//! - Creating and managing stream handles
//! - Reading chunks from streams and converting them to OpenAI format
//! - Managing automatic abort on stream drop (via AbortOnDropStream)
//! - Thread-safe access to streams and response converters
//!
//! # Safety
//!
//! All FFI functions are marked `unsafe` as per Rust FFI conventions. Callers must:
//! - Pass valid pointers
//! - Ensure proper pointer lifetime management
//! - Call corresponding free functions for cleanup
use std::ffi::CString;
use std::os::raw::{c_char, c_int};
use std::ptr;
use std::sync::Arc;
use tokio::runtime::Runtime;
use once_cell::sync::Lazy;
use futures_util::StreamExt;
use smg_grpc_client::{sglang_proto as proto, sglang_scheduler::{SglangSchedulerClient, AbortOnDropStream}};
use super::error::{SglErrorCode, set_error_message};
use super::grpc_converter::{GrpcResponseConverterHandle, convert_proto_chunk_to_openai};
/// Global tokio runtime for async operations
static RUNTIME: Lazy<Runtime> = Lazy::new(|| {
Runtime::new().expect("Failed to create tokio runtime for stream FFI")
});
/// Handle for an active streaming request.
///
/// This struct manages the stream and response converter for a single request.
/// It is wrapped in Arc and Mutex for thread-safe concurrent access.
///
/// # Fields
///
/// * `stream` - The gRPC stream wrapped in AbortOnDropStream for automatic cleanup
/// * `converter` - Response converter that transforms proto messages to OpenAI format
/// * `client` - The underlying gRPC client connection
/// * `prompt_tokens` - Number of prompt tokens from the original request
pub struct SglangStreamHandle {
pub(crate) stream: Arc<tokio::sync::Mutex<AbortOnDropStream>>,
pub(crate) converter: Arc<tokio::sync::Mutex<GrpcResponseConverterHandle>>,
#[allow(dead_code)]
pub(crate) client: Arc<SglangSchedulerClient>,
#[allow(dead_code)]
pub(crate) prompt_tokens: i32, // Number of prompt tokens for this request
}
/// Read next chunk from stream and convert to OpenAI format.
///
/// This function reads the next chunk from the gRPC stream, converts it from the
/// internal protocol format to OpenAI-compatible JSON format, and returns it via
/// the output parameters.
///
/// # Arguments
///
/// * `stream_handle` - Mutable pointer to the stream handle
/// * `response_json_out` - Pointer to receive OpenAI format JSON string
/// - Caller must free this with `sgl_free_string`
/// - May be NULL if no data available
/// * `is_done_out` - Pointer to receive completion status
/// - 0 = stream has more data
/// - 1 = stream is complete
/// * `error_out` - Optional pointer to receive error message
/// - Only set if function returns an error code
/// - Must be freed with `sgl_free_string` if not NULL
///
/// # Returns
///
/// * `SglErrorCode::Success` - Successfully read a chunk or reached end of stream
/// * Other error codes - See `SglErrorCode` for details
///
/// # Safety
///
/// - All pointers must be valid and properly aligned
/// - `stream_handle` must point to a valid `SglangStreamHandle`
/// - Output pointers must be writable
///
/// # Notes
///
/// - Complete messages are identified by the presence of `proto::GenerateResponse::Complete`
/// - When is_done=1, this may be the last readable chunk or the stream may be ending
/// - Subsequent calls after is_done=1 will mark the stream as complete internally
#[no_mangle]
pub unsafe extern "C" fn sgl_stream_read_next(
stream_handle: *mut SglangStreamHandle,
response_json_out: *mut *mut c_char,
is_done_out: *mut c_int,
error_out: *mut *mut c_char,
) -> SglErrorCode {
if stream_handle.is_null() || response_json_out.is_null() || is_done_out.is_null() {
set_error_message(error_out, "Invalid arguments: null pointer");
return SglErrorCode::InvalidArgument;
}
let handle_ref = &*stream_handle;
let stream = Arc::clone(&handle_ref.stream);
let converter = Arc::clone(&handle_ref.converter);
// Read next chunk from stream
let chunk_result = RUNTIME.block_on(async {
let mut stream_guard = stream.lock().await;
stream_guard.next().await
});
match chunk_result {
Some(Ok(proto_response)) => {
// Convert proto response to OpenAI format
// We need to get the converter lock first
let conversion_result = RUNTIME.block_on(async {
let mut converter_guard = converter.lock().await;
// Clone necessary fields for conversion
let tokenizer = Arc::clone(&converter_guard.tokenizer);
let model = converter_guard.model.clone();
let request_id = converter_guard.request_id.clone();
let created = converter_guard.created;
let system_fingerprint = converter_guard.system_fingerprint.clone();
// Call the conversion function
convert_proto_chunk_to_openai(
proto_response.clone(),
&mut *converter_guard,
&tokenizer,
&model,
&request_id,
created,
system_fingerprint.as_deref(),
)
.await
});
match conversion_result {
Ok(Some(openai_response)) => {
// Serialize to JSON
let result_str = match serde_json::to_string(&openai_response) {
Ok(s) => s,
Err(e) => {
set_error_message(error_out, &format!("Failed to serialize response: {}", e));
return SglErrorCode::ParsingError;
}
};
let result_cstr = match CString::new(result_str) {
Ok(s) => s,
Err(e) => {
set_error_message(error_out, &format!("Failed to create result string: {}", e));
return SglErrorCode::MemoryError;
}
};
// Check if this is a complete response (stream done)
let is_complete = matches!(proto_response.response, Some(proto::generate_response::Response::Complete(_)) | Some(proto::generate_response::Response::Error(_)));
*response_json_out = result_cstr.into_raw();
*is_done_out = if is_complete { 1 } else { 0 };
if is_complete {
// Mark stream as completed
// Ensure mark_completed() completes and is visible before returning
// Use yield_now to ensure Release ordering is fully propagated
RUNTIME.block_on(async {
let stream_guard = stream.lock().await;
stream_guard.mark_completed();
// Keep the guard until mark_completed() is fully executed
drop(stream_guard);
// Yield to ensure Release ordering is propagated before returning
// This prevents race condition where Free() is called immediately
// and Drop might not see the mark_completed() write
tokio::task::yield_now().await;
});
}
SglErrorCode::Success
}
Ok(None) => {
// No response to send (e.g., empty chunk)
// Don't mark as completed - stream might continue
// Just return null and let caller read more
*response_json_out = ptr::null_mut();
*is_done_out = 0; // Keep stream open, not done yet
SglErrorCode::Success
}
Err(e) => {
// Conversion error - don't mark as completed
// Let the stream end naturally or return error without stopping stream
set_error_message(error_out, &format!("Conversion error: {}", e));
*response_json_out = ptr::null_mut();
*is_done_out = 0; // Don't mark as done - let caller decide
SglErrorCode::ParsingError
}
}
}
Some(Err(e)) => {
// Stream error - mark as completed to prevent abort
RUNTIME.block_on(async {
let stream_guard = stream.lock().await;
stream_guard.mark_completed();
drop(stream_guard);
// Yield to ensure Release ordering is propagated
tokio::task::yield_now().await;
});
set_error_message(error_out, &format!("Stream error: {}", e));
*is_done_out = 1;
SglErrorCode::UnknownError
}
None => {
// Stream ended naturally (no more chunks)
// Mark stream as completed before returning to prevent abort
RUNTIME.block_on(async {
let stream_guard = stream.lock().await;
stream_guard.mark_completed();
drop(stream_guard);
// Yield to ensure Release ordering is propagated
tokio::task::yield_now().await;
});
*response_json_out = ptr::null_mut();
*is_done_out = 1;
SglErrorCode::Success
}
}
}
/// Free a stream handle and release all associated resources.
///
/// This function must be called exactly once for each stream handle returned by
/// `sgl_client_chat_completion_stream`. It marks the stream as completed internally
/// to prevent abort signals from being sent when resources are cleaned up.
///
/// # Arguments
///
/// * `handle` - Mutable pointer to the stream handle to free
/// - If NULL, this function does nothing
///
/// # Safety
///
/// - Must be called only once per handle
/// - Handle must not be used after calling this function
/// - After this call, the stream is no longer valid
///
/// # Notes
///
/// - This function internally calls `mark_completed()` before freeing to ensure
/// the stream cleanup doesn't trigger an abort RPC to the server
/// - Memory fences are used to ensure visibility across threads
#[no_mangle]
pub unsafe extern "C" fn sgl_stream_free(handle: *mut SglangStreamHandle) {
if !handle.is_null() {
let handle_ref = Box::from_raw(handle);
// Mark stream as completed to prevent abort on drop
// By this point, the stream should already be completed by ReadNext()
// but we call it again to be safe
RUNTIME.block_on(async {
let stream_guard = handle_ref.stream.lock().await;
stream_guard.mark_completed();
// Keep guard alive to ensure mark_completed() write completes
drop(stream_guard);
// Yield to ensure the atomic write is visible
tokio::task::yield_now().await;
});
// Use a strong memory fence to ensure mark_completed()'s Release write
// is visible before we drop the last Arc reference
std::sync::atomic::fence(std::sync::atomic::Ordering::SeqCst);
// Now drop all references - if mark_completed() was called successfully,
// the drop won't send an abort
drop(handle_ref.stream);
// Free converter
let converter = Arc::try_unwrap(handle_ref.converter)
.ok()
.map(|m| m.into_inner());
if let Some(conv) = converter {
super::grpc_converter::sgl_grpc_response_converter_free(Box::into_raw(Box::new(conv)));
}
}
}

388
third_party/sglang/sgl-model-gateway/bindings/golang/src/tokenizer.rs vendored Normal file
View File

@@ -0,0 +1,388 @@
//! Tokenizer FFI functions
use std::ffi::{CStr, CString};
use std::os::raw::{c_char, c_int};
use std::ptr;
use std::sync::Arc;
use serde_json::Value;
use smg::tokenizer::{
create_tokenizer_from_file,
traits::Tokenizer as TokenizerTrait,
chat_template::ChatTemplateParams,
huggingface::HuggingFaceTokenizer,
};
use super::error::{SglErrorCode, set_error_message, clear_error_message};
#[cfg(target_os = "macos")]
type BooleanT = libc::boolean_t;
#[cfg(not(target_os = "macos"))]
type BooleanT = libc::c_int;
/// Opaque handle for a tokenizer instance
#[repr(C)]
pub struct TokenizerHandle {
pub(crate) tokenizer: Arc<dyn TokenizerTrait>,
}
/// Create a tokenizer from a file path
///
/// # Arguments
/// * `path` - Path to tokenizer.json file (null-terminated C string)
/// * `error_out` - Optional pointer to receive error message (must be freed with sgl_free_string)
///
/// # Returns
/// * Pointer to TokenizerHandle on success, null on failure
///
/// # Safety
/// The returned handle must be freed with `sgl_tokenizer_free`.
#[no_mangle]
pub unsafe extern "C" fn sgl_tokenizer_create_from_file(
path: *const c_char,
error_out: *mut *mut c_char,
) -> *mut TokenizerHandle {
if path.is_null() {
set_error_message(error_out, "path cannot be null");
return ptr::null_mut();
}
let path_str = match CStr::from_ptr(path).to_str() {
Ok(s) => s,
Err(e) => {
set_error_message(error_out, &format!("Invalid UTF-8 in path: {}", e));
return ptr::null_mut();
}
};
match create_tokenizer_from_file(path_str) {
Ok(tokenizer) => {
clear_error_message(error_out);
Box::into_raw(Box::new(TokenizerHandle {
tokenizer,
}))
}
Err(e) => {
set_error_message(error_out, &e.to_string());
ptr::null_mut()
}
}
}
/// Encode text to token IDs
///
/// # Arguments
/// * `handle` - Tokenizer handle (must not be null)
/// * `text` - Input text (null-terminated C string)
/// * `add_special_tokens` - Whether to add special tokens
/// * `token_ids_out` - Pointer to receive array of token IDs (must be freed with sgl_free_token_ids)
/// * `token_count_out` - Pointer to receive token count
/// * `error_out` - Optional pointer to receive error message
///
/// # Returns
/// * SglErrorCode::Success on success, error code on failure
///
/// # Safety
/// The token_ids_out array must be freed with sgl_free_token_ids() after use.
#[no_mangle]
pub unsafe extern "C" fn sgl_tokenizer_encode(
handle: *mut TokenizerHandle,
text: *const c_char,
add_special_tokens: BooleanT,
token_ids_out: *mut *mut u32,
token_count_out: *mut usize,
error_out: *mut *mut c_char,
) -> SglErrorCode {
if handle.is_null() || text.is_null() || token_ids_out.is_null() || token_count_out.is_null() {
set_error_message(error_out, "Invalid arguments: null pointer");
return SglErrorCode::InvalidArgument;
}
let text_str = match CStr::from_ptr(text).to_str() {
Ok(s) => s,
Err(_) => {
set_error_message(error_out, "Invalid UTF-8 in text");
return SglErrorCode::InvalidArgument;
}
};
let add_special_tokens_bool = add_special_tokens != 0;
let tokenizer = &(*handle).tokenizer;
match tokenizer.encode(text_str, add_special_tokens_bool) {
Ok(encoding) => {
let token_ids = encoding.token_ids();
let count = token_ids.len();
// Allocate memory for token IDs using Vec, then leak to give ownership to C
let vec = token_ids.to_vec();
let ptr = vec.as_ptr() as *mut u32;
let _ = std::mem::ManuallyDrop::new(vec);
*token_ids_out = ptr;
*token_count_out = count;
clear_error_message(error_out);
SglErrorCode::Success
}
Err(e) => {
set_error_message(error_out, &e.to_string());
SglErrorCode::TokenizationError
}
}
}
/// Apply chat template to messages with tools support
///
/// # Arguments
/// * `handle` - Tokenizer handle
/// * `messages_json` - JSON string of messages array
/// * `tools_json` - Optional JSON string of tools array (null or empty string for no tools)
/// * `result_out` - Pointer to receive result string (must be freed with sgl_free_string)
/// * `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_tokenizer_apply_chat_template_with_tools(
handle: *mut TokenizerHandle,
messages_json: *const c_char,
tools_json: *const c_char,
result_out: *mut *mut c_char,
error_out: *mut *mut c_char,
) -> SglErrorCode {
if handle.is_null() || messages_json.is_null() || result_out.is_null() {
set_error_message(error_out, "Invalid arguments: null pointer");
return SglErrorCode::InvalidArgument;
}
let messages_str = match CStr::from_ptr(messages_json).to_str() {
Ok(s) => s,
Err(_) => {
set_error_message(error_out, "Invalid UTF-8 in messages_json");
return SglErrorCode::InvalidArgument;
}
};
// Parse JSON messages
let messages: Vec<Value> = match serde_json::from_str(messages_str) {
Ok(msgs) => msgs,
Err(e) => {
set_error_message(error_out, &format!("Failed to parse messages JSON: {}", e));
return SglErrorCode::InvalidArgument;
}
};
// Parse tools JSON if provided
let tools: Option<Vec<Value>> = if tools_json.is_null() {
None
} else {
let tools_str = match CStr::from_ptr(tools_json).to_str() {
Ok(s) => {
if s.is_empty() {
None
} else {
match serde_json::from_str::<Vec<Value>>(s) {
Ok(t) => Some(t),
Err(e) => {
set_error_message(error_out, &format!("Failed to parse tools JSON: {}", e));
return SglErrorCode::InvalidArgument;
}
}
}
}
Err(_) => {
set_error_message(error_out, "Invalid UTF-8 in tools_json");
return SglErrorCode::InvalidArgument;
}
};
tools_str
};
// Get the tokenizer from handle
let handle_ref = &*handle;
let tokenizer = &handle_ref.tokenizer;
// Try to downcast to HuggingFaceTokenizer
if let Some(hf_tokenizer) = tokenizer.as_any().downcast_ref::<HuggingFaceTokenizer>() {
// Apply chat template with tools
let empty_docs: [Value; 0] = [];
let tools_slice = tools.as_ref().map(|t| t.as_slice());
let params = ChatTemplateParams {
add_generation_prompt: true,
tools: tools_slice,
documents: Some(&empty_docs),
template_kwargs: None,
};
match hf_tokenizer.apply_chat_template(&messages, params) {
Ok(result) => {
let result_cstr = match CString::new(result) {
Ok(s) => s,
Err(e) => {
set_error_message(error_out, &format!("Failed to create result string: {}", e));
return SglErrorCode::MemoryError;
}
};
*result_out = result_cstr.into_raw();
clear_error_message(error_out);
SglErrorCode::Success
}
Err(e) => {
set_error_message(error_out, &format!("Failed to apply chat template: {}", e));
SglErrorCode::TokenizationError
}
}
} else {
set_error_message(error_out, "Chat template is only supported for HuggingFace tokenizers");
SglErrorCode::TokenizationError
}
}
/// Apply chat template to messages
///
/// # Arguments
/// * `handle` - Tokenizer handle
/// * `messages_json` - JSON string of messages array
/// * `result_out` - Pointer to receive result string (must be freed with sgl_free_string)
/// * `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_tokenizer_apply_chat_template(
handle: *mut TokenizerHandle,
messages_json: *const c_char,
result_out: *mut *mut c_char,
error_out: *mut *mut c_char,
) -> SglErrorCode {
if handle.is_null() || messages_json.is_null() || result_out.is_null() {
set_error_message(error_out, "Invalid arguments: null pointer");
return SglErrorCode::InvalidArgument;
}
let messages_str = match CStr::from_ptr(messages_json).to_str() {
Ok(s) => s,
Err(_) => {
set_error_message(error_out, "Invalid UTF-8 in messages_json");
return SglErrorCode::InvalidArgument;
}
};
// Parse JSON messages
let messages: Vec<Value> = match serde_json::from_str(messages_str) {
Ok(msgs) => msgs,
Err(e) => {
set_error_message(error_out, &format!("Failed to parse messages JSON: {}", e));
return SglErrorCode::InvalidArgument;
}
};
// Get the tokenizer from handle
let handle_ref = &*handle;
let tokenizer = &handle_ref.tokenizer;
// Try to downcast to HuggingFaceTokenizer
if let Some(hf_tokenizer) = tokenizer.as_any().downcast_ref::<HuggingFaceTokenizer>() {
// Apply chat template with default parameters
// Use empty arrays instead of None to avoid template errors
// Set add_generation_prompt to true so the model knows to start generating
let empty_tools: [Value; 0] = [];
let empty_docs: [Value; 0] = [];
let params = ChatTemplateParams {
add_generation_prompt: true, // Important: tells the model to start generating
tools: Some(&empty_tools),
documents: Some(&empty_docs),
template_kwargs: None,
};
match hf_tokenizer.apply_chat_template(&messages, params) {
Ok(result) => {
let result_cstr = match CString::new(result) {
Ok(s) => s,
Err(e) => {
set_error_message(error_out, &format!("Failed to create result string: {}", e));
return SglErrorCode::MemoryError;
}
};
*result_out = result_cstr.into_raw();
clear_error_message(error_out);
SglErrorCode::Success
}
Err(e) => {
set_error_message(error_out, &format!("Failed to apply chat template: {}", e));
SglErrorCode::TokenizationError
}
}
} else {
set_error_message(error_out, "Chat template is only supported for HuggingFace tokenizers");
SglErrorCode::TokenizationError
}
}
/// Decode token IDs to text
///
/// # Arguments
/// * `handle` - Tokenizer handle
/// * `token_ids` - Array of token IDs
/// * `token_count` - Number of tokens
/// * `skip_special_tokens` - Whether to skip special tokens
/// * `result_out` - Pointer to receive result string (must be freed with sgl_free_string)
/// * `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_tokenizer_decode(
handle: *mut TokenizerHandle,
token_ids: *const u32,
token_count: usize,
skip_special_tokens: c_int,
result_out: *mut *mut c_char,
error_out: *mut *mut c_char,
) -> SglErrorCode {
if handle.is_null() || token_ids.is_null() || result_out.is_null() {
set_error_message(error_out, "Invalid arguments: null pointer");
return SglErrorCode::InvalidArgument;
}
if token_count == 0 {
let empty = CString::new("").unwrap();
*result_out = empty.into_raw();
clear_error_message(error_out);
return SglErrorCode::Success;
}
// Convert C array to Rust slice
let token_slice = std::slice::from_raw_parts(token_ids, token_count);
let tokenizer = &(*handle).tokenizer;
match tokenizer.decode(token_slice, skip_special_tokens != 0) {
Ok(text) => {
let result_cstr = match CString::new(text) {
Ok(s) => s,
Err(e) => {
set_error_message(error_out, &format!("Failed to create result string: {}", e));
return SglErrorCode::MemoryError;
}
};
*result_out = result_cstr.into_raw();
clear_error_message(error_out);
SglErrorCode::Success
}
Err(e) => {
set_error_message(error_out, &e.to_string());
SglErrorCode::TokenizationError
}
}
}
/// Free a tokenizer handle
///
/// # Safety
/// This function must only be called once per handle, and the handle must not be used after calling.
#[no_mangle]
pub unsafe extern "C" fn sgl_tokenizer_free(handle: *mut TokenizerHandle) {
if !handle.is_null() {
let _ = Box::from_raw(handle);
}
}

329
third_party/sglang/sgl-model-gateway/bindings/golang/src/tool_parser.rs vendored Normal file
View File

@@ -0,0 +1,329 @@
//! Tool parser FFI functions
use std::ffi::{CStr, CString};
use std::os::raw::{c_char};
use std::ptr;
use std::sync::Arc;
use std::collections::HashMap;
use serde_json::{json, Value};
use tokio::runtime::Runtime;
use once_cell::sync::Lazy;
use smg::tool_parser::{ParserFactory, ToolParser};
use smg::protocols::common::Tool;
use super::error::{SglErrorCode, set_error_message, clear_error_message};
use super::utils::generate_tool_call_id;
/// Global parser factory (initialized once)
static PARSER_FACTORY: Lazy<ParserFactory> = Lazy::new(|| ParserFactory::new());
/// Global tokio runtime for async operations
static RUNTIME: Lazy<Runtime> = Lazy::new(|| {
Runtime::new().expect("Failed to create tokio runtime for tool parser FFI")
});
/// Opaque handle for a tool parser instance
/// Note: For streaming, we need mutable access, so we use Arc<Mutex<>> internally
/// Note: This is an opaque handle, C code doesn't access fields directly
pub struct ToolParserHandle {
parser: Arc<tokio::sync::Mutex<Box<dyn ToolParser>>>,
model: String, // Store model name for ID generation
history_tool_calls_count: usize, // Track tool call count for ID generation
tool_index_to_id: HashMap<usize, String>, // Map tool_index to ID for incremental updates
}
/// Create a tool parser
///
/// # Arguments
/// * `parser_type` - Parser type name (e.g., "json", "llama", "mistral") or model name (e.g., "gpt-4")
/// * `error_out` - Optional pointer to receive error message
///
/// # Returns
/// * Pointer to ToolParserHandle on success, null on failure
#[no_mangle]
pub unsafe extern "C" fn sgl_tool_parser_create(
parser_type: *const c_char,
error_out: *mut *mut c_char,
) -> *mut ToolParserHandle {
if parser_type.is_null() {
set_error_message(error_out, "parser_type cannot be null");
return ptr::null_mut();
}
let type_str = match CStr::from_ptr(parser_type).to_str() {
Ok(s) => s,
Err(_) => {
set_error_message(error_out, "Invalid UTF-8 in parser_type");
return ptr::null_mut();
}
};
// Create parser using factory
// The factory will determine the parser type based on model name or use the provided type
let parser = if let Some(parser_box) = PARSER_FACTORY.registry().create_for_model(type_str) {
parser_box
} else if let Some(parser_box) = PARSER_FACTORY.registry().create_parser(type_str) {
parser_box
} else {
set_error_message(error_out, &format!("Unknown parser type: {}", type_str));
return ptr::null_mut();
};
Box::into_raw(Box::new(ToolParserHandle {
parser: Arc::new(tokio::sync::Mutex::new(parser)),
model: type_str.to_string(),
history_tool_calls_count: 0,
tool_index_to_id: HashMap::new(),
}))
}
/// Parse complete tool calls from text
///
/// # Arguments
/// * `handle` - Tool parser handle
/// * `text` - Input text to parse
/// * `result_json_out` - Pointer to receive JSON result (must be freed with sgl_free_string)
/// * `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_tool_parser_parse_complete(
handle: *mut ToolParserHandle,
text: *const c_char,
result_json_out: *mut *mut c_char,
error_out: *mut *mut c_char,
) -> SglErrorCode {
if handle.is_null() || text.is_null() || result_json_out.is_null() {
set_error_message(error_out, "Invalid arguments: null pointer");
return SglErrorCode::InvalidArgument;
}
let text_str = match CStr::from_ptr(text).to_str() {
Ok(s) => s,
Err(_) => {
set_error_message(error_out, "Invalid UTF-8 in text");
return SglErrorCode::InvalidArgument;
}
};
let handle_ref = &*handle;
let parser = Arc::clone(&handle_ref.parser);
let model = handle_ref.model.clone();
let history_count = handle_ref.history_tool_calls_count;
// Use tokio runtime to run async code
let result = RUNTIME.block_on(async {
let parser_guard = parser.lock().await;
parser_guard.parse_complete(text_str).await
});
match result {
Ok((normal_text, tool_calls)) => {
// Convert Rust ToolCall to OpenAI format
let openai_tool_calls: Vec<Value> = tool_calls
.into_iter()
.enumerate()
.map(|(index, tc)| {
// Generate ID for this tool call
let id = generate_tool_call_id(&model, &tc.function.name, index, history_count);
json!({
"id": id,
"type": "function",
"function": {
"name": tc.function.name,
"arguments": tc.function.arguments
}
})
})
.collect();
// Build result JSON
let result_json = json!({
"normal_text": normal_text,
"tool_calls": openai_tool_calls
});
let result_str = match serde_json::to_string(&result_json) {
Ok(s) => s,
Err(e) => {
set_error_message(error_out, &format!("Failed to serialize JSON: {}", e));
return SglErrorCode::ParsingError;
}
};
let result_cstr = match CString::new(result_str) {
Ok(s) => s,
Err(e) => {
set_error_message(error_out, &format!("Failed to create result string: {}", e));
return SglErrorCode::MemoryError;
}
};
*result_json_out = result_cstr.into_raw();
clear_error_message(error_out);
SglErrorCode::Success
}
Err(e) => {
set_error_message(error_out, &format!("Parse error: {}", e));
SglErrorCode::ParsingError
}
}
}
/// Parse tool calls incrementally from streaming chunks
///
/// # Arguments
/// * `handle` - Tool parser handle
/// * `chunk` - New text chunk from stream
/// * `tools_json` - JSON array of available tools (for validation, can be null/empty)
/// * `result_json_out` - Pointer to receive JSON result (must be freed with sgl_free_string)
/// * `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_tool_parser_parse_incremental(
handle: *mut ToolParserHandle,
chunk: *const c_char,
tools_json: *const c_char,
result_json_out: *mut *mut c_char,
error_out: *mut *mut c_char,
) -> SglErrorCode {
if handle.is_null() || chunk.is_null() || result_json_out.is_null() {
set_error_message(error_out, "Invalid arguments: null pointer");
return SglErrorCode::InvalidArgument;
}
let chunk_str = match CStr::from_ptr(chunk).to_str() {
Ok(s) => s,
Err(_) => {
set_error_message(error_out, "Invalid UTF-8 in chunk");
return SglErrorCode::InvalidArgument;
}
};
// Parse tools JSON if provided
let tools: Vec<Tool> = if !tools_json.is_null() {
let tools_str = match CStr::from_ptr(tools_json).to_str() {
Ok(s) => s,
Err(_) => {
set_error_message(error_out, "Invalid UTF-8 in tools_json");
return SglErrorCode::InvalidArgument;
}
};
match serde_json::from_str::<Vec<Tool>>(tools_str) {
Ok(t) => t,
Err(_) => vec![], // If parsing fails, use empty tools
}
} else {
vec![]
};
let handle_ref = &*handle;
let parser = Arc::clone(&handle_ref.parser);
let model = handle_ref.model.clone();
let history_count = handle_ref.history_tool_calls_count;
// Use tokio runtime to run async code
let result = RUNTIME.block_on(async {
let mut parser_guard = parser.lock().await;
parser_guard.parse_incremental(chunk_str, &tools).await
});
match result {
Ok(streaming_result) => {
// Convert StreamingParseResult to OpenAI format
let handle_mut = &mut *handle;
let openai_tool_calls: Vec<Value> = streaming_result
.calls
.into_iter()
.map(|item| {
// For incremental parsing, we may not have complete tool calls yet
// Generate or reuse ID based on tool_index
let id = if let Some(ref name) = item.name {
// New tool call with name - generate ID and store it
let id = generate_tool_call_id(&model, name, item.tool_index, history_count);
handle_mut.tool_index_to_id.insert(item.tool_index, id.clone());
id
} else {
// Parameter update - reuse existing ID for this tool_index
handle_mut.tool_index_to_id
.get(&item.tool_index)
.cloned()
.unwrap_or_else(|| format!("call_{}", item.tool_index))
};
json!({
"id": id,
"type": "function",
"function": {
"name": item.name.unwrap_or_default(),
"arguments": item.parameters
}
})
})
.collect();
// Build result JSON
let result_json = json!({
"normal_text": streaming_result.normal_text,
"tool_calls": openai_tool_calls
});
let result_str = match serde_json::to_string(&result_json) {
Ok(s) => s,
Err(e) => {
set_error_message(error_out, &format!("Failed to serialize JSON: {}", e));
return SglErrorCode::ParsingError;
}
};
let result_cstr = match CString::new(result_str) {
Ok(s) => s,
Err(e) => {
set_error_message(error_out, &format!("Failed to create result string: {}", e));
return SglErrorCode::MemoryError;
}
};
*result_json_out = result_cstr.into_raw();
clear_error_message(error_out);
SglErrorCode::Success
}
Err(e) => {
set_error_message(error_out, &format!("Parse incremental error: {}", e));
SglErrorCode::ParsingError
}
}
}
/// Reset the parser state for reuse
#[no_mangle]
pub unsafe extern "C" fn sgl_tool_parser_reset(handle: *mut ToolParserHandle) {
if handle.is_null() {
return;
}
let handle_ref = &mut *handle;
let parser = Arc::clone(&handle_ref.parser);
// Reset parser state
RUNTIME.block_on(async {
let mut parser_guard = parser.lock().await;
parser_guard.reset();
});
// Reset history count and tool index mapping
handle_ref.history_tool_calls_count = 0;
handle_ref.tool_index_to_id.clear();
}
/// Free a tool parser handle
#[no_mangle]
pub unsafe extern "C" fn sgl_tool_parser_free(handle: *mut ToolParserHandle) {
if !handle.is_null() {
let _ = Box::from_raw(handle);
}
}

44
third_party/sglang/sgl-model-gateway/bindings/golang/src/utils.rs vendored Normal file
View File

@@ -0,0 +1,44 @@
//! Utility functions for FFI
use uuid::Uuid;
/// Helper function to generate tool call ID (matches router implementation)
pub fn generate_tool_call_id(
model: &str,
function_name: &str,
index: usize,
history_tool_calls_count: usize,
) -> String {
if model.to_lowercase().contains("kimi") {
// KimiK2 format: functions.{name}:{global_index}
format!("functions.{}:{}", function_name, history_tool_calls_count + index)
} else {
// Standard OpenAI format: call_{24-char-uuid}
format!("call_{}", &Uuid::new_v4().simple().to_string()[..24])
}
}
/// Generate tool constraints (placeholder implementation)
///
/// # Arguments
/// * `tools_json` - JSON array of tools
/// * `tool_choice_json` - JSON object representing tool_choice
/// * `constraint_type_out` - Pointer to receive constraint type (e.g., "json_schema")
/// * `constraint_schema_out` - Pointer to receive constraint schema JSON
/// * `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_generate_tool_constraints(
_tools_json: *const std::os::raw::c_char,
_tool_choice_json: *const std::os::raw::c_char,
_constraint_type_out: *mut *mut std::os::raw::c_char,
_constraint_schema_out: *mut *mut std::os::raw::c_char,
error_out: *mut *mut std::os::raw::c_char,
) -> super::error::SglErrorCode {
// Implementation would parse JSON and call generate_tool_constraints
// This is a placeholder
super::error::set_error_message(error_out, "Tool constraint generation not yet implemented in FFI");
super::error::SglErrorCode::UnknownError
}