From a6cfe7f00b13e36b0b293cf17c2583674eb1adfc Mon Sep 17 00:00:00 2001
From: Josh Yan <jyan00017@gmail.com>
Date: Tue, 2 Jul 2024 14:53:54 -0700
Subject: [PATCH] benchmark

---
 cmd/cmd.go      | 34 +++++++++++++++-------------------
 cmd/cmd_test.go | 13 +++++++++++++
 2 files changed, 28 insertions(+), 19 deletions(-)
 create mode 100644 cmd/cmd_test.go

diff --git a/cmd/cmd.go b/cmd/cmd.go
index 3dce7214b..5d0804ec7 100644
--- a/cmd/cmd.go
+++ b/cmd/cmd.go
@@ -290,7 +290,18 @@ func createBlob(cmd *cobra.Command, client *api.Client, path string) (string, er
 	// Resolve server to IP
 	// Check if server is local
 	if client.IsLocal() {
-		err := createBlobLocal(cmd.Context(), client, path, digest)
+		config, err := client.ServerConfig(cmd.Context())
+		if err != nil {
+			return "", err
+		}
+
+		modelDir := config.ModelDir
+
+		// Get blob destination
+		digest = strings.ReplaceAll(digest, ":", "-")
+		dest := filepath.Join(modelDir, "blobs", digest)
+
+		err = createBlobLocal(path, dest)
 		if err == nil {
 			return digest, nil
 		}
@@ -302,32 +313,17 @@ func createBlob(cmd *cobra.Command, client *api.Client, path string) (string, er
 	return digest, nil
 }
 
-func createBlobLocal(ctx context.Context, client *api.Client, path string, digest string) error {
+func createBlobLocal(path string, dest string) error {
 	// This function should be called if the server is local
 	// It should find the model directory, copy the blob over, and return the digest
-
-	// Get the model directory
-	config, err := client.ServerConfig(ctx)
-	if err != nil {
-		return err
-	}
-
-	modelDir := config.ModelDir
-
-	// Get blob destination
-	digest = strings.ReplaceAll(digest, ":", "-")
-	dest := filepath.Join(modelDir, "blobs", digest)
 	dirPath := filepath.Dir(dest)
-	if digest == "" {
-		dirPath = dest
-	}
 
 	if err := os.MkdirAll(dirPath, 0o755); err != nil {
 		return err
 	}
 
 	// Check blob exists
-	_, err = os.Stat(dest)
+	_, err := os.Stat(dest)
 	switch {
 	case errors.Is(err, os.ErrNotExist):
 		// noop
@@ -351,7 +347,7 @@ func createBlobLocal(ctx context.Context, client *api.Client, path string, diges
 	}
 	defer destFile.Close()
 
-	_, err = io.Copy(destFile, sourceFile)
+	_, err = io.CopyBuffer(destFile, sourceFile, make([]byte, 4*1024*1024))
 	if err != nil {
 		return fmt.Errorf("error copying file: %v", err)
 	}
diff --git a/cmd/cmd_test.go b/cmd/cmd_test.go
new file mode 100644
index 000000000..bbdd8d17a
--- /dev/null
+++ b/cmd/cmd_test.go
@@ -0,0 +1,13 @@
+package cmd
+
+import (
+	"testing"
+)
+
+func BenchmarkCreateLocalBlob(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		dest := b.TempDir() + "/hi"
+
+		createBlobLocal("/Users/joshyan/.ollama/models/blobs/sha256-edd739ebd0b09f4e9345e8dc76d06ec37d08a080246560e57f7f1443fa3e57af", dest)
+	}
+}