Init dataset when sim starts and show downloading speed of url

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

Author: pancak3

pkg/common/config.go

@@ -186,7 +186,7 @@ type Dataset struct {
 	// Url is the URL to download the sqlite db file if set, default is empty
 	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, "~/.llmd/dataset.db" will be used
+	// 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"`
 }
 

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

@@ -27,8 +27,10 @@ import (
 	"os/signal"
 	"path/filepath"
 	"syscall"
+	"time"
 
 	"github.com/go-logr/logr"
+	_ "github.com/mattn/go-sqlite3"
 )
 
 type Dataset struct {
@@ -79,10 +81,12 @@ 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,
+		Reader:    resp.Body,
+		total:     resp.ContentLength,
+		logger:    d.logger,
+		ctx:       ctx,
+		startTime: time.Now(),
+		hasShownSpeed: false,
 	}
 
 	written, err := io.Copy(out, pr)
@@ -124,9 +128,11 @@ type progressReader struct {
 	io.Reader
 	total      int64
 	downloaded int64
+	startTime  time.Time
 	lastPct    int
 	logger     logr.Logger
 	ctx        context.Context
+	hasShownSpeed bool
 }
 
 func (pr *progressReader) Read(p []byte) (int, error) {
@@ -139,13 +145,30 @@ func (pr *progressReader) Read(p []byte) (int, error) {
 	pr.downloaded += int64(n)
 	if pr.total > 0 {
 		pct := int(float64(pr.downloaded) * 100 / float64(pr.total))
-		if pct != pr.lastPct && pct%10 == 0 { // log every 10%
-			pr.logger.Info(fmt.Sprintf("Download progress: %d%%", pct))
+		if !pr.hasShownSpeed && time.Since(pr.startTime).Seconds() > 2 {
+			pr.hasShownSpeed = true
+			pr.logProgress(pct)
+			pr.lastPct = pct
+		}
+		if pct != pr.lastPct && pct%10 == 0 {
+			pr.logProgress(pct)
 			pr.lastPct = pct
 		}
 	}
 	return n, err
 }
+
+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) 
+	if pct != 100 {
+		pr.logger.Info(fmt.Sprintf("Download progress: %d%%, Speed: %.2f MB/s, Remaining time: %.2fs", pct, speed, remainingTime))
+	} else {
+		pr.logger.Info(fmt.Sprintf("Download completed: 100%%, Average Speed: %.2f MB/s, Total time: %.2fs", speed, elapsedTime))
+	}
+}
+
 func (d *Dataset) connectToDB(path string) error {
 	// check if file exists
 	_, err := os.Stat(path)
@@ -156,6 +179,8 @@ func (d *Dataset) connectToDB(path string) error {
 	if err != nil {
 		return fmt.Errorf("failed to open database: %w", err)
 	}
+	// Test the connection
+
 	return nil
 }
 
@@ -165,7 +190,11 @@ func (d *Dataset) Init(path string, url string, savePath string) error {
 	}
 	if url != "" {
 		if savePath == "" {
-			savePath = "~/.llmd/dataset.sqlite3"
+			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")
 		}
 
 		_, err := os.Stat(savePath)

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

@@ -116,6 +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
 }
 
 // New creates a new VllmSimulator instance with the given logger
@@ -153,6 +155,15 @@ 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
+	}
+	s.dataset = dataset
+
 	// For Data Parallel, start data-parallel-size - 1 additional simulators
 	g, ctx := errgroup.WithContext(ctx)
 	if s.config.DPSize > 1 {