279 lines
6.9 KiB
Go
279 lines
6.9 KiB
Go
// Package vllm provides an Adapter for OpenAI-compatible inference engines
|
|
// such as vLLM and SGLang via the /v1/chat/completions SSE endpoint.
|
|
package vllm
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/url"
|
|
"strings"
|
|
"time"
|
|
|
|
"go.uber.org/zap"
|
|
|
|
"iop/apps/node/internal/runtime"
|
|
"iop/packages/go/config"
|
|
)
|
|
|
|
const Name = "vllm"
|
|
|
|
type Vllm struct {
|
|
instanceName string
|
|
endpoint string
|
|
client *http.Client
|
|
logger *zap.Logger
|
|
}
|
|
|
|
// New creates a vLLM adapter. The optional instanceName is the registry
|
|
// instance key used to disambiguate multiple vLLM instances on the same node.
|
|
func New(cfg config.VllmConf, logger *zap.Logger, instanceName ...string) *Vllm {
|
|
endpoint := strings.TrimRight(cfg.Endpoint, "/")
|
|
name := Name
|
|
if len(instanceName) > 0 && instanceName[0] != "" {
|
|
name = instanceName[0]
|
|
}
|
|
return &Vllm{
|
|
instanceName: name,
|
|
endpoint: endpoint,
|
|
client: &http.Client{},
|
|
logger: logger,
|
|
}
|
|
}
|
|
|
|
func (v *Vllm) Name() string { return Name }
|
|
|
|
func (v *Vllm) Capabilities(_ context.Context) (runtime.Capabilities, error) {
|
|
return runtime.Capabilities{
|
|
AdapterName: Name,
|
|
InstanceKey: v.instanceName,
|
|
Targets: v.fetchTargets(),
|
|
MaxConcurrency: 8,
|
|
}, nil
|
|
}
|
|
|
|
func (v *Vllm) Execute(ctx context.Context, spec runtime.ExecutionSpec, sink runtime.EventSink) error {
|
|
if v.endpoint == "" {
|
|
return fmt.Errorf("vllm adapter: endpoint is required")
|
|
}
|
|
model := strings.TrimSpace(spec.Target)
|
|
if model == "" {
|
|
model = stringInput(spec.Input, "model")
|
|
}
|
|
if model == "" {
|
|
return fmt.Errorf("vllm adapter: target/model is required")
|
|
}
|
|
messages := messagesFromInput(spec.Input)
|
|
if len(messages) == 0 {
|
|
return fmt.Errorf("vllm adapter: messages are required")
|
|
}
|
|
|
|
if err := sink.Emit(ctx, runtime.RuntimeEvent{
|
|
RunID: spec.RunID,
|
|
Type: runtime.EventTypeStart,
|
|
Timestamp: time.Now(),
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
|
|
chatReq := vllmChatRequest{
|
|
Model: model,
|
|
Messages: messages,
|
|
Stream: true,
|
|
}
|
|
body, err := json.Marshal(chatReq)
|
|
if err != nil {
|
|
return fmt.Errorf("vllm adapter: marshal request: %w", err)
|
|
}
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, joinURL(v.endpoint, "/v1/chat/completions"), bytes.NewReader(body))
|
|
if err != nil {
|
|
return fmt.Errorf("vllm adapter: build request: %w", err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
|
|
v.logger.Info("vllm adapter executing",
|
|
zap.String("run_id", spec.RunID),
|
|
zap.String("target", model),
|
|
zap.String("endpoint", v.endpoint),
|
|
)
|
|
resp, err := v.client.Do(req)
|
|
if err != nil {
|
|
_ = emitError(ctx, sink, spec.RunID, fmt.Sprintf("vllm request failed: %v", err))
|
|
return fmt.Errorf("vllm adapter: request: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
msg := readLimited(resp.Body, 4096)
|
|
if msg == "" {
|
|
msg = resp.Status
|
|
}
|
|
_ = emitError(ctx, sink, spec.RunID, fmt.Sprintf("vllm returned %s: %s", resp.Status, msg))
|
|
return fmt.Errorf("vllm adapter: non-2xx response: %s", resp.Status)
|
|
}
|
|
|
|
scanner := bufio.NewScanner(resp.Body)
|
|
outputTokens := 0
|
|
for scanner.Scan() {
|
|
line := scanner.Text()
|
|
if !strings.HasPrefix(line, "data: ") {
|
|
continue
|
|
}
|
|
payload := strings.TrimPrefix(line, "data: ")
|
|
if payload == "[DONE]" {
|
|
return sink.Emit(ctx, runtime.RuntimeEvent{
|
|
RunID: spec.RunID,
|
|
Type: runtime.EventTypeComplete,
|
|
Message: "vllm chat complete",
|
|
Usage: &runtime.UsageStats{
|
|
OutputTokens: outputTokens,
|
|
},
|
|
Timestamp: time.Now(),
|
|
})
|
|
}
|
|
var chunk vllmChatChunk
|
|
if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
|
|
continue
|
|
}
|
|
if len(chunk.Choices) == 0 {
|
|
continue
|
|
}
|
|
content := chunk.Choices[0].Delta.Content
|
|
if content != "" {
|
|
outputTokens += len(strings.Fields(content))
|
|
if err := sink.Emit(ctx, runtime.RuntimeEvent{
|
|
RunID: spec.RunID,
|
|
Type: runtime.EventTypeDelta,
|
|
Delta: content,
|
|
Timestamp: time.Now(),
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
if err := scanner.Err(); err != nil {
|
|
_ = emitError(ctx, sink, spec.RunID, fmt.Sprintf("vllm stream scan failed: %v", err))
|
|
return fmt.Errorf("vllm adapter: scan stream: %w", err)
|
|
}
|
|
_ = emitError(ctx, sink, spec.RunID, "vllm stream ended without [DONE]")
|
|
return fmt.Errorf("vllm adapter: stream ended without [DONE]")
|
|
}
|
|
|
|
func (v *Vllm) fetchTargets() []string {
|
|
if v.endpoint == "" {
|
|
return nil
|
|
}
|
|
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
|
defer cancel()
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, joinURL(v.endpoint, "/v1/models"), nil)
|
|
if err != nil {
|
|
return nil
|
|
}
|
|
resp, err := v.client.Do(req)
|
|
if err != nil {
|
|
return nil
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
return nil
|
|
}
|
|
var models vllmModelsResponse
|
|
if err := json.NewDecoder(resp.Body).Decode(&models); err != nil {
|
|
return nil
|
|
}
|
|
out := make([]string, 0, len(models.Data))
|
|
for _, model := range models.Data {
|
|
if model.ID != "" {
|
|
out = append(out, model.ID)
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
func messagesFromInput(input map[string]any) []vllmMessage {
|
|
if raw, ok := input["messages"].([]any); ok {
|
|
out := make([]vllmMessage, 0, len(raw))
|
|
for _, item := range raw {
|
|
m, ok := item.(map[string]any)
|
|
if !ok {
|
|
continue
|
|
}
|
|
role, _ := m["role"].(string)
|
|
content, _ := m["content"].(string)
|
|
role = strings.TrimSpace(role)
|
|
content = strings.TrimSpace(content)
|
|
if role == "" || content == "" {
|
|
continue
|
|
}
|
|
out = append(out, vllmMessage{Role: role, Content: content})
|
|
}
|
|
if len(out) > 0 {
|
|
return out
|
|
}
|
|
}
|
|
if prompt := strings.TrimSpace(stringInput(input, "prompt")); prompt != "" {
|
|
return []vllmMessage{{Role: "user", Content: prompt}}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func emitError(ctx context.Context, sink runtime.EventSink, runID, msg string) error {
|
|
return sink.Emit(ctx, runtime.RuntimeEvent{
|
|
RunID: runID,
|
|
Type: runtime.EventTypeError,
|
|
Error: msg,
|
|
Timestamp: time.Now(),
|
|
})
|
|
}
|
|
|
|
func stringInput(input map[string]any, key string) string {
|
|
if input == nil {
|
|
return ""
|
|
}
|
|
if v, ok := input[key].(string); ok {
|
|
return v
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func joinURL(baseURL, path string) string {
|
|
u, err := url.Parse(baseURL)
|
|
if err != nil {
|
|
return strings.TrimRight(baseURL, "/") + path
|
|
}
|
|
u.Path = strings.TrimRight(u.Path, "/") + path
|
|
return u.String()
|
|
}
|
|
|
|
func readLimited(r io.Reader, limit int64) string {
|
|
b, _ := io.ReadAll(io.LimitReader(r, limit))
|
|
return strings.TrimSpace(string(b))
|
|
}
|
|
|
|
type vllmChatRequest struct {
|
|
Model string `json:"model"`
|
|
Messages []vllmMessage `json:"messages"`
|
|
Stream bool `json:"stream"`
|
|
}
|
|
|
|
type vllmMessage struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content"`
|
|
}
|
|
|
|
type vllmChatChunk struct {
|
|
Choices []struct {
|
|
Delta struct {
|
|
Content string `json:"content"`
|
|
} `json:"delta"`
|
|
} `json:"choices"`
|
|
}
|
|
|
|
type vllmModelsResponse struct {
|
|
Data []struct {
|
|
ID string `json:"id"`
|
|
} `json:"data"`
|
|
}
|