fix dataset tests

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

Author: pancak3

pkg/dataset/custom_dataset.go

@@ -28,16 +28,18 @@ import (
 	"os"
 	"os/signal"
 	"path/filepath"
+	"strconv"
 	"syscall"
 	"time"
 
 	"github.com/go-logr/logr"
 	"github.com/google/uuid"
+	openaiserverapi "github.com/llm-d/llm-d-inference-sim/pkg/openai-server-api"
 	_ "github.com/mattn/go-sqlite3"
 )
 
 type CustomDataset struct {
-	Dataset
+	BaseDataset
 	db        *sql.DB
 	hasWarned bool
 }
@@ -53,7 +55,7 @@ const (
 	nGenTokensColType = "INTEGER"
 )
 
-func (d CustomDataset) downloadDataset(url string, savePath string) error {
+func (d *CustomDataset) downloadDataset(url string, savePath string) error {
 	// Set up signal handling for Ctrl+C (SIGINT)
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
@@ -184,7 +186,7 @@ func (pr *progressReader) logProgress(pct int) {
 	}
 }
 
-func (d CustomDataset) verifyDB() error {
+func (d *CustomDataset) verifyDB() error {
 	rows, err := d.db.Query("PRAGMA table_info(" + tableName + ");")
 	if err != nil {
 		return fmt.Errorf("failed to query table info for `%s`: %w", tableName, err)
@@ -234,7 +236,7 @@ func (d CustomDataset) verifyDB() error {
 	return nil
 }
 
-func (d CustomDataset) getRecordsCount() (int, error) {
+func (d *CustomDataset) getRecordsCount() (int, error) {
 	var count int
 	err := d.db.QueryRow("SELECT COUNT(" + promptHashCol + ") FROM " + tableName + ";").Scan(&count)
 	if err != nil {
@@ -243,7 +245,7 @@ func (d CustomDataset) getRecordsCount() (int, error) {
 	return count, nil
 }
 
-func (d CustomDataset) connectToDB(path string) error {
+func (d *CustomDataset) connectToDB(path string) error {
 	if d.db != nil {
 		err := d.db.Close()
 		if err != nil {
@@ -277,7 +279,7 @@ func (d CustomDataset) connectToDB(path string) error {
 	return nil
 }
 
-func (d CustomDataset) Init(path string, url string, savePath string) error {
+func (d *CustomDataset) Init(path string, url string, savePath string) error {
 	d.hasWarned = false
 	if path != "" {
 		return d.connectToDB(path)
@@ -312,7 +314,7 @@ func (d CustomDataset) Init(path string, url string, savePath string) error {
 	return errors.New("no dataset path or url provided")
 }
 
-func (d CustomDataset) Close() error {
+func (d *CustomDataset) Close() error {
 	if d.db != nil {
 		return d.db.Close()
 	}
@@ -336,11 +338,11 @@ func unmarshalAllRecords(rows *sql.Rows) ([][]string, error) {
 	return tokensList, nil
 }
 
-func (d CustomDataset) getRandomTokens(n_gen_tokens int) []string {
-	return nil
+func (d *CustomDataset) getRandomTokens(n_gen_tokens int) []string {
+	return []string{"<|random_tokens|>", strconv.Itoa(n_gen_tokens)}
 }
 
-func (d *CustomDataset) GetTokens(prompt string, n_gen_tokens int) []string {
+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)
 	if err != nil {
@@ -369,3 +371,8 @@ func (d *CustomDataset) GetTokens(prompt string, n_gen_tokens int) []string {
 	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
+}

pkg/dataset/custom_dataset_test.go

@@ -27,7 +27,7 @@ import (
 	_ "github.com/mattn/go-sqlite3"
 )
 
-var _ = Describe("Dataset", func() {
+var _ = Describe("CustomDataset", func() {
 	var (
 		dataset               *CustomDataset
 		file_folder           string
@@ -42,7 +42,7 @@ var _ = Describe("Dataset", func() {
 
 	BeforeEach(func() {
 		dataset = &CustomDataset{
-			Dataset: Dataset{
+			BaseDataset: BaseDataset{
 				Logger: logr.Discard(),
 			},
 		}

pkg/dataset/dataset.go

@@ -68,6 +68,15 @@ var chatCompletionFakeResponses = []string{
 	`Give a man a fish and you feed him for a day; teach a man to fish and you feed him for a lifetime`,
 }
 
+type Dataset interface {
+	// Init initializes the dataset using configs
+	Init(path string, url string, savePath string) error
+	// Close closes the dataset
+	Close() error
+	// GetTokens returns tokens for the given request and mode (echo or random)
+	GetTokens(req openaiserverapi.CompletionRequest, mode string) ([]string, string, error)
+}
+
 func init() {
 	cumulativeBucketsProbabilities = make([]float64, len(respLenBucketsProbabilities))
 	sum := 0.0
@@ -267,19 +276,20 @@ func EchoResponseTokens(maxCompletionTokens *int64, text string) ([]string, stri
 	return tokens[0:*maxCompletionTokens], LengthFinishReason
 }
 
-type Dataset struct {
+type BaseDataset struct {
 	Logger logr.Logger
 }
 
-func (d *Dataset) Init(path string, url string, savePath string) error {
+func (d *BaseDataset) Init(path string, url string, savePath string) error {
 	return nil
 }
 
-func (d *Dataset) Close() error {
+func (d *BaseDataset) Close() error {
 	return nil
 }
 
-func (d *Dataset) GetTokens(req openaiserverapi.CompletionRequest, mode string) ([]string, string, error) {
+// 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)
@@ -295,7 +305,10 @@ func (d *Dataset) GetTokens(req openaiserverapi.CompletionRequest, mode string)
 	return tokens, finishReason, err
 }
 
-func (d *Dataset) extractMaxTokens(req openaiserverapi.CompletionRequest) *int64 {
+// extractMaxTokens extracts the max tokens from the request
+// for chat completion - max_completion_tokens field is used
+// for text completion - max_tokens field is used
+func (d *BaseDataset) extractMaxTokens(req openaiserverapi.CompletionRequest) *int64 {
 	if chatReq, ok := req.(*openaiserverapi.ChatCompletionRequest); ok {
 		return chatReq.GetMaxCompletionTokens()
 	} else if textReq, ok := req.(*openaiserverapi.TextCompletionRequest); ok {
@@ -304,7 +317,10 @@ func (d *Dataset) extractMaxTokens(req openaiserverapi.CompletionRequest) *int64
 	return nil
 }
 
-func (d *Dataset) extractPrompt(req openaiserverapi.CompletionRequest) (string, error) {
+// extractPrompt extracts the prompt from the request
+// for chat completion - the last user message is used as the prompt
+// for text completion - the prompt field is used
+func (d *BaseDataset) extractPrompt(req openaiserverapi.CompletionRequest) (string, error) {
 	if chatReq, ok := req.(*openaiserverapi.ChatCompletionRequest); ok {
 		return chatReq.GetLastUserMsg(), nil
 	} else if textReq, ok := req.(*openaiserverapi.TextCompletionRequest); ok {
@@ -313,7 +329,9 @@ func (d *Dataset) extractPrompt(req openaiserverapi.CompletionRequest) (string,
 	return "", errors.New("unknown request type")
 }
 
-func (d *Dataset) GenerateTokens(req openaiserverapi.CompletionRequest, nTokens int) ([]string, error) {
+// 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_suite_test.go

@@ -0,0 +1,13 @@
+package dataset_test
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+func TestDataset(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "Dataset Suite")
+}

pkg/dataset/dataset_test.go

@@ -19,6 +19,7 @@ package dataset
 import (
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/go-logr/logr"
 	"github.com/llm-d/llm-d-inference-sim/pkg/common"
@@ -27,30 +28,31 @@ import (
 	. "github.com/onsi/gomega"
 )
 
-var _ = Describe("Utils", Ordered, func() {
+var _ = Describe("Dataset", Ordered, func() {
 	var (
-		dataset *Dataset
+		dataset *BaseDataset
 	)
 
+	BeforeAll(func() {
+		common.InitRandom(time.Now().UnixNano())
+	})
 	BeforeEach(func() {
-		dataset = &Dataset{
+		dataset = &BaseDataset{
 			Logger: logr.Discard(),
 		}
 	})
 
 	Context("GetRandomTokens", func() {
+
 		It("should return complete text", func() {
-			var n int64
-			req := &openaiserverapi.ChatCompletionRequest{
-				MaxTokens:           &n,
-				MaxCompletionTokens: &n,
-			}
+			req := &openaiserverapi.ChatCompletionRequest{}
 			tokens, finishReason, err := dataset.GetTokens(req, common.ModeRandom)
 			Expect(err).ShouldNot(HaveOccurred())
 			text := strings.Join(tokens, "")
 			Expect(IsValidText(text)).To(BeTrue())
 			Expect(finishReason).Should(Equal(StopFinishReason))
 		})
+
 		It("should return short text", func() {
 			maxCompletionTokens := int64(2)
 			req := &openaiserverapi.ChatCompletionRequest{
@@ -67,6 +69,7 @@ var _ = Describe("Utils", 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/llm-d-inference-sim/simulator.go

@@ -116,8 +116,8 @@ type VllmSimulator struct {
 	pod string
 	// tokenizer is currently used in kv-cache and in /tokenize
 	tokenizer tokenization.Tokenizer
-	// dataset is used for managing dataset files
-	dataset *dataset.Dataset
+	// dataset is used for token generation in responses
+	dataset dataset.Dataset
 }
 
 // New creates a new VllmSimulator instance with the given logger
@@ -216,18 +216,9 @@ func (s *VllmSimulator) startSim(ctx context.Context) error {
 		go s.kvcacheHelper.Run(ctx)
 	}
 
-	if s.config.Dataset.Path == "" && s.config.Dataset.Url == "" && s.config.Dataset.SavePath == "" {
-		s.dataset = nil
-		s.logger.Info("No dataset provided, will generate random responses")
-	} else {
-		dataset := &dataset.Dataset{
-			Logger: s.logger,
-		}
-		err = dataset.Init(s.config.Dataset.Path, s.config.Dataset.Url, s.config.Dataset.SavePath)
-		if err != nil {
-			return err
-		}
-		s.dataset = dataset
+	err = s.initDataset()
+	if err != nil {
+		return fmt.Errorf("dataset initialization error: %w", err)
 	}
 
 	// run request processing workers
@@ -239,13 +230,98 @@ func (s *VllmSimulator) startSim(ctx context.Context) error {
 
 	listener, err := s.newListener()
 	if err != nil {
-		return err
+		s.logger.Error(err, "Failed to create listener")
+		return fmt.Errorf("listener creation error: %w", err)
 	}
 
 	// start the http server with context support
 	return s.startServer(ctx, listener)
 }
 
+func (s *VllmSimulator) initDataset() error {
+	randDataset := &dataset.BaseDataset{
+		Logger: s.logger,
+	}
+
+	if s.config.Dataset.Path == "" && s.config.Dataset.Url == "" && s.config.Dataset.SavePath == "" {
+		s.logger.Info("No dataset provided, will generate random responses")
+		s.dataset = randDataset
+	} else {
+		s.logger.Info("Custom dataset configuration detected")
+		s.dataset = &dataset.CustomDataset{
+			BaseDataset: *randDataset,
+		}
+	}
+
+	if err := s.dataset.Init(s.config.Dataset.Path, s.config.Dataset.Url, s.config.Dataset.SavePath); err != nil {
+		return fmt.Errorf("dataset initialization error: %w", err)
+	}
+	return nil
+}
+
+func (s *VllmSimulator) newListener() (net.Listener, error) {
+	s.logger.Info("Server starting", "port", s.config.Port)
+	listener, err := net.Listen("tcp4", fmt.Sprintf(":%d", s.config.Port))
+	if err != nil {
+		return nil, err
+	}
+	return listener, nil
+}
+
+// startServer starts http server on port defined in command line
+func (s *VllmSimulator) startServer(ctx context.Context, listener net.Listener) error {
+	r := fasthttprouter.New()
+
+	// support completion APIs
+	r.POST("/v1/chat/completions", s.HandleChatCompletions)
+	r.POST("/v1/completions", s.HandleTextCompletions)
+	// supports /models API
+	r.GET("/v1/models", s.HandleModels)
+	// support load/unload of lora adapter
+	r.POST("/v1/load_lora_adapter", s.HandleLoadLora)
+	r.POST("/v1/unload_lora_adapter", s.HandleUnloadLora)
+	// supports /metrics prometheus API
+	r.GET("/metrics", fasthttpadaptor.NewFastHTTPHandler(promhttp.HandlerFor(s.registry, promhttp.HandlerOpts{})))
+	// supports standard Kubernetes health and readiness checks
+	r.GET("/health", s.HandleHealth)
+	r.GET("/ready", s.HandleReady)
+	r.POST("/tokenize", s.HandleTokenize)
+
+	server := fasthttp.Server{
+		ErrorHandler: s.HandleError,
+		Handler:      r.Handler,
+		Logger:       s,
+	}
+
+	// Start server in a goroutine
+	serverErr := make(chan error, 1)
+	go func() {
+		s.logger.Info("HTTP server starting")
+		serverErr <- server.Serve(listener)
+	}()
+
+	// Wait for either context cancellation or server error
+	select {
+	case <-ctx.Done():
+		s.logger.Info("Shutdown signal received, shutting down HTTP server gracefully")
+
+		// Gracefully shutdown the server
+		if err := server.Shutdown(); err != nil {
+			s.logger.Error(err, "Error during server shutdown")
+			return err
+		}
+
+		s.logger.Info("HTTP server stopped")
+		return nil
+
+	case err := <-serverErr:
+		if err != nil {
+			s.logger.Error(err, "HTTP server failed")
+		}
+		return err
+	}
+}
+
 // Print prints to a log, implementation of fasthttp.Logger
 func (s *VllmSimulator) Printf(format string, args ...interface{}) {
 	s.logger.Info("Server error", "msg", fmt.Sprintf(format, args...))