diff --git a/cmd/cli/commands/package.go b/cmd/cli/commands/package.go index e793150b0..4f98ea1c1 100644 --- a/cmd/cli/commands/package.go +++ b/cmd/cli/commands/package.go @@ -11,6 +11,7 @@ import ( "path/filepath" "github.com/docker/model-runner/pkg/distribution/builder" + "github.com/docker/model-runner/pkg/distribution/distribution" "github.com/docker/model-runner/pkg/distribution/packaging" "github.com/docker/model-runner/pkg/distribution/registry" "github.com/docker/model-runner/pkg/distribution/tarball" @@ -26,11 +27,12 @@ func newPackagedCmd() *cobra.Command { var opts packageOptions c := &cobra.Command{ - Use: "package (--gguf | --safetensors-dir ) [--license ...] [--context-size ] [--push] MODEL", - Short: "Package a GGUF file or Safetensors directory into a Docker model OCI artifact.", - Long: "Package a GGUF file or Safetensors directory into a Docker model OCI artifact, with optional licenses. The package is sent to the model-runner, unless --push is specified.\n" + + Use: "package (--gguf | --safetensors-dir | --from ) [--license ...] [--context-size ] [--push] MODEL", + Short: "Package a GGUF file, Safetensors directory, or existing model into a Docker model OCI artifact.", + Long: "Package a GGUF file, Safetensors directory, or existing model into a Docker model OCI artifact, with optional licenses. The package is sent to the model-runner, unless --push is specified.\n" + "When packaging a sharded GGUF model, --gguf should point to the first shard. All shard files should be siblings and should include the index in the file name (e.g. model-00001-of-00015.gguf).\n" + - "When packaging a Safetensors model, --safetensors-dir should point to a directory containing .safetensors files and config files (*.json, merges.txt). All files will be auto-discovered and config files will be packaged into a tar archive.", + "When packaging a Safetensors model, --safetensors-dir should point to a directory containing .safetensors files and config files (*.json, merges.txt). All files will be auto-discovered and config files will be packaged into a tar archive.\n" + + "When packaging from an existing model using --from, you can modify properties like context size to create a variant of the original model.", Args: func(cmd *cobra.Command, args []string) error { if len(args) != 1 { return fmt.Errorf( @@ -41,16 +43,27 @@ func newPackagedCmd() *cobra.Command { ) } - // Validate that either --gguf or --safetensors-dir is provided (mutually exclusive) - if opts.ggufPath == "" && opts.safetensorsDir == "" { + // Validate that exactly one of --gguf, --safetensors-dir, or --from is provided (mutually exclusive) + sourcesProvided := 0 + if opts.ggufPath != "" { + sourcesProvided++ + } + if opts.safetensorsDir != "" { + sourcesProvided++ + } + if opts.fromModel != "" { + sourcesProvided++ + } + + if sourcesProvided == 0 { return fmt.Errorf( - "Either --gguf or --safetensors-dir path is required.\n\n" + + "One of --gguf, --safetensors-dir, or --from is required.\n\n" + "See 'docker model package --help' for more information", ) } - if opts.ggufPath != "" && opts.safetensorsDir != "" { + if sourcesProvided > 1 { return fmt.Errorf( - "Cannot specify both --gguf and --safetensors-dir. Please use only one format.\n\n" + + "Cannot specify more than one of --gguf, --safetensors-dir, or --from. Please use only one source.\n\n" + "See 'docker model package --help' for more information", ) } @@ -133,6 +146,7 @@ func newPackagedCmd() *cobra.Command { c.Flags().StringVar(&opts.ggufPath, "gguf", "", "absolute path to gguf file") c.Flags().StringVar(&opts.safetensorsDir, "safetensors-dir", "", "absolute path to directory containing safetensors files and config") + c.Flags().StringVar(&opts.fromModel, "from", "", "reference to an existing model to repackage") c.Flags().StringVar(&opts.chatTemplatePath, "chat-template", "", "absolute path to chat template file (must be Jinja format)") c.Flags().StringArrayVarP(&opts.licensePaths, "license", "l", nil, "absolute path to a license file") c.Flags().StringArrayVar(&opts.dirTarPaths, "dir-tar", nil, "relative path to directory to package as tar (can be specified multiple times)") @@ -146,53 +160,88 @@ type packageOptions struct { contextSize uint64 ggufPath string safetensorsDir string + fromModel string licensePaths []string dirTarPaths []string push bool tag string } -func packageModel(cmd *cobra.Command, opts packageOptions) error { - var ( - target builder.Target - err error - ) - if opts.push { - target, err = registry.NewClient( - registry.WithUserAgent("docker-model-cli/" + desktop.Version), - ).NewTarget(opts.tag) - } else { - target, err = newModelRunnerTarget(desktopClient, opts.tag) - } - if err != nil { - return err - } +// builderInitResult contains the result of initializing a builder from various sources +type builderInitResult struct { + builder *builder.Builder + distClient *distribution.Client // Only set when building from existing model + cleanupFunc func() // Optional cleanup function for temporary files +} + +// initializeBuilder creates a package builder from GGUF, Safetensors, or existing model +func initializeBuilder(cmd *cobra.Command, opts packageOptions) (*builderInitResult, error) { + result := &builderInitResult{} + + if opts.fromModel != "" { + // Get the model store path + userHomeDir, err := os.UserHomeDir() + if err != nil { + return nil, fmt.Errorf("get user home directory: %w", err) + } + modelStorePath := filepath.Join(userHomeDir, ".docker", "models") + if envPath := os.Getenv("MODELS_PATH"); envPath != "" { + modelStorePath = envPath + } - // Create package builder based on model format - var pkg *builder.Builder - if opts.ggufPath != "" { + // Create a distribution client to access the model store + distClient, err := distribution.NewClient(distribution.WithStoreRootPath(modelStorePath)) + if err != nil { + return nil, fmt.Errorf("create distribution client: %w", err) + } + result.distClient = distClient + + // Package from existing model + cmd.PrintErrf("Reading model from store: %q\n", opts.fromModel) + + // Get the model from the local store + mdl, err := distClient.GetModel(opts.fromModel) + if err != nil { + return nil, fmt.Errorf("get model from store: %w", err) + } + + // Type assert to ModelArtifact - the Model from store implements both interfaces + modelArtifact, ok := mdl.(types.ModelArtifact) + if !ok { + return nil, fmt.Errorf("model does not implement ModelArtifact interface") + } + + cmd.PrintErrf("Creating builder from existing model\n") + result.builder, err = builder.FromModel(modelArtifact) + if err != nil { + return nil, fmt.Errorf("create builder from model: %w", err) + } + } else if opts.ggufPath != "" { cmd.PrintErrf("Adding GGUF file from %q\n", opts.ggufPath) - pkg, err = builder.FromGGUF(opts.ggufPath) + pkg, err := builder.FromGGUF(opts.ggufPath) if err != nil { - return fmt.Errorf("add gguf file: %w", err) + return nil, fmt.Errorf("add gguf file: %w", err) } + result.builder = pkg } else { // Safetensors model from directory cmd.PrintErrf("Scanning directory %q for safetensors model...\n", opts.safetensorsDir) safetensorsPaths, tempConfigArchive, err := packaging.PackageFromDirectory(opts.safetensorsDir) if err != nil { - return fmt.Errorf("scan safetensors directory: %w", err) + return nil, fmt.Errorf("scan safetensors directory: %w", err) } - // Clean up temp config archive when done + // Set up cleanup for temp config archive if tempConfigArchive != "" { - defer os.Remove(tempConfigArchive) + result.cleanupFunc = func() { + os.Remove(tempConfigArchive) + } } cmd.PrintErrf("Found %d safetensors file(s)\n", len(safetensorsPaths)) - pkg, err = builder.FromSafetensors(safetensorsPaths) + pkg, err := builder.FromSafetensors(safetensorsPaths) if err != nil { - return fmt.Errorf("create safetensors model: %w", err) + return nil, fmt.Errorf("create safetensors model: %w", err) } // Add config archive if it was created @@ -200,11 +249,44 @@ func packageModel(cmd *cobra.Command, opts packageOptions) error { cmd.PrintErrf("Adding config archive from directory\n") pkg, err = pkg.WithConfigArchive(tempConfigArchive) if err != nil { - return fmt.Errorf("add config archive: %w", err) + return nil, fmt.Errorf("add config archive: %w", err) } } + result.builder = pkg + } + + return result, nil +} + +func packageModel(cmd *cobra.Command, opts packageOptions) error { + var ( + target builder.Target + err error + ) + if opts.push { + target, err = registry.NewClient( + registry.WithUserAgent("docker-model-cli/" + desktop.Version), + ).NewTarget(opts.tag) + } else { + target, err = newModelRunnerTarget(desktopClient, opts.tag) + } + if err != nil { + return err } + // Initialize the package builder based on model format + initResult, err := initializeBuilder(cmd, opts) + if err != nil { + return err + } + // Clean up any temporary files when done + if initResult.cleanupFunc != nil { + defer initResult.cleanupFunc() + } + + pkg := initResult.builder + distClient := initResult.distClient + // Set context size if opts.contextSize > 0 { cmd.PrintErrf("Setting context size %d\n", opts.contextSize) @@ -227,32 +309,48 @@ func packageModel(cmd *cobra.Command, opts packageOptions) error { } } - // Process directory tar archives - if len(opts.dirTarPaths) > 0 { - // Determine base directory for resolving relative paths - var baseDir string - if opts.safetensorsDir != "" { - baseDir = opts.safetensorsDir - } else { - // For GGUF, use the directory containing the GGUF file - baseDir = filepath.Dir(opts.ggufPath) - } + // Check if we can use lightweight repackaging (config-only changes from existing model) + useLightweight := opts.fromModel != "" && pkg.HasOnlyConfigChanges() - processor := packaging.NewDirTarProcessor(opts.dirTarPaths, baseDir) - tarPaths, cleanup, err := processor.Process() - if err != nil { - return err + if useLightweight { + cmd.PrintErrln("Creating lightweight model variant...") + + // Get the model artifact with new config + builtModel := pkg.Model() + + // Write using lightweight method + if err := distClient.WriteLightweightModel(builtModel, []string{opts.tag}); err != nil { + return fmt.Errorf("failed to create lightweight model: %w", err) } - defer cleanup() - for _, tarPath := range tarPaths { - pkg, err = pkg.WithDirTar(tarPath) + cmd.PrintErrln("Model variant created successfully") + } else { + // Process directory tar archives + if len(opts.dirTarPaths) > 0 { + // Determine base directory for resolving relative paths + var baseDir string + if opts.safetensorsDir != "" { + baseDir = opts.safetensorsDir + } else { + // For GGUF, use the directory containing the GGUF file + baseDir = filepath.Dir(opts.ggufPath) + } + + processor := packaging.NewDirTarProcessor(opts.dirTarPaths, baseDir) + tarPaths, cleanup, err := processor.Process() if err != nil { - return fmt.Errorf("add directory tar: %w", err) + return err + } + defer cleanup() + + for _, tarPath := range tarPaths { + pkg, err = pkg.WithDirTar(tarPath) + if err != nil { + return fmt.Errorf("add directory tar: %w", err) + } } } } - if opts.push { cmd.PrintErrln("Pushing model to registry...") } else { diff --git a/cmd/cli/docs/reference/docker_model_package.yaml b/cmd/cli/docs/reference/docker_model_package.yaml index b30710331..b0166ee3a 100644 --- a/cmd/cli/docs/reference/docker_model_package.yaml +++ b/cmd/cli/docs/reference/docker_model_package.yaml @@ -1,11 +1,12 @@ command: docker model package short: | - Package a GGUF file or Safetensors directory into a Docker model OCI artifact. + Package a GGUF file, Safetensors directory, or existing model into a Docker model OCI artifact. long: |- - Package a GGUF file or Safetensors directory into a Docker model OCI artifact, with optional licenses. The package is sent to the model-runner, unless --push is specified. + Package a GGUF file, Safetensors directory, or existing model into a Docker model OCI artifact, with optional licenses. The package is sent to the model-runner, unless --push is specified. When packaging a sharded GGUF model, --gguf should point to the first shard. All shard files should be siblings and should include the index in the file name (e.g. model-00001-of-00015.gguf). When packaging a Safetensors model, --safetensors-dir should point to a directory containing .safetensors files and config files (*.json, merges.txt). All files will be auto-discovered and config files will be packaged into a tar archive. -usage: docker model package (--gguf | --safetensors-dir ) [--license ...] [--context-size ] [--push] MODEL + When packaging from an existing model using --from, you can modify properties like context size to create a variant of the original model. +usage: docker model package (--gguf | --safetensors-dir | --from ) [--license ...] [--context-size ] [--push] MODEL pname: docker model plink: docker_model.yaml options: @@ -39,6 +40,15 @@ options: experimentalcli: false kubernetes: false swarm: false + - option: from + value_type: string + description: reference to an existing model to repackage + deprecated: false + hidden: false + experimental: false + experimentalcli: false + kubernetes: false + swarm: false - option: gguf value_type: string description: absolute path to gguf file diff --git a/cmd/cli/docs/reference/model.md b/cmd/cli/docs/reference/model.md index 499b4fad7..6f71d4378 100644 --- a/cmd/cli/docs/reference/model.md +++ b/cmd/cli/docs/reference/model.md @@ -5,29 +5,29 @@ Docker Model Runner ### Subcommands -| Name | Description | -|:------------------------------------------------|:-------------------------------------------------------------------------------| -| [`df`](model_df.md) | Show Docker Model Runner disk usage | -| [`inspect`](model_inspect.md) | Display detailed information on one model | -| [`install-runner`](model_install-runner.md) | Install Docker Model Runner (Docker Engine only) | -| [`list`](model_list.md) | List the models pulled to your local environment | -| [`logs`](model_logs.md) | Fetch the Docker Model Runner logs | -| [`package`](model_package.md) | Package a GGUF file or Safetensors directory into a Docker model OCI artifact. | -| [`ps`](model_ps.md) | List running models | -| [`pull`](model_pull.md) | Pull a model from Docker Hub or HuggingFace to your local environment | -| [`push`](model_push.md) | Push a model to Docker Hub | -| [`reinstall-runner`](model_reinstall-runner.md) | Reinstall Docker Model Runner (Docker Engine only) | -| [`requests`](model_requests.md) | Fetch requests+responses from Docker Model Runner | -| [`restart-runner`](model_restart-runner.md) | Restart Docker Model Runner (Docker Engine only) | -| [`rm`](model_rm.md) | Remove local models downloaded from Docker Hub | -| [`run`](model_run.md) | Run a model and interact with it using a submitted prompt or chat mode | -| [`start-runner`](model_start-runner.md) | Start Docker Model Runner (Docker Engine only) | -| [`status`](model_status.md) | Check if the Docker Model Runner is running | -| [`stop-runner`](model_stop-runner.md) | Stop Docker Model Runner (Docker Engine only) | -| [`tag`](model_tag.md) | Tag a model | -| [`uninstall-runner`](model_uninstall-runner.md) | Uninstall Docker Model Runner (Docker Engine only) | -| [`unload`](model_unload.md) | Unload running models | -| [`version`](model_version.md) | Show the Docker Model Runner version | +| Name | Description | +|:------------------------------------------------|:------------------------------------------------------------------------------------------------| +| [`df`](model_df.md) | Show Docker Model Runner disk usage | +| [`inspect`](model_inspect.md) | Display detailed information on one model | +| [`install-runner`](model_install-runner.md) | Install Docker Model Runner (Docker Engine only) | +| [`list`](model_list.md) | List the models pulled to your local environment | +| [`logs`](model_logs.md) | Fetch the Docker Model Runner logs | +| [`package`](model_package.md) | Package a GGUF file, Safetensors directory, or existing model into a Docker model OCI artifact. | +| [`ps`](model_ps.md) | List running models | +| [`pull`](model_pull.md) | Pull a model from Docker Hub or HuggingFace to your local environment | +| [`push`](model_push.md) | Push a model to Docker Hub | +| [`reinstall-runner`](model_reinstall-runner.md) | Reinstall Docker Model Runner (Docker Engine only) | +| [`requests`](model_requests.md) | Fetch requests+responses from Docker Model Runner | +| [`restart-runner`](model_restart-runner.md) | Restart Docker Model Runner (Docker Engine only) | +| [`rm`](model_rm.md) | Remove local models downloaded from Docker Hub | +| [`run`](model_run.md) | Run a model and interact with it using a submitted prompt or chat mode | +| [`start-runner`](model_start-runner.md) | Start Docker Model Runner (Docker Engine only) | +| [`status`](model_status.md) | Check if the Docker Model Runner is running | +| [`stop-runner`](model_stop-runner.md) | Stop Docker Model Runner (Docker Engine only) | +| [`tag`](model_tag.md) | Tag a model | +| [`uninstall-runner`](model_uninstall-runner.md) | Uninstall Docker Model Runner (Docker Engine only) | +| [`unload`](model_unload.md) | Unload running models | +| [`version`](model_version.md) | Show the Docker Model Runner version | diff --git a/cmd/cli/docs/reference/model_package.md b/cmd/cli/docs/reference/model_package.md index 44e7a4e32..49234cee6 100644 --- a/cmd/cli/docs/reference/model_package.md +++ b/cmd/cli/docs/reference/model_package.md @@ -1,9 +1,10 @@ # docker model package -Package a GGUF file or Safetensors directory into a Docker model OCI artifact, with optional licenses. The package is sent to the model-runner, unless --push is specified. +Package a GGUF file, Safetensors directory, or existing model into a Docker model OCI artifact, with optional licenses. The package is sent to the model-runner, unless --push is specified. When packaging a sharded GGUF model, --gguf should point to the first shard. All shard files should be siblings and should include the index in the file name (e.g. model-00001-of-00015.gguf). When packaging a Safetensors model, --safetensors-dir should point to a directory containing .safetensors files and config files (*.json, merges.txt). All files will be auto-discovered and config files will be packaged into a tar archive. +When packaging from an existing model using --from, you can modify properties like context size to create a variant of the original model. ### Options @@ -12,6 +13,7 @@ When packaging a Safetensors model, --safetensors-dir should point to a director | `--chat-template` | `string` | | absolute path to chat template file (must be Jinja format) | | `--context-size` | `uint64` | `0` | context size in tokens | | `--dir-tar` | `stringArray` | | relative path to directory to package as tar (can be specified multiple times) | +| `--from` | `string` | | reference to an existing model to repackage | | `--gguf` | `string` | | absolute path to gguf file | | `-l`, `--license` | `stringArray` | | absolute path to a license file | | `--push` | `bool` | | push to registry (if not set, the model is loaded into the Model Runner content store) | diff --git a/pkg/distribution/builder/builder.go b/pkg/distribution/builder/builder.go index f64dd7ab4..cb1d4d576 100644 --- a/pkg/distribution/builder/builder.go +++ b/pkg/distribution/builder/builder.go @@ -5,6 +5,8 @@ import ( "fmt" "io" + v1 "github.com/google/go-containerregistry/pkg/v1" + "github.com/docker/model-runner/pkg/distribution/internal/gguf" "github.com/docker/model-runner/pkg/distribution/internal/mutate" "github.com/docker/model-runner/pkg/distribution/internal/partial" @@ -14,7 +16,8 @@ import ( // Builder builds a model artifact type Builder struct { - model types.ModelArtifact + model types.ModelArtifact + originalLayers []v1.Layer // Snapshot of layers when created from existing model } // FromGGUF returns a *Builder that builds a model artifacts from a GGUF file @@ -39,6 +42,19 @@ func FromSafetensors(safetensorsPaths []string) (*Builder, error) { }, nil } +// FromModel returns a *Builder that builds model artifacts from an existing model artifact +func FromModel(mdl types.ModelArtifact) (*Builder, error) { + // Capture original layers for comparison + layers, err := mdl.Layers() + if err != nil { + return nil, fmt.Errorf("getting model layers: %w", err) + } + return &Builder{ + model: mdl, + originalLayers: layers, + }, nil +} + // WithLicense adds a license file to the artifact func (b *Builder) WithLicense(path string) (*Builder, error) { licenseLayer, err := partial.NewLayer(path, types.MediaTypeLicense) @@ -46,13 +62,15 @@ func (b *Builder) WithLicense(path string) (*Builder, error) { return nil, fmt.Errorf("license layer from %q: %w", path, err) } return &Builder{ - model: mutate.AppendLayers(b.model, licenseLayer), + model: mutate.AppendLayers(b.model, licenseLayer), + originalLayers: b.originalLayers, }, nil } func (b *Builder) WithContextSize(size uint64) *Builder { return &Builder{ - model: mutate.ContextSize(b.model, size), + model: mutate.ContextSize(b.model, size), + originalLayers: b.originalLayers, } } @@ -63,7 +81,8 @@ func (b *Builder) WithMultimodalProjector(path string) (*Builder, error) { return nil, fmt.Errorf("mmproj layer from %q: %w", path, err) } return &Builder{ - model: mutate.AppendLayers(b.model, mmprojLayer), + model: mutate.AppendLayers(b.model, mmprojLayer), + originalLayers: b.originalLayers, }, nil } @@ -74,7 +93,8 @@ func (b *Builder) WithChatTemplateFile(path string) (*Builder, error) { return nil, fmt.Errorf("chat template layer from %q: %w", path, err) } return &Builder{ - model: mutate.AppendLayers(b.model, templateLayer), + model: mutate.AppendLayers(b.model, templateLayer), + originalLayers: b.originalLayers, }, nil } @@ -98,7 +118,8 @@ func (b *Builder) WithConfigArchive(path string) (*Builder, error) { return nil, fmt.Errorf("config archive layer from %q: %w", path, err) } return &Builder{ - model: mutate.AppendLayers(b.model, configLayer), + model: mutate.AppendLayers(b.model, configLayer), + originalLayers: b.originalLayers, }, nil } @@ -110,7 +131,8 @@ func (b *Builder) WithDirTar(path string) (*Builder, error) { return nil, fmt.Errorf("dir tar layer from %q: %w", path, err) } return &Builder{ - model: mutate.AppendLayers(b.model, dirTarLayer), + model: mutate.AppendLayers(b.model, dirTarLayer), + originalLayers: b.originalLayers, }, nil } @@ -119,7 +141,50 @@ type Target interface { Write(context.Context, types.ModelArtifact, io.Writer) error } +// Model returns the underlying model artifact +func (b *Builder) Model() types.ModelArtifact { + return b.model +} + // Build finalizes the artifact and writes it to the given target, reporting progress to the given writer func (b *Builder) Build(ctx context.Context, target Target, pw io.Writer) error { return target.Write(ctx, b.model, pw) } + +// HasOnlyConfigChanges returns true if the builder was created from an existing model +// and only configuration changes were made (no layers added or removed). +// This is useful for determining if lightweight repackaging optimizations can be used. +func (b *Builder) HasOnlyConfigChanges() bool { + // If not created from an existing model, return false + if b.originalLayers == nil { + return false + } + + // Get current layers + currentLayers, err := b.model.Layers() + if err != nil { + return false + } + + // If layer count changed, files were added or removed + if len(currentLayers) != len(b.originalLayers) { + return false + } + + // Verify layer digests match to ensure no layer content changed + for i, origLayer := range b.originalLayers { + origDigest, err := origLayer.Digest() + if err != nil { + return false + } + currDigest, err := currentLayers[i].Digest() + if err != nil { + return false + } + if origDigest != currDigest { + return false + } + } + + return true +} diff --git a/pkg/distribution/builder/builder_test.go b/pkg/distribution/builder/builder_test.go index 24e10c77e..9a7bbb3b2 100644 --- a/pkg/distribution/builder/builder_test.go +++ b/pkg/distribution/builder/builder_test.go @@ -2,12 +2,15 @@ package builder_test import ( "context" + "fmt" "io" "path/filepath" + "strings" "testing" "github.com/docker/model-runner/pkg/distribution/builder" "github.com/docker/model-runner/pkg/distribution/types" + v1 "github.com/google/go-containerregistry/pkg/v1" ) func TestBuilder(t *testing.T) { @@ -142,6 +145,178 @@ func TestWithMultimodalProjectorChaining(t *testing.T) { // but we can verify the layers were added with correct media types above } +func TestFromModel(t *testing.T) { + // Step 1: Create an initial model from GGUF with context size 2048 + initialBuilder, err := builder.FromGGUF(filepath.Join("..", "assets", "dummy.gguf")) + if err != nil { + t.Fatalf("Failed to create initial builder from GGUF: %v", err) + } + + // Add license to the initial model + initialBuilder, err = initialBuilder.WithLicense(filepath.Join("..", "assets", "license.txt")) + if err != nil { + t.Fatalf("Failed to add license: %v", err) + } + + // Set initial context size + initialBuilder = initialBuilder.WithContextSize(2048) + + // Build the initial model + initialTarget := &fakeTarget{} + if err := initialBuilder.Build(t.Context(), initialTarget, nil); err != nil { + t.Fatalf("Failed to build initial model: %v", err) + } + + // Verify initial model properties + initialConfig, err := initialTarget.artifact.Config() + if err != nil { + t.Fatalf("Failed to get initial config: %v", err) + } + if initialConfig.ContextSize == nil || *initialConfig.ContextSize != 2048 { + t.Fatalf("Expected initial context size 2048, got %v", initialConfig.ContextSize) + } + + // Step 2: Use FromModel() to create a new builder from the existing model + repackagedBuilder, err := builder.FromModel(initialTarget.artifact) + if err != nil { + t.Fatalf("Failed to create builder from model: %v", err) + } + + // Step 3: Modify the context size to 4096 + repackagedBuilder = repackagedBuilder.WithContextSize(4096) + + // Step 4: Build the repackaged model + repackagedTarget := &fakeTarget{} + if err := repackagedBuilder.Build(t.Context(), repackagedTarget, nil); err != nil { + t.Fatalf("Failed to build repackaged model: %v", err) + } + + // Step 5: Verify the repackaged model has the new context size + repackagedConfig, err := repackagedTarget.artifact.Config() + if err != nil { + t.Fatalf("Failed to get repackaged config: %v", err) + } + + if repackagedConfig.ContextSize == nil || *repackagedConfig.ContextSize != 4096 { + t.Errorf("Expected repackaged context size 4096, got %v", repackagedConfig.ContextSize) + } + + // Step 6: Verify the original layers are preserved + initialManifest, err := initialTarget.artifact.Manifest() + if err != nil { + t.Fatalf("Failed to get initial manifest: %v", err) + } + + repackagedManifest, err := repackagedTarget.artifact.Manifest() + if err != nil { + t.Fatalf("Failed to get repackaged manifest: %v", err) + } + + // Should have the same number of layers (GGUF + license) + if len(repackagedManifest.Layers) != len(initialManifest.Layers) { + t.Errorf("Expected %d layers in repackaged model, got %d", len(initialManifest.Layers), len(repackagedManifest.Layers)) + } + + // Verify layer media types are preserved + for i, initialLayer := range initialManifest.Layers { + if i >= len(repackagedManifest.Layers) { + break + } + if initialLayer.MediaType != repackagedManifest.Layers[i].MediaType { + t.Errorf("Layer %d media type mismatch: expected %s, got %s", i, initialLayer.MediaType, repackagedManifest.Layers[i].MediaType) + } + } +} + +func TestFromModelWithAdditionalLayers(t *testing.T) { + // Create an initial model from GGUF + initialBuilder, err := builder.FromGGUF(filepath.Join("..", "assets", "dummy.gguf")) + if err != nil { + t.Fatalf("Failed to create initial builder from GGUF: %v", err) + } + + // Build the initial model + initialTarget := &fakeTarget{} + if err := initialBuilder.Build(t.Context(), initialTarget, nil); err != nil { + t.Fatalf("Failed to build initial model: %v", err) + } + + // Use FromModel() and add additional layers + repackagedBuilder, err := builder.FromModel(initialTarget.artifact) + if err != nil { + t.Fatalf("Failed to create builder from model: %v", err) + } + repackagedBuilder, err = repackagedBuilder.WithLicense(filepath.Join("..", "assets", "license.txt")) + if err != nil { + t.Fatalf("Failed to add license to repackaged model: %v", err) + } + + repackagedBuilder, err = repackagedBuilder.WithMultimodalProjector(filepath.Join("..", "assets", "dummy.mmproj")) + if err != nil { + t.Fatalf("Failed to add multimodal projector to repackaged model: %v", err) + } + + // Build the repackaged model + repackagedTarget := &fakeTarget{} + if err := repackagedBuilder.Build(t.Context(), repackagedTarget, nil); err != nil { + t.Fatalf("Failed to build repackaged model: %v", err) + } + + // Verify the repackaged model has all layers + initialManifest, err := initialTarget.artifact.Manifest() + if err != nil { + t.Fatalf("Failed to get initial manifest: %v", err) + } + + repackagedManifest, err := repackagedTarget.artifact.Manifest() + if err != nil { + t.Fatalf("Failed to get repackaged manifest: %v", err) + } + + // Should have original layers plus license and mmproj (2 additional layers) + expectedLayers := len(initialManifest.Layers) + 2 + if len(repackagedManifest.Layers) != expectedLayers { + t.Errorf("Expected %d layers in repackaged model, got %d", expectedLayers, len(repackagedManifest.Layers)) + } + + // Verify the new layers were added + hasLicense := false + hasMMProj := false + for _, layer := range repackagedManifest.Layers { + if layer.MediaType == types.MediaTypeLicense { + hasLicense = true + } + if layer.MediaType == types.MediaTypeMultimodalProjector { + hasMMProj = true + } + } + + if !hasLicense { + t.Error("Expected repackaged model to have license layer") + } + if !hasMMProj { + t.Error("Expected repackaged model to have multimodal projector layer") + } +} + +// TestFromModelErrorHandling tests that FromModel properly handles and surfaces errors from mdl.Layers() +func TestFromModelErrorHandling(t *testing.T) { + // Create a mock model that fails when Layers() is called + mockModel := &mockFailingModel{} + + // Attempt to create a builder from the failing model + _, err := builder.FromModel(mockModel) + if err == nil { + t.Fatal("Expected error when model.Layers() fails, got nil") + } + + // Verify the error message indicates the issue + expectedErrMsg := "getting model layers" + if !strings.Contains(err.Error(), expectedErrMsg) { + t.Errorf("Expected error message to contain %q, got: %v", expectedErrMsg, err) + } +} + var _ builder.Target = &fakeTarget{} type fakeTarget struct { @@ -152,3 +327,12 @@ func (ft *fakeTarget) Write(ctx context.Context, artifact types.ModelArtifact, w ft.artifact = artifact return nil } + +// mockFailingModel is a mock that fails when Layers() is called +type mockFailingModel struct { + types.ModelArtifact +} + +func (m *mockFailingModel) Layers() ([]v1.Layer, error) { + return nil, fmt.Errorf("simulated layers error") +} diff --git a/pkg/distribution/distribution/client.go b/pkg/distribution/distribution/client.go index 07360efa6..8b51c9010 100644 --- a/pkg/distribution/distribution/client.go +++ b/pkg/distribution/distribution/client.go @@ -386,6 +386,14 @@ func (c *Client) PushModel(ctx context.Context, tag string, progressWriter io.Wr return nil } +// WriteLightweightModel writes a model to the store without transferring layer data. +// This is used for config-only modifications where the layer data hasn't changed. +// The layers must already exist in the store. +func (c *Client) WriteLightweightModel(mdl types.ModelArtifact, tags []string) error { + c.log.Infoln("Writing lightweight model variant") + return c.store.WriteLightweight(mdl, tags) +} + func (c *Client) ResetStore() error { c.log.Infoln("Resetting store") if err := c.store.Reset(); err != nil { diff --git a/pkg/distribution/internal/store/store.go b/pkg/distribution/internal/store/store.go index bae3b6a37..ae4da969b 100644 --- a/pkg/distribution/internal/store/store.go +++ b/pkg/distribution/internal/store/store.go @@ -259,6 +259,52 @@ func (s *LocalStore) Write(mdl v1.Image, tags []string, w io.Writer) error { return err } +// WriteLightweight writes only the manifest and config for a model, assuming layers already exist in the store. +// This is used for config-only modifications where the layer data hasn't changed. +func (s *LocalStore) WriteLightweight(mdl v1.Image, tags []string) error { + // Verify that all layers already exist in the store + layers, err := mdl.Layers() + if err != nil { + return fmt.Errorf("getting layers: %w", err) + } + + for _, layer := range layers { + digest, err := layer.Digest() + if err != nil { + return fmt.Errorf("getting layer digest: %w", err) + } + hasBlob, err := s.hasBlob(digest) + if err != nil { + return fmt.Errorf("checking if layer %s exists: %w", digest, err) + } + if !hasBlob { + return fmt.Errorf("layer %s not found in store, cannot use lightweight write", digest) + } + } + + // Write the config JSON file + if err := s.writeConfigFile(mdl); err != nil { + return fmt.Errorf("writing config file: %w", err) + } + + // Write the manifest + digest, err := mdl.Digest() + if err != nil { + return fmt.Errorf("get digest: %w", err) + } + rm, err := mdl.RawManifest() + if err != nil { + return fmt.Errorf("get raw manifest: %w", err) + } + if err := s.WriteManifest(digest, rm); err != nil { + return fmt.Errorf("write manifest: %w", err) + } + if err := s.AddTags(digest.String(), tags); err != nil { + return fmt.Errorf("adding tags: %w", err) + } + return nil +} + // Read reads a model from the store by reference (either tag or ID) func (s *LocalStore) Read(reference string) (*Model, error) { models, err := s.List() diff --git a/pkg/distribution/internal/store/store_test.go b/pkg/distribution/internal/store/store_test.go index 31c75c46a..a63d71783 100644 --- a/pkg/distribution/internal/store/store_test.go +++ b/pkg/distribution/internal/store/store_test.go @@ -600,3 +600,357 @@ func newTestModelWithMultimodalProjector(t *testing.T) types.ModelArtifact { mdl = mutate.AppendLayers(mdl, licenseLayer, mmprojLayer) return mdl } + +// TestWriteLightweight tests the WriteLightweight method +func TestWriteLightweight(t *testing.T) { + // Create a temporary directory for the test store + tempDir, err := os.MkdirTemp("", "lightweight-write-test") + if err != nil { + t.Fatalf("Failed to create temp directory: %v", err) + } + defer os.RemoveAll(tempDir) + + // Create store + storePath := filepath.Join(tempDir, "lightweight-model-store") + s, err := store.New(store.Options{ + RootPath: storePath, + }) + if err != nil { + t.Fatalf("Failed to create store: %v", err) + } + + t.Run("SuccessfulLightweightWrite", func(t *testing.T) { + // Create and write a model normally to populate the store with blobs + baseModel := newTestModel(t) + if err := s.Write(baseModel, []string{"base-model:v1"}, nil); err != nil { + t.Fatalf("Write base model failed: %v", err) + } + + // Get original digest + originalDigest, err := baseModel.Digest() + if err != nil { + t.Fatalf("Failed to get original digest: %v", err) + } + + // Modify the model's config by changing context size + newContextSize := uint64(4096) + modifiedModel := mutate.ContextSize(baseModel, newContextSize) + + // Use WriteLightweight to write the modified model + if err := s.WriteLightweight(modifiedModel, []string{"base-model:v2"}); err != nil { + t.Fatalf("WriteLightweight failed: %v", err) + } + + // Verify the new model can be read + readModel, err := s.Read("base-model:v2") + if err != nil { + t.Fatalf("Read modified model failed: %v", err) + } + + // Verify the new model has a different digest + newDigest, err := readModel.Digest() + if err != nil { + t.Fatalf("Failed to get new digest: %v", err) + } + + if originalDigest.String() == newDigest.String() { + t.Error("Expected different digest for modified model") + } + + // Verify both models exist in the index + models, err := s.List() + if err != nil { + t.Fatalf("List failed: %v", err) + } + + if len(models) != 2 { + t.Fatalf("Expected 2 models in index, got %d", len(models)) + } + + // Verify both models reference the same layer blobs + var baseModelFiles []string + var modifiedModelFiles []string + for _, m := range models { + if containsTag(m.Tags, "base-model:v1") { + baseModelFiles = m.Files + } + if containsTag(m.Tags, "base-model:v2") { + modifiedModelFiles = m.Files + } + } + + if len(baseModelFiles) == 0 || len(modifiedModelFiles) == 0 { + t.Fatal("Failed to find both models in index") + } + + // The first file should be the same (gguf blob), but config should differ + if baseModelFiles[0] != modifiedModelFiles[0] { + t.Error("Expected models to share the same layer blobs") + } + + // The license blob should also be shared + if len(baseModelFiles) >= 2 && len(modifiedModelFiles) >= 2 { + if baseModelFiles[1] != modifiedModelFiles[1] { + t.Error("Expected models to share the same license blob") + } + } + }) + + t.Run("FailureWhenLayerMissing", func(t *testing.T) { + // Create a separate store to ensure no blobs from previous tests exist + freshStorePath := filepath.Join(tempDir, "fresh-model-store") + freshStore, err := store.New(store.Options{ + RootPath: freshStorePath, + }) + if err != nil { + t.Fatalf("Failed to create fresh store: %v", err) + } + + // Create a new model without writing its blobs first + freshModel := newTestModel(t) + + // Attempt to use WriteLightweight without the blobs existing + err = freshStore.WriteLightweight(freshModel, []string{"missing-blobs:latest"}) + if err == nil { + t.Fatal("Expected error when layer blobs are missing, got nil") + } + + if !strings.Contains(err.Error(), "not found in store") { + t.Errorf("Expected error about missing layer, got: %v", err) + } + }) + + t.Run("MultipleTags", func(t *testing.T) { + // Create and write a model normally + baseModel := newTestModel(t) + if err := s.Write(baseModel, []string{"multi-tag:base"}, nil); err != nil { + t.Fatalf("Write base model failed: %v", err) + } + + // Create a variant with different config + newContextSize := uint64(8192) + variant := mutate.ContextSize(baseModel, newContextSize) + + // Use WriteLightweight with multiple tags + if err := s.WriteLightweight(variant, []string{"multi-tag:variant1", "multi-tag:variant2"}); err != nil { + t.Fatalf("WriteLightweight with multiple tags failed: %v", err) + } + + // Verify both tags point to the same model + model1, err := s.Read("multi-tag:variant1") + if err != nil { + t.Fatalf("Read variant1 failed: %v", err) + } + + model2, err := s.Read("multi-tag:variant2") + if err != nil { + t.Fatalf("Read variant2 failed: %v", err) + } + + digest1, err := model1.Digest() + if err != nil { + t.Fatalf("Failed to get digest1: %v", err) + } + + digest2, err := model2.Digest() + if err != nil { + t.Fatalf("Failed to get digest2: %v", err) + } + + if digest1.String() != digest2.String() { + t.Error("Expected both tags to point to the same model") + } + + // Verify they share the same blobs as the base model + baseModelRead, err := s.Read("multi-tag:base") + if err != nil { + t.Fatalf("Read base model failed: %v", err) + } + + baseLayers, err := baseModelRead.Layers() + if err != nil { + t.Fatalf("Failed to get base layers: %v", err) + } + + variantLayers, err := model1.Layers() + if err != nil { + t.Fatalf("Failed to get variant layers: %v", err) + } + + // Verify they have the same number of layers + if len(baseLayers) != len(variantLayers) { + t.Fatalf("Expected same number of layers, got base=%d variant=%d", len(baseLayers), len(variantLayers)) + } + + // Verify layer digests match (same blobs) + for i := range baseLayers { + baseDigest, err := baseLayers[i].Digest() + if err != nil { + t.Fatalf("Failed to get base layer digest: %v", err) + } + variantDigest, err := variantLayers[i].Digest() + if err != nil { + t.Fatalf("Failed to get variant layer digest: %v", err) + } + if baseDigest.String() != variantDigest.String() { + t.Errorf("Layer %d digest mismatch", i) + } + } + }) + + t.Run("WithMultimodalProjector", func(t *testing.T) { + // Create and write a model with multimodal projector + baseModel := newTestModelWithMultimodalProjector(t) + if err := s.Write(baseModel, []string{"mmproj-model:base"}, nil); err != nil { + t.Fatalf("Write base model with mmproj failed: %v", err) + } + + // Create a variant with different context size + newContextSize := uint64(2048) + variant := mutate.ContextSize(baseModel, newContextSize) + + // Use WriteLightweight for the variant + if err := s.WriteLightweight(variant, []string{"mmproj-model:variant"}); err != nil { + t.Fatalf("WriteLightweight with mmproj failed: %v", err) + } + + // Read the variant back + readVariant, err := s.Read("mmproj-model:variant") + if err != nil { + t.Fatalf("Read variant failed: %v", err) + } + + // Verify multimodal projector path exists + mmprojPath, err := readVariant.MMPROJPath() + if err != nil { + t.Fatalf("Failed to get mmproj path: %v", err) + } + + if mmprojPath == "" { + t.Error("Expected non-empty multimodal projector path") + } + + // Verify the manifest has all layer types + manifest, err := readVariant.Manifest() + if err != nil { + t.Fatalf("Failed to get manifest: %v", err) + } + + // Should have 3 layers: GGUF + license + multimodal projector + if len(manifest.Layers) != 3 { + t.Fatalf("Expected 3 layers, got %d", len(manifest.Layers)) + } + + // Verify multimodal projector layer exists + foundMMProj := false + for _, layer := range manifest.Layers { + if layer.MediaType == types.MediaTypeMultimodalProjector { + foundMMProj = true + break + } + } + + if !foundMMProj { + t.Error("Expected to find multimodal projector layer") + } + }) + + t.Run("IndexIntegrity", func(t *testing.T) { + // Create and write a base model + baseModel := newTestModel(t) + if err := s.Write(baseModel, []string{"integrity-test:base"}, nil); err != nil { + t.Fatalf("Write base model failed: %v", err) + } + + // Create multiple variants using WriteLightweight + for i := 1; i <= 3; i++ { + contextSize := uint64(1024 * i) + variant := mutate.ContextSize(baseModel, contextSize) + tag := fmt.Sprintf("integrity-test:variant%d", i) + if err := s.WriteLightweight(variant, []string{tag}); err != nil { + t.Fatalf("WriteLightweight variant%d failed: %v", i, err) + } + } + + // Verify the index has all models + models, err := s.List() + if err != nil { + t.Fatalf("List failed: %v", err) + } + + // Should have base + 3 variants + any models from previous tests + integrityTestCount := 0 + for _, m := range models { + for _, tag := range m.Tags { + if strings.HasPrefix(tag, "integrity-test:") { + integrityTestCount++ + break + } + } + } + + if integrityTestCount != 4 { + t.Fatalf("Expected 4 integrity-test models, got %d", integrityTestCount) + } + + // Verify blob reference counts by checking that the GGUF blob + // is listed in all 4 models' Files + var ggufBlobHash string + blobRefCount := 0 + + for _, m := range models { + hasIntegrityTag := false + for _, tag := range m.Tags { + if strings.HasPrefix(tag, "integrity-test:") { + hasIntegrityTag = true + break + } + } + + if hasIntegrityTag { + if len(m.Files) > 0 { + if ggufBlobHash == "" { + ggufBlobHash = m.Files[0] + } + if m.Files[0] == ggufBlobHash { + blobRefCount++ + } + } + } + } + + if blobRefCount != 4 { + t.Errorf("Expected GGUF blob to be referenced by 4 models, got %d", blobRefCount) + } + + // Verify the blob file exists only once on disk + if ggufBlobHash != "" { + // Remove the "sha256:" prefix if present + hashStr := strings.TrimPrefix(ggufBlobHash, "sha256:") + blobPath := filepath.Join(storePath, "blobs", "sha256", hashStr) + + if _, err := os.Stat(blobPath); os.IsNotExist(err) { + t.Errorf("Shared blob file doesn't exist: %s", blobPath) + } + + // Verify there's only one copy by checking file count in blobs directory + blobsDir := filepath.Join(storePath, "blobs", "sha256") + entries, err := os.ReadDir(blobsDir) + if err != nil { + t.Fatalf("Failed to read blobs directory: %v", err) + } + + // Count non-config blobs (config blobs will be different for each variant) + ggufCount := 0 + for _, entry := range entries { + if !entry.IsDir() && entry.Name() == hashStr { + ggufCount++ + } + } + + if ggufCount != 1 { + t.Errorf("Expected exactly 1 GGUF blob file, found %d", ggufCount) + } + } + }) +}