Remove unnecessary config

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

Author: pancak3

pkg/common/config.go

@@ -190,10 +190,9 @@ type DatasetConf struct {
 	// when path is empty Url will be checked
 	Path string `yaml:"path" json:"path"`
 	// Url is the URL to download the sqlite db file if set, default is empty
+	// if Path is not provided and Url is provided, the file will be downloaded
+	// to "USER_HOME/.llm-d/dataset.db"
 	Url string `yaml:"url" json:"url"`
-	// SavePath is the local path to save the downloaded sqlite db file
-	// if Url is set but SavePath is not, "USER_HOME/.llm-d/dataset.db" will be used
-	SavePath string `yaml:"save-path" json:"save-path"`
 }
 
 type Metrics struct {

pkg/dataset/custom_dataset.go

@@ -55,7 +55,7 @@ const (
 	nGenTokensColType = "INTEGER"
 )
 
-func (d *CustomDataset) downloadDataset(url string, savePath string) error {
+func (d *CustomDataset) downloadDataset(url string, path string) error {
 	// Set up signal handling for Ctrl+C (SIGINT)
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
@@ -70,7 +70,7 @@ func (d *CustomDataset) downloadDataset(url string, savePath string) error {
 		cancel()
 	}()
 
-	out, err := os.Create(savePath)
+	out, err := os.Create(path)
 	if err != nil {
 		return err
 	}
@@ -109,7 +109,7 @@ func (d *CustomDataset) downloadDataset(url string, savePath string) error {
 	written, err := io.Copy(out, pr)
 	if err != nil {
 		// Remove incomplete file
-		cerr := os.Remove(savePath)
+		cerr := os.Remove(path)
 		if cerr != nil {
 			d.Logger.Error(cerr, "failed to remove incomplete file after download")
 		}
@@ -121,7 +121,7 @@ func (d *CustomDataset) downloadDataset(url string, savePath string) error {
 	}
 	// Check if file size is zero or suspiciously small
 	if written == 0 {
-		cerr := os.Remove(savePath)
+		cerr := os.Remove(path)
 		if cerr != nil {
 			d.Logger.Error(cerr, "failed to remove empty file after download")
 		}
@@ -130,7 +130,7 @@ func (d *CustomDataset) downloadDataset(url string, savePath string) error {
 
 	// Ensure file is fully flushed and closed before returning success
 	if err := out.Sync(); err != nil {
-		cerr := os.Remove(savePath)
+		cerr := os.Remove(path)
 		if cerr != nil {
 			d.Logger.Error(cerr, "failed to remove incomplete file after download")
 		}
@@ -279,37 +279,39 @@ 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) error {
 	d.hasWarned = false
-	if path != "" {
+	if path != "" && url == "" {
+		d.Logger.Info("Using dataset from", "path", path)
 		return d.connectToDB(path)
 	}
 	if url != "" {
-		if savePath == "" {
+		d.Logger.Info("Url detected", "url", url)
+		if path == "" {
 			user, err := os.UserHomeDir()
 			if err != nil {
 				return fmt.Errorf("failed to get user home directory: %w", err)
 			}
-			savePath = filepath.Join(user, ".llm-d", "dataset.sqlite3")
+			path = filepath.Join(user, ".llm-d", "dataset.sqlite3")
+			d.Logger.Info("Using default for dataset", "path", path)
+		} else {
+			d.Logger.Info("Using provided path for dataset", "path", path)
 		}
 
-		_, err := os.Stat(savePath)
+		_, err := os.Stat(path)
 		if err != nil {
 			// file does not exist, download it
-			folder := filepath.Dir(savePath)
+			folder := filepath.Dir(path)
 			err := os.MkdirAll(folder, 0755)
 			if err != nil {
 				return fmt.Errorf("failed to create parent directory: %w", err)
 			}
-			d.Logger.Info("Downloading dataset from URL", "url", url, "to", savePath)
-			err = d.downloadDataset(url, savePath)
+			err = d.downloadDataset(url, path)
 			if err != nil {
 				return fmt.Errorf("failed to download dataset: %w", err)
 			}
 		}
-		d.Logger.Info("Using dataset from", "path", savePath)
-
-		return d.connectToDB(savePath)
+		return d.connectToDB(path)
 	}
 	return errors.New("no dataset path or url provided")
 }

pkg/dataset/custom_dataset_test.go

@@ -38,7 +38,7 @@ var _ = Describe("CustomDataset", Ordered, func() {
 	var (
 		dataset               *CustomDataset
 		file_folder           string
-		savePath              string
+		path                  string
 		validDBPath           string
 		pathToInvalidDB       string
 		pathNotExist          string
@@ -58,7 +58,7 @@ var _ = Describe("CustomDataset", Ordered, func() {
 			},
 		}
 		file_folder = ".llm-d"
-		savePath = file_folder + "/test.sqlite3"
+		path = file_folder + "/test.sqlite3"
 		err := os.MkdirAll(file_folder, os.ModePerm)
 		Expect(err).NotTo(HaveOccurred())
 		validDBPath = file_folder + "/test.valid.sqlite3"
@@ -83,22 +83,22 @@ var _ = Describe("CustomDataset", Ordered, func() {
 
 	It("should download file from url", func() {
 		url := "https://llm-d.ai"
-		err := dataset.downloadDataset(url, savePath)
+		err := dataset.downloadDataset(url, path)
 		Expect(err).NotTo(HaveOccurred())
-		_, err = os.Stat(savePath)
+		_, err = os.Stat(path)
 		Expect(err).NotTo(HaveOccurred())
-		err = os.Remove(savePath)
+		err = os.Remove(path)
 		Expect(err).NotTo(HaveOccurred())
 	})
 
 	It("should not download file from url", func() {
 		url := "https://256.256.256.256" // invalid url
-		err := dataset.downloadDataset(url, savePath)
+		err := dataset.downloadDataset(url, path)
 		Expect(err).To(HaveOccurred())
 	})
 
 	It("should successfully init dataset", func() {
-		err := dataset.Init(validDBPath, "", "")
+		err := dataset.Init(validDBPath, "")
 		Expect(err).NotTo(HaveOccurred())
 
 		row := dataset.db.QueryRow("SELECT n_gen_tokens FROM llmd WHERE prompt_hash=X'74bf14c09c038321cba39717dae1dc732823ae4abd8e155959367629a3c109a8';")
@@ -173,7 +173,7 @@ var _ = Describe("CustomDataset", Ordered, func() {
 	})
 
 	It("should return tokens for existing prompt", func() {
-		err := dataset.Init(validDBPath, "", "")
+		err := dataset.Init(validDBPath, "")
 		Expect(err).NotTo(HaveOccurred())
 
 		req := &openaiserverapi.TextCompletionRequest{

pkg/dataset/dataset.go

@@ -70,7 +70,7 @@ var chatCompletionFakeResponses = []string{
 
 type Dataset interface {
 	// Init initializes the dataset using configs
-	Init(path string, url string, savePath string) error
+	Init(path string, url string) error
 	// Close closes the dataset
 	Close() error
 	// GetTokens returns tokens for the given request and mode (echo or random)
@@ -280,7 +280,7 @@ type BaseDataset struct {
 	Logger logr.Logger
 }
 
-func (d *BaseDataset) Init(path string, url string, savePath string) error {
+func (d *BaseDataset) Init(path string, url string) error {
 	return nil
 }
 

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

@@ -243,7 +243,7 @@ func (s *VllmSimulator) initDataset() error {
 		Logger: s.logger,
 	}
 
-	if s.config.Dataset.Path == "" && s.config.Dataset.Url == "" && s.config.Dataset.SavePath == "" {
+	if s.config.Dataset.Path == "" && s.config.Dataset.Url == "" {
 		s.logger.Info("No dataset provided, will generate random responses from preset text")
 		s.dataset = randDataset
 	} else {
@@ -253,7 +253,7 @@ func (s *VllmSimulator) initDataset() error {
 		}
 	}
 
-	if err := s.dataset.Init(s.config.Dataset.Path, s.config.Dataset.Url, s.config.Dataset.SavePath); err != nil {
+	if err := s.dataset.Init(s.config.Dataset.Path, s.config.Dataset.Url); err != nil {
 		return fmt.Errorf("dataset initialization error: %w", err)
 	}
 	return nil