Fix tests and init dataset when loading sim

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

Author: pancak3

pkg/llm-d-inference-sim/.llm-d/test.valid.sqlite3

@@ -81,11 +81,11 @@ func (d *Dataset) downloadDataset(url string, savePath string) error {
 
 	// Progress reader with context
 	pr := &progressReader{
-		Reader:    resp.Body,
-		total:     resp.ContentLength,
-		logger:    d.logger,
-		ctx:       ctx,
-		startTime: time.Now(),
+		Reader:        resp.Body,
+		total:         resp.ContentLength,
+		logger:        d.logger,
+		ctx:           ctx,
+		startTime:     time.Now(),
 		hasShownSpeed: false,
 	}
 
@@ -126,12 +126,12 @@ func (d *Dataset) downloadDataset(url string, savePath string) error {
 // progressReader wraps an io.Reader and logs download progress.
 type progressReader struct {
 	io.Reader
-	total      int64
-	downloaded int64
-	startTime  time.Time
-	lastPct    int
-	logger     logr.Logger
-	ctx        context.Context
+	total         int64
+	downloaded    int64
+	startTime     time.Time
+	lastPct       int
+	logger        logr.Logger
+	ctx           context.Context
 	hasShownSpeed bool
 }
 
@@ -161,7 +161,7 @@ func (pr *progressReader) Read(p []byte) (int, error) {
 func (pr *progressReader) logProgress(pct int) {
 	elapsedTime := time.Since(pr.startTime).Seconds()
 	speed := float64(pr.downloaded) / (1024 * 1024 * elapsedTime)
-	remainingTime := float64(pr.total-pr.downloaded) / (float64(pr.downloaded) / elapsedTime) 
+	remainingTime := float64(pr.total-pr.downloaded) / (float64(pr.downloaded) / elapsedTime)
 	if pct != 100 {
 		pr.logger.Info(fmt.Sprintf("Download progress: %d%%, Speed: %.2f MB/s, Remaining time: %.2fs", pct, speed, remainingTime))
 	} else {
@@ -179,7 +179,17 @@ func (d *Dataset) connectToDB(path string) error {
 	if err != nil {
 		return fmt.Errorf("failed to open database: %w", err)
 	}
-	// Test the connection
+
+	var count int
+	err = d.db.QueryRow("SELECT COUNT(generated) FROM llmd;").Scan(&count)
+	if err != nil {
+		err := d.db.Close()
+		if err != nil {
+			d.logger.Error(err, "failed to close database after query failure")
+		}
+		return fmt.Errorf("failed to query database: %w", err)
+	}
+	d.logger.Info("Database connected successfully", "path", path, "records count", count)
 
 	return nil
 }

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

@@ -76,20 +76,18 @@ var _ = Describe("Dataset", func() {
 		err := dataset.Init(validDBPath, "", "")
 		Expect(err).NotTo(HaveOccurred())
 
-		// read from the db to verify it's valid
-		row := dataset.db.QueryRow("SELECT * FROM t;")
+		row := dataset.db.QueryRow("SELECT generated FROM llmd WHERE prompt_hash=X'b94d27b9934d041c52e5b721d7373f13a07ed5e79179d63c5d8a0c102a9d00b2';")
 		var value string
 		err = row.Scan(&value)
 		Expect(err).NotTo(HaveOccurred())
-		Expect(value).To(Equal("llm-d"))
+		Expect(value).To(Equal("world!"))
 	})
 
 	It("should raise err with invalid DB content", func() {
 		err := dataset.connectToDB(file_folder)
-		Expect(err).NotTo(HaveOccurred())
-
+		Expect(err).To(HaveOccurred())
 		// read from the db to verify it's not valid
-		row := dataset.db.QueryRow("SELECT * FROM t;")
+		row := dataset.db.QueryRow("SELECT * FROM llmd;")
 		var value string
 		err = row.Scan(&value)
 		Expect(err).To(HaveOccurred())

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

@@ -155,14 +155,18 @@ func (s *VllmSimulator) Start(ctx context.Context) error {
 		return err
 	}
 
-	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
+	if s.config.Dataset.Path == "" && s.config.Dataset.Url == "" && s.config.Dataset.SavePath == "" {
+		s.dataset = nil
+	} else {
+		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
 	}
-	s.dataset = dataset
 
 	// For Data Parallel, start data-parallel-size - 1 additional simulators
 	g, ctx := errgroup.WithContext(ctx)