diff --git a/core/gallery/backends.go b/core/gallery/backends.go index 55f998a74..cbf99ac8e 100644 --- a/core/gallery/backends.go +++ b/core/gallery/backends.go @@ -436,7 +436,7 @@ func InstallBackend(ctx context.Context, systemState *system.SystemState, modelL // Record the OCI digest for upgrade detection (non-fatal on failure) if uri.LooksLikeOCI() { - digest, digestErr := oci.GetImageDigest(string(uri), "", nil, nil) + digest, digestErr := oci.GetImageDigest(uri.OCIReference(), "", nil, nil) if digestErr != nil { xlog.Warn("Failed to get OCI image digest for backend", "uri", string(uri), "error", digestErr) } else { diff --git a/core/gallery/backends_oci_uri_test.go b/core/gallery/backends_oci_uri_test.go new file mode 100644 index 000000000..fcb8a930a --- /dev/null +++ b/core/gallery/backends_oci_uri_test.go @@ -0,0 +1,121 @@ +package gallery + +import ( + "context" + "os" + "path/filepath" + "strings" + + "github.com/google/go-containerregistry/pkg/crane" + "github.com/google/go-containerregistry/pkg/name" + "github.com/google/go-containerregistry/pkg/v1/empty" + "github.com/google/go-containerregistry/pkg/v1/mutate" + "github.com/google/go-containerregistry/pkg/v1/remote" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/pkg/system" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gopkg.in/yaml.v3" +) + +// pushBackendImage publishes a minimal backend image (just a run.sh) to the +// in-process registry and returns the registry host and the manifest digest. +func pushBackendImage(serverURL, repoPath, tag, runScript string) (string, string) { + GinkgoHelper() + host := strings.TrimPrefix(serverURL, "http://") + + layer, err := crane.Layer(map[string][]byte{"run.sh": []byte(runScript)}) + Expect(err).ToNot(HaveOccurred()) + img, err := mutate.AppendLayers(empty.Image, layer) + Expect(err).ToNot(HaveOccurred()) + + ref, err := name.ParseReference(host + "/" + repoPath + ":" + tag) + Expect(err).ToNot(HaveOccurred()) + Expect(remote.Write(ref, img)).To(Succeed()) + + digest, err := img.Digest() + Expect(err).ToNot(HaveOccurred()) + return host, digest.String() +} + +var _ = Describe("backends installed from an oci:// URI", func() { + var ( + registryURL string + systemState *system.SystemState + ml *model.ModelLoader + backendsDir string + ) + + BeforeEach(func() { + srv, _, _ := ociRegistry() + registryURL = srv.URL + + backendsDir = GinkgoT().TempDir() + var err error + systemState, err = system.GetSystemState(system.WithBackendPath(backendsDir)) + Expect(err).ToNot(HaveOccurred()) + ml = model.NewModelLoader(systemState) + }) + + installedDigest := func(name string) string { + GinkgoHelper() + meta, err := readBackendMetadata(filepath.Join(backendsDir, name)) + Expect(err).ToNot(HaveOccurred()) + Expect(meta).ToNot(BeNil()) + return meta.Digest + } + + // The digest lookup that follows an install used to hand the raw oci:// + // URI to the registry client, which read "oci" as the registry host and + // queried https://oci/v2/. The install still succeeded, so the only trace + // was an empty digest in metadata.json and a warning in the log. + DescribeTable("records the image digest from the registry the URI names", + func(form string) { + host, digest := pushBackendImage(registryURL, "acme/backend", "v1", "#!/bin/sh\necho v1\n") + uri := "oci://" + host + "/acme/backend:v1" + if form == "digest" { + uri = "oci://" + host + "/acme/backend@" + digest + } + + Expect(InstallBackend(context.Background(), systemState, ml, &GalleryBackend{ + Metadata: Metadata{Name: "acme-backend"}, + URI: uri, + }, nil, false)).To(Succeed()) + + Expect(filepath.Join(backendsDir, "acme-backend", "run.sh")).To(BeARegularFile()) + Expect(installedDigest("acme-backend")).To(Equal(digest)) + }, + Entry("tag form", "tag"), + Entry("digest form", "digest"), + ) + + It("records the new image digest after an upgrade", func() { + host, oldDigest := pushBackendImage(registryURL, "acme/backend", "v1", "#!/bin/sh\necho v1\n") + _, newDigest := pushBackendImage(registryURL, "acme/backend", "v2", "#!/bin/sh\necho v2\n") + Expect(newDigest).ToNot(Equal(oldDigest)) + + Expect(InstallBackend(context.Background(), systemState, ml, &GalleryBackend{ + Metadata: Metadata{Name: "acme-backend"}, + URI: "oci://" + host + "/acme/backend:v1", + }, nil, false)).To(Succeed()) + Expect(installedDigest("acme-backend")).To(Equal(oldDigest)) + + galleryFile := filepath.Join(backendsDir, "gallery.yaml") + data, err := yaml.Marshal([]GalleryBackend{{ + Metadata: Metadata{Name: "acme-backend"}, + URI: "oci://" + host + "/acme/backend@" + newDigest, + }}) + Expect(err).ToNot(HaveOccurred()) + Expect(os.WriteFile(galleryFile, data, 0o644)).To(Succeed()) + galleries := []config.Gallery{{Name: "acme", URL: "file://" + galleryFile}} + + upgrades, err := CheckBackendUpgrades(context.Background(), galleries, systemState) + Expect(err).ToNot(HaveOccurred()) + Expect(upgrades).To(HaveKey("acme-backend")) + Expect(upgrades["acme-backend"].AvailableDigest).To(Equal(newDigest)) + + Expect(UpgradeBackend(context.Background(), systemState, ml, galleries, "acme-backend", nil, false)).To(Succeed()) + Expect(installedDigest("acme-backend")).To(Equal(newDigest)) + }) +}) diff --git a/core/gallery/gallery_oci.go b/core/gallery/gallery_oci.go index 27915dc6f..620ccc2cb 100644 --- a/core/gallery/gallery_oci.go +++ b/core/gallery/gallery_oci.go @@ -143,7 +143,7 @@ func fetchOCIGalleryIndex(ctx context.Context, g config.Gallery, candidate, base return body, nil } - pullRef := strings.TrimPrefix(candidate, downloader.OCIPrefix) + pullRef := downloader.URI(candidate).OCIReference() if g.Verification != nil { // Resolve first, verify the digest, then pull that same digest. diff --git a/core/gallery/importers/llama-cpp.go b/core/gallery/importers/llama-cpp.go index a1cbb6d1b..c803f8ae2 100644 --- a/core/gallery/importers/llama-cpp.go +++ b/core/gallery/importers/llama-cpp.go @@ -182,7 +182,7 @@ func (i *LlamaCPPImporter) Import(details Details) (gallery.ModelConfig, error) switch { case uri.LooksLikeOCI(): - ociName := strings.TrimPrefix(string(uri), downloader.OCIPrefix) + ociName := uri.OCIReference() ociName = strings.TrimPrefix(ociName, downloader.OllamaPrefix) ociName = strings.ReplaceAll(ociName, "/", "__") ociName = strings.ReplaceAll(ociName, ":", "__") diff --git a/core/gallery/upgrade.go b/core/gallery/upgrade.go index 4af56e256..fb6bc52c8 100644 --- a/core/gallery/upgrade.go +++ b/core/gallery/upgrade.go @@ -169,13 +169,8 @@ func CheckUpgradesAgainst(ctx context.Context, galleries []config.Gallery, syste } // Fall back to OCI digest comparison when versions are unavailable. - if downloader.URI(galleryEntry.URI).LooksLikeOCI() { - // Strip the oci:// scheme — name.ParseReference (called by - // GetImageDigest) cannot parse it. Self-hosted registries must - // set the scheme to be recognised at all, so without stripping - // they silently lose upgrade detection. - remoteDigest, err := oci.GetImageDigest( - strings.TrimPrefix(galleryEntry.URI, downloader.OCIPrefix), "", nil, nil) + if galleryURI := downloader.URI(galleryEntry.URI); galleryURI.LooksLikeOCI() { + remoteDigest, err := oci.GetImageDigest(galleryURI.OCIReference(), "", nil, nil) if err != nil { xlog.Warn("Failed to get remote OCI digest for upgrade check", "backend", installed.Metadata.Name, "error", err) continue @@ -359,7 +354,7 @@ func UpgradeBackend(ctx context.Context, systemState *system.SystemState, modelL // Record OCI digest if applicable (non-fatal on failure) if uri.LooksLikeOCI() { - digest, digestErr := oci.GetImageDigest(galleryEntry.URI, "", nil, nil) + digest, digestErr := oci.GetImageDigest(uri.OCIReference(), "", nil, nil) if digestErr != nil { xlog.Warn("Failed to get OCI image digest after upgrade", "uri", galleryEntry.URI, "error", digestErr) } else { diff --git a/pkg/downloader/uri.go b/pkg/downloader/uri.go index 8da23931d..8db9d6dc7 100644 --- a/pkg/downloader/uri.go +++ b/pkg/downloader/uri.go @@ -315,6 +315,14 @@ func (s URI) LooksLikeOCI() bool { strings.HasPrefix(string(s), "docker.io") } +// OCIReference returns the registry reference an OCI URI names, without the +// oci:// scheme. The scheme is LocalAI's own marker for "this is an image": +// registry clients do not know it and read "oci" as the registry host, so +// every consumer that hands a URI to a registry client goes through here. +func (s URI) OCIReference() string { + return strings.TrimPrefix(string(s), OCIPrefix) +} + func (s URI) LooksLikeOCIFile() bool { return strings.HasPrefix(string(s), OCIFilePrefix) } @@ -606,7 +614,7 @@ func (uri URI) DownloadFileWithContext(ctx context.Context, filePath, sha string return oci.ExtractOCIImage(ctx, img, url, filePath, downloadStatus) } - url = strings.TrimPrefix(url, OCIPrefix) + url = URI(url).OCIReference() img, err := oci.GetImage(url, "", nil, nil) if err != nil { return fmt.Errorf("failed to get image %q: %v", url, err) diff --git a/pkg/downloader/uri_test.go b/pkg/downloader/uri_test.go index 9cb667b57..bbc202c9a 100644 --- a/pkg/downloader/uri_test.go +++ b/pkg/downloader/uri_test.go @@ -87,6 +87,17 @@ var _ = Describe("Gallery API tests", func() { }) }) +var _ = Describe("OCIReference", func() { + DescribeTable("returns the registry reference without the oci:// scheme", + func(uri, want string) { + Expect(URI(uri).OCIReference()).To(Equal(want)) + }, + Entry("tag form", "oci://registry.example.com/acme/backend:v1", "registry.example.com/acme/backend:v1"), + Entry("digest form", "oci://registry.example.com:5000/acme/backend@sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", "registry.example.com:5000/acme/backend@sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"), + Entry("reference without a scheme", "quay.io/acme/backend:latest", "quay.io/acme/backend:latest"), + ) +}) + var _ = Describe("ContentLength", func() { Context("local file", func() { It("returns file size for existing file", func() {