diff --git a/go.mod b/go.mod index 146a7c5b..cd57291b 100644 --- a/go.mod +++ b/go.mod @@ -38,7 +38,6 @@ require ( github.com/wk8/go-ordered-map/v2 v2.1.8 golang.org/x/image v0.22.0 golang.org/x/mod v0.30.0 - golang.org/x/tools v0.38.0 gonum.org/v1/gonum v0.15.0 ) diff --git a/go.sum b/go.sum index 56a7e298..5c57a64b 100644 --- a/go.sum +++ b/go.sum @@ -392,8 +392,6 @@ golang.org/x/tools v0.0.0-20200130002326-2f3ba24bd6e7/go.mod h1:TB2adYChydJhpapK golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE= golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= golang.org/x/tools v0.1.4/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk= -golang.org/x/tools v0.38.0 h1:Hx2Xv8hISq8Lm16jvBZ2VQf+RLmbd7wVUsALibYI/IQ= -golang.org/x/tools v0.38.0/go.mod h1:yEsQ/d/YK8cjh0L6rZlY8tgtlKiBNTL14pGDJPJpYQs= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= diff --git a/server/internal/client/ollama/registry.go b/server/internal/client/ollama/registry.go deleted file mode 100644 index eae130bf..00000000 --- a/server/internal/client/ollama/registry.go +++ /dev/null @@ -1,1197 +0,0 @@ -// Package ollama provides a client for interacting with an Ollama registry -// which pushes and pulls model manifests and layers as defined by the -// [ollama.com/manifest]. -package ollama - -import ( - "bufio" - "bytes" - "cmp" - "context" - "crypto" - "crypto/ed25519" - "crypto/sha256" - "crypto/tls" - "encoding/base64" - "encoding/hex" - "encoding/json" - "errors" - "fmt" - "io" - "io/fs" - "iter" - "log/slog" - "net/http" - "os" - "path/filepath" - "runtime" - "runtime/debug" - "slices" - "strconv" - "strings" - "sync" - "sync/atomic" - "time" - - "golang.org/x/crypto/ssh" - "golang.org/x/sync/errgroup" - - "github.com/ollama/ollama/server/internal/cache/blob" - "github.com/ollama/ollama/server/internal/internal/names" - - _ "embed" -) - -// Errors -var ( - // ErrModelNotFound is returned when a manifest is not found in the - // cache or registry. - ErrModelNotFound = errors.New("model not found") - - // ErrManifestInvalid is returned when a manifest found in a local or - // remote cache is invalid. - ErrManifestInvalid = errors.New("invalid manifest") - - // ErrMissingModel is returned when the model part of a name is missing - // or invalid. - ErrNameInvalid = errors.New("invalid or missing name") - - // ErrCached is passed to [Trace.PushUpdate] when a layer already - // exists. It is a non-fatal error and is never returned by [Registry.Push]. - ErrCached = errors.New("cached") - - // ErrIncomplete is returned by [Registry.Pull] when a model pull was - // incomplete due to one or more layer download failures. Users that - // want specific errors should use [WithTrace]. - ErrIncomplete = errors.New("incomplete") -) - -// Defaults -const ( - // DefaultChunkingThreshold is the threshold at which a layer should be - // split up into chunks when downloading. - DefaultChunkingThreshold = 64 << 20 -) - -var defaultCache = sync.OnceValues(func() (*blob.DiskCache, error) { - dir := os.Getenv("OLLAMA_MODELS") - if dir == "" { - home, _ := os.UserHomeDir() - home = cmp.Or(home, ".") - dir = filepath.Join(home, ".ollama", "models") - } - return blob.Open(dir) -}) - -// DefaultCache returns the default cache used by the registry. It is -// configured from the OLLAMA_MODELS environment variable, or defaults to -// $HOME/.ollama/models, or, if an error occurs obtaining the home directory, -// it uses the current working directory. -func DefaultCache() (*blob.DiskCache, error) { - return defaultCache() -} - -// Error is the standard error returned by Ollama APIs. It can represent a -// single or multiple error response. -// -// Single error responses have the following format: -// -// {"code": "optional_code","error":"error message"} -// -// Multiple error responses have the following format: -// -// {"errors": [{"code": "optional_code","message":"error message"}]} -// -// Note, that the error field is used in single error responses, while the -// message field is used in multiple error responses. -// -// In both cases, the code field is optional and may be empty. -type Error struct { - status int `json:"-"` // TODO(bmizerany): remove this - Code string `json:"code"` - Message string `json:"message"` -} - -// Temporary reports if the error is temporary (e.g. 5xx status code). -func (e *Error) Temporary() bool { - return e.status >= 500 -} - -func (e *Error) Error() string { - var b strings.Builder - b.WriteString("registry responded with status ") - b.WriteString(strconv.Itoa(e.status)) - if e.Code != "" { - b.WriteString(": code ") - b.WriteString(e.Code) - } - if e.Message != "" { - b.WriteString(": ") - b.WriteString(e.Message) - } - return b.String() -} - -func (e *Error) LogValue() slog.Value { - return slog.GroupValue( - slog.Int("status", e.status), - slog.String("code", e.Code), - slog.String("message", e.Message), - ) -} - -// UnmarshalJSON implements json.Unmarshaler. -func (e *Error) UnmarshalJSON(b []byte) error { - type E Error - var v struct { - // Single error - Code string - Error string - - // Multiple errors - Errors []E - } - if err := json.Unmarshal(b, &v); err != nil { - return err - } - if v.Error != "" { - // Single error case - e.Code = v.Code - e.Message = v.Error - return nil - } - if len(v.Errors) == 0 { - return fmt.Errorf("no messages in error response: %s", string(b)) - } - *e = Error(v.Errors[0]) // our registry only returns one error. - return nil -} - -const DefaultMask = "registry.ollama.ai/library/_:latest" - -var defaultMask = func() names.Name { - n := names.Parse(DefaultMask) - if !n.IsFullyQualified() { - panic("default mask is not fully qualified") - } - return n -}() - -// CompleteName returns a fully qualified name by merging the given name with -// the default mask. If the name is already fully qualified, it is returned -// unchanged. -func CompleteName(name string) string { - return names.Merge(names.Parse(name), defaultMask).String() -} - -// Registry is a client for performing push and pull operations against an -// Ollama registry. -type Registry struct { - // Cache is the cache used to store models. If nil, [DefaultCache] is - // used. - Cache *blob.DiskCache - - // UserAgent is the User-Agent header to send with requests to the - // registry. If empty, the User-Agent is determined by HTTPClient. - UserAgent string - - // Key is the key used to authenticate with the registry. - // - // Currently, only Ed25519 keys are supported. - Key crypto.PrivateKey - - // HTTPClient is the HTTP client used to make requests to the registry. - // - // If nil, [http.DefaultClient] is used. - // - // As a quick note: If a Registry function that makes a call to a URL - // with the "https+insecure" scheme, the client will be cloned and the - // transport will be set to skip TLS verification, unless the client's - // Transport done not have a Clone method with the same signature as - // [http.Transport.Clone], which case, the call will fail. - HTTPClient *http.Client - - // MaxStreams is the maximum number of concurrent streams to use when - // pushing or pulling models. If zero, the number of streams is - // determined by [runtime.GOMAXPROCS]. - // - // A negative value means no limit. - MaxStreams int - - // ChunkingThreshold is the maximum size of a layer to download in a single - // request. If zero, [DefaultChunkingThreshold] is used. - ChunkingThreshold int64 - - // Mask, if set, is the name used to convert non-fully qualified names - // to fully qualified names. - // If empty, [DefaultMask] is used. - Mask string - - // ReadTimeout is the maximum duration for reading the entire request, - // including the body. - // A zero or negative value means there will be no timeout. - ReadTimeout time.Duration -} - -func (r *Registry) readTimeout() time.Duration { - if r.ReadTimeout > 0 { - return r.ReadTimeout - } - return 1<<63 - 1 // no timeout, max int -} - -func (r *Registry) cache() (*blob.DiskCache, error) { - if r.Cache != nil { - return r.Cache, nil - } - return defaultCache() -} - -func (r *Registry) parseName(name string) (names.Name, error) { - mask := defaultMask - if r.Mask != "" { - mask = names.Parse(r.Mask) - } - n := names.Merge(names.Parse(name), mask) - if !n.IsFullyQualified() { - return names.Name{}, fmt.Errorf("%w: %q", ErrNameInvalid, name) - } - return n, nil -} - -// DefaultRegistry returns a new Registry configured from the environment. The -// key is read from $HOME/.ollama/id_ed25519, MaxStreams is set to the -// value of OLLAMA_REGISTRY_MAXSTREAMS, and ReadTimeout is set to 30 seconds. -// -// It returns an error if any configuration in the environment is invalid. -func DefaultRegistry() (*Registry, error) { - home, err := os.UserHomeDir() - if err != nil { - return nil, err - } - keyPEM, err := os.ReadFile(filepath.Join(home, ".ollama/id_ed25519")) - if err != nil && errors.Is(err, fs.ErrNotExist) { - return nil, err - } - - var rc Registry - rc.ReadTimeout = 30 * time.Second - rc.UserAgent = UserAgent() - rc.Key, err = ssh.ParseRawPrivateKey(keyPEM) - if err != nil { - return nil, err - } - maxStreams := os.Getenv("OLLAMA_REGISTRY_MAXSTREAMS") - if maxStreams != "" { - var err error - rc.MaxStreams, err = strconv.Atoi(maxStreams) - if err != nil { - return nil, fmt.Errorf("invalid OLLAMA_REGISTRY_MAXSTREAMS: %w", err) - } - } - return &rc, nil -} - -func UserAgent() string { - buildinfo, _ := debug.ReadBuildInfo() - - version := buildinfo.Main.Version - if version == "(devel)" { - // When using `go run .` the version is "(devel)". This is seen - // as an invalid version by ollama.com and so it defaults to - // "needs upgrade" for some requests, such as pulls. These - // checks can be skipped by using the special version "v0.0.0", - // so we set it to that here. - version = "v0.0.0" - } - - return fmt.Sprintf("ollama/%s (%s %s) Go/%s", - version, - runtime.GOARCH, - runtime.GOOS, - runtime.Version(), - ) -} - -func (r *Registry) maxStreams() int { - return cmp.Or(r.MaxStreams, runtime.GOMAXPROCS(0)) -} - -func (r *Registry) maxChunkingThreshold() int64 { - return cmp.Or(r.ChunkingThreshold, DefaultChunkingThreshold) -} - -type PushParams struct { - // From is an optional destination name for the model. If empty, the - // destination name is the same as the source name. - From string -} - -// Push pushes the model with the name in the cache to the remote registry. -func (r *Registry) Push(ctx context.Context, name string, p *PushParams) error { - if p == nil { - p = &PushParams{} - } - - c, err := r.cache() - if err != nil { - return err - } - - m, err := r.ResolveLocal(cmp.Or(p.From, name)) - if err != nil { - return err - } - - // Before much else happens, check layers at not null, the blobs exist, - // and the sizes match. This prevents long uploads followed by - // disappointment. - for _, l := range m.Layers { - if l == nil { - return fmt.Errorf("%w: null layer", ErrManifestInvalid) - } - info, err := c.Get(l.Digest) - if err != nil { - return fmt.Errorf("error getting %s: %w", l.Digest.Short(), err) - } - if info.Size != l.Size { - return fmt.Errorf("size mismatch for %s: %d != %d", l.Digest.Short(), info.Size, l.Size) - } - } - - t := traceFromContext(ctx) - - scheme, n, _, err := r.parseNameExtended(name) - if err != nil { - // This should never happen since ResolveLocal should have - // already validated the name. - panic(err) - } - - ctx, cancel := context.WithCancel(ctx) - defer cancel() - - var g errgroup.Group - g.SetLimit(r.maxStreams()) - for _, l := range m.Layers { - var progress atomic.Int64 - g.Go(func() (err error) { - defer func() { t.update(l, progress.Load(), err) }() - - t.update(l, 0, nil) - - startURL := fmt.Sprintf("%s://%s/v2/%s/%s/blobs/uploads/?digest=%s", - scheme, - n.Host(), - n.Namespace(), - n.Model(), - l.Digest, - ) - res, err := r.send(ctx, "POST", startURL, nil) - if err != nil { - return err - } - res.Body.Close() - - f, err := os.Open(c.GetFile(l.Digest)) - if err != nil { - return err - } - defer f.Close() - - uploadURL := res.Header.Get("Location") - if uploadURL == "" { - t.update(l, l.Size, ErrCached) - return nil - } - - req, err := r.newRequest(ctx, "PUT", uploadURL, f) - if err != nil { - return fmt.Errorf("invalid upload URL returned from registry: %q: %w", uploadURL, err) - } - req.ContentLength = l.Size - - res, err = sendRequest(r.client(), req) - if err == nil { - res.Body.Close() - } - return err - }) - } - - if err := g.Wait(); err != nil { - return err - } - - // Commit - path := fmt.Sprintf("%s://%s/v2/%s/%s/manifests/%s", - scheme, - n.Host(), - n.Namespace(), - n.Model(), - n.Tag(), - ) - res, err := r.send(ctx, "PUT", path, bytes.NewReader(m.Data)) - if err == nil { - res.Body.Close() - } - // TODO(bmizerany): add a "commit" trace event - return err -} - -// trackingReader is an io.Reader that tracks the number of bytes read and -// calls the update function with the layer, the number of bytes read. -// -// It always calls update with a nil error. -type trackingReader struct { - r io.Reader - update func(n int64, err error) // err is always nil -} - -func (r *trackingReader) Read(p []byte) (n int, err error) { - n, err = r.r.Read(p) - r.update(int64(n), nil) - return -} - -// Pull pulls the model with the given name from the remote registry into the -// cache. -// -// For layers larger then [Registry.MaxChunkSize], the layer is downloaded in -// chunks of the specified size, and then reassembled and verified. This is -// typically slower than splitting the model up across layers, and is mostly -// utilized for layers of type equal to "application/vnd.ollama.image". -func (r *Registry) Pull(ctx context.Context, name string) error { - m, err := r.Resolve(ctx, name) - if err != nil { - return err - } - - // TODO(bmizerany): decide if this should be considered valid. Maybe - // server-side we special case '{}' to have some special meaning? Maybe - // "archiving" a tag (which is how we reason about it in the registry - // already, just with a different twist). - if len(m.Layers) == 0 { - return fmt.Errorf("%w: no layers", ErrManifestInvalid) - } - - c, err := r.cache() - if err != nil { - return err - } - - // TODO(bmizerany): work to remove the need to do this - layers := m.Layers - if m.Config != nil && m.Config.Digest.IsValid() { - layers = append(layers, m.Config) - } - - // Send initial layer trace events to allow clients to have an - // understanding of work to be done before work starts. - var expected int64 - t := traceFromContext(ctx) - for _, l := range layers { - t.update(l, 0, nil) - expected += l.Size - } - - var g errgroup.Group - g.SetLimit(r.maxStreams()) - - var completed atomic.Int64 - for _, l := range layers { - var received atomic.Int64 - update := func(n int64, err error) { - if n == 0 && err == nil { - // Clients expect an update with no progress and no error to mean "starting download". - // This is not the case here, - // so we don't want to send an update in this case. - return - } - completed.Add(n) - t.update(l, received.Add(n), err) - } - - info, err := c.Get(l.Digest) - if err == nil && info.Size == l.Size { - update(l.Size, ErrCached) - continue - } - - func() (err error) { - defer func() { - if err != nil { - update(0, err) - } - }() - - var wg sync.WaitGroup - chunked, err := c.Chunked(l.Digest, l.Size) - if err != nil { - return err - } - defer func() { - // Close the chunked writer when all chunks are - // downloaded. - // - // This is done as a background task in the - // group to allow the next layer to start while - // we wait for the final chunk in this layer to - // complete. It also ensures this is done - // before we exit Pull. - g.Go(func() error { - wg.Wait() - chunked.Close() - return nil - }) - }() - - for cs, err := range r.chunksums(ctx, name, l) { - if err != nil { - // Note the chunksum stream - // interruption, but do not cancel - // in-flight downloads. We can still - // make progress on them. Once they are - // done, ErrIncomplete will be returned - // below. - update(0, err) - break - } - - cacheKey := fmt.Sprintf( - "v1 pull chunksum %s %s %d-%d", - l.Digest, - cs.Digest, - cs.Chunk.Start, - cs.Chunk.End, - ) - cacheKeyDigest := blob.DigestFromBytes(cacheKey) - _, err := c.Get(cacheKeyDigest) - if err == nil { - update(cs.Chunk.Size(), ErrCached) - continue - } - - wg.Add(1) - g.Go(func() (err error) { - defer func() { - defer wg.Done() - if err != nil { - update(0, err) - } - }() - - ctx, cancel := context.WithCancelCause(ctx) - defer cancel(nil) - - timer := time.AfterFunc(r.readTimeout(), func() { - cancel(fmt.Errorf("%w: downloading %s %d-%d/%d", - context.DeadlineExceeded, - cs.Digest.Short(), - cs.Chunk.Start, - cs.Chunk.End, - l.Size, - )) - }) - defer timer.Stop() - - req, err := http.NewRequestWithContext(ctx, "GET", cs.URL, nil) - if err != nil { - return err - } - 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 - } - defer res.Body.Close() - - tr := &trackingReader{ - r: res.Body, - update: func(n int64, err error) { - timer.Reset(r.readTimeout()) - update(n, err) - }, - } - if err := chunked.Put(cs.Chunk, cs.Digest, tr); err != nil { - return err - } - - // Record the downloading of this chunk. - return blob.PutBytes(c, cacheKeyDigest, cacheKey) - }) - } - - return nil - }() - } - if err := g.Wait(); err != nil { - return err - } - if recv := completed.Load(); recv != expected { - return fmt.Errorf("%w: received %d/%d bytes", ErrIncomplete, recv, expected) - } - - md := blob.DigestFromBytes(m.Data) - if err := blob.PutBytes(c, md, m.Data); err != nil { - return err - } - return c.Link(m.Name, md) -} - -// Unlink is like [blob.DiskCache.Unlink], but makes name fully qualified -// before attempting to unlink the model. -func (r *Registry) Unlink(name string) (ok bool, _ error) { - n, err := r.parseName(name) - if err != nil { - return false, err - } - c, err := r.cache() - if err != nil { - return false, err - } - return c.Unlink(n.String()) -} - -// Manifest represents a [ollama.com/manifest]. -type Manifest struct { - Name string `json:"-"` // the canonical name of the model - Data []byte `json:"-"` // the raw data of the manifest - Layers []*Layer `json:"layers"` - - // For legacy reasons, we still have to download the config layer. - Config *Layer `json:"config"` -} - -// Layer returns the layer with the given -// digest, or nil if not found. -func (m *Manifest) Layer(d blob.Digest) *Layer { - for _, l := range m.Layers { - if l.Digest == d { - return l - } - } - return nil -} - -func (m *Manifest) All() iter.Seq[*Layer] { - return func(yield func(*Layer) bool) { - if !yield(m.Config) { - return - } - for _, l := range m.Layers { - if !yield(l) { - return - } - } - } -} - -func (m *Manifest) Size() int64 { - var size int64 - if m.Config != nil { - size += m.Config.Size - } - for _, l := range m.Layers { - size += l.Size - } - return size -} - -// MarshalJSON implements json.Marshaler. -// -// NOTE: It adds an empty config object to the manifest, which is required by -// the registry, but not used by the client. In the future, the config object -// will not be required by the registry and this will should be removed. -func (m Manifest) MarshalJSON() ([]byte, error) { - type M Manifest - v := struct { - M - - // This is ignored, mostly, by the registry But, if not - // present, it will cause an error to be returned during the - // last phase of the commit which expects it, but does nothing - // with it. This will be fixed in a future release of - // ollama.com. - Config Layer `json:"config"` - }{ - M: M(m), - } - return json.Marshal(v) -} - -// unmarshalManifest unmarshals the data into a manifest, and sets the name -// field to the string representation of the name. -// -// It panics if the name is not fully qualified. Callers should ensure the name -// is fully qualified before calling this function. -func unmarshalManifest(n names.Name, data []byte) (*Manifest, error) { - if !n.IsFullyQualified() { - panic(fmt.Sprintf("unmarshalManifest: name is not fully qualified: %s", n.String())) - } - var m Manifest - if err := json.Unmarshal(data, &m); err != nil { - return nil, err - } - m.Name = n.String() - m.Data = data - return &m, nil -} - -// Layer is a layer in a model. -type Layer struct { - Digest blob.Digest `json:"digest"` - MediaType string `json:"mediaType"` - Size int64 `json:"size"` -} - -// ResolveLocal resolves a name to a Manifest in the local cache. -func (r *Registry) ResolveLocal(name string) (*Manifest, error) { - _, n, d, err := r.parseNameExtended(name) - if err != nil { - return nil, err - } - c, err := r.cache() - if err != nil { - return nil, err - } - if !d.IsValid() { - // No digest, so resolve the manifest by name. - d, err = c.Resolve(n.String()) - if err != nil { - return nil, err - } - } - data, err := os.ReadFile(c.GetFile(d)) - if err != nil { - if errors.Is(err, fs.ErrNotExist) { - return nil, fmt.Errorf("%w: %s", ErrModelNotFound, name) - } - return nil, err - } - m, err := unmarshalManifest(n, data) - if err != nil { - return nil, fmt.Errorf("%s: %w", name, errors.Join(ErrManifestInvalid, err)) - } - return m, nil -} - -// Resolve resolves a name to a Manifest in the remote registry. -func (r *Registry) Resolve(ctx context.Context, name string) (*Manifest, error) { - scheme, n, d, err := r.parseNameExtended(name) - if err != nil { - return nil, err - } - - manifestURL := fmt.Sprintf("%s://%s/v2/%s/%s/manifests/%s", scheme, n.Host(), n.Namespace(), n.Model(), n.Tag()) - if d.IsValid() { - manifestURL = fmt.Sprintf("%s://%s/v2/%s/%s/blobs/%s", scheme, n.Host(), n.Namespace(), n.Model(), d) - } - - res, err := r.send(ctx, "GET", manifestURL, nil) - if err != nil { - return nil, err - } - defer res.Body.Close() - data, err := io.ReadAll(res.Body) - if err != nil { - return nil, err - } - // TODO(bmizerany): return digest here - m, err := unmarshalManifest(n, data) - if err != nil { - return nil, fmt.Errorf("%s: %w", name, errors.Join(ErrManifestInvalid, err)) - } - return m, nil -} - -type chunksum struct { - URL string - Chunk blob.Chunk - Digest blob.Digest -} - -// chunksums returns a sequence of chunksums for the given layer. If the layer is under the -// chunking threshold, a single chunksum is returned that covers the entire layer. If the layer -// is over the chunking threshold, the chunksums are read from the chunksums endpoint. -func (r *Registry) chunksums(ctx context.Context, name string, l *Layer) iter.Seq2[chunksum, error] { - return func(yield func(chunksum, error) bool) { - scheme, n, _, err := r.parseNameExtended(name) - if err != nil { - yield(chunksum{}, err) - return - } - - if l.Size < r.maxChunkingThreshold() { - // any layer under the threshold should be downloaded - // in one go. - cs := chunksum{ - URL: fmt.Sprintf("%s://%s/v2/%s/%s/blobs/%s", - scheme, - n.Host(), - n.Namespace(), - n.Model(), - l.Digest, - ), - Chunk: blob.Chunk{Start: 0, End: l.Size - 1}, - Digest: l.Digest, - } - yield(cs, nil) - return - } - - // The response is a sequence of chunksums. - // - // Chunksums are chunks of a larger blob that can be - // downloaded and verified independently. - // - // The chunksums endpoint is a GET request that returns a - // sequence of chunksums in the following format: - // - // > GET /v2///chunksums/ - // - // < HTTP/1.1 200 OK - // < Content-Location: - // < - // < - - // < ... - // - // The is the URL to download the chunks from and - // each is the digest of the chunk, and - - // is the range the chunk in the blob. - // - // Ranges may be used directly in Range headers like - // "bytes=-". - // - // The chunksums returned are guaranteed to be contiguous and - // include all bytes of the layer. If the stream is cut short, - // clients should retry. - - chunksumsURL := fmt.Sprintf("%s://%s/v2/%s/%s/chunksums/%s", - scheme, - n.Host(), - n.Namespace(), - n.Model(), - l.Digest, - ) - - req, err := r.newRequest(ctx, "GET", chunksumsURL, nil) - if err != nil { - yield(chunksum{}, err) - return - } - res, err := sendRequest(r.client(), req) - if err != nil { - yield(chunksum{}, err) - return - } - defer res.Body.Close() - if res.StatusCode != 200 { - err := fmt.Errorf("chunksums: unexpected status code %d", res.StatusCode) - yield(chunksum{}, err) - return - } - blobURL := res.Header.Get("Content-Location") - - s := bufio.NewScanner(res.Body) - s.Split(bufio.ScanWords) - for { - if !s.Scan() { - if s.Err() != nil { - yield(chunksum{}, s.Err()) - } - return - } - d, err := blob.ParseDigest(s.Bytes()) - if err != nil { - yield(chunksum{}, fmt.Errorf("invalid digest: %q", s.Bytes())) - return - } - - if !s.Scan() { - err := s.Err() - if err == nil { - err = fmt.Errorf("missing chunk range for digest %s", d) - } - yield(chunksum{}, err) - return - } - chunk, err := parseChunk(s.Bytes()) - if err != nil { - yield(chunksum{}, fmt.Errorf("invalid chunk range for digest %s: %q", d, s.Bytes())) - return - } - - cs := chunksum{ - URL: blobURL, - Chunk: chunk, - Digest: d, - } - if !yield(cs, nil) { - return - } - } - } -} - -func (r *Registry) client() *http.Client { - if r.HTTPClient != nil { - return r.HTTPClient - } - return http.DefaultClient -} - -// newRequest constructs a new request, ready to use, with the given method, -// url, and body, pre-signed with client [Key] and [UserAgent]. -func (r *Registry) newRequest(ctx context.Context, method, url string, body io.Reader) (*http.Request, error) { - req, err := http.NewRequestWithContext(ctx, method, url, body) - if err != nil { - return nil, err - } - if r.UserAgent != "" { - req.Header.Set("User-Agent", r.UserAgent) - } - if r.Key != nil { - token, err := makeAuthToken(r.Key) - if err != nil { - return nil, err - } - req.Header.Set("Authorization", "Bearer "+token) - } - return req, nil -} - -// sendRequest makes a request with the given client and request, and returns the -// response if the status code is 200. If the status code is not 200, an Error -// is parsed from the response body and returned. If any other error occurs, it -// is returned. -func sendRequest(c *http.Client, r *http.Request) (_ *http.Response, err error) { - if r.URL.Scheme == "https+insecure" { - // TODO(bmizerany): clone client.Transport, set - // InsecureSkipVerify, etc. - - type cloner interface { - Clone() *http.Transport - } - - // Attempt to configure the transport to skip TLS verification - // if we can clone it, otherwise fall through and let the http - // client complain and the scheme being invalid. - x, ok := cmp.Or(c.Transport, http.DefaultTransport).(cloner) - if ok { - tr := x.Clone() - tr.TLSClientConfig = cmp.Or(tr.TLSClientConfig, &tls.Config{}) - tr.TLSClientConfig.InsecureSkipVerify = true - - cc := *c // shallow copy - cc.Transport = tr - c = &cc - - r = r.Clone(r.Context()) - r.URL.Scheme = "https" - - // fall through - } - } - - res, err := c.Do(r) - if err != nil { - return nil, err - } - if res.StatusCode/100 != 2 { - out, err := io.ReadAll(res.Body) - if err != nil { - return nil, err - } - var re Error - if err := json.Unmarshal(out, &re); err != nil { - // Use the raw body if we can't parse it as an error object. - re.Message = string(out) - } - - // coerce MANIFEST_UNKNOWN to ErrManifestNotFound - if strings.EqualFold(re.Code, "MANIFEST_UNKNOWN") { - return nil, ErrModelNotFound - } - - re.status = res.StatusCode - return nil, &re - } - return res, nil -} - -// send is a convenience method for making a request with newRequest and -// passing it to send with r.client(). -func (r *Registry) send(ctx context.Context, method, path string, body io.Reader) (*http.Response, error) { - req, err := r.newRequest(ctx, method, path, body) - if err != nil { - return nil, err - } - return sendRequest(r.client(), req) -} - -// makeAuthToken creates an Ollama auth token for the given private key. -// -// NOTE: This format is OLD, overly complex, and should be replaced. We're -// inheriting it from the original Ollama client and ollama.com -// implementations, so we need to support it for now. -func makeAuthToken(key crypto.PrivateKey) (string, error) { - privKey, _ := key.(*ed25519.PrivateKey) - if privKey == nil { - return "", fmt.Errorf("unsupported private key type: %T", key) - } - - url := fmt.Sprintf("https://ollama.com?ts=%d", time.Now().Unix()) - // Part 1: the checkData (e.g. the URL with a timestamp) - - // Part 2: the public key - pubKeyShort, err := func() ([]byte, error) { - sshPubKey, err := ssh.NewPublicKey(privKey.Public()) - if err != nil { - return nil, err - } - pubKeyParts := bytes.Fields(ssh.MarshalAuthorizedKey(sshPubKey)) - if len(pubKeyParts) < 2 { - return nil, fmt.Errorf("malformed public key: %q", pubKeyParts) - } - pubKeyShort := pubKeyParts[1] - return pubKeyShort, nil - }() - if err != nil { - return "", err - } - - // Part 3: the signature - sig := ed25519.Sign(*privKey, []byte(checkData(url))) - - // Assemble the token: :: - var b strings.Builder - io.WriteString(&b, base64.StdEncoding.EncodeToString([]byte(url))) - b.WriteByte(':') - b.Write(pubKeyShort) - b.WriteByte(':') - io.WriteString(&b, base64.StdEncoding.EncodeToString(sig)) - - return b.String(), nil -} - -// The original spec for Ollama tokens was to use the SHA256 of the zero -// string as part of the signature. I'm not sure why that was, but we still -// need it to verify the signature. -var zeroSum = func() string { - sha256sum := sha256.Sum256(nil) - x := base64.StdEncoding.EncodeToString([]byte(hex.EncodeToString(sha256sum[:]))) - return x -}() - -// checkData takes a URL and creates the original string format of the -// data signature that is used by the ollama client to sign requests -func checkData(url string) string { - return fmt.Sprintf("GET,%s,%s", url, zeroSum) -} - -type publicError struct { - wrapped error - message string -} - -func withPublicMessagef(err error, message string, args ...any) error { - return publicError{wrapped: err, message: fmt.Sprintf(message, args...)} -} - -func (e publicError) Error() string { return e.message } -func (e publicError) Unwrap() error { return e.wrapped } - -var supportedSchemes = []string{ - "http", - "https", - "https+insecure", -} - -var supportedSchemesMessage = fmt.Sprintf("supported schemes are %v", strings.Join(supportedSchemes, ", ")) - -// parseNameExtended parses and validates an extended name, returning the scheme, name, -// and digest. -// -// If the scheme is empty, scheme will be "https". If an unsupported scheme is -// given, [ErrNameInvalid] wrapped with a display friendly message is returned. -// -// If the digest is invalid, [ErrNameInvalid] wrapped with a display friendly -// message is returned. -// -// If the name is not, once merged with the mask, fully qualified, -// [ErrNameInvalid] wrapped with a display friendly message is returned. -func (r *Registry) parseNameExtended(s string) (scheme string, _ names.Name, _ blob.Digest, _ error) { - scheme, name, digest := splitExtended(s) - scheme = cmp.Or(scheme, "https") - if !slices.Contains(supportedSchemes, scheme) { - err := withPublicMessagef(ErrNameInvalid, "unsupported scheme: %q: %s", scheme, supportedSchemesMessage) - return "", names.Name{}, blob.Digest{}, err - } - - var d blob.Digest - if digest != "" { - var err error - d, err = blob.ParseDigest(digest) - if err != nil { - err = withPublicMessagef(ErrNameInvalid, "invalid digest: %q", digest) - return "", names.Name{}, blob.Digest{}, err - } - if name == "" { - // We have can resolve a manifest from a digest only, - // so skip name validation and return the scheme and - // digest. - return scheme, names.Name{}, d, nil - } - } - - n, err := r.parseName(name) - if err != nil { - return "", names.Name{}, blob.Digest{}, err - } - return scheme, n, d, nil -} - -// splitExtended splits an extended name string into its scheme, name, and digest -// parts. -// -// Examples: -// -// http://ollama.com/bmizerany/smol:latest@digest -// https://ollama.com/bmizerany/smol:latest -// ollama.com/bmizerany/smol:latest@digest // returns "https" scheme. -// model@digest -// @digest -func splitExtended(s string) (scheme, name, digest string) { - i := strings.Index(s, "://") - if i >= 0 { - scheme = s[:i] - s = s[i+3:] - } - i = strings.LastIndex(s, "@") - if i >= 0 { - digest = s[i+1:] - s = s[:i] - } - 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 -} diff --git a/server/internal/client/ollama/registry_synctest_test.go b/server/internal/client/ollama/registry_synctest_test.go deleted file mode 100644 index 2b454337..00000000 --- a/server/internal/client/ollama/registry_synctest_test.go +++ /dev/null @@ -1,51 +0,0 @@ -// TODO: go:build goexperiment.synctest - -package ollama - -import ( - "context" - "errors" - "io" - "net/http" - "strings" - "testing" - "time" -) - -func TestPullDownloadTimeout(t *testing.T) { - rc, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) { - defer t.Log("upstream", r.Method, r.URL.Path) - switch { - case strings.HasPrefix(r.URL.Path, "/v2/library/smol/manifests/"): - io.WriteString(w, `{ - "layers": [{"digest": "sha256:1111111111111111111111111111111111111111111111111111111111111111", "size": 3}] - }`) - case strings.HasPrefix(r.URL.Path, "/v2/library/smol/blobs/sha256:1111111111111111111111111111111111111111111111111111111111111111"): - // Get headers out to client and then hang on the response - w.WriteHeader(200) - w.(http.Flusher).Flush() - - // Hang on the response and unblock when the client - // gives up - <-r.Context().Done() - default: - t.Fatalf("unexpected request: %s", r.URL.Path) - } - }) - rc.ReadTimeout = 100 * time.Millisecond - - done := make(chan error, 1) - go func() { - done <- rc.Pull(ctx, "http://example.com/library/smol") - }() - - select { - case err := <-done: - want := context.DeadlineExceeded - if !errors.Is(err, want) { - t.Errorf("err = %v, want %v", err, want) - } - case <-time.After(3 * time.Second): - t.Error("timeout waiting for Pull to finish") - } -} diff --git a/server/internal/client/ollama/registry_test.go b/server/internal/client/ollama/registry_test.go deleted file mode 100644 index c0bd2863..00000000 --- a/server/internal/client/ollama/registry_test.go +++ /dev/null @@ -1,953 +0,0 @@ -package ollama - -import ( - "bytes" - "cmp" - "context" - "encoding/json" - "errors" - "fmt" - "io" - "io/fs" - "net" - "net/http" - "net/http/httptest" - "os" - "reflect" - "strings" - "sync/atomic" - "testing" - - "github.com/ollama/ollama/server/internal/cache/blob" - "github.com/ollama/ollama/server/internal/testutil" -) - -func ExampleRegistry_cancelOnFirstError() { - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - ctx = WithTrace(ctx, &Trace{ - Update: func(l *Layer, n int64, err error) { - if err != nil { - // Discontinue pulling layers if there is an - // error instead of continuing to pull more - // data. - cancel() - } - }, - }) - - var r Registry - if err := r.Pull(ctx, "model"); err != nil { - // panic for demo purposes - panic(err) - } -} - -func TestManifestMarshalJSON(t *testing.T) { - // All manifests should contain an "empty" config object. - var m Manifest - data, err := json.Marshal(m) - if err != nil { - t.Fatal(err) - } - if !bytes.Contains(data, []byte(`"config":{"digest":"sha256:`)) { - t.Error("expected manifest to contain empty config") - t.Fatalf("got:\n%s", string(data)) - } -} - -var errRoundTrip = errors.New("forced roundtrip error") - -type recordRoundTripper http.HandlerFunc - -func (rr recordRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { - w := httptest.NewRecorder() - rr(w, req) - if w.Code == 499 { - return nil, errRoundTrip - } - resp := w.Result() - // For some reason, Response.Request is not set by httptest.NewRecorder, so we - // set it manually. - resp.Request = req - return w.Result(), nil -} - -// newClient constructs a cache with predefined manifests for testing. The manifests are: -// -// empty: no data -// zero: no layers -// single: one layer with the contents "exists" -// multiple: two layers with the contents "exists" and "here" -// notfound: a layer that does not exist in the cache -// null: one null layer (e.g. [null]) -// sizemismatch: one valid layer, and one with a size mismatch (file size is less than the reported size) -// invalid: a layer with invalid JSON data -// -// Tests that want to ensure the client does not communicate with the upstream -// registry should pass a nil handler, which will cause a panic if -// communication is attempted. -// -// To simulate a network error, pass a handler that returns a 499 status code. -func newClient(t *testing.T, upstreamRegistry http.HandlerFunc) (*Registry, *blob.DiskCache) { - t.Helper() - - c, err := blob.Open(t.TempDir()) - if err != nil { - t.Fatal(err) - } - - mklayer := func(data string) *Layer { - return &Layer{ - Digest: importBytes(t, c, data), - Size: int64(len(data)), - } - } - - r := &Registry{ - Cache: c, - HTTPClient: &http.Client{ - Transport: recordRoundTripper(upstreamRegistry), - }, - } - - link := func(name string, manifest string) { - n, err := r.parseName(name) - if err != nil { - panic(err) - } - d, err := c.Import(bytes.NewReader([]byte(manifest)), int64(len(manifest))) - if err != nil { - panic(err) - } - if err := c.Link(n.String(), d); err != nil { - panic(err) - } - } - - commit := func(name string, layers ...*Layer) { - t.Helper() - data, err := json.Marshal(&Manifest{Layers: layers}) - if err != nil { - t.Fatal(err) - } - link(name, string(data)) - } - - link("empty", "") - commit("zero") - commit("single", mklayer("exists")) - commit("multiple", mklayer("exists"), mklayer("present")) - commit("notfound", &Layer{Digest: blob.DigestFromBytes("notfound"), Size: int64(len("notfound"))}) - commit("null", nil) - commit("sizemismatch", mklayer("exists"), &Layer{Digest: blob.DigestFromBytes("present"), Size: 499}) - link("invalid", "!!!!!") - - return r, c -} - -func okHandler(w http.ResponseWriter, r *http.Request) { - w.WriteHeader(http.StatusOK) -} - -func checkErrCode(t *testing.T, err error, status int, code string) { - t.Helper() - var e *Error - if !errors.As(err, &e) || e.status != status || e.Code != code { - t.Errorf("err = %v; want %v %v", err, status, code) - } -} - -func importBytes(t *testing.T, c *blob.DiskCache, data string) blob.Digest { - d, err := c.Import(strings.NewReader(data), int64(len(data))) - if err != nil { - t.Fatal(err) - } - return d -} - -func withTraceUnexpected(ctx context.Context) (context.Context, *Trace) { - t := &Trace{Update: func(*Layer, int64, error) { panic("unexpected") }} - return WithTrace(ctx, t), t -} - -func TestPushZero(t *testing.T) { - rc, _ := newClient(t, okHandler) - err := rc.Push(t.Context(), "empty", nil) - if !errors.Is(err, ErrManifestInvalid) { - t.Errorf("err = %v; want %v", err, ErrManifestInvalid) - } -} - -func TestPushSingle(t *testing.T) { - rc, _ := newClient(t, okHandler) - err := rc.Push(t.Context(), "single", nil) - testutil.Check(t, err) -} - -func TestPushMultiple(t *testing.T) { - rc, _ := newClient(t, okHandler) - err := rc.Push(t.Context(), "multiple", nil) - testutil.Check(t, err) -} - -func TestPushNotFound(t *testing.T) { - rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) { - t.Errorf("unexpected request: %v", r) - }) - err := rc.Push(t.Context(), "notfound", nil) - if !errors.Is(err, fs.ErrNotExist) { - t.Errorf("err = %v; want %v", err, fs.ErrNotExist) - } -} - -func TestPushNullLayer(t *testing.T) { - rc, _ := newClient(t, nil) - err := rc.Push(t.Context(), "null", nil) - if err == nil || !strings.Contains(err.Error(), "invalid manifest") { - t.Errorf("err = %v; want invalid manifest", err) - } -} - -func TestPushSizeMismatch(t *testing.T) { - rc, _ := newClient(t, nil) - ctx, _ := withTraceUnexpected(t.Context()) - got := rc.Push(ctx, "sizemismatch", nil) - if got == nil || !strings.Contains(got.Error(), "size mismatch") { - t.Errorf("err = %v; want size mismatch", got) - } -} - -func TestPushInvalid(t *testing.T) { - rc, _ := newClient(t, nil) - err := rc.Push(t.Context(), "invalid", nil) - if err == nil || !strings.Contains(err.Error(), "invalid manifest") { - t.Errorf("err = %v; want invalid manifest", err) - } -} - -func TestPushExistsAtRemote(t *testing.T) { - var pushed bool - rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) { - if strings.Contains(r.URL.Path, "/uploads/") { - if !pushed { - // First push. Return an uploadURL. - pushed = true - w.Header().Set("Location", "http://blob.store/blobs/123") - return - } - w.WriteHeader(http.StatusAccepted) - return - } - - io.Copy(io.Discard, r.Body) - w.WriteHeader(http.StatusOK) - }) - - rc.MaxStreams = 1 // prevent concurrent uploads - - var errs []error - ctx := WithTrace(t.Context(), &Trace{ - Update: func(_ *Layer, n int64, err error) { - // uploading one at a time so no need to lock - errs = append(errs, err) - }, - }) - - check := testutil.Checker(t) - - err := rc.Push(ctx, "single", nil) - check(err) - - if !errors.Is(errors.Join(errs...), nil) { - t.Errorf("errs = %v; want %v", errs, []error{ErrCached}) - } - - err = rc.Push(ctx, "single", nil) - check(err) -} - -func TestPushRemoteError(t *testing.T) { - rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) { - if strings.Contains(r.URL.Path, "/blobs/") { - w.WriteHeader(500) - io.WriteString(w, `{"errors":[{"code":"blob_error"}]}`) - return - } - }) - got := rc.Push(t.Context(), "single", nil) - checkErrCode(t, got, 500, "blob_error") -} - -func TestPushLocationError(t *testing.T) { - rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Location", ":///x") - w.WriteHeader(http.StatusAccepted) - }) - got := rc.Push(t.Context(), "single", nil) - wantContains := "invalid upload URL" - if got == nil || !strings.Contains(got.Error(), wantContains) { - t.Errorf("err = %v; want to contain %v", got, wantContains) - } -} - -func TestPushUploadRoundtripError(t *testing.T) { - rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) { - if r.Host == "blob.store" { - w.WriteHeader(499) // force RoundTrip error on upload - return - } - w.Header().Set("Location", "http://blob.store/blobs/123") - }) - got := rc.Push(t.Context(), "single", nil) - if !errors.Is(got, errRoundTrip) { - t.Errorf("got = %v; want %v", got, errRoundTrip) - } -} - -func TestPushUploadFileOpenError(t *testing.T) { - rc, c := newClient(t, okHandler) - ctx := WithTrace(t.Context(), &Trace{ - Update: func(l *Layer, _ int64, err error) { - // Remove the file just before it is opened for upload, - // but after the initial Stat that happens before the - // upload starts - os.Remove(c.GetFile(l.Digest)) - }, - }) - got := rc.Push(ctx, "single", nil) - if !errors.Is(got, fs.ErrNotExist) { - t.Errorf("got = %v; want fs.ErrNotExist", got) - } -} - -func TestPushCommitRoundtripError(t *testing.T) { - rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) { - if strings.Contains(r.URL.Path, "/blobs/") { - panic("unexpected") - } - w.WriteHeader(499) // force RoundTrip error - }) - err := rc.Push(t.Context(), "zero", nil) - if !errors.Is(err, errRoundTrip) { - t.Errorf("err = %v; want %v", err, errRoundTrip) - } -} - -func TestRegistryPullInvalidName(t *testing.T) { - rc, _ := newRegistryClient(t, nil) - err := rc.Pull(t.Context(), "://") - if !errors.Is(err, ErrNameInvalid) { - t.Errorf("err = %v; want %v", err, ErrNameInvalid) - } -} - -func TestRegistryPullInvalidManifest(t *testing.T) { - cases := []string{ - "", - "null", - "!!!", - `{"layers":[]}`, - } - - for _, resp := range cases { - rc, _ := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) { - io.WriteString(w, resp) - }) - err := rc.Pull(t.Context(), "http://example.com/a/b") - if !errors.Is(err, ErrManifestInvalid) { - t.Errorf("err = %v; want invalid manifest", err) - } - } -} - -func TestRegistryResolveByDigest(t *testing.T) { - check := testutil.Checker(t) - - exists := blob.DigestFromBytes("exists") - rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path != "/v2/alice/palace/blobs/"+exists.String() { - w.WriteHeader(499) // should not hit manifest endpoint - } - fmt.Fprintf(w, `{"layers":[{"digest":%q,"size":5}]}`, exists) - }) - - _, err := rc.Resolve(t.Context(), "alice/palace@"+exists.String()) - check(err) -} - -func TestInsecureSkipVerify(t *testing.T) { - exists := blob.DigestFromBytes("exists") - - s := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - fmt.Fprintf(w, `{"layers":[{"digest":%q,"size":5}]}`, exists) - })) - defer s.Close() - - const name = "library/insecure" - - var rc Registry - url := fmt.Sprintf("https://%s/%s", s.Listener.Addr(), name) - _, err := rc.Resolve(t.Context(), url) - if err == nil || !strings.Contains(err.Error(), "failed to verify") { - t.Errorf("err = %v; want cert verification failure", err) - } - - url = fmt.Sprintf("https+insecure://%s/%s", s.Listener.Addr(), name) - _, err = rc.Resolve(t.Context(), url) - testutil.Check(t, err) -} - -func TestErrorUnmarshal(t *testing.T) { - cases := []struct { - name string - data string - want *Error - wantErr bool - }{ - { - name: "errors empty", - data: `{"errors":[]}`, - wantErr: true, - }, - { - name: "errors empty", - data: `{"errors":[]}`, - wantErr: true, - }, - { - name: "errors single", - data: `{"errors":[{"code":"blob_unknown"}]}`, - want: &Error{Code: "blob_unknown", Message: ""}, - }, - { - name: "errors multiple", - data: `{"errors":[{"code":"blob_unknown"},{"code":"blob_error"}]}`, - want: &Error{Code: "blob_unknown", Message: ""}, - }, - { - name: "error empty", - data: `{"error":""}`, - wantErr: true, - }, - { - name: "error very empty", - data: `{}`, - wantErr: true, - }, - { - name: "error message", - data: `{"error":"message", "code":"code"}`, - want: &Error{Code: "code", Message: "message"}, - }, - { - name: "invalid value", - data: `{"error": 1}`, - wantErr: true, - }, - } - for _, tt := range cases { - t.Run(tt.name, func(t *testing.T) { - var got Error - err := json.Unmarshal([]byte(tt.data), &got) - if err != nil { - if tt.wantErr { - return - } - t.Errorf("Unmarshal() error = %v", err) - // fallthrough and check got - } - if tt.want == nil { - tt.want = &Error{} - } - if !reflect.DeepEqual(got, *tt.want) { - t.Errorf("got = %v; want %v", got, *tt.want) - } - }) - } -} - -// TestParseNameExtendedErrors tests that parseName returns errors messages with enough -// detail for users to debug naming issues they may encounter. Previous to this -// test, the error messages were not very helpful and each problem was reported -// as the same message. -// -// It is only for testing error messages, not that all invalids and valids are -// covered. Those are in other tests for names.Name and blob.Digest. -func TestParseNameExtendedErrors(t *testing.T) { - cases := []struct { - name string - err error - want string - }{} - - var r Registry - for _, tt := range cases { - _, _, _, err := r.parseNameExtended(tt.name) - if !errors.Is(err, tt.err) { - t.Errorf("[%s]: err = %v; want %v", tt.name, err, tt.err) - } - if err != nil && !strings.Contains(err.Error(), tt.want) { - t.Errorf("[%s]: err =\n\t%v\nwant\n\t%v", tt.name, err, tt.want) - } - } -} - -func TestParseNameExtended(t *testing.T) { - cases := []struct { - in string - scheme string - name string - digest string - err string - }{ - {in: "http://m", scheme: "http", name: "m"}, - {in: "https+insecure://m", scheme: "https+insecure", name: "m"}, - {in: "http+insecure://m", err: "unsupported scheme"}, - - {in: "http://m@sha256:1111111111111111111111111111111111111111111111111111111111111111", scheme: "http", name: "m", digest: "sha256:1111111111111111111111111111111111111111111111111111111111111111"}, - - {in: "", err: "invalid or missing name"}, - {in: "m", scheme: "https", name: "m"}, - {in: "://", err: "invalid or missing name"}, - {in: "@sha256:deadbeef", err: "invalid digest"}, - {in: "@sha256:deadbeef@sha256:deadbeef", err: "invalid digest"}, - } - for _, tt := range cases { - t.Run(tt.in, func(t *testing.T) { - var r Registry - scheme, n, digest, err := r.parseNameExtended(tt.in) - if err != nil { - if tt.err == "" { - t.Errorf("err = %v; want nil", err) - } else if !strings.Contains(err.Error(), tt.err) { - t.Errorf("err = %v; want %q", err, tt.err) - } - } else if tt.err != "" { - t.Errorf("err = nil; want %q", tt.err) - } - if err == nil && !n.IsFullyQualified() { - t.Errorf("name = %q; want fully qualified", n) - } - - if scheme != tt.scheme { - t.Errorf("scheme = %q; want %q", scheme, tt.scheme) - } - - // smoke-test name is superset of tt.name - if !strings.Contains(n.String(), tt.name) { - t.Errorf("name = %q; want %q", n, tt.name) - } - - tt.digest = cmp.Or(tt.digest, (&blob.Digest{}).String()) - if digest.String() != tt.digest { - t.Errorf("digest = %q; want %q", digest, tt.digest) - } - }) - } -} - -func TestUnlink(t *testing.T) { - t.Run("found by name", func(t *testing.T) { - check := testutil.Checker(t) - - rc, _ := newRegistryClient(t, nil) - // make a blob and link it - d := blob.DigestFromBytes("{}") - err := blob.PutBytes(rc.Cache, d, "{}") - check(err) - err = rc.Cache.Link("registry.ollama.ai/library/single:latest", d) - check(err) - - // confirm linked - _, err = rc.ResolveLocal("single") - check(err) - - // unlink - _, err = rc.Unlink("single") - check(err) - - // confirm unlinked - _, err = rc.ResolveLocal("single") - if !errors.Is(err, fs.ErrNotExist) { - t.Errorf("err = %v; want fs.ErrNotExist", err) - } - }) - t.Run("not found by name", func(t *testing.T) { - rc, _ := newRegistryClient(t, nil) - ok, err := rc.Unlink("manifestNotFound") - if err != nil { - t.Fatal(err) - } - if ok { - t.Error("expected not found") - } - }) -} - -// Many tests from here out, in this file are based on a single blob, "abc", -// with the checksum of its sha256 hash. The checksum is: -// -// "abc" -> sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad -// -// Using the literal value instead of a constant with fmt.Xprintf calls proved -// to be the most readable and maintainable approach. The sum is consistently -// used in the tests and unique so searches do not yield false positives. - -func checkRequest(t *testing.T, req *http.Request, method, path string) { - t.Helper() - if got := req.URL.Path; got != path { - t.Errorf("URL = %q, want %q", got, path) - } - if req.Method != method { - t.Errorf("Method = %q, want %q", req.Method, method) - } -} - -func newRegistryClient(t *testing.T, upstream http.HandlerFunc) (*Registry, context.Context) { - s := httptest.NewServer(upstream) - t.Cleanup(s.Close) - cache, err := blob.Open(t.TempDir()) - if err != nil { - t.Fatal(err) - } - - ctx := WithTrace(t.Context(), &Trace{ - Update: func(l *Layer, n int64, err error) { - t.Log("trace:", l.Digest.Short(), n, err) - }, - }) - - rc := &Registry{ - Cache: cache, - HTTPClient: &http.Client{Transport: &http.Transport{ - Dial: func(network, addr string) (net.Conn, error) { - return net.Dial(network, s.Listener.Addr().String()) - }, - }}, - } - return rc, ctx -} - -func TestPullChunked(t *testing.T) { - var steps atomic.Int64 - c, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) { - switch steps.Add(1) { - case 1: - checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest") - io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`) - case 2: - checkRequest(t, r, "GET", "/v2/library/abc/chunksums/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad") - w.Header().Set("Content-Location", "http://blob.store/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad") - fmt.Fprintf(w, "%s 0-1\n", blob.DigestFromBytes("ab")) - fmt.Fprintf(w, "%s 2-2\n", blob.DigestFromBytes("c")) - case 3, 4: - checkRequest(t, r, "GET", "/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad") - switch rng := r.Header.Get("Range"); rng { - case "bytes=0-1": - io.WriteString(w, "ab") - case "bytes=2-2": - t.Logf("writing c") - io.WriteString(w, "c") - default: - t.Errorf("unexpected range %q", rng) - } - default: - t.Errorf("unexpected steps %d: %v", steps.Load(), r) - http.Error(w, "unexpected steps", http.StatusInternalServerError) - } - }) - - c.ChunkingThreshold = 1 // force chunking - - err := c.Pull(ctx, "http://o.com/library/abc") - testutil.Check(t, err) - - _, err = c.Cache.Resolve("o.com/library/abc:latest") - testutil.Check(t, err) - - if g := steps.Load(); g != 4 { - t.Fatalf("got %d steps, want 4", g) - } -} - -func TestPullCached(t *testing.T) { - c, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) { - checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest") - io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`) - }) - - check := testutil.Checker(t) - - // Premeptively cache the blob - d, err := blob.ParseDigest("sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad") - check(err) - err = blob.PutBytes(c.Cache, d, []byte("abc")) - check(err) - - // Pull only the manifest, which should be enough to resolve the cached blob - err = c.Pull(ctx, "http://o.com/library/abc") - check(err) -} - -func TestPullManifestError(t *testing.T) { - c, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) { - checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest") - w.WriteHeader(http.StatusNotFound) - io.WriteString(w, `{"errors":[{"code":"MANIFEST_UNKNOWN"}]}`) - }) - - err := c.Pull(ctx, "http://o.com/library/abc") - if err == nil { - t.Fatalf("expected error") - } - var got *Error - if !errors.Is(err, ErrModelNotFound) { - t.Fatalf("err = %v, want %v", got, ErrModelNotFound) - } -} - -func TestPullLayerError(t *testing.T) { - c, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) { - checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest") - io.WriteString(w, `!`) - }) - - err := c.Pull(ctx, "http://o.com/library/abc") - if err == nil { - t.Fatalf("expected error") - } - var want *json.SyntaxError - if !errors.As(err, &want) { - t.Fatalf("err = %T, want %T", err, want) - } -} - -func TestPullLayerChecksumError(t *testing.T) { - var step atomic.Int64 - c, _ := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) { - switch step.Add(1) { - case 1: - checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest") - io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`) - case 2: - checkRequest(t, r, "GET", "/v2/library/abc/chunksums/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad") - w.Header().Set("Content-Location", "http://blob.store/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad") - fmt.Fprintf(w, "%s 0-1\n", blob.DigestFromBytes("ab")) - fmt.Fprintf(w, "%s 2-2\n", blob.DigestFromBytes("c")) - case 3: - w.WriteHeader(http.StatusNotFound) - io.WriteString(w, `{"errors":[{"code":"BLOB_UNKNOWN"}]}`) - case 4: - io.WriteString(w, "c") - default: - t.Errorf("unexpected steps %d: %v", step.Load(), r) - http.Error(w, "unexpected steps", http.StatusInternalServerError) - } - }) - - c.MaxStreams = 1 - c.ChunkingThreshold = 1 // force chunking - - var written atomic.Int64 - ctx := WithTrace(t.Context(), &Trace{ - Update: func(l *Layer, n int64, err error) { - t.Log("trace:", l.Digest.Short(), n, err) - written.Add(n) - }, - }) - - err := c.Pull(ctx, "http://o.com/library/abc") - var got *Error - if !errors.As(err, &got) || got.Code != "BLOB_UNKNOWN" { - t.Fatalf("err = %v, want %v", err, got) - } - - if g := written.Load(); g != 1 { - t.Fatalf("wrote %d bytes, want 1", g) - } -} - -func TestPullChunksumStreamError(t *testing.T) { - var step atomic.Int64 - c, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) { - switch step.Add(1) { - case 1: - checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest") - io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`) - case 2: - w.Header().Set("Content-Location", "http://blob.store/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad") - - // Write one valid chunksum and one invalid chunksum - fmt.Fprintf(w, "%s 0-1\n", blob.DigestFromBytes("ab")) // valid - fmt.Fprint(w, "sha256:!") // invalid - case 3: - io.WriteString(w, "ab") - default: - t.Errorf("unexpected steps %d: %v", step.Load(), r) - http.Error(w, "unexpected steps", http.StatusInternalServerError) - } - }) - - c.ChunkingThreshold = 1 // force chunking - - got := c.Pull(ctx, "http://o.com/library/abc") - if !errors.Is(got, ErrIncomplete) { - t.Fatalf("err = %v, want %v", got, ErrIncomplete) - } -} - -type flushAfterWriter struct { - w io.Writer -} - -func (f *flushAfterWriter) Write(p []byte) (n int, err error) { - n, err = f.w.Write(p) - f.w.(http.Flusher).Flush() // panic if not a flusher - return -} - -func TestPullChunksumStreaming(t *testing.T) { - csr, csw := io.Pipe() - defer csw.Close() - - var step atomic.Int64 - c, _ := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) { - switch step.Add(1) { - case 1: - checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest") - io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`) - case 2: - w.Header().Set("Content-Location", "http://blob.store/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad") - fw := &flushAfterWriter{w} // ensure client gets data as it arrives by aggressively flushing - _, err := io.Copy(fw, csr) - if err != nil { - t.Errorf("copy: %v", err) - } - case 3: - io.WriteString(w, "ab") - case 4: - io.WriteString(w, "c") - default: - t.Errorf("unexpected steps %d: %v", step.Load(), r) - http.Error(w, "unexpected steps", http.StatusInternalServerError) - } - }) - - c.ChunkingThreshold = 1 // force chunking - - update := make(chan int64, 1) - ctx := WithTrace(t.Context(), &Trace{ - Update: func(l *Layer, n int64, err error) { - t.Log("trace:", l.Digest.Short(), n, err) - if n > 0 { - update <- n - } - }, - }) - - errc := make(chan error, 1) - go func() { - errc <- c.Pull(ctx, "http://o.com/library/abc") - }() - - // Send first chunksum and ensure it kicks off work immediately - fmt.Fprintf(csw, "%s 0-1\n", blob.DigestFromBytes("ab")) - if g := <-update; g != 2 { - t.Fatalf("got %d, want 2", g) - } - - // now send the second chunksum and ensure it kicks off work immediately - fmt.Fprintf(csw, "%s 2-2\n", blob.DigestFromBytes("c")) - if g := <-update; g != 3 { - t.Fatalf("got %d, want 3", g) - } - csw.Close() - testutil.Check(t, <-errc) -} - -func TestPullChunksumsCached(t *testing.T) { - var step atomic.Int64 - c, _ := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) { - switch step.Add(1) { - case 1: - checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest") - io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`) - case 2: - w.Header().Set("Content-Location", "http://blob.store/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad") - fmt.Fprintf(w, "%s 0-1\n", blob.DigestFromBytes("ab")) - fmt.Fprintf(w, "%s 2-2\n", blob.DigestFromBytes("c")) - case 3, 4: - switch rng := r.Header.Get("Range"); rng { - case "bytes=0-1": - io.WriteString(w, "ab") - case "bytes=2-2": - io.WriteString(w, "c") - default: - t.Errorf("unexpected range %q", rng) - } - default: - t.Errorf("unexpected steps %d: %v", step.Load(), r) - http.Error(w, "unexpected steps", http.StatusInternalServerError) - } - }) - - c.MaxStreams = 1 // force serial processing of chunksums - c.ChunkingThreshold = 1 // force chunking - - ctx, cancel := context.WithCancel(t.Context()) - defer cancel() - - // Cancel the pull after the first chunksum is processed, but before - // the second chunksum is processed (which is waiting because - // MaxStreams=1). This should cause the second chunksum to error out - // leaving the blob incomplete. - ctx = WithTrace(ctx, &Trace{ - Update: func(l *Layer, n int64, err error) { - if n > 0 { - cancel() - } - }, - }) - err := c.Pull(ctx, "http://o.com/library/abc") - if !errors.Is(err, context.Canceled) { - t.Fatalf("err = %v, want %v", err, context.Canceled) - } - - _, err = c.Cache.Resolve("o.com/library/abc:latest") - if !errors.Is(err, fs.ErrNotExist) { - t.Fatalf("err = %v, want nil", err) - } - - // Reset state and pull again to ensure the blob chunks that should - // have been cached are, and the remaining chunk was downloaded, making - // the blob complete. - step.Store(0) - var written atomic.Int64 - var cached atomic.Int64 - ctx = WithTrace(t.Context(), &Trace{ - Update: func(l *Layer, n int64, err error) { - t.Log("trace:", l.Digest.Short(), n, err) - if errors.Is(err, ErrCached) { - cached.Add(n) - } - written.Add(n) - }, - }) - - check := testutil.Checker(t) - - err = c.Pull(ctx, "http://o.com/library/abc") - check(err) - - _, err = c.Cache.Resolve("o.com/library/abc:latest") - check(err) - - if g := written.Load(); g != 5 { - t.Fatalf("wrote %d bytes, want 3", g) - } - if g := cached.Load(); g != 2 { // "ab" should have been cached - t.Fatalf("cached %d bytes, want 5", g) - } -} diff --git a/server/internal/client/ollama/trace.go b/server/internal/client/ollama/trace.go deleted file mode 100644 index a7cac0d5..00000000 --- a/server/internal/client/ollama/trace.go +++ /dev/null @@ -1,72 +0,0 @@ -package ollama - -import ( - "context" -) - -// Trace is a set of functions that are called to report progress during blob -// downloads and uploads. -// -// Use [WithTrace] to attach a Trace to a context for use with [Registry.Push] -// and [Registry.Pull]. -type Trace struct { - // Update is called during [Registry.Push] and [Registry.Pull] to - // report the progress of blob uploads and downloads. - // - // The n argument is the number of bytes transferred so far, and err is - // any error that has occurred. If n == 0, and err is nil, the download - // or upload has just started. If err is [ErrCached], the download or - // upload has been skipped because the blob is already present in the - // local cache or remote registry, respectively. Otherwise, if err is - // non-nil, the download or upload has failed. When l.Size == n, and - // err is nil, the download or upload has completed. - // - // A function assigned must be safe for concurrent use. The function is - // called synchronously and so should not block or take long to run. - Update func(_ *Layer, n int64, _ error) -} - -func (t *Trace) update(l *Layer, n int64, err error) { - if t.Update != nil { - t.Update(l, n, err) - } -} - -type traceKey struct{} - -// WithTrace adds a trace to the context for transfer progress reporting. -func WithTrace(ctx context.Context, t *Trace) context.Context { - old := traceFromContext(ctx) - if old == t { - // No change, return the original context. This also prevents - // infinite recursion below, if the caller passes the same - // Trace. - return ctx - } - - // Create a new Trace that wraps the old one, if any. If we used the - // same pointer t, we end up with a recursive structure. - composed := &Trace{ - Update: func(l *Layer, n int64, err error) { - if old != nil { - old.update(l, n, err) - } - t.update(l, n, err) - }, - } - return context.WithValue(ctx, traceKey{}, composed) -} - -var emptyTrace = &Trace{} - -// traceFromContext returns the Trace associated with ctx, or an empty Trace if -// none is found. -// -// It never returns nil. -func traceFromContext(ctx context.Context) *Trace { - t, _ := ctx.Value(traceKey{}).(*Trace) - if t == nil { - return emptyTrace - } - return t -} diff --git a/server/internal/registry/server.go b/server/internal/registry/server.go deleted file mode 100644 index f62a622a..00000000 --- a/server/internal/registry/server.go +++ /dev/null @@ -1,417 +0,0 @@ -// Package registry implements an http.Handler for handling local Ollama API -// model management requests. See [Local] for details. -package registry - -import ( - "cmp" - "context" - "encoding/json" - "errors" - "fmt" - "io" - "log/slog" - "net/http" - "slices" - "strings" - "sync" - "time" - - "github.com/ollama/ollama/server/internal/cache/blob" - "github.com/ollama/ollama/server/internal/client/ollama" - "github.com/ollama/ollama/server/internal/internal/backoff" -) - -// Local implements an http.Handler for handling local Ollama API model -// management requests, such as pushing, pulling, and deleting models. -// -// It can be arranged for all unknown requests to be passed through to a -// fallback handler, if one is provided. -type Local struct { - Client *ollama.Registry // required - Logger *slog.Logger // required - - // Fallback, if set, is used to handle requests that are not handled by - // this handler. - Fallback http.Handler - - // Prune, if set, is called to prune the local disk cache after a model - // is deleted. - Prune func() error // optional -} - -// serverError is like ollama.Error, but with a Status field for the HTTP -// response code. We want to avoid adding that field to ollama.Error because it -// would always be 0 to clients (we don't want to leak the status code in -// errors), and so it would be confusing to have a field that is always 0. -type serverError struct { - Status int `json:"-"` - - // TODO(bmizerany): Decide if we want to keep this and maybe - // bring back later. - Code string `json:"code"` - - Message string `json:"error"` -} - -func (e serverError) Error() string { - return e.Message -} - -// Common API errors -var ( - errMethodNotAllowed = &serverError{405, "method_not_allowed", "method not allowed"} - errNotFound = &serverError{404, "not_found", "not found"} - errModelNotFound = &serverError{404, "not_found", "model not found"} - errInternalError = &serverError{500, "internal_error", "internal server error"} -) - -type statusCodeRecorder struct { - _status int // use status() to get the status code - http.ResponseWriter -} - -func (r *statusCodeRecorder) WriteHeader(status int) { - if r._status == 0 { - r._status = status - r.ResponseWriter.WriteHeader(status) - } -} - -func (r *statusCodeRecorder) Write(b []byte) (int, error) { - r._status = r.status() - return r.ResponseWriter.Write(b) -} - -var ( - _ http.ResponseWriter = (*statusCodeRecorder)(nil) - _ http.CloseNotifier = (*statusCodeRecorder)(nil) - _ http.Flusher = (*statusCodeRecorder)(nil) -) - -// CloseNotify implements the http.CloseNotifier interface, for Gin. Remove with Gin. -// -// It panics if the underlying ResponseWriter is not a CloseNotifier. -func (r *statusCodeRecorder) CloseNotify() <-chan bool { - return r.ResponseWriter.(http.CloseNotifier).CloseNotify() -} - -// Flush implements the http.Flusher interface, for Gin. Remove with Gin. -// -// It panics if the underlying ResponseWriter is not a Flusher. -func (r *statusCodeRecorder) Flush() { - r.ResponseWriter.(http.Flusher).Flush() -} - -func (r *statusCodeRecorder) status() int { - return cmp.Or(r._status, 200) -} - -func (s *Local) ServeHTTP(w http.ResponseWriter, r *http.Request) { - rec := &statusCodeRecorder{ResponseWriter: w} - s.serveHTTP(rec, r) -} - -func (s *Local) serveHTTP(rec *statusCodeRecorder, r *http.Request) { - var errattr slog.Attr - proxied, err := func() (bool, error) { - 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) - return true, nil - } - return false, errNotFound - } - }() - if err != nil { - // We always log the error, so fill in the error log attribute - errattr = slog.String("error", err.Error()) - - var e *serverError - switch { - case errors.As(err, &e): - case errors.Is(err, ollama.ErrNameInvalid): - e = &serverError{400, "bad_request", err.Error()} - default: - e = errInternalError - } - - data, err := json.Marshal(e) - if err != nil { - // unreachable - panic(err) - } - rec.Header().Set("Content-Type", "application/json") - rec.WriteHeader(e.Status) - rec.Write(data) - - // fallthrough to log - } - - if !proxied { - // we're only responsible for logging if we handled the request - var level slog.Level - if rec.status() >= 500 { - level = slog.LevelError - } else if rec.status() >= 400 { - level = slog.LevelWarn - } - - s.Logger.LogAttrs(r.Context(), level, "http", - errattr, // report first in line to make it easy to find - - // TODO(bmizerany): Write a test to ensure that we are logging - // all of this correctly. That also goes for the level+error - // logic above. - slog.Int("status", rec.status()), - slog.String("method", r.Method), - slog.String("path", r.URL.Path), - slog.Int64("content-length", r.ContentLength), - slog.String("remote", r.RemoteAddr), - slog.String("proto", r.Proto), - slog.String("query", r.URL.RawQuery), - ) - } -} - -type params struct { - // DeprecatedName is the name of the model to push, pull, or delete, - // but is deprecated. New clients should use [Model] instead. - // - // Use [model()] to get the model name for both old and new API requests. - DeprecatedName string `json:"name"` - - // Model is the name of the model to push, pull, or delete. - // - // Use [model()] to get the model name for both old and new API requests. - Model string `json:"model"` - - // AllowNonTLS is a flag that indicates a client using HTTP - // is doing so, deliberately. - // - // Deprecated: This field is ignored and only present for this - // deprecation message. It should be removed in a future release. - // - // Users can just use http or https+insecure to show intent to - // communicate they want to do insecure things, without awkward and - // confusing flags such as this. - AllowNonTLS bool `json:"insecure"` - - // Stream, if true, will make the server send progress updates in a - // streaming of JSON objects. If false, the server will send a single - // JSON object with the final status as "success", or an error object - // if an error occurred. - // - // Unfortunately, this API was designed to be a bit awkward. Stream is - // defined to default to true if not present, so we need a way to check - // if the client decisively set it to false. So, we use a pointer to a - // bool. Gross. - // - // Use [stream()] to get the correct value for this field. - Stream *bool `json:"stream"` -} - -// model returns the model name for both old and new API requests. -func (p params) model() string { - return cmp.Or(p.Model, p.DeprecatedName) -} - -func (p params) stream() bool { - if p.Stream == nil { - return true - } - return *p.Stream -} - -func (s *Local) handleDelete(_ http.ResponseWriter, r *http.Request) error { - if r.Method != "DELETE" { - return errMethodNotAllowed - } - p, err := decodeUserJSON[*params](r.Body) - if err != nil { - return err - } - ok, err := s.Client.Unlink(p.model()) - if err != nil { - return err - } - if !ok { - return errModelNotFound - } - if s.Prune != nil { - return s.Prune() - } - return nil -} - -type progressUpdateJSON struct { - Error string `json:"error,omitempty,omitzero"` - Status string `json:"status,omitempty,omitzero"` - 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 - } - - enc := json.NewEncoder(w) - if !p.stream() { - if err := s.Client.Pull(r.Context(), p.model()); err != nil { - if errors.Is(err, ollama.ErrModelNotFound) { - return errModelNotFound - } - return err - } - enc.Encode(progressUpdateJSON{Status: "success"}) - return nil - } - - var mu sync.Mutex - var progress []progressUpdateJSON - flushProgress := func() { - mu.Lock() - progress := slices.Clone(progress) // make a copy and release lock before encoding to the wire - mu.Unlock() - for _, p := range progress { - enc.Encode(p) - } - fl, _ := w.(http.Flusher) - if fl != nil { - fl.Flush() - } - } - - t := time.NewTicker(1<<63 - 1) // "unstarted" timer - start := sync.OnceFunc(func() { - flushProgress() // flush initial state - t.Reset(100 * time.Millisecond) - }) - ctx := ollama.WithTrace(r.Context(), &ollama.Trace{ - Update: func(l *ollama.Layer, n int64, err error) { - if err != nil && !errors.Is(err, ollama.ErrCached) { - s.Logger.Error("pulling", "model", p.model(), "error", err) - return - } - - func() { - mu.Lock() - defer mu.Unlock() - for i, p := range progress { - if p.Digest == l.Digest { - progress[i].Completed = n - return - } - } - progress = append(progress, progressUpdateJSON{ - Digest: l.Digest, - Total: l.Size, - }) - }() - - // Block flushing progress updates until every - // layer is accounted for. Clients depend on a - // complete model size to calculate progress - // correctly; if they use an incomplete total, - // progress indicators would erratically jump - // as new layers are registered. - start() - }, - }) - - done := make(chan error, 1) - go func() (err error) { - defer func() { done <- err }() - for _, err := range backoff.Loop(ctx, 3*time.Second) { - if err != nil { - return err - } - err := s.Client.Pull(ctx, p.model()) - if canRetry(err) { - continue - } - return err - } - return nil - }() - - enc.Encode(progressUpdateJSON{Status: "pulling manifest"}) - for { - select { - case <-t.C: - flushProgress() - case err := <-done: - flushProgress() - if err != nil { - if errors.Is(err, ollama.ErrModelNotFound) { - return &serverError{ - Status: 404, - Code: "not_found", - Message: fmt.Sprintf("model %q not found", p.model()), - } - } else { - return err - } - } - - // Emulate old client pull progress (for now): - enc.Encode(progressUpdateJSON{Status: "verifying sha256 digest"}) - enc.Encode(progressUpdateJSON{Status: "writing manifest"}) - enc.Encode(progressUpdateJSON{Status: "success"}) - return nil - } - } -} - -func decodeUserJSON[T any](r io.Reader) (T, error) { - var v T - err := json.NewDecoder(r).Decode(&v) - if err == nil { - return v, nil - } - var zero T - - // Not sure why, but I can't seem to be able to use: - // - // errors.As(err, &json.UnmarshalTypeError{}) - // - // This is working fine in stdlib, so I'm not sure what rules changed - // and why this no longer works here. So, we do it the verbose way. - var a *json.UnmarshalTypeError - var b *json.SyntaxError - if errors.As(err, &a) || errors.As(err, &b) { - err = &serverError{Status: 400, Message: err.Error(), Code: "bad_request"} - } - if errors.Is(err, io.EOF) { - err = &serverError{Status: 400, Message: "empty request body", Code: "bad_request"} - } - return zero, err -} - -func canRetry(err error) bool { - if err == nil { - return false - } - var oe *ollama.Error - if errors.As(err, &oe) { - return oe.Temporary() - } - s := err.Error() - return cmp.Or( - errors.Is(err, context.DeadlineExceeded), - strings.Contains(s, "unreachable"), - strings.Contains(s, "no route to host"), - strings.Contains(s, "connection reset by peer"), - ) -} diff --git a/server/internal/registry/server_test.go b/server/internal/registry/server_test.go deleted file mode 100644 index 15d8d828..00000000 --- a/server/internal/registry/server_test.go +++ /dev/null @@ -1,302 +0,0 @@ -package registry - -import ( - "bytes" - "context" - "encoding/json" - "io" - "io/fs" - "net" - "net/http" - "net/http/httptest" - "os" - "regexp" - "strings" - "sync" - "testing" - - "github.com/ollama/ollama/server/internal/cache/blob" - "github.com/ollama/ollama/server/internal/client/ollama" - "github.com/ollama/ollama/server/internal/testutil" - "golang.org/x/tools/txtar" - - _ "embed" -) - -type panicTransport struct{} - -func (t *panicTransport) RoundTrip(r *http.Request) (*http.Response, error) { - panic("unexpected RoundTrip call") -} - -var panicOnRoundTrip = &http.Client{Transport: &panicTransport{}} - -// bytesResetter is an interface for types that can be reset and return a byte -// slice, only. This is to prevent inadvertent use of bytes.Buffer.Read/Write -// etc for the purpose of checking logs. -type bytesResetter interface { - Bytes() []byte - Reset() -} - -func newTestServer(t *testing.T, upstreamRegistry http.HandlerFunc) *Local { - t.Helper() - dir := t.TempDir() - err := os.CopyFS(dir, os.DirFS("testdata/models")) - if err != nil { - t.Fatal(err) - } - c, err := blob.Open(dir) - if err != nil { - t.Fatal(err) - } - - client := panicOnRoundTrip - if upstreamRegistry != nil { - s := httptest.NewTLSServer(upstreamRegistry) - t.Cleanup(s.Close) - tr := s.Client().Transport.(*http.Transport).Clone() - tr.DialContext = func(ctx context.Context, _, _ string) (net.Conn, error) { - var d net.Dialer - return d.DialContext(ctx, "tcp", s.Listener.Addr().String()) - } - client = &http.Client{Transport: tr} - } - - rc := &ollama.Registry{ - Cache: c, - HTTPClient: client, - Mask: "example.com/library/_:latest", - } - - l := &Local{ - Client: rc, - Logger: testutil.Slogger(t), - } - return l -} - -func (s *Local) send(t *testing.T, method, path, body string) *httptest.ResponseRecorder { - t.Helper() - ctx := ollama.WithTrace(t.Context(), &ollama.Trace{ - Update: func(l *ollama.Layer, n int64, err error) { - t.Logf("update: %s %d %v", l.Digest, n, err) - }, - }) - req := httptest.NewRequestWithContext(ctx, method, path, strings.NewReader(body)) - return s.sendRequest(t, req) -} - -func (s *Local) sendRequest(t *testing.T, req *http.Request) *httptest.ResponseRecorder { - t.Helper() - w := httptest.NewRecorder() - s.ServeHTTP(w, req) - return w -} - -type invalidReader struct{} - -func (r *invalidReader) Read(p []byte) (int, error) { - return 0, os.ErrInvalid -} - -// captureLogs is a helper to capture logs from the server. It returns a -// shallow copy of the server with a new logger and a bytesResetter for the -// logs. -func captureLogs(t *testing.T, s *Local) (*Local, bytesResetter) { - t.Helper() - log, logs := testutil.SlogBuffer() - l := *s // shallow copy - l.Logger = log - return &l, logs -} - -func TestServerDelete(t *testing.T) { - check := testutil.Checker(t) - - s := newTestServer(t, nil) - - _, err := s.Client.ResolveLocal("smol") - check(err) - - got := s.send(t, "DELETE", "/api/delete", `{"model": "smol"}`) - if got.Code != 200 { - t.Fatalf("Code = %d; want 200", got.Code) - } - - _, err = s.Client.ResolveLocal("smol") - if err == nil { - t.Fatal("expected smol to have been deleted") - } - - got = s.send(t, "DELETE", "/api/delete", `!`) - checkErrorResponse(t, got, 400, "bad_request", "invalid character '!' looking for beginning of value") - - got = s.send(t, "GET", "/api/delete", `{"model": "smol"}`) - checkErrorResponse(t, got, 405, "method_not_allowed", "method not allowed") - - got = s.send(t, "DELETE", "/api/delete", ``) - checkErrorResponse(t, got, 400, "bad_request", "empty request body") - - got = s.send(t, "DELETE", "/api/delete", `{"model": "://"}`) - checkErrorResponse(t, got, 400, "bad_request", "invalid or missing name") - - got = s.send(t, "DELETE", "/unknown_path", `{}`) // valid body - checkErrorResponse(t, got, 404, "not_found", "not found") - - s, logs := captureLogs(t, s) - req := httptest.NewRequestWithContext(t.Context(), "DELETE", "/api/delete", &invalidReader{}) - got = s.sendRequest(t, req) - checkErrorResponse(t, got, 500, "internal_error", "internal server error") - ok, err := regexp.Match(`ERROR.*error="invalid argument"`, logs.Bytes()) - check(err) - if !ok { - t.Logf("logs:\n%s", logs) - t.Fatalf("expected log to contain ERROR with invalid argument") - } -} - -//go:embed testdata/registry.txt -var registryTXT []byte - -var registryFS = sync.OnceValue(func() fs.FS { - // Txtar gets hung up on \r\n line endings, so we need to convert them - // to \n when parsing the txtar on Windows. - data := bytes.ReplaceAll(registryTXT, []byte("\r\n"), []byte("\n")) - a := txtar.Parse(data) - fsys, err := txtar.FS(a) - if err != nil { - panic(err) - } - return fsys -}) - -func TestServerPull(t *testing.T) { - modelsHandler := http.FileServerFS(registryFS()) - s := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case "/v2/library/BOOM/manifests/latest": - w.WriteHeader(999) - io.WriteString(w, `{"error": "boom"}`) - case "/v2/library/unknown/manifests/latest": - w.WriteHeader(404) - io.WriteString(w, `{"errors": [{"code": "MANIFEST_UNKNOWN", "message": "manifest unknown"}]}`) - default: - t.Logf("serving blob: %s", r.URL.Path) - modelsHandler.ServeHTTP(w, r) - } - }) - - checkResponse := func(got *httptest.ResponseRecorder, wantlines string) { - t.Helper() - if got.Code != 200 { - t.Errorf("Code = %d; want 200", got.Code) - } - gotlines := got.Body.String() - if strings.TrimSpace(gotlines) == "" { - gotlines = "" - } - t.Logf("got:\n%s", gotlines) - for want := range strings.Lines(wantlines) { - want = strings.TrimSpace(want) - want, unwanted := strings.CutPrefix(want, "!") - want = strings.TrimSpace(want) - if !unwanted && !strings.Contains(gotlines, want) { - t.Errorf("\t! missing %q in body", want) - } - if unwanted && strings.Contains(gotlines, want) { - t.Errorf("\t! unexpected %q in body", want) - } - } - } - - got := s.send(t, "POST", "/api/pull", `{"model": "smol"}`) - checkResponse(got, ` - {"status":"pulling manifest"} - {"digest":"sha256:68e0ec597aee59d35f8dc44942d7b17d471ade10d3aca07a5bb7177713950312","total":5,"completed":5} - {"status":"verifying sha256 digest"} - {"status":"writing manifest"} - {"status":"success"} - `) - - got = s.send(t, "POST", "/api/pull", `{"model": "unknown"}`) - checkResponse(got, ` - {"code":"not_found","error":"model \"unknown\" not found"} - `) - - got = s.send(t, "DELETE", "/api/pull", `{"model": "smol"}`) - checkErrorResponse(t, got, 405, "method_not_allowed", "method not allowed") - - got = s.send(t, "POST", "/api/pull", `!`) - checkErrorResponse(t, got, 400, "bad_request", "invalid character '!' looking for beginning of value") - - got = s.send(t, "POST", "/api/pull", ``) - checkErrorResponse(t, got, 400, "bad_request", "empty request body") - - got = s.send(t, "POST", "/api/pull", `{"model": "://"}`) - checkResponse(got, ` - {"code":"bad_request","error":"invalid or missing name: \"\""} - `) - - // Non-streaming pulls - got = s.send(t, "POST", "/api/pull", `{"model": "://", "stream": false}`) - checkErrorResponse(t, got, 400, "bad_request", "invalid or missing name") - got = s.send(t, "POST", "/api/pull", `{"model": "smol", "stream": false}`) - checkResponse(got, ` - {"status":"success"} - !digest - !total - !completed - `) - got = s.send(t, "POST", "/api/pull", `{"model": "unknown", "stream": false}`) - checkErrorResponse(t, got, 404, "not_found", "model not found") -} - -func TestServerUnknownPath(t *testing.T) { - s := newTestServer(t, nil) - got := s.send(t, "DELETE", "/api/unknown", `{}`) - checkErrorResponse(t, got, 404, "not_found", "not found") - - var fellback bool - s.Fallback = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - fellback = true - }) - got = s.send(t, "DELETE", "/api/unknown", `{}`) - if !fellback { - t.Fatal("expected Fallback to be called") - } - if got.Code != 200 { - t.Fatalf("Code = %d; want 200", got.Code) - } -} - -func checkErrorResponse(t *testing.T, got *httptest.ResponseRecorder, status int, code, msg string) { - t.Helper() - - var printedBody bool - errorf := func(format string, args ...any) { - t.Helper() - if !printedBody { - t.Logf("BODY:\n%s", got.Body.String()) - printedBody = true - } - t.Errorf(format, args...) - } - - if got.Code != status { - errorf("Code = %d; want %d", got.Code, status) - } - - // unmarshal the error as *ollama.Error (proving *serverError is an *ollama.Error) - var e *ollama.Error - if err := json.Unmarshal(got.Body.Bytes(), &e); err != nil { - errorf("unmarshal error: %v", err) - t.FailNow() - } - if e.Code != code { - errorf("Code = %q; want %q", e.Code, code) - } - if !strings.Contains(e.Message, msg) { - errorf("Message = %q; want to contain %q", e.Message, msg) - } -} diff --git a/server/internal/registry/testdata/models/blobs/sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99 b/server/internal/registry/testdata/models/blobs/sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99 deleted file mode 100644 index def4dffc..00000000 Binary files a/server/internal/registry/testdata/models/blobs/sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99 and /dev/null differ diff --git a/server/internal/registry/testdata/models/blobs/sha256-ecfb1acfca9c76444d622fcdc3840217bd502124a9d3687d438c19b3cb9c3cb1 b/server/internal/registry/testdata/models/blobs/sha256-ecfb1acfca9c76444d622fcdc3840217bd502124a9d3687d438c19b3cb9c3cb1 deleted file mode 100644 index 62114d06..00000000 --- a/server/internal/registry/testdata/models/blobs/sha256-ecfb1acfca9c76444d622fcdc3840217bd502124a9d3687d438c19b3cb9c3cb1 +++ /dev/null @@ -1 +0,0 @@ -{"schemaVersion":2,"mediaType":"application/vnd.docker.distribution.manifest.v2+json","config":{"mediaType":"application/vnd.docker.container.image.v1+json","digest":"sha256:ca239d7bd8ea90e4a5d2e6bf88f8d74a47b14336e73eb4e18bed4dd325018116","size":267},"layers":[{"mediaType":"application/vnd.ollama.image.model","digest":"sha256:a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99","size":24}]} \ No newline at end of file diff --git a/server/internal/registry/testdata/models/manifests/example.com/library/smol/latest b/server/internal/registry/testdata/models/manifests/example.com/library/smol/latest deleted file mode 100644 index 62114d06..00000000 --- a/server/internal/registry/testdata/models/manifests/example.com/library/smol/latest +++ /dev/null @@ -1 +0,0 @@ -{"schemaVersion":2,"mediaType":"application/vnd.docker.distribution.manifest.v2+json","config":{"mediaType":"application/vnd.docker.container.image.v1+json","digest":"sha256:ca239d7bd8ea90e4a5d2e6bf88f8d74a47b14336e73eb4e18bed4dd325018116","size":267},"layers":[{"mediaType":"application/vnd.ollama.image.model","digest":"sha256:a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99","size":24}]} \ No newline at end of file diff --git a/server/internal/registry/testdata/registry.txt b/server/internal/registry/testdata/registry.txt deleted file mode 100644 index 2fc363fc..00000000 --- a/server/internal/registry/testdata/registry.txt +++ /dev/null @@ -1,22 +0,0 @@ --- v2/library/smol/manifests/latest -- -{ - "schemaVersion": 2, - "mediaType": "application/vnd.docker.distribution.manifest.v2+json", - "config": { - "mediaType": "application/vnd.docker.container.image.v1+json", - "digest": "sha256:ca3d163bab055381827226140568f3bef7eaac187cebd76878e0b63e9e442356", - "size": 3 - }, - "layers": [ - { - "mediaType": "application/vnd.ollama.image.model", - "digest": "sha256:68e0ec597aee59d35f8dc44942d7b17d471ade10d3aca07a5bb7177713950312", - "size": 5 - } - ] -} - --- v2/library/smol/blobs/sha256:68e0ec597aee59d35f8dc44942d7b17d471ade10d3aca07a5bb7177713950312 -- -GGUF --- v2/library/smol/blobs/sha256:ca3d163bab055381827226140568f3bef7eaac187cebd76878e0b63e9e442356 -- -{} diff --git a/server/model_recommendations_test.go b/server/model_recommendations_test.go index a918e517..67e0d6f4 100644 --- a/server/model_recommendations_test.go +++ b/server/model_recommendations_test.go @@ -360,7 +360,7 @@ func TestModelRecommendationsRouteRegistration(t *testing.T) { cache.set([]api.ModelRecommendation{{Model: "route-model", Description: "route description"}}) s := &Server{modelCaches: &modelCaches{recommendations: cache}} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatalf("GenerateRoutes failed: %v", err) } diff --git a/server/model_show_cache_test.go b/server/model_show_cache_test.go index 04bab53b..70312d64 100644 --- a/server/model_show_cache_test.go +++ b/server/model_show_cache_test.go @@ -323,7 +323,7 @@ func TestModelShowCacheCloudColdMissFallsBackToProxy(t *testing.T) { withCloudProxyBaseURL(t, upstream.URL) s := &Server{modelCaches: &modelCaches{show: newModelShowCache()}} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatalf("GenerateRoutes failed: %v", err) } diff --git a/server/routes.go b/server/routes.go index 432dc408..e9d34c76 100644 --- a/server/routes.go +++ b/server/routes.go @@ -44,8 +44,6 @@ import ( "github.com/ollama/ollama/middleware" "github.com/ollama/ollama/model/parsers" "github.com/ollama/ollama/model/renderers" - "github.com/ollama/ollama/server/internal/client/ollama" - "github.com/ollama/ollama/server/internal/registry" "github.com/ollama/ollama/template" "github.com/ollama/ollama/thinking" "github.com/ollama/ollama/tools" @@ -1833,7 +1831,7 @@ func allowedHostsMiddleware(addr net.Addr) gin.HandlerFunc { } } -func (s *Server) GenerateRoutes(rc *ollama.Registry) (http.Handler, error) { +func (s *Server) GenerateRoutes() (http.Handler, error) { corsConfig := cors.DefaultConfig() corsConfig.AllowWildcard = true corsConfig.AllowBrowserExtensions = true @@ -1923,18 +1921,6 @@ func (s *Server) GenerateRoutes(rc *ollama.Registry) (http.Handler, error) { // Inference (Anthropic compatibility) r.POST("/v1/messages", s.withInferenceRequestLogging("/v1/messages", cloudPassthroughMiddleware(cloudErrRemoteInferenceUnavailable), middleware.AnthropicMessagesMiddleware(), s.ChatHandler)...) - if rc != nil { - // wrap old with new - rs := ®istry.Local{ - Client: rc, - Logger: slog.Default(), // TODO(bmizerany): Take a logger, do not use slog.Default() - Fallback: r, - - Prune: PruneLayers, - } - return rs, nil - } - return r, nil } @@ -1998,16 +1984,11 @@ func Serve(ln net.Listener) error { return err } - var rc *ollama.Registry if useClient2 { - var err error - rc, err = ollama.DefaultRegistry() - if err != nil { - return err - } + slog.Warn("OLLAMA_EXPERIMENT=client2 is no longer available. Please remove this environment.") } - h, err := s.GenerateRoutes(rc) + h, err := s.GenerateRoutes() if err != nil { return err } diff --git a/server/routes_cloud_test.go b/server/routes_cloud_test.go index aaaf5b73..3ae8069a 100644 --- a/server/routes_cloud_test.go +++ b/server/routes_cloud_test.go @@ -171,7 +171,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -222,7 +222,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -265,7 +265,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -308,7 +308,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -351,7 +351,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -394,7 +394,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -450,7 +450,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -506,7 +506,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -557,7 +557,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -608,7 +608,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -675,7 +675,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -722,7 +722,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -768,7 +768,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -803,7 +803,7 @@ func TestCloudDisabledBlocksExplicitCloudPassthrough(t *testing.T) { t.Setenv("OLLAMA_NO_CLOUD", "1") s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -864,7 +864,7 @@ func TestCloudPassthroughStreamsPromptly(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -1048,7 +1048,7 @@ func TestCloudPassthroughSigningFailureReturnsUnauthorized(t *testing.T) { }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -1106,7 +1106,7 @@ func TestCloudPassthroughSigningFailureWithoutSigninURL(t *testing.T) { }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } diff --git a/server/routes_list_test.go b/server/routes_list_test.go index f7b8bd6e..eb409c8b 100644 --- a/server/routes_list_test.go +++ b/server/routes_list_test.go @@ -133,7 +133,7 @@ func TestOpenAIListMatchesTagsModels(t *testing.T) { setManifestTime("older-model:latest", older) setManifestTime("newer-model:latest", newer) - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } diff --git a/server/routes_test.go b/server/routes_test.go index 045f8a31..894035e6 100644 --- a/server/routes_test.go +++ b/server/routes_test.go @@ -28,7 +28,6 @@ import ( "github.com/ollama/ollama/fs/ggml" "github.com/ollama/ollama/manifest" "github.com/ollama/ollama/openai" - "github.com/ollama/ollama/server/internal/client/ollama" "github.com/ollama/ollama/types/model" "github.com/ollama/ollama/version" ) @@ -506,22 +505,7 @@ func TestRoutes(t *testing.T) { }, } - rc := &ollama.Registry{ - // This is a temporary measure to allow us to move forward, - // surfacing any code contacting ollama.com we do not intended - // to. - // - // Currently, this only handles DELETE /api/delete, which - // should not make any contact with the ollama.com registry, so - // be clear about that. - // - // Tests that do need to contact the registry here, will be - // consumed into our new server/api code packages and removed - // from here. - HTTPClient: panicOnRoundTrip, - } - - router, err := s.GenerateRoutes(rc) + router, err := s.GenerateRoutes() if err != nil { t.Fatalf("failed to generate routes: %v", err) } @@ -877,7 +861,7 @@ func TestShowCopilotUserAgentOverwritesExistingBasename(t *testing.T) { t.Fatalf("expected status code 200 creating model, actual %d", w.Code) } - h, err := s.GenerateRoutes(nil) + h, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -934,7 +918,7 @@ func TestShowCopilotUserAgentSetsBasenameWhenModelInfoIsEmpty(t *testing.T) { t.Fatalf("expected status code 200 creating model, actual %d", w.Code) } - h, err := s.GenerateRoutes(nil) + h, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } diff --git a/server/routes_web_experimental_test.go b/server/routes_web_experimental_test.go index 60be5f38..8bac9758 100644 --- a/server/routes_web_experimental_test.go +++ b/server/routes_web_experimental_test.go @@ -79,7 +79,7 @@ func TestExperimentalWebEndpointsPassthrough(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -129,7 +129,7 @@ func TestExperimentalWebEndpointsMissingBody(t *testing.T) { setTestHome(t, t.TempDir()) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -173,7 +173,7 @@ func TestExperimentalWebEndpointsCloudDisabled(t *testing.T) { t.Setenv("OLLAMA_NO_CLOUD", "1") s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -249,7 +249,7 @@ func TestExperimentalWebEndpointSigningFailureReturnsUnauthorized(t *testing.T) }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -303,7 +303,7 @@ func TestExperimentalWebEndpointSigningFailureWithoutSigninURL(t *testing.T) { }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) }