Josh Yan 10 miesięcy temu
rodzic
commit
461c964941
2 zmienionych plików z 28 dodań i 19 usunięć
  1. 15 19
      cmd/cmd.go
  2. 13 0
      cmd/cmd_test.go

+ 15 - 19
cmd/cmd.go

@@ -287,7 +287,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
 		}
@@ -299,32 +310,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
@@ -348,7 +344,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)
 	}

+ 13 - 0
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)
+	}
+}