visor/test/gpu/sglang/sglang.go

// Copyright 2023 The gVisor Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package sglang provides an SGLang API client.
package sglang

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"gvisor.dev/gvisor/pkg/test/dockerutil"
	"gvisor.dev/gvisor/pkg/test/llmutil"
	"gvisor.dev/gvisor/pkg/test/testutil"
)

const (
	// Port is the port used by the sglang server.
	Port = 30000

	// curtQuery is a query that should result in a very curt response.
	curtQuery = `Please reply with the single word: "Hello". Do not reply with any other word.`
)

// SGLang is an sglang client.
type SGLang struct {
	// server is used to perform requests against the server.
	server llmutil.Server

	// logger is used to log.
	logger testutil.Logger
}

// New starts a new SGLang server in the given container,
// then waits for it to serve and returns the client.
func New(ctx context.Context, server llmutil.Server, logger testutil.Logger) (*SGLang, error) {
	started := time.Now()
	llm := &SGLang{
		logger: logger,
		server: server,
	}

	// Wait until serving.
	if err := llm.WaitUntilServing(ctx); err != nil {
		return nil, fmt.Errorf("sglang did not come up for serving: %w", err)
	}
	logger.Logf("SGLang serving API requests after %v", time.Since(started))

	// Run a warmup query to force the model to load.
	_, err := llm.WarmModel(ctx)
	if err != nil {
		return nil, fmt.Errorf("could not warmup the model: %w", err)
	}
	logger.Logf("Loaded sglang model. (%v since container start)", time.Since(started))

	logger.Logf("SGLang successfully initialized in a total of %v", time.Since(started))
	return llm, nil
}

// ModelLoadStats holds metrics about the model loading process.
type ModelLoadStats struct {
	// ClientReportedDuration is the duration to load the model as perceived
	// by the client, measured by HTTP client metrics.
	ClientReportedDuration time.Duration
}

// WarmModel pre-warms a model in memory and keeps it warm for `keepWarmFor`.
// If `unloadFirst` is true, another model will be loaded before loading the
// requested model. This ensures that the model was loaded from a cold state.
func (llm *SGLang) WarmModel(ctx context.Context) (*ModelLoadStats, error) {
	prompt := ZeroTemperaturePrompt(curtQuery)
	prompt.SetMaxNewTokens(1)
	resp, err := llm.Prompt(ctx, prompt)
	if err != nil {
		return nil, llm.withServerLogsErr(ctx, fmt.Errorf("warmup prompt (%s) failed: %w", prompt.Text, err))
	}
	return &ModelLoadStats{
		ClientReportedDuration: resp.metrics.TimeToFirstByte(),
	}, nil
}

// NewDocker returns a new SGLang client talking to an SGLang server that runs
// in a local Docker container.
func NewDocker(ctx context.Context, cont *dockerutil.Container, logger testutil.Logger) (*SGLang, error) {
	opts, err := dockerutil.GPURunOpts(dockerutil.SniffGPUOpts{})
	if err != nil {
		return nil, fmt.Errorf("failed to get GPU run options: %w", err)
	}
	opts.Image = "gpu/sglang"
	started := time.Now()
	if err := cont.Spawn(ctx, opts); err != nil {
		return nil, fmt.Errorf("could not start sglang: %v", err)
	}
	logger.Logf("SGLang container started after %v", time.Since(started))
	ds := &llmutil.DockerServer{
		Container:   cont,
		Logger:      logger,
		Port:        Port,
		ClientImage: "gpu/ollama/client",
	}
	return New(ctx, ds, logger)
}

// instrumentedRequest makes an HTTP request to the sglang API.
// It returns the raw bytestream from the instrumented request logs.
func (llm *SGLang) instrumentedRequest(ctx context.Context, method, endpoint string, data []byte) ([]byte, error) {
	if endpoint != "" && !strings.HasPrefix(endpoint, "/") {
		return nil, fmt.Errorf("endpoint must be empty or start with '/', got %q", endpoint)
	}
	argvFn := func(hostPort string) []string {
		argv := []string{
			"httpclient",
			fmt.Sprintf("--method=%s", method),
			fmt.Sprintf("--url=%s%s", hostPort, endpoint),
			"--strip_prefix=data: ",
		}
		if data != nil {
			argv = append(argv, fmt.Sprintf("--post_base64=%s", base64.StdEncoding.EncodeToString(data)))
		}
		if ctxDeadline, hasDeadline := ctx.Deadline(); hasDeadline {
			argv = append(argv, fmt.Sprintf("--timeout=%v", time.Until(ctxDeadline)))
		}
		return argv
	}
	rawResponse, err := llm.server.InstrumentedRequest(ctx, argvFn)
	if err != nil {
		return nil, fmt.Errorf("%s: %w", endpoint, err)
	}
	return rawResponse, nil
}

// jsonGet performs a JSON HTTP GET request.
func jsonGet[Out any](ctx context.Context, llm *SGLang, endpoint string) (*llmutil.APIResponse[Out], error) {
	out, err := llm.instrumentedRequest(ctx, "GET", endpoint, nil)
	if err != nil {
		return nil, fmt.Errorf("GET %q failed: %w", endpoint, err)
	}
	return llmutil.MakeAPIResponse[Out](out)
}

// jsonPost performs a JSON HTTP POST request.
func jsonPost[In, Out any](ctx context.Context, llm *SGLang, endpoint string, input In) (*llmutil.APIResponse[Out], error) {
	query, err := json.Marshal(input)
	if err != nil {
		return nil, fmt.Errorf("could not marshal input %v: %w", input, err)
	}
	out, err := llm.instrumentedRequest(ctx, "POST", endpoint, query)
	if err != nil {
		return nil, fmt.Errorf("POST %q %v failed: %w", endpoint, string(query), err)
	}

	return llmutil.MakeAPIResponse[Out](out)
}

// WaitUntilServing waits until sglang is serving, or the context expires.
func (llm *SGLang) WaitUntilServing(ctx context.Context) error {
	for ctx.Err() == nil {
		out, err := llm.instrumentedRequest(ctx, "GET", "/health", nil)
		if err != nil {
			continue
		}
		if strings.Contains(string(out), "uvicorn") {
			return nil
		}
	}
	return fmt.Errorf("sglang did not respond: %w", ctx.Err())
}

// temperatureOption is the temperature option that most models have
// which controls how free they are from deviating from their most-likely
// token chain.
const temperatureOption = "temperature"
const maxNewTokensOption = "max_new_tokens"

// RaiseTemperature increases the "temperature" option of the model,
// if any.
func (p *Prompt) RaiseTemperature() {
	temp, ok := p.Options[temperatureOption]
	if !ok {
		temp = float64(0.0)
	}
	if p.Options == nil {
		p.Options = map[string]any{}
	}
	p.Options[temperatureOption] = min(1.0, temp.(float64)*2+.025)
}

// Copy returns a copy of the prompt.
func (p *Prompt) Copy() *Prompt {
	promptCopy := *p
	promptCopy.Options = make(map[string]any, len(p.Options))
	for k, v := range p.Options {
		promptCopy.Options[k] = v
	}
	return &promptCopy
}

// SetTemperature sets the "temperature" option of the prompt to the given
// value.
func (p *Prompt) SetTemperature(temperature float64) {
	if p.Options == nil {
		p.Options = map[string]any{}
	}
	p.Options[temperatureOption] = temperature
}

// SetMaxNewTokens sets the "max_new_tokens" option of the prompt to the given
// value.
func (p *Prompt) SetMaxNewTokens(maxNewTokens int) {
	if p.Options == nil {
		p.Options = map[string]any{}
	}
	p.Options[maxNewTokensOption] = maxNewTokens
}

// ZeroTemperaturePrompt returns a Prompt with the given text and an initial
// temperature setting of zero. This setting allows for consistent settings.
func ZeroTemperaturePrompt(text string) *Prompt {
	return &Prompt{
		Text: text,
		Options: map[string]any{
			temperatureOption: 0.0,
		},
	}
}

// Prompt is an sglang prompt.
type Prompt struct {

	// Text is the prompt string.
	// Common leading whitespace will be removed.
	Text string

	// images is a set of attached images.
	// Use AddImage to add an image.
	images [][]byte

	// Options maps parameter names to JSON-compatible values.
	Options map[string]any
}

// AddImage attaches an image to the prompt.
// Returns itself for chainability.
func (p *Prompt) AddImage(data []byte) *Prompt {
	p.images = append(p.images, data)
	return p
}

// CleanQuery removes common whitespace from query lines, and all
// leading/ending whitespace-only lines.
// It is useful to be able to specify query string as indented strings
// without breaking visual continuity in Go code.
// For example (where dots are spaces):
//
// """\n
// ..The Quick Brown Fox\n
// ..Jumps Over\n
// ....The Lazy Dog\n
// ."""
//
// becomes:
// Jumps Over\n
// ..The Lazy Dog"""
func (p *Prompt) CleanQuery() string {
	lines := strings.Split(p.Text, "\n")

	// Trim lines at the beginning and end that are only whitespace.
	trimmedLines := make([]string, 0, len(lines))
	startedNonWhitespace := false
	var block []string
	for _, line := range lines {
		trimmedLine := strings.TrimSpace(line)
		if !startedNonWhitespace && trimmedLine != "" {
			startedNonWhitespace = true
		}
		if startedNonWhitespace {
			block = append(block, line)
		}
		if trimmedLine != "" {
			trimmedLines = append(trimmedLines, block...)
			block = block[:0]
		}
	}

	// Find longest common whitespace prefix.
	if len(trimmedLines) == 0 {
		return ""
	}
	trimmedFirstLine := strings.TrimSpace(trimmedLines[0])
	common := []rune(trimmedLines[0][:strings.Index(trimmedLines[0], trimmedFirstLine)])
	for ; len(common) > 0; common = common[:len(common)-1] {
		allMatch := true
		for _, line := range trimmedLines[1:] {
			if strings.TrimSpace(line) == "" {
				continue // Ignore whitespace-only or empty lines.
			}
			if !strings.HasPrefix(line, string(common)) {
				allMatch = false
				break
			}
		}
		if allMatch {
			break
		}
	}

	// Remove it.
	if len(common) > 0 {
		for i, line := range trimmedLines {
			trimmedLines[i] = strings.TrimPrefix(line, string(common))
		}
	}

	return strings.Join(trimmedLines, "\n")
}

// WithHotterModel returns a copy of this prompt with the same model having
// a higher temperature.
func (p *Prompt) WithHotterModel() *Prompt {
	promptCopy := p.Copy()
	promptCopy.RaiseTemperature()
	return promptCopy
}

// Request defines the structure for the JSON payload.
// https://docs.sglang.ai/basic_usage/sampling_params.html
type promptJSON struct {
	Prompt  string         `json:"text,omitempty"`
	Options map[string]any `json:"sampling_params"`
	Images  []string       `json:"image_data"`
	Stream  bool           `json:"stream"`
}

// json encodes this prompt to the JSON format expected by SGLang.
func (p *Prompt) json() promptJSON {
	images := make([]string, len(p.images))
	for i, image := range p.images {
		images[i] = base64.StdEncoding.EncodeToString(image)
	}
	return promptJSON{
		Prompt:  p.CleanQuery(),
		Images:  images,
		Stream:  true,
		Options: p.Options,
	}
}

// MetaInfo contains the meta information about the response.
// Every token has a meta info.
type MetaInfo struct {
	// E2ELatency only exists in the last token.
	E2ELatency float64 `json:"e2e_latency"`
}

// responseJSON is the JSON-format response from sglang about a prompt.
// Note that in `streamed` mode, the `Response` field contains a single token.
// To recover the whole response, all `Response` fields must be concatenated
// until the last `responseJSON`, identified as such by the `Done` field.
type responseJSON struct {
	Text     string   `json:"text"`
	MetaInfo MetaInfo `json:"meta_info"`
}

// Response represents a response to a query from SGLang.
type Response struct {
	data    []*responseJSON
	metrics llmutil.ResponseMetrics
}

// Done returns whether the response was completely generated.
func (r *Response) Done() bool {
	if len(r.data) == 0 {
		return false
	}
	return r.data[len(r.data)-1].MetaInfo.E2ELatency > 0
}

// NumTokens returns the number of tokens in the response.
func (r *Response) NumTokens() int {
	return len(r.data)
}

// TimeToFirstToken returns the time it took between the request starting
// and the first token being received by the client.
func (r *Response) TimeToFirstToken() time.Duration {
	if !r.Done() {
		return -1
	}
	return r.metrics.FirstByteRead.Sub(r.metrics.RequestSent)
}

// TimeToLastToken returns the time it took between the request starting
// and the last token being received by the client.
func (r *Response) TimeToLastToken() time.Duration {
	if !r.Done() {
		return -1
	}
	return r.metrics.LastByteRead.Sub(r.metrics.RequestSent)
}

// OutputTokensPerSecond computes the average number of output tokens
// generated per second.
func (r *Response) OutputTokensPerSecond() float64 {
	if !r.Done() {
		return -1
	}
	return float64(len(r.data)-1) / r.data[len(r.data)-1].MetaInfo.E2ELatency
}

// E2ELatency returns the elapsed time between when start_time was recorded and
// the current moment in seconds.
// See https://github.com/sgl-project/sglang/blob/4a2768a86b2905b9b7f19d415261b9d4af639e19/sgl-router/src/routers/grpc/regular/streaming.rs#L904
func (r *Response) E2ELatency() float64 {
	if len(r.data) == 0 {
		return 0
	}
	return r.data[len(r.data)-1].MetaInfo.E2ELatency
}

// String returns the response text, if it is done.
func (r *Response) String() string {
	if len(r.data) == 0 {
		return "<EMPTY>"
	}
	response := r.data[len(r.data)-1]
	if response.MetaInfo.E2ELatency > 0 {
		return response.Text
	}
	return "<NOT DONE>"
}

// Text returns the body of the response, if it is done.
func (r *Response) Text() string {
	if !r.Done() {
		return ""
	}
	return r.String()
}

// withServerLogsErr adds server logs to `err` if possible.
func (llm *SGLang) withServerLogsErr(ctx context.Context, err error) error {
	if err == nil {
		return nil
	}
	if ctx.Err() != nil {
		return fmt.Errorf("%w (+ context err: %v)", err, ctx.Err())
	}
	serverLogs, logsErr := llm.server.Logs(ctx)
	if logsErr != nil {
		return fmt.Errorf("%w (could not get server logs: %v)", err, logsErr)
	}
	if serverLogs != "" {
		return fmt.Errorf("%w; sglang server logs:\n%v\n(end of sglang server logs)", err, serverLogs)
	}
	return fmt.Errorf("%w (server logs are empty)", err)
}

// Prompt returns the result of prompting the given `model` with `prompt`.
func (llm *SGLang) Prompt(ctx context.Context, prompt *Prompt) (*Response, error) {
	resp, err := jsonPost[promptJSON, responseJSON](ctx, llm, "/generate", prompt.json())
	if err != nil {
		return nil, llm.withServerLogsErr(ctx, fmt.Errorf("prompt (%q) request failed: %w", prompt.CleanQuery(), err))
	}
	return &Response{data: resp.Objects, metrics: resp.Metrics}, nil
}

// PromptUntil repeatedly issues a prompt until `iterate` returns a nil error.
// `iterate` may optionally return an updated `Prompt` which will be used to
// follow up. This is useful to work around the flakiness of LLMs in tests.
func (llm *SGLang) PromptUntil(ctx context.Context, prompt *Prompt, iterate func(*Prompt, *Response) (*Prompt, error)) (*Response, error) {
	var lastResponse *Response
	var lastError error
	attempts := 0
	for ctx.Err() == nil {
		response, err := llm.Prompt(ctx, prompt)
		if err != nil {
			return nil, fmt.Errorf("prompt request failed: %w", err)
		}
		attempts++
		newPrompt, err := iterate(prompt, response)
		if err == nil {
			return response, nil
		}
		if newPrompt != nil {
			prompt = newPrompt
		}
		lastResponse = response
		lastError = err
	}
	return nil, fmt.Errorf("response %q (attempt #%d with prompt %v) did not match predicate: %v", lastResponse, attempts, prompt, lastError)
}