server/internal/registry: take over pulls from server package (#9485)

This commit replaces the old pull implementation in the server package
with the new, faster, more robust pull implementation in the registry
package.

The new endpoint, and now the remove endpoint too, are behind the
feature gate "client2" enabled only by setting the OLLAMA_EXPERIMENT
environment variable include "client2".

Currently, the progress indication is wired to perform the same as the
previous implementation to avoid making changes to the CLI, and because
the status reports happen at the start of the download, and the end of
the write to disk, the progress indication is not as smooth as it could
be. This is a known issue and will be addressed in a future change.

This implementation may be ~0.5-1.0% slower in rare cases, depending on
network and disk speed, but is generally MUCH faster and more robust
than the its predecessor in all other cases.

server/internal/registry/server.go

@@ -7,10 +7,14 @@ import (
 	"cmp"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"io"
 	"log/slog"
 	"net/http"
+	"sync"
+	"time"
 
+	"github.com/ollama/ollama/server/internal/cache/blob"
 	"github.com/ollama/ollama/server/internal/client/ollama"
 )
 
@@ -109,6 +113,8 @@ func (s *Local) serveHTTP(rec *statusCodeRecorder, r *http.Request) {
 		switch r.URL.Path {
 		case "/api/delete":
 			return false, s.handleDelete(rec, r)
+		case "/api/pull":
+			return false, s.handlePull(rec, r)
 		default:
 			if s.Fallback != nil {
 				s.Fallback.ServeHTTP(rec, r)
@@ -214,6 +220,97 @@ func (s *Local) handleDelete(_ http.ResponseWriter, r *http.Request) error {
 	return s.Prune()
 }
 
+type progressUpdateJSON struct {
+	Status    string      `json:"status"`
+	Digest    blob.Digest `json:"digest,omitempty,omitzero"`
+	Total     int64       `json:"total,omitempty,omitzero"`
+	Completed int64       `json:"completed,omitempty,omitzero"`
+}
+
+func (s *Local) handlePull(w http.ResponseWriter, r *http.Request) error {
+	if r.Method != "POST" {
+		return errMethodNotAllowed
+	}
+
+	p, err := decodeUserJSON[*params](r.Body)
+	if err != nil {
+		return err
+	}
+
+	maybeFlush := func() {
+		fl, _ := w.(http.Flusher)
+		if fl != nil {
+			fl.Flush()
+		}
+	}
+	defer maybeFlush()
+
+	var mu sync.Mutex
+	enc := json.NewEncoder(w)
+	enc.Encode(progressUpdateJSON{Status: "pulling manifest"})
+
+	ctx := ollama.WithTrace(r.Context(), &ollama.Trace{
+		Update: func(l *ollama.Layer, n int64, err error) {
+			mu.Lock()
+			defer mu.Unlock()
+
+			// TODO(bmizerany): coalesce these updates; writing per
+			// update is expensive
+			enc.Encode(progressUpdateJSON{
+				Digest:    l.Digest,
+				Status:    "pulling",
+				Total:     l.Size,
+				Completed: n,
+			})
+		},
+	})
+
+	done := make(chan error, 1)
+	go func() {
+		// TODO(bmizerany): continue to support non-streaming responses
+		done <- s.Client.Pull(ctx, p.model())
+	}()
+
+	func() {
+		t := time.NewTicker(100 * time.Millisecond)
+		defer t.Stop()
+		for {
+			select {
+			case <-t.C:
+				mu.Lock()
+				maybeFlush()
+				mu.Unlock()
+			case err := <-done:
+				if err != nil {
+					var status string
+					if errors.Is(err, ollama.ErrModelNotFound) {
+						status = fmt.Sprintf("error: model %q not found", p.model())
+						enc.Encode(progressUpdateJSON{Status: status})
+					} else {
+						status = fmt.Sprintf("error: %v", err)
+						enc.Encode(progressUpdateJSON{Status: status})
+					}
+					return
+				}
+
+				// These final updates are not strictly necessary, because they have
+				// already happened at this point. Our pull handler code used to do
+				// these steps after, not during, the pull, and they were slow, so we
+				// wanted to provide feedback to users what was happening. For now, we
+				// keep them to not jar users who are used to seeing them. We can phase
+				// them out with a new and nicer UX later. One without progress bars
+				// and digests that no one cares about.
+				enc.Encode(progressUpdateJSON{Status: "verifying layers"})
+				enc.Encode(progressUpdateJSON{Status: "writing manifest"})
+				enc.Encode(progressUpdateJSON{Status: "success"})
+				return
+			}
+		}
+	}()
+
+	return nil
+}
+
 func decodeUserJSON[T any](r io.Reader) (T, error) {
 	var v T
 	err := json.NewDecoder(r).Decode(&v)