From df295a73077379a34e83dcb9311720ae483a3455 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20L=C3=B3pez=20Luna?= Date: Mon, 6 Oct 2025 12:09:20 +0200 Subject: [PATCH 1/5] bump model-runner --- cmd/cli/go.mod | 4 +++- cmd/cli/go.sum | 2 -- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/cmd/cli/go.mod b/cmd/cli/go.mod index 43d0e5596..aede61d10 100644 --- a/cmd/cli/go.mod +++ b/cmd/cli/go.mod @@ -13,7 +13,7 @@ require ( github.com/docker/go-connections v0.5.0 github.com/docker/go-units v0.5.0 github.com/docker/model-distribution v0.0.0-20250918153037-7d9fc7b72b57 - github.com/docker/model-runner v0.0.0-20250911130340-38bb0171c947 + github.com/docker/model-runner v0.0.0 github.com/fatih/color v1.18.0 github.com/google/go-containerregistry v0.20.6 github.com/mattn/go-isatty v0.0.20 @@ -139,3 +139,5 @@ require ( ) replace github.com/kolesnikovae/go-winjob => github.com/docker/go-winjob v0.0.0-20250829235554-57b487ebcbc5 + +replace github.com/docker/model-runner => ../.. diff --git a/cmd/cli/go.sum b/cmd/cli/go.sum index c17f5b4b0..564232898 100644 --- a/cmd/cli/go.sum +++ b/cmd/cli/go.sum @@ -112,8 +112,6 @@ github.com/docker/go-winjob v0.0.0-20250829235554-57b487ebcbc5/go.mod h1:ICOGmIX github.com/docker/libtrust v0.0.0-20160708172513-aabc10ec26b7/go.mod h1:cyGadeNEkKy96OOhEzfZl+yxihPEzKnqJwvfuSUqbZE= github.com/docker/model-distribution v0.0.0-20250918153037-7d9fc7b72b57 h1:WHiPO9UmO5v97T3ksQUA2SbYVkTdUCSFobznegL97kk= github.com/docker/model-distribution v0.0.0-20250918153037-7d9fc7b72b57/go.mod h1:bV1RH2e79nTwOW38GoMU9UO8gpZVLH9+cZeEeR4wSeE= -github.com/docker/model-runner v0.0.0-20250911130340-38bb0171c947 h1:6Dz1SFZONEd8tlKetn2Gu6v5HDJI/YtUFwkqHGwrsV0= -github.com/docker/model-runner v0.0.0-20250911130340-38bb0171c947/go.mod h1:cl7panafjkSHllYCCGYAzty2aUvbwk55Gi35v06XL80= github.com/dvsekhvalnov/jose2go v0.0.0-20170216131308-f21a8cedbbae/go.mod h1:7BvyPhdbLxMXIYTFPLsyJRFMsKmOZnQmzh6Gb+uquuM= github.com/elastic/go-sysinfo v1.15.3 h1:W+RnmhKFkqPTCRoFq2VCTmsT4p/fwpo+3gKNQsn1XU0= github.com/elastic/go-sysinfo v1.15.3/go.mod h1:K/cNrqYTDrSoMh2oDkYEMS2+a72GRxMvNP+GC+vRIlo= From 5b6d4509992afd5e22a606b7778992faf503957a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20L=C3=B3pez=20Luna?= Date: Mon, 6 Oct 2025 12:11:10 +0200 Subject: [PATCH 2/5] feat: add support for packaging Safetensors models alongside GGUF files --- cmd/cli/commands/package.go | 258 +++++++++++++++++++++++++++++++++--- 1 file changed, 239 insertions(+), 19 deletions(-) diff --git a/cmd/cli/commands/package.go b/cmd/cli/commands/package.go index d9a41c6b0..6ccbaca9f 100644 --- a/cmd/cli/commands/package.go +++ b/cmd/cli/commands/package.go @@ -1,18 +1,21 @@ package commands import ( + "archive/tar" "bufio" "context" "encoding/json" "fmt" "html" "io" + "os" "path/filepath" + "sort" - "github.com/docker/model-distribution/builder" - "github.com/docker/model-distribution/registry" - "github.com/docker/model-distribution/tarball" - "github.com/docker/model-distribution/types" + "github.com/docker/model-runner/pkg/distribution/builder" + "github.com/docker/model-runner/pkg/distribution/registry" + "github.com/docker/model-runner/pkg/distribution/tarball" + "github.com/docker/model-runner/pkg/distribution/types" "github.com/google/go-containerregistry/pkg/name" "github.com/spf13/cobra" @@ -24,10 +27,11 @@ func newPackagedCmd() *cobra.Command { var opts packageOptions c := &cobra.Command{ - Use: "package --gguf [--license ...] [--context-size ] [--push] MODEL", - Short: "Package a GGUF file into a Docker model OCI artifact, with optional licenses.", - Long: "Package a GGUF file 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 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).", + 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" + + "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.", Args: func(cmd *cobra.Command, args []string) error { if len(args) != 1 { return fmt.Errorf( @@ -37,19 +41,59 @@ func newPackagedCmd() *cobra.Command { cmd.Use, ) } - if opts.ggufPath == "" { + + // Validate that either --gguf or --safetensors-dir is provided (mutually exclusive) + if opts.ggufPath == "" && opts.safetensorsDir == "" { return fmt.Errorf( - "GGUF path is required.\n\n" + + "Either --gguf or --safetensors-dir path is required.\n\n" + "See 'docker model package --help' for more information", ) } - if !filepath.IsAbs(opts.ggufPath) { + if opts.ggufPath != "" && opts.safetensorsDir != "" { return fmt.Errorf( - "GGUF path must be absolute.\n\n" + + "Cannot specify both --gguf and --safetensors-dir. Please use only one format.\n\n" + "See 'docker model package --help' for more information", ) } - opts.ggufPath = filepath.Clean(opts.ggufPath) + + // Validate GGUF path if provided + if opts.ggufPath != "" { + if !filepath.IsAbs(opts.ggufPath) { + return fmt.Errorf( + "GGUF path must be absolute.\n\n" + + "See 'docker model package --help' for more information", + ) + } + opts.ggufPath = filepath.Clean(opts.ggufPath) + } + + // Validate safetensors directory if provided + if opts.safetensorsDir != "" { + if !filepath.IsAbs(opts.safetensorsDir) { + return fmt.Errorf( + "Safetensors directory path must be absolute.\n\n" + + "See 'docker model package --help' for more information", + ) + } + opts.safetensorsDir = filepath.Clean(opts.safetensorsDir) + + // Check if it's a directory + info, err := os.Stat(opts.safetensorsDir) + if err != nil { + return fmt.Errorf( + "Safetensors directory does not exist: %s\n\n"+ + "See 'docker model package --help' for more information", + opts.safetensorsDir, + ) + } + if !info.IsDir() { + return fmt.Errorf( + "Safetensors path must be a directory: %s\n\n"+ + "See 'docker model package --help' for more information", + opts.safetensorsDir, + ) + } + } for i, l := range opts.licensePaths { if !filepath.IsAbs(l) { @@ -73,7 +117,8 @@ func newPackagedCmd() *cobra.Command { ValidArgsFunction: completion.NoComplete, } - c.Flags().StringVar(&opts.ggufPath, "gguf", "", "absolute path to gguf file (required)") + 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.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().BoolVar(&opts.push, "push", false, "push to registry (if not set, the model is loaded into the Model Runner content store)") @@ -85,6 +130,7 @@ type packageOptions struct { chatTemplatePath string contextSize uint64 ggufPath string + safetensorsDir string licensePaths []string push bool tag string @@ -106,11 +152,41 @@ func packageModel(cmd *cobra.Command, opts packageOptions) error { return err } - // Create package builder with GGUF file - cmd.PrintErrf("Adding GGUF file from %q\n", opts.ggufPath) - pkg, err := builder.FromGGUF(opts.ggufPath) - if err != nil { - return fmt.Errorf("add gguf file: %w", err) + // Create package builder based on model format + var pkg *builder.Builder + if opts.ggufPath != "" { + cmd.PrintErrf("Adding GGUF file from %q\n", opts.ggufPath) + pkg, err = builder.FromGGUF(opts.ggufPath) + if err != nil { + return fmt.Errorf("add gguf file: %w", err) + } + } else { + // Safetensors model from directory + cmd.PrintErrf("Scanning directory %q for safetensors model...\n", opts.safetensorsDir) + safetensorsPaths, tempConfigArchive, err := packageFromDirectory(opts.safetensorsDir) + if err != nil { + return fmt.Errorf("scan safetensors directory: %w", err) + } + + // Clean up temp config archive when done + if tempConfigArchive != "" { + defer os.Remove(tempConfigArchive) + } + + cmd.PrintErrf("Found %d safetensors file(s)\n", len(safetensorsPaths)) + pkg, err = builder.FromSafetensors(safetensorsPaths) + if err != nil { + return fmt.Errorf("create safetensors model: %w", err) + } + + // Add config archive if it was created + if tempConfigArchive != "" { + cmd.PrintErrf("Adding config archive from directory\n") + pkg, err = pkg.WithConfigArchive(tempConfigArchive) + if err != nil { + return fmt.Errorf("add config archive: %w", err) + } + } } // Set context size @@ -236,3 +312,147 @@ func (t *modelRunnerTarget) Write(ctx context.Context, mdl types.ModelArtifact, } return nil } + +// packageFromDirectory scans a directory for safetensors files and config files, +// creating a temporary tar archive of the config files +func packageFromDirectory(dirPath string) (safetensorsPaths []string, tempConfigArchive string, err error) { + // Read directory contents (only top level, no subdirectories) + entries, err := os.ReadDir(dirPath) + if err != nil { + return nil, "", fmt.Errorf("read directory: %w", err) + } + + var configFiles []string + + for _, entry := range entries { + if entry.IsDir() { + continue // Skip subdirectories + } + + name := entry.Name() + fullPath := filepath.Join(dirPath, name) + + // Collect safetensors files + if filepath.Ext(name) == ".safetensors" { + safetensorsPaths = append(safetensorsPaths, fullPath) + } + + // Collect config files: *.json, merges.txt + if filepath.Ext(name) == ".json" || name == "merges.txt" { + configFiles = append(configFiles, fullPath) + } + } + + if len(safetensorsPaths) == 0 { + return nil, "", fmt.Errorf("no safetensors files found in directory: %s", dirPath) + } + + // Sort to ensure reproducible artifacts + sortStrings(safetensorsPaths) + + // Create temporary tar archive with config files if any exist + if len(configFiles) > 0 { + // Sort config files for reproducible tar archive + sortStrings(configFiles) + + tempConfigArchive, err = createTempConfigArchive(configFiles) + if err != nil { + return nil, "", fmt.Errorf("create config archive: %w", err) + } + } + + return safetensorsPaths, tempConfigArchive, nil +} + +// createTempConfigArchive creates a temporary tar archive containing the specified config files +func createTempConfigArchive(configFiles []string) (string, error) { + // Create temp file + tmpFile, err := os.CreateTemp("", "vllm-config-*.tar") + if err != nil { + return "", fmt.Errorf("create temp file: %w", err) + } + tmpPath := tmpFile.Name() + + // Import archive/tar at the top if not already imported + // We'll use the tar package here + tw := newTarWriter(tmpFile) + + // Add each config file to tar (preserving just filename, not full path) + for _, filePath := range configFiles { + // Open the file + file, err := os.Open(filePath) + if err != nil { + tw.Close() + tmpFile.Close() + os.Remove(tmpPath) + return "", fmt.Errorf("open config file %s: %w", filePath, err) + } + + // Get file info for tar header + fileInfo, err := file.Stat() + if err != nil { + file.Close() + tw.Close() + tmpFile.Close() + os.Remove(tmpPath) + return "", fmt.Errorf("stat config file %s: %w", filePath, err) + } + + // Create tar header (use only basename, not full path) + header := newTarHeader(filepath.Base(filePath), fileInfo) + + // Write header + if err := tw.WriteHeader(header); err != nil { + file.Close() + tw.Close() + tmpFile.Close() + os.Remove(tmpPath) + return "", fmt.Errorf("write tar header for %s: %w", filePath, err) + } + + // Copy file contents + if _, err := io.Copy(tw, file); err != nil { + file.Close() + tw.Close() + tmpFile.Close() + os.Remove(tmpPath) + return "", fmt.Errorf("write tar content for %s: %w", filePath, err) + } + + file.Close() + } + + // Close tar writer and file + if err := tw.Close(); err != nil { + tmpFile.Close() + os.Remove(tmpPath) + return "", fmt.Errorf("close tar writer: %w", err) + } + + if err := tmpFile.Close(); err != nil { + os.Remove(tmpPath) + return "", fmt.Errorf("close temp file: %w", err) + } + + return tmpPath, nil +} + +// sortStrings sorts a slice of strings in place +func sortStrings(s []string) { + sort.Strings(s) +} + +// newTarWriter creates a new tar writer +func newTarWriter(w io.Writer) *tar.Writer { + return tar.NewWriter(w) +} + +// newTarHeader creates a tar header from file info +func newTarHeader(name string, fileInfo os.FileInfo) *tar.Header { + return &tar.Header{ + Name: name, + Size: fileInfo.Size(), + Mode: int64(fileInfo.Mode()), + ModTime: fileInfo.ModTime(), + } +} From 983916dd60fbc39f147f47daa8308620878daaeb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20L=C3=B3pez=20Luna?= Date: Mon, 6 Oct 2025 12:25:46 +0200 Subject: [PATCH 3/5] feat: refactor packaging logic for Safetensors models into a dedicated package --- cmd/cli/commands/package.go | 149 +--------------------- cmd/mdltool/main.go | 135 +------------------- pkg/distribution/packaging/safetensors.go | 143 +++++++++++++++++++++ 3 files changed, 147 insertions(+), 280 deletions(-) create mode 100644 pkg/distribution/packaging/safetensors.go diff --git a/cmd/cli/commands/package.go b/cmd/cli/commands/package.go index 6ccbaca9f..a8eac7ce2 100644 --- a/cmd/cli/commands/package.go +++ b/cmd/cli/commands/package.go @@ -1,7 +1,6 @@ package commands import ( - "archive/tar" "bufio" "context" "encoding/json" @@ -10,9 +9,9 @@ import ( "io" "os" "path/filepath" - "sort" "github.com/docker/model-runner/pkg/distribution/builder" + "github.com/docker/model-runner/pkg/distribution/packaging" "github.com/docker/model-runner/pkg/distribution/registry" "github.com/docker/model-runner/pkg/distribution/tarball" "github.com/docker/model-runner/pkg/distribution/types" @@ -163,7 +162,7 @@ func packageModel(cmd *cobra.Command, opts packageOptions) error { } else { // Safetensors model from directory cmd.PrintErrf("Scanning directory %q for safetensors model...\n", opts.safetensorsDir) - safetensorsPaths, tempConfigArchive, err := packageFromDirectory(opts.safetensorsDir) + safetensorsPaths, tempConfigArchive, err := packaging.PackageFromDirectory(opts.safetensorsDir) if err != nil { return fmt.Errorf("scan safetensors directory: %w", err) } @@ -312,147 +311,3 @@ func (t *modelRunnerTarget) Write(ctx context.Context, mdl types.ModelArtifact, } return nil } - -// packageFromDirectory scans a directory for safetensors files and config files, -// creating a temporary tar archive of the config files -func packageFromDirectory(dirPath string) (safetensorsPaths []string, tempConfigArchive string, err error) { - // Read directory contents (only top level, no subdirectories) - entries, err := os.ReadDir(dirPath) - if err != nil { - return nil, "", fmt.Errorf("read directory: %w", err) - } - - var configFiles []string - - for _, entry := range entries { - if entry.IsDir() { - continue // Skip subdirectories - } - - name := entry.Name() - fullPath := filepath.Join(dirPath, name) - - // Collect safetensors files - if filepath.Ext(name) == ".safetensors" { - safetensorsPaths = append(safetensorsPaths, fullPath) - } - - // Collect config files: *.json, merges.txt - if filepath.Ext(name) == ".json" || name == "merges.txt" { - configFiles = append(configFiles, fullPath) - } - } - - if len(safetensorsPaths) == 0 { - return nil, "", fmt.Errorf("no safetensors files found in directory: %s", dirPath) - } - - // Sort to ensure reproducible artifacts - sortStrings(safetensorsPaths) - - // Create temporary tar archive with config files if any exist - if len(configFiles) > 0 { - // Sort config files for reproducible tar archive - sortStrings(configFiles) - - tempConfigArchive, err = createTempConfigArchive(configFiles) - if err != nil { - return nil, "", fmt.Errorf("create config archive: %w", err) - } - } - - return safetensorsPaths, tempConfigArchive, nil -} - -// createTempConfigArchive creates a temporary tar archive containing the specified config files -func createTempConfigArchive(configFiles []string) (string, error) { - // Create temp file - tmpFile, err := os.CreateTemp("", "vllm-config-*.tar") - if err != nil { - return "", fmt.Errorf("create temp file: %w", err) - } - tmpPath := tmpFile.Name() - - // Import archive/tar at the top if not already imported - // We'll use the tar package here - tw := newTarWriter(tmpFile) - - // Add each config file to tar (preserving just filename, not full path) - for _, filePath := range configFiles { - // Open the file - file, err := os.Open(filePath) - if err != nil { - tw.Close() - tmpFile.Close() - os.Remove(tmpPath) - return "", fmt.Errorf("open config file %s: %w", filePath, err) - } - - // Get file info for tar header - fileInfo, err := file.Stat() - if err != nil { - file.Close() - tw.Close() - tmpFile.Close() - os.Remove(tmpPath) - return "", fmt.Errorf("stat config file %s: %w", filePath, err) - } - - // Create tar header (use only basename, not full path) - header := newTarHeader(filepath.Base(filePath), fileInfo) - - // Write header - if err := tw.WriteHeader(header); err != nil { - file.Close() - tw.Close() - tmpFile.Close() - os.Remove(tmpPath) - return "", fmt.Errorf("write tar header for %s: %w", filePath, err) - } - - // Copy file contents - if _, err := io.Copy(tw, file); err != nil { - file.Close() - tw.Close() - tmpFile.Close() - os.Remove(tmpPath) - return "", fmt.Errorf("write tar content for %s: %w", filePath, err) - } - - file.Close() - } - - // Close tar writer and file - if err := tw.Close(); err != nil { - tmpFile.Close() - os.Remove(tmpPath) - return "", fmt.Errorf("close tar writer: %w", err) - } - - if err := tmpFile.Close(); err != nil { - os.Remove(tmpPath) - return "", fmt.Errorf("close temp file: %w", err) - } - - return tmpPath, nil -} - -// sortStrings sorts a slice of strings in place -func sortStrings(s []string) { - sort.Strings(s) -} - -// newTarWriter creates a new tar writer -func newTarWriter(w io.Writer) *tar.Writer { - return tar.NewWriter(w) -} - -// newTarHeader creates a tar header from file info -func newTarHeader(name string, fileInfo os.FileInfo) *tar.Header { - return &tar.Header{ - Name: name, - Size: fileInfo.Size(), - Mode: int64(fileInfo.Mode()), - ModTime: fileInfo.ModTime(), - } -} diff --git a/cmd/mdltool/main.go b/cmd/mdltool/main.go index 490551a52..79faccf62 100644 --- a/cmd/mdltool/main.go +++ b/cmd/mdltool/main.go @@ -1,18 +1,16 @@ package main import ( - "archive/tar" "context" "flag" "fmt" - "io" "os" "path/filepath" - "sort" "strings" "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" ) @@ -220,7 +218,7 @@ func cmdPackage(args []string) int { if sourceInfo.IsDir() { fmt.Printf("Detected directory, scanning for safetensors model...\n") var err error - safetensorsPaths, configArchive, err = packageFromDirectory(source) + safetensorsPaths, configArchive, err = packaging.PackageFromDirectory(source) if err != nil { fmt.Fprintf(os.Stderr, "Error scanning directory: %v\n", err) return 1 @@ -581,132 +579,3 @@ func cmdBundle(client *distribution.Client, args []string) int { fmt.Fprint(os.Stdout, bundle.RootDir()) return 0 } - -// packageFromDirectory scans a directory for safetensors files and config files, -// creating a temporary tar archive of the config files -func packageFromDirectory(dirPath string) (safetensorsPaths []string, tempConfigArchive string, err error) { - // Read directory contents (only top level, no subdirectories) - entries, err := os.ReadDir(dirPath) - if err != nil { - return nil, "", fmt.Errorf("read directory: %w", err) - } - - var configFiles []string - - for _, entry := range entries { - if entry.IsDir() { - continue // Skip subdirectories - } - - name := entry.Name() - fullPath := filepath.Join(dirPath, name) - - // Collect safetensors files - if strings.HasSuffix(strings.ToLower(name), ".safetensors") { - safetensorsPaths = append(safetensorsPaths, fullPath) - } - - // Collect config files: *.json, merges.txt - if strings.HasSuffix(strings.ToLower(name), ".json") || - name == "merges.txt" { - configFiles = append(configFiles, fullPath) - } - } - - if len(safetensorsPaths) == 0 { - return nil, "", fmt.Errorf("no safetensors files found in directory: %s", dirPath) - } - - // Sort to ensure reproducible artifacts - sort.Strings(safetensorsPaths) - - // Create temporary tar archive with config files if any exist - if len(configFiles) > 0 { - // Sort config files for reproducible tar archive - sort.Strings(configFiles) - - tempConfigArchive, err = createTempConfigArchive(configFiles) - if err != nil { - return nil, "", fmt.Errorf("create config archive: %w", err) - } - } - - return safetensorsPaths, tempConfigArchive, nil -} - -// createTempConfigArchive creates a temporary tar archive containing the specified config files -func createTempConfigArchive(configFiles []string) (string, error) { - // Create temp file - tmpFile, err := os.CreateTemp("", "vllm-config-*.tar") - if err != nil { - return "", fmt.Errorf("create temp file: %w", err) - } - tmpPath := tmpFile.Name() - - // Create tar writer - tw := tar.NewWriter(tmpFile) - - // Add each config file to tar (preserving just filename, not full path) - for _, filePath := range configFiles { - // Open the file - file, err := os.Open(filePath) - if err != nil { - tw.Close() - tmpFile.Close() - os.Remove(tmpPath) - return "", fmt.Errorf("open config file %s: %w", filePath, err) - } - - // Get file info for tar header - fileInfo, err := file.Stat() - if err != nil { - file.Close() - tw.Close() - tmpFile.Close() - os.Remove(tmpPath) - return "", fmt.Errorf("stat config file %s: %w", filePath, err) - } - - // Create tar header (use only basename, not full path) - header := &tar.Header{ - Name: filepath.Base(filePath), - Size: fileInfo.Size(), - Mode: int64(fileInfo.Mode()), - ModTime: fileInfo.ModTime(), - } - - // Write header - if err := tw.WriteHeader(header); err != nil { - file.Close() - tw.Close() - tmpFile.Close() - os.Remove(tmpPath) - return "", fmt.Errorf("write tar header for %s: %w", filePath, err) - } - - // Copy file contents - if _, err := io.Copy(tw, file); err != nil { - file.Close() - tw.Close() - tmpFile.Close() - os.Remove(tmpPath) - return "", fmt.Errorf("write tar content for %s: %w", filePath, err) - } - - file.Close() - } - - // Close tar writer and file - if err := tw.Close(); err != nil { - tmpFile.Close() - os.Remove(tmpPath) - return "", fmt.Errorf("close tar writer: %w", err) - } - - if err := tmpFile.Close(); err != nil { - os.Remove(tmpPath) - return "", fmt.Errorf("close temp file: %w", err) - } - - return tmpPath, nil -} diff --git a/pkg/distribution/packaging/safetensors.go b/pkg/distribution/packaging/safetensors.go new file mode 100644 index 000000000..a536d1974 --- /dev/null +++ b/pkg/distribution/packaging/safetensors.go @@ -0,0 +1,143 @@ +package packaging + +import ( + "archive/tar" + "fmt" + "io" + "os" + "path/filepath" + "sort" + "strings" +) + +// PackageFromDirectory scans a directory for safetensors files and config files, +// creating a temporary tar archive of the config files. +// It returns the paths to safetensors files, path to temporary config archive (if created), +// and any error encountered. +func PackageFromDirectory(dirPath string) (safetensorsPaths []string, tempConfigArchive string, err error) { + // Read directory contents (only top level, no subdirectories) + entries, err := os.ReadDir(dirPath) + if err != nil { + return nil, "", fmt.Errorf("read directory: %w", err) + } + + var configFiles []string + + for _, entry := range entries { + if entry.IsDir() { + continue // Skip subdirectories + } + + name := entry.Name() + fullPath := filepath.Join(dirPath, name) + + // Collect safetensors files + if strings.HasSuffix(strings.ToLower(name), ".safetensors") { + safetensorsPaths = append(safetensorsPaths, fullPath) + } + + // Collect config files: *.json, merges.txt + if strings.HasSuffix(strings.ToLower(name), ".json") || strings.EqualFold(name, "merges.txt") { + configFiles = append(configFiles, fullPath) + } + } + + if len(safetensorsPaths) == 0 { + return nil, "", fmt.Errorf("no safetensors files found in directory: %s", dirPath) + } + + // Sort to ensure reproducible artifacts + sort.Strings(safetensorsPaths) + + // Create temporary tar archive with config files if any exist + if len(configFiles) > 0 { + // Sort config files for reproducible tar archive + sort.Strings(configFiles) + + tempConfigArchive, err = CreateTempConfigArchive(configFiles) + if err != nil { + return nil, "", fmt.Errorf("create config archive: %w", err) + } + } + + return safetensorsPaths, tempConfigArchive, nil +} + +// CreateTempConfigArchive creates a temporary tar archive containing the specified config files. +// It returns the path to the temporary tar file and any error encountered. +// The caller is responsible for removing the temporary file when done. +func CreateTempConfigArchive(configFiles []string) (string, error) { + // Create temp file + tmpFile, err := os.CreateTemp("", "vllm-config-*.tar") + if err != nil { + return "", fmt.Errorf("create temp file: %w", err) + } + tmpPath := tmpFile.Name() + + // Create tar writer + tw := tar.NewWriter(tmpFile) + + // Add each config file to tar (preserving just filename, not full path) + for _, filePath := range configFiles { + // Open the file + file, err := os.Open(filePath) + if err != nil { + tw.Close() + tmpFile.Close() + os.Remove(tmpPath) + return "", fmt.Errorf("open config file %s: %w", filePath, err) + } + + // Get file info for tar header + fileInfo, err := file.Stat() + if err != nil { + file.Close() + tw.Close() + tmpFile.Close() + os.Remove(tmpPath) + return "", fmt.Errorf("stat config file %s: %w", filePath, err) + } + + // Create tar header (use only basename, not full path) + header := &tar.Header{ + Name: filepath.Base(filePath), + Size: fileInfo.Size(), + Mode: int64(fileInfo.Mode()), + ModTime: fileInfo.ModTime(), + } + + // Write header + if err := tw.WriteHeader(header); err != nil { + file.Close() + tw.Close() + tmpFile.Close() + os.Remove(tmpPath) + return "", fmt.Errorf("write tar header for %s: %w", filePath, err) + } + + // Copy file contents + if _, err := io.Copy(tw, file); err != nil { + file.Close() + tw.Close() + tmpFile.Close() + os.Remove(tmpPath) + return "", fmt.Errorf("write tar content for %s: %w", filePath, err) + } + + file.Close() + } + + // Close tar writer and file + if err := tw.Close(); err != nil { + tmpFile.Close() + os.Remove(tmpPath) + return "", fmt.Errorf("close tar writer: %w", err) + } + + if err := tmpFile.Close(); err != nil { + os.Remove(tmpPath) + return "", fmt.Errorf("close temp file: %w", err) + } + + return tmpPath, nil +} From 213dfc8ceabdc2ce3b51e16be51f3eddd8d1d5ef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20L=C3=B3pez=20Luna?= Date: Tue, 7 Oct 2025 14:20:16 +0200 Subject: [PATCH 4/5] fix: improve error handling for non-existent Safetensors directory --- cmd/cli/commands/package.go | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/cmd/cli/commands/package.go b/cmd/cli/commands/package.go index 0c98a2a70..fa1683f30 100644 --- a/cmd/cli/commands/package.go +++ b/cmd/cli/commands/package.go @@ -79,11 +79,14 @@ func newPackagedCmd() *cobra.Command { // Check if it's a directory info, err := os.Stat(opts.safetensorsDir) if err != nil { - return fmt.Errorf( - "Safetensors directory does not exist: %s\n\n"+ - "See 'docker model package --help' for more information", - opts.safetensorsDir, - ) + if os.IsNotExist(err) { + return fmt.Errorf( + "Safetensors directory does not exist: %s\n\n"+ + "See 'docker model package --help' for more information", + opts.safetensorsDir, + ) + } + return fmt.Errorf("could not access safetensors directory %q: %w", opts.safetensorsDir, err) } if !info.IsDir() { return fmt.Errorf( From 6311b2b252458605825b836d5ec380bf2253a750 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20L=C3=B3pez=20Luna?= Date: Tue, 7 Oct 2025 16:26:11 +0200 Subject: [PATCH 5/5] refactor: streamline tar file creation and error handling --- pkg/distribution/packaging/safetensors.go | 95 ++++++++++++----------- 1 file changed, 49 insertions(+), 46 deletions(-) diff --git a/pkg/distribution/packaging/safetensors.go b/pkg/distribution/packaging/safetensors.go index a536d1974..46c27c7ec 100644 --- a/pkg/distribution/packaging/safetensors.go +++ b/pkg/distribution/packaging/safetensors.go @@ -74,70 +74,73 @@ func CreateTempConfigArchive(configFiles []string) (string, error) { } tmpPath := tmpFile.Name() + // Track success to determine if we should clean up the temp file + shouldKeepTempFile := false + defer func() { + if !shouldKeepTempFile { + os.Remove(tmpPath) + } + }() + // Create tar writer tw := tar.NewWriter(tmpFile) // Add each config file to tar (preserving just filename, not full path) for _, filePath := range configFiles { - // Open the file - file, err := os.Open(filePath) - if err != nil { - tw.Close() - tmpFile.Close() - os.Remove(tmpPath) - return "", fmt.Errorf("open config file %s: %w", filePath, err) - } - - // Get file info for tar header - fileInfo, err := file.Stat() - if err != nil { - file.Close() - tw.Close() - tmpFile.Close() - os.Remove(tmpPath) - return "", fmt.Errorf("stat config file %s: %w", filePath, err) - } - - // Create tar header (use only basename, not full path) - header := &tar.Header{ - Name: filepath.Base(filePath), - Size: fileInfo.Size(), - Mode: int64(fileInfo.Mode()), - ModTime: fileInfo.ModTime(), - } - - // Write header - if err := tw.WriteHeader(header); err != nil { - file.Close() - tw.Close() - tmpFile.Close() - os.Remove(tmpPath) - return "", fmt.Errorf("write tar header for %s: %w", filePath, err) - } - - // Copy file contents - if _, err := io.Copy(tw, file); err != nil { - file.Close() + if err := addFileToTar(tw, filePath); err != nil { tw.Close() tmpFile.Close() - os.Remove(tmpPath) - return "", fmt.Errorf("write tar content for %s: %w", filePath, err) + return "", err } - - file.Close() } - // Close tar writer and file + // Close tar writer first if err := tw.Close(); err != nil { tmpFile.Close() - os.Remove(tmpPath) return "", fmt.Errorf("close tar writer: %w", err) } + // Close temp file if err := tmpFile.Close(); err != nil { - os.Remove(tmpPath) return "", fmt.Errorf("close temp file: %w", err) } + shouldKeepTempFile = true return tmpPath, nil } + +// addFileToTar adds a single file to the tar archive with only its basename (not full path) +func addFileToTar(tw *tar.Writer, filePath string) error { + // Open the file + file, err := os.Open(filePath) + if err != nil { + return fmt.Errorf("open file %s: %w", filePath, err) + } + defer file.Close() + + // Get file info for tar header + fileInfo, err := file.Stat() + if err != nil { + return fmt.Errorf("stat file %s: %w", filePath, err) + } + + // Create tar header (use only basename, not full path) + header := &tar.Header{ + Name: filepath.Base(filePath), + Size: fileInfo.Size(), + Mode: int64(fileInfo.Mode()), + ModTime: fileInfo.ModTime(), + } + + // Write header + if err := tw.WriteHeader(header); err != nil { + return fmt.Errorf("write tar header for %s: %w", filePath, err) + } + + // Copy file contents + if _, err := io.Copy(tw, file); err != nil { + return fmt.Errorf("write tar content for %s: %w", filePath, err) + } + + return nil +}