From 37a5544fd72eb41bf8d8dbd8181ec489cc3636dd Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Mon, 26 Aug 2024 13:19:56 +0800 Subject: [PATCH] Support ollama image media types Signed-off-by: Xiaodong Ye --- internal/manifest/manifest.go | 21 +++++++++++++++++++++ manifest/manifest.go | 23 +++++++++++++++++++++++ storage/storage_dest.go | 26 +++++++++++++++++++------- 3 files changed, 63 insertions(+), 7 deletions(-) diff --git a/internal/manifest/manifest.go b/internal/manifest/manifest.go index 3fb52104a6..5744a3e869 100644 --- a/internal/manifest/manifest.go +++ b/internal/manifest/manifest.go @@ -32,6 +32,27 @@ const ( DockerV2Schema2ForeignLayerMediaType = "application/vnd.docker.image.rootfs.foreign.diff.tar" // DockerV2Schema2ForeignLayerMediaType is the MIME type used for gzipped schema 2 foreign layers. DockerV2Schema2ForeignLayerMediaTypeGzip = "application/vnd.docker.image.rootfs.foreign.diff.tar.gzip" + + // OllamaImageLayerMediaTypePrefix is the prefix for Ollama image layer media types. + OllamaImageLayerMediaTypePrefix = "application/vnd.ollama.image." + // OllamaImageLayerMediaType is the media type used for Ollama image model layers. + OllamaImageModelLayerMediaType = OllamaImageLayerMediaTypePrefix + "model" + // OllamaImageAdapterLayerMediaType is the media type used for Ollama image adapter layers. + OllamaImageAdapterLayerMediaType = OllamaImageLayerMediaTypePrefix + "adapter" + // OllamaImageProjectorLayerMediaType is the media type used for Ollama image projector layers. + OllamaImageProjectorLayerMediaType = OllamaImageLayerMediaTypePrefix + "projector" + // OllamaImagePromptLayerMediaType is the media type used for Ollama image prompt layers. + OllamaImagePromptLayerMediaType = OllamaImageLayerMediaTypePrefix + "prompt" + // OllamaImageTemplateLayerMediaType is the media type used for Ollama image template layers. + OllamaImageTemplateLayerMediaType = OllamaImageLayerMediaTypePrefix + "template" + // OllamaImageSystemLayerMediaType is the media type used for Ollama image system layers. + OllamaImageSystemLayerMediaType = OllamaImageLayerMediaTypePrefix + "system" + // OllamaImageParamsLayerMediaType is the media type used for Ollama image params layers. + OllamaImageParamsLayerMediaType = OllamaImageLayerMediaTypePrefix + "params" + // OllamaImageMessagesLayerMediaType is the media type used for Ollama image messages layers. + OllamaImageMessagesLayerMediaType = OllamaImageLayerMediaTypePrefix + "messages" + // OllamaImageLicenseLayerMediaType is the media type used for Ollama image license layers. + OllamaImageLicenseLayerMediaType = OllamaImageLayerMediaTypePrefix + "license" ) // GuessMIMEType guesses MIME type of a manifest and returns it _if it is recognized_, or "" if unknown or unrecognized. diff --git a/manifest/manifest.go b/manifest/manifest.go index d8f37eb45d..839d83e3ae 100644 --- a/manifest/manifest.go +++ b/manifest/manifest.go @@ -32,6 +32,27 @@ const ( DockerV2Schema2ForeignLayerMediaType = manifest.DockerV2Schema2ForeignLayerMediaType // DockerV2Schema2ForeignLayerMediaType is the MIME type used for gzipped schema 2 foreign layers. DockerV2Schema2ForeignLayerMediaTypeGzip = manifest.DockerV2Schema2ForeignLayerMediaTypeGzip + + // OllamaImageLayerMediaTypePrefix is the prefix for all Ollama image layer media types. + OllamaImageLayerMediaTypePrefix = manifest.OllamaImageLayerMediaTypePrefix + // OllamaImageModelLayerMediaType is the media type used for model layers. + OllamaImageModelLayerMediaType = manifest.OllamaImageModelLayerMediaType + // OllamaImageAdapterLayerMediaType is the media type used for adapter layers. + OllamaImageAdapterLayerMediaType = manifest.OllamaImageAdapterLayerMediaType + // OllamaImageProjectorLayerMediaType is the media type used for projector layers. + OllamaImageProjectorLayerMediaType = manifest.OllamaImageProjectorLayerMediaType + // OllamaImagePromptLayerMediaType is the media type used for prompt layers. + OllamaImagePromptLayerMediaType = manifest.OllamaImagePromptLayerMediaType + // OllamaImageTemplateLayerMediaType is the media type used for template layers. + OllamaImageTemplateLayerMediaType = manifest.OllamaImageTemplateLayerMediaType + // OllamaImageSystemLayerMediaType is the media type used for system layers. + OllamaImageSystemLayerMediaType = manifest.OllamaImageSystemLayerMediaType + // OllamaImageParamsLayerMediaType is the media type used for params layers. + OllamaImageParamsLayerMediaType = manifest.OllamaImageParamsLayerMediaType + // OllamaImageMessagesLayerMediaType is the media type used for messages layers. + OllamaImageMessagesLayerMediaType = manifest.OllamaImageMessagesLayerMediaType + // OllamaImageLicenseLayerMediaType is the media type used for license layers. + OllamaImageLicenseLayerMediaType = manifest.OllamaImageLicenseLayerMediaType ) // NonImageArtifactError (detected via errors.As) is used when asking for an image-specific operation @@ -43,6 +64,8 @@ func SupportedSchema2MediaType(m string) error { switch m { case DockerV2ListMediaType, DockerV2Schema1MediaType, DockerV2Schema1SignedMediaType, DockerV2Schema2ConfigMediaType, DockerV2Schema2ForeignLayerMediaType, DockerV2Schema2ForeignLayerMediaTypeGzip, DockerV2Schema2LayerMediaType, DockerV2Schema2MediaType, DockerV2SchemaLayerMediaTypeUncompressed: return nil + case OllamaImageModelLayerMediaType, OllamaImageAdapterLayerMediaType, OllamaImageProjectorLayerMediaType, OllamaImagePromptLayerMediaType, OllamaImageTemplateLayerMediaType, OllamaImageSystemLayerMediaType, OllamaImageParamsLayerMediaType, OllamaImageMessagesLayerMediaType, OllamaImageLicenseLayerMediaType: + return nil default: return fmt.Errorf("unsupported docker v2s2 media type: %q", m) } diff --git a/storage/storage_dest.go b/storage/storage_dest.go index 842a3ab068..a5d19423f1 100644 --- a/storage/storage_dest.go +++ b/storage/storage_dest.go @@ -13,6 +13,7 @@ import ( "os" "path/filepath" "slices" + "strings" "sync" "sync/atomic" @@ -125,8 +126,9 @@ type storageImageDestinationLockProtected struct { // addedLayerInfo records data about a layer to use in this image. type addedLayerInfo struct { - digest digest.Digest // Mandatory, the digest of the layer. - emptyLayer bool // The layer is an “empty”/“throwaway” one, and may or may not be physically represented in various transport / storage systems. false if the manifest type does not have the concept. + digest digest.Digest // Mandatory, the digest of the layer. + emptyLayer bool // The layer is an “empty”/“throwaway” one, and may or may not be physically represented in various transport / storage systems. false if the manifest type does not have the concept. + layerFilename *string } // newImageDestination sets us up to write a new image, caching blobs in a temporary directory until @@ -218,9 +220,18 @@ func (s *storageImageDestination) PutBlobWithOptions(ctx context.Context, stream return info, nil } + layerFilename := func() *string { + if strings.HasPrefix(blobinfo.MediaType, manifest.OllamaImageLayerMediaTypePrefix) { + filename := strings.TrimPrefix(blobinfo.MediaType, manifest.OllamaImageLayerMediaTypePrefix) + return &filename + } + return nil + }() + return info, s.queueOrCommit(*options.LayerIndex, addedLayerInfo{ - digest: info.Digest, - emptyLayer: options.EmptyLayer, + digest: info.Digest, + emptyLayer: options.EmptyLayer, + layerFilename: layerFilename, }) } @@ -817,7 +828,7 @@ func (s *storageImageDestination) commitLayer(index int, info addedLayerInfo, si if index != 0 { prev, ok := s.indexToStorageID[index-1] if !ok { - return false, fmt.Errorf("Internal error: commitLayer called with previous layer %d not committed yet", index-1) + return false, fmt.Errorf("internal error: commitLayer called with previous layer %d not committed yet", index-1) } parentLayer = prev } @@ -873,7 +884,7 @@ func (s *storageImageDestination) commitLayer(index int, info addedLayerInfo, si return false, nil } - layer, err := s.createNewLayer(index, info.digest, parentLayer, id) + layer, err := s.createNewLayer(index, info.digest, parentLayer, id, info.layerFilename) if err != nil { return false, err } @@ -886,7 +897,7 @@ func (s *storageImageDestination) commitLayer(index int, info addedLayerInfo, si // createNewLayer creates a new layer newLayerID for (index, layerDigest) on top of parentLayer (which may be ""). // If the layer cannot be committed yet, the function returns (nil, nil). -func (s *storageImageDestination) createNewLayer(index int, layerDigest digest.Digest, parentLayer, newLayerID string) (*storage.Layer, error) { +func (s *storageImageDestination) createNewLayer(index int, layerDigest digest.Digest, parentLayer, newLayerID string, layerFilename *string) (*storage.Layer, error) { s.lock.Lock() diffOutput, ok := s.lockProtected.diffOutputs[index] s.lock.Unlock() @@ -1061,6 +1072,7 @@ func (s *storageImageDestination) createNewLayer(index int, layerDigest digest.D OriginalDigest: trustedOriginalDigest, // This might be "" if trusted.layerIdentifiedByTOC; in that case PutLayer will compute the value from the stream. UncompressedDigest: trusted.diffID, + LayerFilename: layerFilename, }, file) if err != nil && !errors.Is(err, storage.ErrDuplicateID) { return nil, fmt.Errorf("adding layer with blob %q/%q/%q: %w", trusted.blobDigest, trusted.tocDigest, trusted.diffID, err)