add tests for custom dataset

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

Author: pancak3

pkg/dataset/.llm-d/test.valid.sqlite3

@@ -18,7 +18,9 @@ package dataset
 
 import (
 	"context"
+	"crypto/sha256"
 	"database/sql"
+	"encoding/hex"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -28,12 +30,11 @@ import (
 	"os"
 	"os/signal"
 	"path/filepath"
-	"strconv"
 	"syscall"
 	"time"
 
 	"github.com/go-logr/logr"
-	"github.com/google/uuid"
+	"github.com/llm-d/llm-d-inference-sim/pkg/common"
 	openaiserverapi "github.com/llm-d/llm-d-inference-sim/pkg/openai-server-api"
 	_ "github.com/mattn/go-sqlite3"
 )
@@ -338,19 +339,35 @@ func unmarshalAllRecords(rows *sql.Rows) ([][]string, error) {
 	return tokensList, nil
 }
 
-func (d *CustomDataset) getRandomTokens(n_gen_tokens int) []string {
-	return []string{"<|random_tokens|>", strconv.Itoa(n_gen_tokens)}
+func (d *CustomDataset) GetPromptHash(req openaiserverapi.CompletionRequest) []byte {
+	hashArray := sha256.Sum256([]byte(req.GetFullPrompt()))
+	return hashArray[:]
 }
 
-func (d *CustomDataset) readTokensFromDB(prompt string, n_gen_tokens int) []string {
-	promptHash := uuid.NewSHA1(uuid.NameSpaceOID, []byte(prompt)).NodeID()
-	rows, err := d.db.Query("SELECT "+genTokensCol+" FROM "+tableName+" WHERE "+promptHashCol+" = ?;", promptHash)
+func (d *CustomDataset) GetPromptHashHex(hashBytes []byte) string {
+	return hex.EncodeToString(hashBytes)
+}
+
+// GetTokens returns tokens and finishReason for the given request and mode (echo or random)
+func (d *CustomDataset) GetTokens(req openaiserverapi.CompletionRequest, mode string) ([]string, string, error) {
+	if mode == common.ModeEcho {
+		return d.echo(req)
+	}
+	nTokensToGen, finishReason := howManyTokensToGen(d.extractMaxTokens(req), req.GetIgnoreEOS())
+	tokens, err := d.GenerateTokens(req, nTokensToGen)
+	return tokens, finishReason, err
+}
+
+func (d *CustomDataset) GenerateTokens(req openaiserverapi.CompletionRequest, nTokens int) ([]string, error) {
+	promptHash := d.GetPromptHash(req)
+	promptHashHex := d.GetPromptHashHex(promptHash)
+	rows, err := d.db.Query("SELECT " + genTokensCol + " FROM " + tableName + " WHERE " + promptHashCol + "=X'" + promptHashHex + "';")
 	if err != nil {
 		if !d.hasWarned {
 			d.Logger.Error(err, "failed to query database. Ensure the prompt hash exists in the dataset. Will generate random tokens instead.")
 			d.hasWarned = true
 		}
-		return d.getRandomTokens(n_gen_tokens)
+		return GenPresetRandomTokens(nTokens), nil
 	}
 	defer func() {
 		if cerr := rows.Close(); cerr != nil {
@@ -361,18 +378,13 @@ func (d *CustomDataset) readTokensFromDB(prompt string, n_gen_tokens int) []stri
 	tokensList, err := unmarshalAllRecords(rows)
 	if err != nil {
 		d.Logger.Error(err, "failed to unmarshal records from database")
-		return d.getRandomTokens(n_gen_tokens)
+		return GenPresetRandomTokens(nTokens), nil
 	}
 
 	if len(tokensList) == 0 {
-		return d.getRandomTokens(n_gen_tokens)
+		return GenPresetRandomTokens(nTokens), nil
 	}
 	d.hasWarned = false
 	randIndex := rand.Intn(len(tokensList))
-	return tokensList[randIndex]
-}
-
-func (d *CustomDataset) GenerateTokens(req openaiserverapi.CompletionRequest, nTokens int) ([]string, error) {
-	tokens := d.readTokensFromDB("", nTokens)
-	return tokens, nil
+	return tokensList[randIndex], nil
 }

pkg/dataset/custom_dataset.go

@@ -21,12 +21,18 @@ import (
 	"os"
 
 	"github.com/go-logr/logr"
+	"github.com/llm-d/llm-d-inference-sim/pkg/common"
+	openaiserverapi "github.com/llm-d/llm-d-inference-sim/pkg/openai-server-api"
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
 
 	_ "github.com/mattn/go-sqlite3"
 )
 
+const (
+	testPrompt = "Hello world!"
+)
+
 var _ = Describe("CustomDataset", func() {
 	var (
 		dataset               *CustomDataset
@@ -90,20 +96,20 @@ var _ = Describe("CustomDataset", func() {
 		err := dataset.Init(validDBPath, "", "")
 		Expect(err).NotTo(HaveOccurred())
 
-		row := dataset.db.QueryRow("SELECT n_gen_tokens FROM llmd WHERE prompt_hash=X'b94d27b9934d041c52e5b721d7373f13a07ed5e79179d63c5d8a0c102a9d00b2';")
+		row := dataset.db.QueryRow("SELECT n_gen_tokens FROM llmd WHERE prompt_hash=X'74bf14c09c038321cba39717dae1dc732823ae4abd8e155959367629a3c109a8';")
 		var n_gen_tokens int
 		err = row.Scan(&n_gen_tokens)
 		Expect(err).NotTo(HaveOccurred())
-		Expect(n_gen_tokens).To(Equal(3))
+		Expect(n_gen_tokens).To(Equal(4))
 
 		var jsonStr string
-		row = dataset.db.QueryRow("SELECT gen_tokens FROM llmd WHERE prompt_hash=X'b94d27b9934d041c52e5b721d7373f13a07ed5e79179d63c5d8a0c102a9d00b2';")
+		row = dataset.db.QueryRow("SELECT gen_tokens FROM llmd WHERE prompt_hash=X'74bf14c09c038321cba39717dae1dc732823ae4abd8e155959367629a3c109a8';")
 		err = row.Scan(&jsonStr)
 		Expect(err).NotTo(HaveOccurred())
 		var tokens []string
 		err = json.Unmarshal([]byte(jsonStr), &tokens)
 		Expect(err).NotTo(HaveOccurred())
-		Expect(tokens).To(Equal([]string{"Hello", "world", "!"}))
+		Expect(tokens).To(Equal([]string{"Hello", " llm-d ", "world", "!"}))
 
 	})
 
@@ -136,4 +142,41 @@ var _ = Describe("CustomDataset", func() {
 		Expect(err).To(HaveOccurred())
 		Expect(err.Error()).To(ContainSubstring("incorrect type"))
 	})
+
+	It("should return correct prompt hash in bytes", func() {
+		// b't\xbf\x14\xc0\x9c\x03\x83!\xcb\xa3\x97\x17\xda\xe1\xdcs(#\xaeJ\xbd\x8e\x15YY6v)\xa3\xc1\t\xa8'
+		expectedHashBytes := []byte{0x74, 0xbf, 0x14, 0xc0, 0x9c, 0x03, 0x83, 0x21, 0xcb, 0xa3, 0x97, 0x17, 0xda, 0xe1, 0xdc, 0x73, 0x28, 0x23, 0xae, 0x4a, 0xbd, 0x8e, 0x15, 0x59, 0x59, 0x36, 0x76, 0x29, 0xa3, 0xc1, 0x09, 0xa8}
+
+		req := &openaiserverapi.TextCompletionRequest{
+			Prompt: testPrompt,
+		}
+
+		hashBytes := dataset.GetPromptHash(req)
+		Expect(hashBytes).To(Equal(expectedHashBytes))
+	})
+
+	It("should return correct prompt hash in hex", func() {
+		expectedHashHex := "74bf14c09c038321cba39717dae1dc732823ae4abd8e155959367629a3c109a8"
+
+		req := &openaiserverapi.TextCompletionRequest{
+			Prompt: testPrompt,
+		}
+
+		hashBytes := dataset.GetPromptHash(req)
+		hashHex := dataset.GetPromptHashHex(hashBytes)
+		Expect(hashHex).To(Equal(expectedHashHex))
+	})
+
+	It("should return tokens for existing prompt", func() {
+		err := dataset.Init(validDBPath, "", "")
+		Expect(err).NotTo(HaveOccurred())
+
+		req := &openaiserverapi.TextCompletionRequest{
+			Prompt: testPrompt,
+		}
+		tokens, finishReason, err := dataset.GetTokens(req, common.ModeRandom)
+		Expect(err).NotTo(HaveOccurred())
+		Expect(finishReason).To(Equal(StopFinishReason))
+		Expect(tokens).To(Equal([]string{"Hello", " llm-d ", "world", "!"}))
+	})
 })

pkg/dataset/custom_dataset_test.go

@@ -288,21 +288,23 @@ func (d *BaseDataset) Close() error {
 	return nil
 }
 
+func (d *BaseDataset) echo(req openaiserverapi.CompletionRequest) ([]string, string, error) {
+	nMaxTokens := d.extractMaxTokens(req)
+	prompt, err := d.extractPrompt(req)
+	if err != nil {
+		return nil, "", err
+	}
+	tokens, finishReason := EchoResponseTokens(nMaxTokens, prompt)
+	return tokens, finishReason, nil
+}
+
 // GetTokens returns tokens and finishReason for the given request and mode (echo or random)
 func (d *BaseDataset) GetTokens(req openaiserverapi.CompletionRequest, mode string) ([]string, string, error) {
-	nMaxTokens := d.extractMaxTokens(req)
 	if mode == common.ModeEcho {
-		prompt, err := d.extractPrompt(req)
-		if err != nil {
-			return nil, "", err
-		}
-		tokens, finishReason := EchoResponseTokens(nMaxTokens, prompt)
-		return tokens, finishReason, nil
+		return d.echo(req)
 	}
-
-	nTokensToGen, finishReason := howManyTokensToGen(nMaxTokens, req.GetIgnoreEOS())
-	tokens, err := d.GenerateTokens(req, nTokensToGen)
-	return tokens, finishReason, err
+	nTokensToGen, finishReason := howManyTokensToGen(d.extractMaxTokens(req), req.GetIgnoreEOS())
+	return GenPresetRandomTokens(nTokensToGen), finishReason, nil
 }
 
 // extractMaxTokens extracts the max tokens from the request
@@ -328,10 +330,3 @@ func (d *BaseDataset) extractPrompt(req openaiserverapi.CompletionRequest) (stri
 	}
 	return "", errors.New("unknown request type")
 }
-
-// GenerateTokens generates random tokens for the required number of tokens
-// other dataset types should override this function
-func (d *BaseDataset) GenerateTokens(req openaiserverapi.CompletionRequest, nTokens int) ([]string, error) {
-	tokens := GenPresetRandomTokens(nTokens)
-	return tokens, nil
-}

pkg/dataset/dataset.go

@@ -69,7 +69,7 @@ var _ = Describe("Dataset", Ordered, func() {
 				Expect(finishReason).To(Equal(StopFinishReason))
 			}
 		})
-		
+
 		It("should return long text", func() {
 			// return required number of tokens although it is higher than ResponseLenMax
 			maxCompletionTokens := int64(ResponseLenMax * 5)

pkg/dataset/dataset_test.go

@@ -244,7 +244,7 @@ func (s *VllmSimulator) initDataset() error {
 	}
 
 	if s.config.Dataset.Path == "" && s.config.Dataset.Url == "" && s.config.Dataset.SavePath == "" {
-		s.logger.Info("No dataset provided, will generate random responses")
+		s.logger.Info("No dataset provided, will generate random responses from preset text")
 		s.dataset = randDataset
 	} else {
 		s.logger.Info("Custom dataset configuration detected")

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

@@ -65,6 +65,8 @@ type CompletionRequest interface {
 	// when the field is true, the prefill phase should be done on remote pod,
 	// whereas decode phase is done on local pod, thus this is a decode request
 	IsDoRemotePrefill() bool
+	// GetFullPrompt returns the full prompt including system and user prompts
+	GetFullPrompt() string
 }
 
 // BaseCompletionRequest contains base completion request related information
@@ -236,6 +238,21 @@ func (req *ChatCompletionRequest) GetLastUserMsg() string {
 	return ""
 }
 
+func (req *ChatCompletionRequest) GetFullPrompt() string {
+	prompt := ""
+	for _, msg := range req.Messages {
+		switch msg.Role {
+		case RoleUser:
+			prompt += "### user:\n" + msg.Content.Raw + "\n"
+		case RoleAssistant:
+			prompt += "### assistant:\n" + msg.Content.Raw + "\n"
+		default:
+			prompt += "### unknown:\n" + msg.Content.Raw + "\n"
+		}
+	}
+	return prompt
+}
+
 // v1/completion
 // TextCompletionRequest defines structure of /completion request
 type TextCompletionRequest struct {
@@ -270,3 +287,7 @@ func (c *TextCompletionRequest) GetToolChoice() string {
 func (c *TextCompletionRequest) GetMaxCompletionTokens() *int64 {
 	return c.MaxTokens
 }
+
+func (t *TextCompletionRequest) GetFullPrompt() string {
+	return "### user:\n" + t.Prompt + "\n"
+}