server/internal/chunks: remove chunks package (ollama#9755)

Author: bmizerany

server/internal/cache/blob/chunked.go

@@ -5,11 +5,18 @@ import (
 	"errors"
 	"io"
 	"os"
-
-	"github.com/ollama/ollama/server/internal/chunks"
 )
 
-type Chunk = chunks.Chunk // TODO: move chunks here?
+// Chunk represents a range of bytes in a blob.
+type Chunk struct {
+	Start int64
+	End   int64
+}
+
+// Size returns end minus start plus one.
+func (c Chunk) Size() int64 {
+	return c.End - c.Start + 1
+}
 
 // Chunker writes to a blob in chunks.
 // Its zero value is invalid. Use [DiskCache.Chunked] to create a new Chunker.

server/internal/chunks/chunks.go

@@ -36,7 +36,6 @@ import (
 	"golang.org/x/sync/errgroup"
 
 	"github.com/ollama/ollama/server/internal/cache/blob"
-	"github.com/ollama/ollama/server/internal/chunks"
 	"github.com/ollama/ollama/server/internal/internal/backoff"
 	"github.com/ollama/ollama/server/internal/internal/names"
 
@@ -500,7 +499,7 @@ func (r *Registry) Pull(ctx context.Context, name string) error {
 						if err != nil {
 							return err
 						}
-						req.Header.Set("Range", fmt.Sprintf("bytes=%s", cs.Chunk))
+						req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", cs.Chunk.Start, cs.Chunk.End))
 						res, err := sendRequest(r.client(), req)
 						if err != nil {
 							return err
@@ -794,7 +793,7 @@ func (r *Registry) chunksums(ctx context.Context, name string, l *Layer) iter.Se
 				yield(chunksum{}, err)
 				return
 			}
-			chunk, err := chunks.Parse(s.Bytes())
+			chunk, err := parseChunk(s.Bytes())
 			if err != nil {
 				yield(chunksum{}, fmt.Errorf("invalid chunk range for digest %s: %q", d, s.Bytes()))
 				return
@@ -1059,3 +1058,23 @@ func splitExtended(s string) (scheme, name, digest string) {
 	}
 	return scheme, s, digest
 }
+
+// parseChunk parses a string in the form "start-end" and returns the Chunk.
+func parseChunk[S ~string | ~[]byte](s S) (blob.Chunk, error) {
+	startPart, endPart, found := strings.Cut(string(s), "-")
+	if !found {
+		return blob.Chunk{}, fmt.Errorf("chunks: invalid range %q: missing '-'", s)
+	}
+	start, err := strconv.ParseInt(startPart, 10, 64)
+	if err != nil {
+		return blob.Chunk{}, fmt.Errorf("chunks: invalid start to %q: %v", s, err)
+	}
+	end, err := strconv.ParseInt(endPart, 10, 64)
+	if err != nil {
+		return blob.Chunk{}, fmt.Errorf("chunks: invalid end to %q: %v", s, err)
+	}
+	if start > end {
+		return blob.Chunk{}, fmt.Errorf("chunks: invalid range %q: start > end", s)
+	}
+	return blob.Chunk{Start: start, End: end}, nil
+}

server/internal/chunks/chunks_test.go

@@ -21,7 +21,6 @@ import (
 	"time"
 
 	"github.com/ollama/ollama/server/internal/cache/blob"
-	"github.com/ollama/ollama/server/internal/chunks"
 	"github.com/ollama/ollama/server/internal/testutil"
 )
 
@@ -531,56 +530,6 @@ func TestRegistryPullMixedCachedNotCached(t *testing.T) {
 	}
 }
 
-func TestRegistryPullChunking(t *testing.T) {
-	t.Skip("TODO: BRING BACK BEFORE LANDING")
-
-	rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) {
-		t.Log("request:", r.URL.Host, r.Method, r.URL.Path, r.Header.Get("Range"))
-		if r.URL.Host != "blob.store" {
-			// The production registry redirects to the blob store.
-			http.Redirect(w, r, "http://blob.store"+r.URL.Path, http.StatusFound)
-			return
-		}
-		if strings.Contains(r.URL.Path, "/blobs/") {
-			rng := r.Header.Get("Range")
-			if rng == "" {
-				http.Error(w, "missing range", http.StatusBadRequest)
-				return
-			}
-			_, c, err := chunks.ParseRange(r.Header.Get("Range"))
-			if err != nil {
-				panic(err)
-			}
-			io.WriteString(w, "remote"[c.Start:c.End+1])
-			return
-		}
-		fmt.Fprintf(w, `{"layers":[{"digest":%q,"size":6}]}`, blob.DigestFromBytes("remote"))
-	})
-
-	// Force chunking by setting the threshold to less than the size of the
-	// layer.
-	rc.ChunkingThreshold = 3
-	rc.MaxChunkSize = 3
-
-	var reads []int64
-	ctx := WithTrace(t.Context(), &Trace{
-		Update: func(d *Layer, n int64, err error) {
-			if err != nil {
-				t.Errorf("update %v %d %v", d, n, err)
-			}
-			reads = append(reads, n)
-		},
-	})
-
-	err := rc.Pull(ctx, "remote")
-	testutil.Check(t, err)
-
-	want := []int64{0, 3, 6}
-	if !slices.Equal(reads, want) {
-		t.Errorf("reads = %v; want %v", reads, want)
-	}
-}
-
 func TestRegistryResolveByDigest(t *testing.T) {
 	check := testutil.Checker(t)