Move token generation to simulator

Signed-off-by: Qifan Deng <[email protected]>

Author: pancak3

.gitignore

@@ -8,4 +8,4 @@ vendor
 *.test
 manifests/dev-config.yaml
 pkg/llm-d-inference-sim/.llm-d
-.llm-d/
+pkg/llm-d-inference-sim/tests-tmp/

pkg/llm-d-inference-sim/simulator.go

@@ -339,7 +339,7 @@ func (s *VllmSimulator) reqProcessingWorker(ctx context.Context, id int) {
 			if toolCalls == nil && err == nil {
 				// Either no tool calls were defined, or we randomly chose not to create tool calls,
 				// so we generate a response text.
-				responseTokens, finishReason, completionTokens, err = req.CreateResponseText(s.config.Mode)
+				responseTokens, finishReason, completionTokens, err = s.generateTokens(req)
 			}
 			if err != nil {
 				prefix := ""
@@ -514,8 +514,6 @@ func (s *VllmSimulator) createModelsResponse() *vllmapi.ModelsResponse {
 
 	return &modelsResp
 }
-<<<<<<< HEAD
-=======
 
 // HandleHealth http handler for /health
 func (s *VllmSimulator) HandleHealth(ctx *fasthttp.RequestCtx) {
@@ -599,4 +597,38 @@ func (s *VllmSimulator) GetPrefillTimePerToken() int {
 func (s *VllmSimulator) GetInterTokenLatency() int {
 	return int(float64(s.config.InterTokenLatency) * s.getCurrFactor())
 }
->>>>>>> 482434e (Show config in yaml)
+
+// generateTokens creates and returns response payload based on this request,
+// i.e., an array of generated tokens, the finish reason, and the number of created tokens
+func (s *VllmSimulator) generateTokens(req openaiserverapi.CompletionRequest) ([]string, string, int, error) {
+	// if req is ChatCompletionRequest
+	ignoreEOS := req.GetIgnoreEOS()
+	var maxTokens *int64
+	var prompt string
+
+	if chatReq, ok := req.(*openaiserverapi.ChatCompletionRequest); ok {
+		maxTokens = chatReq.GetMaxCompletionTokens()
+		prompt = chatReq.GetLastUserMsg()
+	} else if textReq, ok := req.(*openaiserverapi.TextCompletionRequest); ok {
+		maxTokens = textReq.MaxTokens
+		prompt = textReq.GetPrompt()
+	} else {
+		return nil, "", 0, fmt.Errorf("unknown request type: %T", req)
+	}
+
+	maxTokensValue, err := common.GetMaxTokens(nil, maxTokens)
+	if err != nil {
+		return nil, "", 0, err
+	}
+
+	var text, finishReason string
+	if s.config.Mode == common.ModeEcho {
+		text, finishReason = common.GetResponseText(maxTokensValue, prompt)
+	} else {
+		text, finishReason = common.GetRandomResponseText(maxTokensValue, ignoreEOS)
+	}
+
+	tokens := common.Tokenize(text)
+	return tokens, finishReason, len(tokens), nil
+}
+>>>>>>> 48ec8bc (Move token generation to simulator)

pkg/openai-server-api/request.go

@@ -33,10 +33,6 @@ const (
 type CompletionRequest interface {
 	// GetRequestID returns the unique request id
 	GetRequestID() string
-	// CreateResponseText creates and returns response payload based on this request,
-	// i.e., an array of generated tokens, the finish reason, and the number of created
-	// tokens
-	CreateResponseText(mode string) ([]string, string, int, error)
 	// IsStream returns boolean that defines is response should be streamed
 	IsStream() bool
 	// GetModel returns model name as defined in the request
@@ -230,7 +226,7 @@ func (c *ChatCompletionRequest) GetMaxCompletionTokens() *int64 {
 
 // getLastUserMsg returns last message from this request's messages with user role,
 // if does not exist - returns an empty string
-func (req *ChatCompletionRequest) getLastUserMsg() string {
+func (req *ChatCompletionRequest) GetLastUserMsg() string {
 	for i := len(req.Messages) - 1; i >= 0; i-- {
 		if req.Messages[i].Role == RoleUser {
 			return req.Messages[i].Content.PlainText()
@@ -240,31 +236,6 @@ func (req *ChatCompletionRequest) getLastUserMsg() string {
 	return ""
 }
 
-// CreateResponseText creates and returns response payload based on this request,
-// i.e., an array of generated tokens, the finish reason, and the number of created
-// tokens
-func (req ChatCompletionRequest) CreateResponseText(mode string) ([]string, string, int, error) {
-	return generateResponseText(mode, req.GetMaxCompletionTokens(), req.getLastUserMsg(), req.GetIgnoreEOS())
-}
-
-// Helper function to generate response text
-func generateResponseText(mode string, maxTokens *int64, prompt string, ignoreEOS bool) ([]string, string, int, error) {
-	maxTokensValue, err := common.GetMaxTokens(nil, maxTokens)
-	if err != nil {
-		return nil, "", 0, err
-	}
-
-	var text, finishReason string
-	if mode == common.ModeEcho {
-		text, finishReason = common.GetResponseText(maxTokensValue, prompt)
-	} else {
-		text, finishReason = common.GetRandomResponseText(maxTokensValue, ignoreEOS)
-	}
-
-	tokens := common.Tokenize(text)
-	return tokens, finishReason, len(tokens), nil
-}
-
 // v1/completion
 // TextCompletionRequest defines structure of /completion request
 type TextCompletionRequest struct {
@@ -299,10 +270,3 @@ func (c *TextCompletionRequest) GetToolChoice() string {
 func (c *TextCompletionRequest) GetMaxCompletionTokens() *int64 {
 	return c.MaxTokens
 }
-
-// CreateResponseText creates and returns response payload based on this request,
-// i.e., an array of generated tokens, the finish reason, and the number of created
-// tokens
-func (req TextCompletionRequest) CreateResponseText(mode string) ([]string, string, int, error) {
-	return generateResponseText(mode, req.MaxTokens, req.Prompt, req.GetIgnoreEOS())
-}