From ac57eb7ab3384afdfa587712c6a8bb05674a5129 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20L=C3=B3pez=20Luna?= Date: Thu, 12 Feb 2026 14:18:40 +0100 Subject: [PATCH 1/7] feat(unpack): enhance safetensors model handling with V0.2 packaging support --- pkg/distribution/builder/from_directory.go | 22 +++++-- pkg/distribution/distribution/client.go | 2 +- pkg/distribution/distribution/client_test.go | 2 +- pkg/distribution/huggingface/model.go | 30 +++++++++- pkg/distribution/internal/bundle/parse.go | 34 +++++++++-- .../internal/bundle/parse_test.go | 59 +++++++++++++++++++ 6 files changed, 138 insertions(+), 11 deletions(-) diff --git a/pkg/distribution/builder/from_directory.go b/pkg/distribution/builder/from_directory.go index a769eb900..d499415c5 100644 --- a/pkg/distribution/builder/from_directory.go +++ b/pkg/distribution/builder/from_directory.go @@ -8,6 +8,7 @@ import ( "time" "github.com/docker/model-runner/pkg/distribution/files" + "github.com/docker/model-runner/pkg/distribution/format" "github.com/docker/model-runner/pkg/distribution/internal/mutate" "github.com/docker/model-runner/pkg/distribution/internal/partial" "github.com/docker/model-runner/pkg/distribution/oci" @@ -182,13 +183,26 @@ func FromDirectory(dirPath string, opts ...DirectoryOption) (*Builder, error) { return nil, fmt.Errorf("no weight files (safetensors, GGUF, or DDUF) found in directory: %s", dirPath) } - // Build config + // Extract config metadata from weight files using format-specific logic config := types.Config{ Format: detectedFormat, } - - // TODO: Extract additional metadata from weight files if needed - // For safetensors, we might want to read config.json from the directory + if detectedFormat != "" { + f, fmtErr := format.Get(detectedFormat) + if fmtErr == nil { + extracted, extractErr := f.ExtractConfig(weightFiles) + if extractErr == nil { + // Preserve the detected format, overlay extracted metadata + config.Parameters = extracted.Parameters + config.Quantization = extracted.Quantization + config.Architecture = extracted.Architecture + config.Size = extracted.Size + config.GGUF = extracted.GGUF + config.Safetensors = extracted.Safetensors + config.Diffusers = extracted.Diffusers + } + } + } // Build the model with V0.2 config (layer-per-file with annotations) created := time.Now() diff --git a/pkg/distribution/distribution/client.go b/pkg/distribution/distribution/client.go index d9a2cca5f..e77f8f543 100644 --- a/pkg/distribution/distribution/client.go +++ b/pkg/distribution/distribution/client.go @@ -713,7 +713,7 @@ func checkCompat(image types.ModelArtifact, log *logrus.Entry, reference string, if err != nil { return err } - if manifest.Config.MediaType != types.MediaTypeModelConfigV01 { + if manifest.Config.MediaType != types.MediaTypeModelConfigV01 && manifest.Config.MediaType != types.MediaTypeModelConfigV02 { return fmt.Errorf("config type %q is unsupported: %w", manifest.Config.MediaType, ErrUnsupportedMediaType) } diff --git a/pkg/distribution/distribution/client_test.go b/pkg/distribution/distribution/client_test.go index a6f92af45..cfbfab049 100644 --- a/pkg/distribution/distribution/client_test.go +++ b/pkg/distribution/distribution/client_test.go @@ -395,7 +395,7 @@ func TestClientPullModel(t *testing.T) { }) t.Run("pull unsupported (newer) version", func(t *testing.T) { - newMdl := mutate.ConfigMediaType(model, "application/vnd.docker.ai.model.config.v0.2+json") + newMdl := mutate.ConfigMediaType(model, "application/vnd.docker.ai.model.config.v99.0+json") // Push model to local store testTag := registryHost + "/unsupported-test/model:v1.0.0" ref, err := reference.ParseReference(testTag) diff --git a/pkg/distribution/huggingface/model.go b/pkg/distribution/huggingface/model.go index e70dc3cb6..72cd386a3 100644 --- a/pkg/distribution/huggingface/model.go +++ b/pkg/distribution/huggingface/model.go @@ -87,8 +87,36 @@ func BuildModel(ctx context.Context, client *Client, repo, revision, tag string, return model, nil } -// buildModelFromFiles constructs an OCI model artifact from downloaded files +// buildModelFromFiles constructs an OCI model artifact from downloaded files. +// For safetensors models, it uses the V0.2 layer-per-file packaging (FromDirectory) +// which preserves directory structure and adds each file as an individual layer with +// filepath annotations. For GGUF models, it uses the V0.1 packaging (FromPaths) +// for backward compatibility. func buildModelFromFiles(localPaths map[string]string, weightFiles, configFiles []RepoFile, tempDir string) (types.ModelArtifact, error) { + // Check if this is a safetensors model - use V0.2 packaging + if isSafetensorsModel(weightFiles) { + return buildSafetensorsModelV02(tempDir) + } + + // For GGUF models, use V0.1 packaging (backward compatible) + return buildGGUFModelV01(localPaths, weightFiles, configFiles, tempDir) +} + +// buildSafetensorsModelV02 builds a safetensors model using V0.2 layer-per-file packaging. +// It uses builder.FromDirectory which recursively scans the tempDir and creates one layer +// per file, preserving nested directory structure with filepath annotations. +func buildSafetensorsModelV02(tempDir string) (types.ModelArtifact, error) { + b, err := builder.FromDirectory(tempDir) + if err != nil { + return nil, fmt.Errorf("create builder from directory: %w", err) + } + + return b.Model(), nil +} + +// buildGGUFModelV01 builds a GGUF model using V0.1 packaging (backward compatible). +// GGUF uses FromPaths + config archive approach. +func buildGGUFModelV01(localPaths map[string]string, weightFiles, configFiles []RepoFile, tempDir string) (types.ModelArtifact, error) { // Collect weight file paths (sorted for reproducibility) var weightPaths []string for _, f := range weightFiles { diff --git a/pkg/distribution/internal/bundle/parse.go b/pkg/distribution/internal/bundle/parse.go index 4e2cf49ad..3df746e78 100644 --- a/pkg/distribution/internal/bundle/parse.go +++ b/pkg/distribution/internal/bundle/parse.go @@ -5,6 +5,7 @@ import ( "fmt" "os" "path/filepath" + "strings" "github.com/docker/model-runner/pkg/distribution/modelpack" "github.com/docker/model-runner/pkg/distribution/types" @@ -99,15 +100,40 @@ func findGGUFFile(modelDir string) (string, error) { } func findSafetensorsFile(modelDir string) (string, error) { + // First check top-level directory (most common case) safetensors, err := filepath.Glob(filepath.Join(modelDir, "[^.]*.safetensors")) if err != nil { return "", fmt.Errorf("find safetensors files: %w", err) } - if len(safetensors) == 0 { - // Safetensors files are optional - GGUF models won't have them - return "", nil + if len(safetensors) > 0 { + return filepath.Base(safetensors[0]), nil } - return filepath.Base(safetensors[0]), nil + + // Search recursively for V0.2 models with nested directory structure + // (e.g., text_encoder/model.safetensors) + var firstFound string + walkErr := filepath.Walk(modelDir, func(path string, info os.FileInfo, err error) error { + if err != nil { + return nil // skip errors + } + if info.IsDir() { + return nil + } + if filepath.Ext(path) == ".safetensors" && !strings.HasPrefix(info.Name(), ".") { + rel, relErr := filepath.Rel(modelDir, path) + if relErr == nil { + firstFound = rel + return filepath.SkipAll // found one, stop walking + } + } + return nil + }) + if walkErr != nil { + return "", fmt.Errorf("walk for safetensors files: %w", walkErr) + } + + // Safetensors files are optional - GGUF models won't have them + return firstFound, nil } func findMultiModalProjectorFile(modelDir string) (string, error) { diff --git a/pkg/distribution/internal/bundle/parse_test.go b/pkg/distribution/internal/bundle/parse_test.go index 50d7460cf..625aeba06 100644 --- a/pkg/distribution/internal/bundle/parse_test.go +++ b/pkg/distribution/internal/bundle/parse_test.go @@ -139,6 +139,65 @@ func TestParse_WithSafetensors(t *testing.T) { } } +func TestParse_WithNestedSafetensors(t *testing.T) { + // Create a temporary directory for the test bundle + tempDir := t.TempDir() + + // Create model subdirectory + modelDir := filepath.Join(tempDir, ModelSubdir) + if err := os.MkdirAll(modelDir, 0755); err != nil { + t.Fatalf("Failed to create model directory: %v", err) + } + + // Create nested directory structure (V0.2 layout) + textEncoderDir := filepath.Join(modelDir, "text_encoder") + if err := os.MkdirAll(textEncoderDir, 0755); err != nil { + t.Fatalf("Failed to create text_encoder directory: %v", err) + } + + // Create a safetensors file in the nested directory (no safetensors at top level) + nestedSafetensorsPath := filepath.Join(textEncoderDir, "model.safetensors") + if err := os.WriteFile(nestedSafetensorsPath, []byte("dummy nested safetensors content"), 0644); err != nil { + t.Fatalf("Failed to create nested safetensors file: %v", err) + } + + // Create a valid config.json at bundle root + cfg := types.Config{ + Format: types.FormatSafetensors, + } + configPath := filepath.Join(tempDir, "config.json") + f, err := os.Create(configPath) + if err != nil { + t.Fatalf("Failed to create config.json: %v", err) + } + if err := json.NewEncoder(f).Encode(cfg); err != nil { + f.Close() + t.Fatalf("Failed to encode config: %v", err) + } + f.Close() + + // Parse the bundle - should succeed by finding safetensors recursively + bundle, err := Parse(tempDir) + if err != nil { + t.Fatalf("Expected successful parse with nested safetensors, got error: %v", err) + } + + // The safetensorsFile should include the relative path from modelDir + expectedPath := filepath.Join("text_encoder", "model.safetensors") + if bundle.safetensorsFile != expectedPath { + t.Errorf("Expected safetensorsFile to be %q, got: %s", expectedPath, bundle.safetensorsFile) + } + + // Verify SafetensorsPath() returns the full path + fullPath := bundle.SafetensorsPath() + if fullPath == "" { + t.Error("Expected SafetensorsPath() to return a non-empty path") + } + if !strings.HasSuffix(fullPath, expectedPath) { + t.Errorf("Expected SafetensorsPath() to end with %q, got: %s", expectedPath, fullPath) + } +} + func TestParse_WithBothFormats(t *testing.T) { // Create a temporary directory for the test bundle tempDir := t.TempDir() From f06e39b2860f024341725e6ebc19a8dc4f6a4c21 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20L=C3=B3pez=20Luna?= Date: Thu, 12 Feb 2026 15:02:11 +0100 Subject: [PATCH 2/7] feat(unpack): update Safetensors model packaging to support layer-per-file structure --- cmd/cli/commands/package.go | 71 +--- pkg/distribution/builder/builder.go | 13 - pkg/distribution/huggingface/model.go | 52 --- pkg/distribution/packaging/directory.go | 157 -------- pkg/distribution/packaging/directory_test.go | 297 --------------- pkg/distribution/packaging/dirtar.go | 209 ----------- pkg/distribution/packaging/dirtar_test.go | 362 ------------------- 7 files changed, 4 insertions(+), 1157 deletions(-) delete mode 100644 pkg/distribution/packaging/directory.go delete mode 100644 pkg/distribution/packaging/directory_test.go delete mode 100644 pkg/distribution/packaging/dirtar.go delete mode 100644 pkg/distribution/packaging/dirtar_test.go diff --git a/cmd/cli/commands/package.go b/cmd/cli/commands/package.go index da6b612d7..c919e80d2 100644 --- a/cmd/cli/commands/package.go +++ b/cmd/cli/commands/package.go @@ -16,7 +16,6 @@ import ( "github.com/docker/model-runner/pkg/distribution/distribution" "github.com/docker/model-runner/pkg/distribution/oci" "github.com/docker/model-runner/pkg/distribution/oci/reference" - "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" @@ -43,7 +42,7 @@ func newPackagedCmd() *cobra.Command { Short: "Package a GGUF file, Safetensors directory, DDUF file, or existing model into a Docker model OCI artifact.", Long: "Package a GGUF file, Safetensors directory, DDUF file, or existing model into a Docker model OCI artifact, with optional licenses and multimodal projector. 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.\n" + + "When packaging a Safetensors model, --safetensors-dir should point to a directory containing .safetensors files and config files. All files (including nested subdirectories) will be auto-discovered and each file is packaged as an individual layer.\n" + "When packaging a DDUF file (Diffusers Unified Format), --dduf should point to a .dduf archive file.\n" + "When packaging from an existing model using --from, you can modify properties like context size to create a variant of the original model.\n" + "For multimodal models, use --mmproj to include a multimodal projector file.", @@ -155,17 +154,6 @@ func newPackagedCmd() *cobra.Command { } } - // Validate dir-tar paths are relative (not absolute) - for _, dirPath := range opts.dirTarPaths { - if filepath.IsAbs(dirPath) { - return fmt.Errorf( - "dir-tar path must be relative, got absolute path: %s\n\n"+ - "See 'docker model package --help' for more information", - dirPath, - ) - } - } - return nil }, RunE: func(cmd *cobra.Command, args []string) error { @@ -185,7 +173,6 @@ func newPackagedCmd() *cobra.Command { 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)") c.Flags().StringVar(&opts.mmprojPath, "mmproj", "", "absolute path to multimodal projector file") c.Flags().BoolVar(&opts.push, "push", false, "push to registry (if not set, the model is loaded into the Model Runner content store)") c.Flags().Uint64Var(&opts.contextSize, "context-size", 0, "context size in tokens") @@ -200,7 +187,6 @@ type packageOptions struct { ddufPath string fromModel string licensePaths []string - dirTarPaths []string mmprojPath string push bool tag string @@ -274,33 +260,11 @@ func initializeBuilder(ctx context.Context, cmd *cobra.Command, client *desktop. } result.builder = pkg } else if opts.safetensorsDir != "" { - // Safetensors model from directory + // Safetensors model from directory — uses V0.2 layer-per-file packaging cmd.PrintErrf("Scanning directory %q for safetensors model...\n", opts.safetensorsDir) - safetensorsPaths, tempConfigArchive, err := packaging.PackageFromDirectory(opts.safetensorsDir) - if err != nil { - return nil, fmt.Errorf("scan safetensors directory: %w", err) - } - - // Set up cleanup for temp config archive - if tempConfigArchive != "" { - result.cleanupFunc = func() { - os.Remove(tempConfigArchive) - } - } - - cmd.PrintErrf("Found %d safetensors file(s)\n", len(safetensorsPaths)) - pkg, err := builder.FromPaths(safetensorsPaths) + pkg, err := builder.FromDirectory(opts.safetensorsDir) if err != nil { - return nil, 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 nil, fmt.Errorf("add config archive: %w", err) - } + return nil, fmt.Errorf("create safetensors model from directory: %w", err) } result.builder = pkg } else { @@ -354,7 +318,6 @@ func packageModel(ctx context.Context, cmd *cobra.Command, client *desktop.Clien len(opts.licensePaths) == 0 && opts.chatTemplatePath == "" && opts.mmprojPath == "" && - len(opts.dirTarPaths) == 0 && cmd.Flags().Changed("context-size") if canUseDaemonRepackage { @@ -456,32 +419,6 @@ func packageModel(ctx context.Context, cmd *cobra.Command, client *desktop.Clien cmd.PrintErrln("Model variant created successfully") return nil // Return early to avoid the Build operation in lightweight case - } 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 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...") diff --git a/pkg/distribution/builder/builder.go b/pkg/distribution/builder/builder.go index 1b566c611..4464d39d5 100644 --- a/pkg/distribution/builder/builder.go +++ b/pkg/distribution/builder/builder.go @@ -185,19 +185,6 @@ func (b *Builder) WithConfigArchive(path string) (*Builder, error) { }, nil } -// WithDirTar adds a directory tar archive to the artifact. -// Multiple directory tar archives can be added by calling this method multiple times. -func (b *Builder) WithDirTar(path string) (*Builder, error) { - dirTarLayer, err := partial.NewLayer(path, types.MediaTypeDirTar) - if err != nil { - return nil, fmt.Errorf("dir tar layer from %q: %w", path, err) - } - return &Builder{ - model: mutate.AppendLayers(b.model, dirTarLayer), - originalLayers: b.originalLayers, - }, nil -} - // Target represents a build target type Target interface { Write(context.Context, types.ModelArtifact, io.Writer) error diff --git a/pkg/distribution/huggingface/model.go b/pkg/distribution/huggingface/model.go index 72cd386a3..32f36e68d 100644 --- a/pkg/distribution/huggingface/model.go +++ b/pkg/distribution/huggingface/model.go @@ -11,7 +11,6 @@ import ( "github.com/docker/model-runner/pkg/distribution/builder" "github.com/docker/model-runner/pkg/distribution/internal/progress" - "github.com/docker/model-runner/pkg/distribution/packaging" "github.com/docker/model-runner/pkg/distribution/types" ) @@ -134,25 +133,6 @@ func buildGGUFModelV01(localPaths map[string]string, weightFiles, configFiles [] return nil, fmt.Errorf("create builder: %w", err) } - // Create config archive if we have config files - if len(configFiles) > 0 { - configArchive, configArchiveErr := createConfigArchive(localPaths, configFiles, tempDir) - if configArchiveErr != nil { - return nil, fmt.Errorf("create config archive: %w", configArchiveErr) - } - // Note: configArchive is created inside tempDir and will be cleaned up when - // the caller removes tempDir. The file must exist until after store.Write() - // completes since the model artifact references it lazily. - - if configArchive != "" { - var withConfigErr error - b, withConfigErr = b.WithConfigArchive(configArchive) - if withConfigErr != nil { - return nil, fmt.Errorf("add config archive: %w", withConfigErr) - } - } - } - // Check for chat template and add it for _, f := range configFiles { if isChatTemplate(f.Path) { @@ -170,38 +150,6 @@ func buildGGUFModelV01(localPaths map[string]string, weightFiles, configFiles [] return b.Model(), nil } -// createConfigArchive creates a tar archive of config files in the specified tempDir -func createConfigArchive(localPaths map[string]string, configFiles []RepoFile, tempDir string) (string, error) { - // Collect config file paths (excluding chat templates which are added separately) - var configPaths []string - for _, f := range configFiles { - if isChatTemplate(f.Path) { - continue // Chat templates are added as separate layers - } - localPath, ok := localPaths[f.Path] - if !ok { - return "", fmt.Errorf("internal error: missing local path for downloaded config file %s", f.Path) - } - configPaths = append(configPaths, localPath) - } - - if len(configPaths) == 0 { - // No config files to archive - return "", nil - } - - // Sort for reproducibility - sort.Strings(configPaths) - - // Create the archive in our tempDir so it gets cleaned up with everything else - archivePath, err := packaging.CreateConfigArchiveInDir(configPaths, tempDir) - if err != nil { - return "", fmt.Errorf("create config archive: %w", err) - } - - return archivePath, nil -} - // isChatTemplate checks if a file is a chat template func isChatTemplate(path string) bool { filename := filepath.Base(path) diff --git a/pkg/distribution/packaging/directory.go b/pkg/distribution/packaging/directory.go deleted file mode 100644 index a4348f202..000000000 --- a/pkg/distribution/packaging/directory.go +++ /dev/null @@ -1,157 +0,0 @@ -package packaging - -import ( - "archive/tar" - "fmt" - "io" - "os" - "path/filepath" - "sort" - - "github.com/docker/model-runner/pkg/distribution/files" -) - -// 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) - - // Classify file using centralized files package - fileType := files.Classify(name) - - switch fileType { - case files.FileTypeSafetensors: - safetensorsPaths = append(safetensorsPaths, fullPath) - case files.FileTypeConfig, files.FileTypeChatTemplate: - configFiles = append(configFiles, fullPath) - case files.FileTypeUnknown, files.FileTypeGGUF, files.FileTypeLicense, files.FileTypeDDUF: - // Skip these file types - } - } - - 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) { - return CreateConfigArchiveInDir(configFiles, "") -} - -// CreateConfigArchiveInDir creates a tar archive containing the specified config files in the given directory. -// If dir is empty, the system temp directory is used. -// It returns the path to the tar file and any error encountered. -// The caller is responsible for removing the file when done. -func CreateConfigArchiveInDir(configFiles []string, dir string) (string, error) { - // Create temp file in specified directory - tmpFile, err := os.CreateTemp(dir, "vllm-config-*.tar") - if err != nil { - return "", fmt.Errorf("create temp file: %w", err) - } - 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 { - if err := addFileToTar(tw, filePath); err != nil { - tw.Close() - tmpFile.Close() - return "", err - } - } - - // Close tar writer first - if err := tw.Close(); err != nil { - tmpFile.Close() - return "", fmt.Errorf("close tar writer: %w", err) - } - - // Close temp file - if err := tmpFile.Close(); err != nil { - 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 -} diff --git a/pkg/distribution/packaging/directory_test.go b/pkg/distribution/packaging/directory_test.go deleted file mode 100644 index 1031f6b71..000000000 --- a/pkg/distribution/packaging/directory_test.go +++ /dev/null @@ -1,297 +0,0 @@ -package packaging - -import ( - "archive/tar" - "io" - "os" - "path/filepath" - "sort" - "strings" - "testing" -) - -func TestPackageFromDirectory_WithTokenizerModel(t *testing.T) { - // Create temporary directory - tempDir := t.TempDir() - - // Create test files - files := map[string]string{ - "model.safetensors": "safetensors content", - "config.json": `{"model_type": "test"}`, - "tokenizer.model": "tokenizer model binary content", - "tokenizer_config.json": `{"tokenizer_class": "TestTokenizer"}`, - "not.included": `not included content`, - } - - for name, content := range files { - path := filepath.Join(tempDir, name) - if err := os.WriteFile(path, []byte(content), 0644); err != nil { - t.Fatalf("Failed to create test file %s: %v", name, err) - } - } - - // Call PackageFromDirectory - safetensorsPaths, tempConfigArchive, err := PackageFromDirectory(tempDir) - if err != nil { - t.Fatalf("PackageFromDirectory failed: %v", err) - } - - // Clean up temp archive - if tempConfigArchive != "" { - defer os.Remove(tempConfigArchive) - } - - // Verify safetensors files were found - if len(safetensorsPaths) != 1 { - t.Errorf("Expected 1 safetensors file, got %d", len(safetensorsPaths)) - } - - // Verify config archive was created - if tempConfigArchive == "" { - t.Fatal("Expected config archive to be created") - } - - // Verify tokenizer.model is in the archive - archiveFiles, err := readTarArchive(tempConfigArchive) - if err != nil { - t.Fatalf("Failed to read tar archive: %v", err) - } - - expectedFiles := []string{"config.json", "tokenizer.model", "tokenizer_config.json"} - sort.Strings(expectedFiles) - sort.Strings(archiveFiles) - - if len(archiveFiles) != len(expectedFiles) { - t.Errorf("Expected %d files in archive, got %d", len(expectedFiles), len(archiveFiles)) - } - - for i, expected := range expectedFiles { - if i >= len(archiveFiles) || archiveFiles[i] != expected { - t.Errorf("Expected file %s in archive, got %v", expected, archiveFiles) - } - } -} - -func TestPackageFromDirectory_BasicFunctionality(t *testing.T) { - // Create temporary directory - tempDir := t.TempDir() - - // Create test files - files := map[string]string{ - "model-00001-of-00002.safetensors": "safetensors content 1", - "model-00002-of-00002.safetensors": "safetensors content 2", - "config.json": `{"model_type": "test"}`, - "merges.txt": "merge1 merge2", - "tokenizer.model": "tokenizer content", - "special_tokens_map.json": `{"unk_token": ""}`, - } - - for name, content := range files { - path := filepath.Join(tempDir, name) - if err := os.WriteFile(path, []byte(content), 0644); err != nil { - t.Fatalf("Failed to create test file %s: %v", name, err) - } - } - - // Call PackageFromDirectory - safetensorsPaths, tempConfigArchive, err := PackageFromDirectory(tempDir) - if err != nil { - t.Fatalf("PackageFromDirectory failed: %v", err) - } - - // Clean up temp archive - if tempConfigArchive != "" { - defer os.Remove(tempConfigArchive) - } - - // Verify safetensors files - if len(safetensorsPaths) != 2 { - t.Errorf("Expected 2 safetensors files, got %d", len(safetensorsPaths)) - } - - // Verify files are sorted - for i := 0; i < len(safetensorsPaths)-1; i++ { - if safetensorsPaths[i] > safetensorsPaths[i+1] { - t.Error("Safetensors paths are not sorted") - } - } - - // Verify config archive was created - if tempConfigArchive == "" { - t.Fatal("Expected config archive to be created") - } - - // Verify archive contents - archiveFiles, err := readTarArchive(tempConfigArchive) - if err != nil { - t.Fatalf("Failed to read tar archive: %v", err) - } - - expectedConfigFiles := []string{ - "config.json", - "merges.txt", - "tokenizer.model", - "special_tokens_map.json", - } - sort.Strings(expectedConfigFiles) - sort.Strings(archiveFiles) - - if len(archiveFiles) != len(expectedConfigFiles) { - t.Errorf("Expected %d config files in archive, got %d", len(expectedConfigFiles), len(archiveFiles)) - } - - for i, expected := range expectedConfigFiles { - if i >= len(archiveFiles) || archiveFiles[i] != expected { - t.Errorf("Expected file %s in archive at position %d, got %v", expected, i, archiveFiles) - } - } -} - -func TestPackageFromDirectory_NoSafetensorsFiles(t *testing.T) { - // Create temporary directory - tempDir := t.TempDir() - - // Create only config files (no safetensors) - files := map[string]string{ - "config.json": `{"model_type": "test"}`, - "tokenizer.model": "tokenizer content", - } - - for name, content := range files { - path := filepath.Join(tempDir, name) - if err := os.WriteFile(path, []byte(content), 0644); err != nil { - t.Fatalf("Failed to create test file %s: %v", name, err) - } - } - - // Call PackageFromDirectory - _, _, err := PackageFromDirectory(tempDir) - if err == nil { - t.Fatal("Expected error when no safetensors files found, got nil") - } - - expectedError := "no safetensors files found" - if !strings.Contains(err.Error(), expectedError) { - t.Errorf("Expected error containing %q, got: %v", expectedError, err) - } -} - -func TestPackageFromDirectory_OnlySafetensorsFiles(t *testing.T) { - // Create temporary directory - tempDir := t.TempDir() - - // Create only safetensors files (no config files) - files := map[string]string{ - "model.safetensors": "safetensors content", - } - - for name, content := range files { - path := filepath.Join(tempDir, name) - if err := os.WriteFile(path, []byte(content), 0644); err != nil { - t.Fatalf("Failed to create test file %s: %v", name, err) - } - } - - // Call PackageFromDirectory - safetensorsPaths, tempConfigArchive, err := PackageFromDirectory(tempDir) - if err != nil { - t.Fatalf("PackageFromDirectory failed: %v", err) - } - - // Verify safetensors files were found - if len(safetensorsPaths) != 1 { - t.Errorf("Expected 1 safetensors file, got %d", len(safetensorsPaths)) - } - - // Verify no config archive was created - if tempConfigArchive != "" { - defer os.Remove(tempConfigArchive) - t.Error("Expected no config archive to be created when no config files exist") - } -} - -func TestPackageFromDirectory_SkipsSubdirectories(t *testing.T) { - // Create temporary directory - tempDir := t.TempDir() - - // Create test files in root - rootFiles := map[string]string{ - "model.safetensors": "safetensors content", - "config.json": `{"model_type": "test"}`, - } - - for name, content := range rootFiles { - path := filepath.Join(tempDir, name) - if err := os.WriteFile(path, []byte(content), 0644); err != nil { - t.Fatalf("Failed to create test file %s: %v", name, err) - } - } - - // Create subdirectory with files that should be ignored - subDir := filepath.Join(tempDir, "subdir") - if err := os.Mkdir(subDir, 0755); err != nil { - t.Fatalf("Failed to create subdirectory: %v", err) - } - - subFiles := map[string]string{ - "ignored.safetensors": "should be ignored", - "ignored.json": `{"ignored": true}`, - } - - for name, content := range subFiles { - path := filepath.Join(subDir, name) - if err := os.WriteFile(path, []byte(content), 0644); err != nil { - t.Fatalf("Failed to create test file in subdir %s: %v", name, err) - } - } - - // Call PackageFromDirectory - safetensorsPaths, tempConfigArchive, err := PackageFromDirectory(tempDir) - if err != nil { - t.Fatalf("PackageFromDirectory failed: %v", err) - } - - // Clean up temp archive - if tempConfigArchive != "" { - defer os.Remove(tempConfigArchive) - } - - // Verify only root-level files were processed - if len(safetensorsPaths) != 1 { - t.Errorf("Expected 1 safetensors file from root directory, got %d", len(safetensorsPaths)) - } - - archiveFiles, err := readTarArchive(tempConfigArchive) - if err != nil { - t.Fatalf("Failed to read tar archive: %v", err) - } - - if len(archiveFiles) != 1 || archiveFiles[0] != "config.json" { - t.Errorf("Expected only config.json from root directory, got %v", archiveFiles) - } -} - -// Helper function to read tar archive and return list of file names -func readTarArchive(archivePath string) ([]string, error) { - file, err := os.Open(archivePath) - if err != nil { - return nil, err - } - defer file.Close() - - tr := tar.NewReader(file) - var files []string - - for { - header, err := tr.Next() - if err == io.EOF { - break - } - if err != nil { - return nil, err - } - files = append(files, header.Name) - } - - return files, nil -} diff --git a/pkg/distribution/packaging/dirtar.go b/pkg/distribution/packaging/dirtar.go deleted file mode 100644 index 82c5aa2cf..000000000 --- a/pkg/distribution/packaging/dirtar.go +++ /dev/null @@ -1,209 +0,0 @@ -package packaging - -import ( - "archive/tar" - "fmt" - "io" - "os" - "path/filepath" - "strings" -) - -// CreateDirectoryTarArchive creates a temporary tar archive containing the specified directory -// with its structure preserved. Symlinks encountered in the directory are skipped and will not be included -// in the archive. 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 CreateDirectoryTarArchive(dirPath string) (string, error) { - // Verify directory exists - info, err := os.Stat(dirPath) - if err != nil { - return "", fmt.Errorf("stat directory: %w", err) - } - if !info.IsDir() { - return "", fmt.Errorf("path is not a directory: %s", dirPath) - } - - // Create temp file - tmpFile, err := os.CreateTemp("", "dir-tar-*.tar") - if err != nil { - return "", fmt.Errorf("create temp file: %w", err) - } - 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) - - // Walk the directory tree - err = filepath.Walk(dirPath, func(path string, info os.FileInfo, err error) error { - if err != nil { - return err - } - if info == nil { - return fmt.Errorf("nil FileInfo for path: %s", path) - } - // Skip symlinks - they're not needed for model distribution and are - // skipped during extraction for security reasons - if info.Mode()&os.ModeSymlink != 0 { - return nil - } - - // Create tar header - header, err := tar.FileInfoHeader(info, "") - if err != nil { - return fmt.Errorf("create tar header for %s: %w", path, err) - } - - // Compute relative path from the parent of dirPath - relPath, err := filepath.Rel(filepath.Dir(dirPath), path) - if err != nil { - return fmt.Errorf("compute relative path: %w", err) - } - - // Use forward slashes for tar archive paths - header.Name = filepath.ToSlash(relPath) - - // Write header - if err := tw.WriteHeader(header); err != nil { - return fmt.Errorf("write tar header: %w", err) - } - - // If it's a file, write its contents - if !info.IsDir() { - file, err := os.Open(path) - if err != nil { - return fmt.Errorf("open file %s: %w", path, err) - } - - // Copy file contents - if _, err := io.Copy(tw, file); err != nil { - file.Close() - return fmt.Errorf("write tar content for %s: %w", path, err) - } - if err := file.Close(); err != nil { - return fmt.Errorf("close file %s: %w", path, err) - } - } - - return nil - }) - - if err != nil { - tw.Close() - tmpFile.Close() - return "", fmt.Errorf("walk directory: %w", err) - } - - // Close tar writer - if err := tw.Close(); err != nil { - tmpFile.Close() - return "", fmt.Errorf("close tar writer: %w", err) - } - - // Close temp file - if err := tmpFile.Close(); err != nil { - return "", fmt.Errorf("close temp file: %w", err) - } - - shouldKeepTempFile = true - return tmpPath, nil -} - -// DirTarProcessor handles processing of directory tar paths for packaging -type DirTarProcessor struct { - dirTarPaths []string - baseDir string - tempFiles []string -} - -// NewDirTarProcessor creates a new processor for directory tar paths -func NewDirTarProcessor(dirTarPaths []string, baseDir string) *DirTarProcessor { - return &DirTarProcessor{ - dirTarPaths: dirTarPaths, - baseDir: baseDir, - tempFiles: make([]string, 0), - } -} - -// Process processes all directory tar paths, validates them, and creates temporary tar archives. -// Returns a list of temporary tar file paths, cleanup function, and any error encountered. -// The caller is responsible for adding these tar files to the builder. -func (p *DirTarProcessor) Process() ([]string, func(), error) { - var tarPaths []string - - // Get absolute paths for robust security validation - absBase, err := filepath.Abs(p.baseDir) - if err != nil { - return nil, nil, fmt.Errorf("resolve base directory: %w", err) - } - - // Return cleanup function - cleanup := func() { - for _, tempFile := range p.tempFiles { - os.Remove(tempFile) - } - } - - for _, relDirPath := range p.dirTarPaths { - // Reject absolute paths - if filepath.IsAbs(relDirPath) { - return nil, cleanup, fmt.Errorf("dir-tar path must be relative: %s", relDirPath) - } - - // Resolve the full directory path - fullDirPath := filepath.Join(p.baseDir, relDirPath) - fullDirPath = filepath.Clean(fullDirPath) - - absFull, err := filepath.Abs(fullDirPath) - if err != nil { - return nil, cleanup, fmt.Errorf("resolve full path: %w", err) - } - - // Use filepath.Rel to compute the relative path from base to full - // This is the canonical way to check if a path is within another - relPathCheck, err := filepath.Rel(absBase, absFull) - if err != nil { - return nil, cleanup, fmt.Errorf("dir-tar path %q could not be validated: %w", relDirPath, err) - } - - // If the relative path starts with ".." as a path component (not just as prefix), - // it means absFull is outside absBase. We check for ".." followed by separator - // or as the entire path to avoid false positives with directories like "..data" - if relPathCheck == ".." || strings.HasPrefix(relPathCheck, ".."+string(os.PathSeparator)) { - return nil, cleanup, fmt.Errorf("dir-tar path %q escapes base directory", relDirPath) - } - - // Use Lstat (not Stat) to check if the path itself is a symlink - // Stat would follow the symlink, but we want to detect symlinks themselves - linfo, err := os.Lstat(fullDirPath) - if err != nil { - return nil, cleanup, fmt.Errorf("cannot access directory %q (resolved from %q): %w", fullDirPath, relDirPath, err) - } - - // Reject symlinks to prevent empty tar archives (CreateDirectoryTarArchive skips symlinks) - if linfo.Mode()&os.ModeSymlink != 0 { - return nil, cleanup, fmt.Errorf("path %q is a symlink; symlinked directories are not supported", relDirPath) - } - - // Verify it's a directory - if !linfo.IsDir() { - return nil, cleanup, fmt.Errorf("path %q is not a directory", fullDirPath) - } - - tempTarPath, err := CreateDirectoryTarArchive(fullDirPath) - if err != nil { - return nil, cleanup, fmt.Errorf("create tar archive for directory %q: %w", relDirPath, err) - } - p.tempFiles = append(p.tempFiles, tempTarPath) - tarPaths = append(tarPaths, tempTarPath) - } - - return tarPaths, cleanup, nil -} diff --git a/pkg/distribution/packaging/dirtar_test.go b/pkg/distribution/packaging/dirtar_test.go deleted file mode 100644 index ea764f591..000000000 --- a/pkg/distribution/packaging/dirtar_test.go +++ /dev/null @@ -1,362 +0,0 @@ -package packaging - -import ( - "archive/tar" - "io" - "os" - "path/filepath" - "strings" - "testing" -) - -func TestCreateDirTarArchive(t *testing.T) { - tempDir := t.TempDir() - - // Create test directory structure - testDir := filepath.Join(tempDir, "test_directory") - if err := os.MkdirAll(filepath.Join(testDir, "subdir"), 0755); err != nil { - t.Fatalf("Failed to create test directory: %v", err) - } - - // Create test files - testFiles := map[string]string{ - "file1.txt": "content1", - "subdir/file2.txt": "content2", - } - - for relPath, content := range testFiles { - fullPath := filepath.Join(testDir, relPath) - if err := os.WriteFile(fullPath, []byte(content), 0644); err != nil { - t.Fatalf("Failed to write test file %s: %v", relPath, err) - } - } - - // Create tar archive - tarPath, err := CreateDirectoryTarArchive(testDir) - if err != nil { - t.Fatalf("CreateDirectoryTarArchive failed: %v", err) - } - defer os.Remove(tarPath) - - // Verify tar archive exists - if _, err := os.Stat(tarPath); os.IsNotExist(err) { - t.Fatal("Tar archive was not created") - } - - // Read and verify tar contents - file, err := os.Open(tarPath) - if err != nil { - t.Fatalf("Failed to open tar archive: %v", err) - } - defer file.Close() - - tr := tar.NewReader(file) - foundFiles := make(map[string]bool) - - for { - header, err := tr.Next() - if err == io.EOF { - break - } - if err != nil { - t.Fatalf("Failed to read tar header: %v", err) - } - - foundFiles[header.Name] = true - - // Verify it's within the test_directory structure - if header.Typeflag == tar.TypeReg { - // Read file content - content, err := io.ReadAll(tr) - if err != nil { - t.Fatalf("Failed to read file content: %v", err) - } - - // Verify content matches - expectedPath := header.Name[len("test_directory/"):] - if expectedContent, ok := testFiles[expectedPath]; ok { - if string(content) != expectedContent { - t.Errorf("File %s content mismatch: got %q, want %q", expectedPath, string(content), expectedContent) - } - } - } - } - - // Verify all expected entries are present - expectedEntries := []string{ - "test_directory", - "test_directory/file1.txt", - "test_directory/subdir", - "test_directory/subdir/file2.txt", - } - - for _, entry := range expectedEntries { - if !foundFiles[entry] { - t.Errorf("Expected entry %q not found in tar archive", entry) - } - } -} - -func TestCreateDirTarArchive_NonExistentDir(t *testing.T) { - _, err := CreateDirectoryTarArchive("/nonexistent/directory") - if err == nil { - t.Error("Expected error for non-existent directory, got nil") - } -} - -func TestCreateDirTarArchive_NotADirectory(t *testing.T) { - // Create a temporary file - tempFile, err := os.CreateTemp("", "not-a-dir-*") - if err != nil { - t.Fatalf("Failed to create temp file: %v", err) - } - defer os.Remove(tempFile.Name()) - tempFile.Close() - - _, err = CreateDirectoryTarArchive(tempFile.Name()) - if err == nil { - t.Error("Expected error for file path instead of directory, got nil") - } -} - -func TestDirTarProcessor_ValidRelativePaths(t *testing.T) { - tempDir := t.TempDir() - - // Create test subdirectories - subDir1 := filepath.Join(tempDir, "config") - subDir2 := filepath.Join(tempDir, "templates") - if err := os.MkdirAll(subDir1, 0755); err != nil { - t.Fatalf("Failed to create subdir1: %v", err) - } - if err := os.MkdirAll(subDir2, 0755); err != nil { - t.Fatalf("Failed to create subdir2: %v", err) - } - - // Create test files - if err := os.WriteFile(filepath.Join(subDir1, "config.json"), []byte("{}"), 0644); err != nil { - t.Fatalf("Failed to write config file: %v", err) - } - if err := os.WriteFile(filepath.Join(subDir2, "template.txt"), []byte("template"), 0644); err != nil { - t.Fatalf("Failed to write template file: %v", err) - } - - // Test processing valid relative paths - processor := NewDirTarProcessor([]string{"config", "templates"}, tempDir) - tarPaths, cleanup, err := processor.Process() - if err != nil { - t.Fatalf("Process failed: %v", err) - } - defer cleanup() - - // Verify we got 2 tar files - if len(tarPaths) != 2 { - t.Errorf("Expected 2 tar files, got %d", len(tarPaths)) - } - - // Verify tar files exist - for _, tarPath := range tarPaths { - if _, err := os.Stat(tarPath); os.IsNotExist(err) { - t.Errorf("Tar file does not exist: %s", tarPath) - } - } - - // Test cleanup - cleanup() - for _, tarPath := range tarPaths { - if _, err := os.Stat(tarPath); !os.IsNotExist(err) { - t.Errorf("Tar file was not cleaned up: %s", tarPath) - } - } -} - -func TestDirTarProcessor_DirectoryTraversal_DoubleDot(t *testing.T) { - tempDir := t.TempDir() - - // Test various directory traversal attempts using OS-specific path separator - sep := string(os.PathSeparator) - traversalAttempts := []string{ - "..", - ".." + sep, // "../" on Unix, "..\\" on Windows - ".." + sep + ".." + sep + "etc", // "../../etc" or "..\\..\\etc" - ".." + sep + ".." + sep + ".." + sep + "etc", // "../../../etc" or "..\\..\\..\\etc" - "foo" + sep + ".." + sep + ".." + sep + ".." + sep + "etc", - "subdir" + sep + ".." + sep + ".." + sep + "..", - "." + sep + ".." + sep + ".." + sep + "etc", - } - - for _, attempt := range traversalAttempts { - t.Run(attempt, func(t *testing.T) { - processor := NewDirTarProcessor([]string{attempt}, tempDir) - _, _, err := processor.Process() - if err == nil { - t.Errorf("Expected error for traversal attempt %q, got nil", attempt) - } - if err != nil && !strings.Contains(err.Error(), "escapes base directory") { - t.Errorf("Expected 'escapes base directory' error for %q, got: %v", attempt, err) - } - }) - } -} - -func TestDirTarProcessor_AbsolutePathRejection(t *testing.T) { - tempDir := t.TempDir() - - // Test absolute path rejection - absolutePaths := []string{ - "/etc/passwd", - "/usr/bin", - filepath.Join(tempDir, "subdir"), // Absolute even though within tempDir - } - - for _, absPath := range absolutePaths { - t.Run(absPath, func(t *testing.T) { - processor := NewDirTarProcessor([]string{absPath}, tempDir) - _, _, err := processor.Process() - if err == nil { - t.Errorf("Expected error for absolute path %q, got nil", absPath) - } - if err != nil && !strings.Contains(err.Error(), "must be relative") { - t.Errorf("Expected 'must be relative' error for %q, got: %v", absPath, err) - } - }) - } -} - -func TestDirTarProcessor_NonExistentDirectory(t *testing.T) { - processor := NewDirTarProcessor([]string{"nonexistent"}, t.TempDir()) - _, _, err := processor.Process() - if err == nil { - t.Error("Expected error for non-existent directory, got nil") - } - if err != nil && !strings.Contains(err.Error(), "cannot access directory") { - t.Errorf("Expected 'cannot access directory' error, got: %v", err) - } -} - -func TestDirTarProcessor_FileInsteadOfDirectory(t *testing.T) { - tempDir := t.TempDir() - - // Create a file instead of directory - filePath := filepath.Join(tempDir, "not-a-dir.txt") - if err := os.WriteFile(filePath, []byte("content"), 0644); err != nil { - t.Fatalf("Failed to create test file: %v", err) - } - - processor := NewDirTarProcessor([]string{"not-a-dir.txt"}, tempDir) - _, _, err := processor.Process() - if err == nil { - t.Error("Expected error for file instead of directory, got nil") - } - if err != nil && !strings.Contains(err.Error(), "not a directory") { - t.Errorf("Expected 'not a directory' error, got: %v", err) - } -} - -func TestDirTarProcessor_BaseDirItself(t *testing.T) { - tempDir := t.TempDir() - - // Create a file in the base directory - if err := os.WriteFile(filepath.Join(tempDir, "file.txt"), []byte("content"), 0644); err != nil { - t.Fatalf("Failed to create test file: %v", err) - } - - // Test with "." which should reference the base directory itself - processor := NewDirTarProcessor([]string{"."}, tempDir) - tarPaths, cleanup, err := processor.Process() - if err != nil { - t.Fatalf("Process failed for base directory: %v", err) - } - defer cleanup() - - if len(tarPaths) != 1 { - t.Errorf("Expected 1 tar file, got %d", len(tarPaths)) - } -} - -func TestDirTarProcessor_EmptyList(t *testing.T) { - processor := NewDirTarProcessor([]string{}, t.TempDir()) - tarPaths, cleanup, err := processor.Process() - if err != nil { - t.Fatalf("Process failed for empty list: %v", err) - } - defer cleanup() - - if len(tarPaths) != 0 { - t.Errorf("Expected 0 tar files for empty list, got %d", len(tarPaths)) - } -} - -func TestDirTarProcessor_DoubleDotPrefixedDirectories(t *testing.T) { - tempDir := t.TempDir() - - // Create directories with names that start with ".." but are not traversal attempts - doubleDotDirs := []string{ - "..data", - "..gitkeep", - "..metadata", - "..2024_backup", - } - - for _, dirName := range doubleDotDirs { - dirPath := filepath.Join(tempDir, dirName) - if err := os.MkdirAll(dirPath, 0755); err != nil { - t.Fatalf("Failed to create test directory %s: %v", dirName, err) - } - // Add a file so the tar isn't empty - if err := os.WriteFile(filepath.Join(dirPath, "file.txt"), []byte("content"), 0644); err != nil { - t.Fatalf("Failed to create test file: %v", err) - } - } - - // Test processing these directories - processor := NewDirTarProcessor(doubleDotDirs, tempDir) - tarPaths, cleanup, err := processor.Process() - if err != nil { - t.Fatalf("Process failed for double-dot prefixed directories: %v", err) - } - defer cleanup() - - if len(tarPaths) != len(doubleDotDirs) { - t.Errorf("Expected %d tar files, got %d", len(doubleDotDirs), len(tarPaths)) - } - - // Verify tar files exist and are not empty - for i, tarPath := range tarPaths { - info, err := os.Stat(tarPath) - if err != nil { - t.Errorf("Failed to stat test file %q: %v", tarPath, err) - } - if info.Size() == 0 { - t.Errorf("Tar file is empty for %s", doubleDotDirs[i]) - } - } -} - -func TestDirTarProcessor_SymlinkedDirectory(t *testing.T) { - tempDir := t.TempDir() - - // Create a real directory - realDir := filepath.Join(tempDir, "realdir") - if err := os.MkdirAll(realDir, 0755); err != nil { - t.Fatalf("Failed to create real directory: %v", err) - } - if err := os.WriteFile(filepath.Join(realDir, "file.txt"), []byte("content"), 0644); err != nil { - t.Fatalf("Failed to create test file: %v", err) - } - - // Create a symlink to the directory - symlinkPath := filepath.Join(tempDir, "symlink") - if err := os.Symlink(realDir, symlinkPath); err != nil { - t.Skipf("Cannot create symlink (may not be supported on this system): %v", err) - } - - // Test that the symlink is rejected - processor := NewDirTarProcessor([]string{"symlink"}, tempDir) - _, _, err := processor.Process() - if err == nil { - t.Error("Expected error for symlinked directory, got nil") - } - if err != nil && !strings.Contains(err.Error(), "symlink") { - t.Errorf("Expected error about symlink, got: %v", err) - } -} From ff896b554b2900c44bf7d29111beb202abedee01 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20L=C3=B3pez=20Luna?= Date: Thu, 12 Feb 2026 15:03:38 +0100 Subject: [PATCH 3/7] feat(unpack): update Safetensors model packaging to support individual layer files from nested subdirectories --- cmd/cli/docs/reference/docker_model_package.yaml | 13 +------------ cmd/cli/docs/reference/model_package.md | 3 +-- 2 files changed, 2 insertions(+), 14 deletions(-) diff --git a/cmd/cli/docs/reference/docker_model_package.yaml b/cmd/cli/docs/reference/docker_model_package.yaml index 79d608997..77849d47e 100644 --- a/cmd/cli/docs/reference/docker_model_package.yaml +++ b/cmd/cli/docs/reference/docker_model_package.yaml @@ -4,7 +4,7 @@ short: | long: |- Package a GGUF file, Safetensors directory, DDUF file, or existing model into a Docker model OCI artifact, with optional licenses and multimodal projector. 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 a Safetensors model, --safetensors-dir should point to a directory containing .safetensors files and config files. All files (including nested subdirectories) will be auto-discovered and each file is packaged as an individual layer. When packaging a DDUF file (Diffusers Unified Format), --dduf should point to a .dduf archive file. When packaging from an existing model using --from, you can modify properties like context size to create a variant of the original model. For multimodal models, use --mmproj to include a multimodal projector file. @@ -40,17 +40,6 @@ options: experimentalcli: false kubernetes: false swarm: false - - option: dir-tar - value_type: stringArray - default_value: '[]' - description: | - relative path to directory to package as tar (can be specified multiple times) - deprecated: false - hidden: false - experimental: false - experimentalcli: false - kubernetes: false - swarm: false - option: from value_type: string description: reference to an existing model to repackage diff --git a/cmd/cli/docs/reference/model_package.md b/cmd/cli/docs/reference/model_package.md index 062f15815..ebb1c294f 100644 --- a/cmd/cli/docs/reference/model_package.md +++ b/cmd/cli/docs/reference/model_package.md @@ -3,7 +3,7 @@ Package a GGUF file, Safetensors directory, DDUF file, or existing model into a Docker model OCI artifact, with optional licenses and multimodal projector. 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 a Safetensors model, --safetensors-dir should point to a directory containing .safetensors files and config files. All files (including nested subdirectories) will be auto-discovered and each file is packaged as an individual layer. When packaging a DDUF file (Diffusers Unified Format), --dduf should point to a .dduf archive file. When packaging from an existing model using --from, you can modify properties like context size to create a variant of the original model. For multimodal models, use --mmproj to include a multimodal projector file. @@ -15,7 +15,6 @@ For multimodal models, use --mmproj to include a multimodal projector file. | `--chat-template` | `string` | | absolute path to chat template file (must be Jinja format) | | `--context-size` | `uint64` | `0` | context size in tokens | | `--dduf` | `string` | | absolute path to DDUF archive file (Diffusers Unified Format) | -| `--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 | From 1f61305769ba488b74844b2dc40055a0477a419a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20L=C3=B3pez=20Luna?= Date: Thu, 12 Feb 2026 15:17:56 +0100 Subject: [PATCH 4/7] feat(unpack): refine model packaging documentation and streamline command descriptions --- cmd/cli/commands/package.go | 46 ++++++++++++--- .../docs/reference/docker_model_package.yaml | 47 ++++++++++++--- cmd/cli/docs/reference/model.md | 58 +++++++++---------- cmd/cli/docs/reference/model_package.md | 44 ++++++++++++-- 4 files changed, 145 insertions(+), 50 deletions(-) diff --git a/cmd/cli/commands/package.go b/cmd/cli/commands/package.go index c919e80d2..d1c8dcdab 100644 --- a/cmd/cli/commands/package.go +++ b/cmd/cli/commands/package.go @@ -39,13 +39,45 @@ func newPackagedCmd() *cobra.Command { c := &cobra.Command{ Use: "package (--gguf | --safetensors-dir | --dduf | --from ) [--license ...] [--mmproj ] [--context-size ] [--push] MODEL", - Short: "Package a GGUF file, Safetensors directory, DDUF file, or existing model into a Docker model OCI artifact.", - Long: "Package a GGUF file, Safetensors directory, DDUF file, or existing model into a Docker model OCI artifact, with optional licenses and multimodal projector. 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. All files (including nested subdirectories) will be auto-discovered and each file is packaged as an individual layer.\n" + - "When packaging a DDUF file (Diffusers Unified Format), --dduf should point to a .dduf archive file.\n" + - "When packaging from an existing model using --from, you can modify properties like context size to create a variant of the original model.\n" + - "For multimodal models, use --mmproj to include a multimodal projector file.", + Short: "Package a model into a Docker Model OCI artifact", + Long: `Package a model into a Docker Model OCI artifact. + +The model source must be one of: + --gguf A GGUF file (single file or first shard of a sharded model) + --safetensors-dir A directory containing .safetensors and configuration files + --dduf A .dduf (Diffusers Unified Format) archive + --from An existing packaged model reference + +By default, the packaged artifact is loaded into the local Model Runner content store. +Use --push to publish the model to a registry instead. + +MODEL specifies the target model reference (for example: myorg/llama3:8b). +When using --push, MODEL must be a registry-qualified reference. + +Packaging behavior: + + GGUF + --gguf must point to a .gguf file. + For sharded models, point to the first shard. All shards must: + • reside in the same directory + • follow an indexed naming convention (e.g. model-00001-of-00015.gguf) + All shards are automatically discovered and packaged together. + + Safetensors + --safetensors-dir must point to a directory containing .safetensors files + and required configuration files (e.g. model config, tokenizer files). + All files under the directory (including nested subdirectories) are + automatically discovered. Each file is packaged as a separate OCI layer. + + DDUF + --dduf must point to a .dduf archive file. + + Repackaging + --from repackages an existing model. You may override selected properties + such as --context-size to create a variant of the original model. + + Multimodal models + Use --mmproj to include a multimodal projector file.`, Args: func(cmd *cobra.Command, args []string) error { if err := requireExactArgs(1, "package", "MODEL")(cmd, args); err != nil { return err diff --git a/cmd/cli/docs/reference/docker_model_package.yaml b/cmd/cli/docs/reference/docker_model_package.yaml index 77849d47e..ce868c1e5 100644 --- a/cmd/cli/docs/reference/docker_model_package.yaml +++ b/cmd/cli/docs/reference/docker_model_package.yaml @@ -1,13 +1,44 @@ command: docker model package -short: | - Package a GGUF file, Safetensors directory, DDUF file, or existing model into a Docker model OCI artifact. +short: Package a model into a Docker Model OCI artifact long: |- - Package a GGUF file, Safetensors directory, DDUF file, or existing model into a Docker model OCI artifact, with optional licenses and multimodal projector. 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. All files (including nested subdirectories) will be auto-discovered and each file is packaged as an individual layer. - When packaging a DDUF file (Diffusers Unified Format), --dduf should point to a .dduf archive file. - When packaging from an existing model using --from, you can modify properties like context size to create a variant of the original model. - For multimodal models, use --mmproj to include a multimodal projector file. + Package a model into a Docker Model OCI artifact. + + The model source must be one of: + --gguf A GGUF file (single file or first shard of a sharded model) + --safetensors-dir A directory containing .safetensors and configuration files + --dduf A .dduf (Diffusers Unified Format) archive + --from An existing packaged model reference + + By default, the packaged artifact is loaded into the local Model Runner content store. + Use --push to publish the model to a registry instead. + + MODEL specifies the target model reference (for example: myorg/llama3:8b). + When using --push, MODEL must be a registry-qualified reference. + + Packaging behavior: + + GGUF + --gguf must point to a .gguf file. + For sharded models, point to the first shard. All shards must: + • reside in the same directory + • follow an indexed naming convention (e.g. model-00001-of-00015.gguf) + All shards are automatically discovered and packaged together. + + Safetensors + --safetensors-dir must point to a directory containing .safetensors files + and required configuration files (e.g. model config, tokenizer files). + All files under the directory (including nested subdirectories) are + automatically discovered. Each file is packaged as a separate OCI layer. + + DDUF + --dduf must point to a .dduf archive file. + + Repackaging + --from repackages an existing model. You may override selected properties + such as --context-size to create a variant of the original model. + + Multimodal models + Use --mmproj to include a multimodal projector file. usage: docker model package (--gguf | --safetensors-dir | --dduf | --from ) [--license ...] [--mmproj ] [--context-size ] [--push] MODEL pname: docker model plink: docker_model.yaml diff --git a/cmd/cli/docs/reference/model.md b/cmd/cli/docs/reference/model.md index e2869cd98..b439eeaf2 100644 --- a/cmd/cli/docs/reference/model.md +++ b/cmd/cli/docs/reference/model.md @@ -5,35 +5,35 @@ Docker Model Runner ### Subcommands -| Name | Description | -|:------------------------------------------------|:-----------------------------------------------------------------------------------------------------------| -| [`bench`](model_bench.md) | Benchmark a model's performance at different concurrency levels | -| [`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) | -| [`launch`](model_launch.md) | Launch an app configured to use Docker Model Runner | -| [`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, DDUF file, 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 | -| [`purge`](model_purge.md) | Remove all models | -| [`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 | -| [`search`](model_search.md) | Search for models on Docker Hub and HuggingFace | -| [`show`](model_show.md) | Show information for a model | -| [`skills`](model_skills.md) | Install Docker Model Runner skills for AI coding assistants | -| [`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 | +|:------------------------------------------------|:-----------------------------------------------------------------------| +| [`bench`](model_bench.md) | Benchmark a model's performance at different concurrency levels | +| [`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) | +| [`launch`](model_launch.md) | Launch an app configured to use Docker Model Runner | +| [`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 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 | +| [`purge`](model_purge.md) | Remove all models | +| [`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 | +| [`search`](model_search.md) | Search for models on Docker Hub and HuggingFace | +| [`show`](model_show.md) | Show information for a model | +| [`skills`](model_skills.md) | Install Docker Model Runner skills for AI coding assistants | +| [`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 ebb1c294f..ade44149b 100644 --- a/cmd/cli/docs/reference/model_package.md +++ b/cmd/cli/docs/reference/model_package.md @@ -1,12 +1,44 @@ # docker model package -Package a GGUF file, Safetensors directory, DDUF file, or existing model into a Docker model OCI artifact, with optional licenses and multimodal projector. 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. All files (including nested subdirectories) will be auto-discovered and each file is packaged as an individual layer. -When packaging a DDUF file (Diffusers Unified Format), --dduf should point to a .dduf archive file. -When packaging from an existing model using --from, you can modify properties like context size to create a variant of the original model. -For multimodal models, use --mmproj to include a multimodal projector file. +Package a model into a Docker Model OCI artifact. + +The model source must be one of: + --gguf A GGUF file (single file or first shard of a sharded model) + --safetensors-dir A directory containing .safetensors and configuration files + --dduf A .dduf (Diffusers Unified Format) archive + --from An existing packaged model reference + +By default, the packaged artifact is loaded into the local Model Runner content store. +Use --push to publish the model to a registry instead. + +MODEL specifies the target model reference (for example: myorg/llama3:8b). +When using --push, MODEL must be a registry-qualified reference. + +Packaging behavior: + + GGUF + --gguf must point to a .gguf file. + For sharded models, point to the first shard. All shards must: + • reside in the same directory + • follow an indexed naming convention (e.g. model-00001-of-00015.gguf) + All shards are automatically discovered and packaged together. + + Safetensors + --safetensors-dir must point to a directory containing .safetensors files + and required configuration files (e.g. model config, tokenizer files). + All files under the directory (including nested subdirectories) are + automatically discovered. Each file is packaged as a separate OCI layer. + + DDUF + --dduf must point to a .dduf archive file. + + Repackaging + --from repackages an existing model. You may override selected properties + such as --context-size to create a variant of the original model. + + Multimodal models + Use --mmproj to include a multimodal projector file. ### Options From 72102221479d160cb1f250f6f5d001806ba81cb2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20L=C3=B3pez=20Luna?= Date: Fri, 13 Feb 2026 10:00:48 +0100 Subject: [PATCH 5/7] feat(unpack): add logging for format handler and extraction errors --- pkg/distribution/builder/from_directory.go | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/pkg/distribution/builder/from_directory.go b/pkg/distribution/builder/from_directory.go index d499415c5..a52181681 100644 --- a/pkg/distribution/builder/from_directory.go +++ b/pkg/distribution/builder/from_directory.go @@ -2,6 +2,7 @@ package builder import ( "fmt" + "log" "os" "path/filepath" "strings" @@ -189,9 +190,13 @@ func FromDirectory(dirPath string, opts ...DirectoryOption) (*Builder, error) { } if detectedFormat != "" { f, fmtErr := format.Get(detectedFormat) - if fmtErr == nil { + if fmtErr != nil { + log.Printf("warning: could not get format handler for %q: %v", detectedFormat, fmtErr) + } else { extracted, extractErr := f.ExtractConfig(weightFiles) - if extractErr == nil { + if extractErr != nil { + log.Printf("warning: could not extract config from weight files: %v", extractErr) + } else { // Preserve the detected format, overlay extracted metadata config.Parameters = extracted.Parameters config.Quantization = extracted.Quantization From e2549e3b0cc7dee69daf5a348faa9f2be395e4fd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20L=C3=B3pez=20Luna?= Date: Fri, 13 Feb 2026 12:03:33 +0100 Subject: [PATCH 6/7] fix(parse): propagate filesystem errors in model directory walk --- pkg/distribution/huggingface/model.go | 1 - pkg/distribution/internal/bundle/parse.go | 13 +++++++++---- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/pkg/distribution/huggingface/model.go b/pkg/distribution/huggingface/model.go index b97f7695f..c714bfd67 100644 --- a/pkg/distribution/huggingface/model.go +++ b/pkg/distribution/huggingface/model.go @@ -133,7 +133,6 @@ func buildSafetensorsModelV02(tempDir string, createdTime *time.Time) (types.Mod } // buildGGUFModelV01 builds a GGUF model using V0.1 packaging (backward compatible). -// GGUF uses FromPaths + config archive approach. func buildGGUFModelV01(localPaths map[string]string, weightFiles, configFiles []RepoFile, createdTime *time.Time) (types.ModelArtifact, error) { // Collect weight file paths (sorted for reproducibility) var weightPaths []string diff --git a/pkg/distribution/internal/bundle/parse.go b/pkg/distribution/internal/bundle/parse.go index 3df746e78..239cb8056 100644 --- a/pkg/distribution/internal/bundle/parse.go +++ b/pkg/distribution/internal/bundle/parse.go @@ -114,17 +114,22 @@ func findSafetensorsFile(modelDir string) (string, error) { var firstFound string walkErr := filepath.Walk(modelDir, func(path string, info os.FileInfo, err error) error { if err != nil { - return nil // skip errors + // Propagate filesystem errors so callers can distinguish them from + // the case where no safetensors files are present. + return err } if info.IsDir() { return nil } if filepath.Ext(path) == ".safetensors" && !strings.HasPrefix(info.Name(), ".") { rel, relErr := filepath.Rel(modelDir, path) - if relErr == nil { - firstFound = rel - return filepath.SkipAll // found one, stop walking + if relErr != nil { + // Treat a bad relative path as a real error instead of silently + // ignoring it, so malformed bundles surface to the caller. + return relErr } + firstFound = rel + return filepath.SkipAll // found one, stop walking } return nil }) From d638a65dc16d38bc65ddc0f686c59ef1458aee56 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20L=C3=B3pez=20Luna?= Date: Fri, 13 Feb 2026 15:53:16 +0100 Subject: [PATCH 7/7] fix(parse): update safetensors file walk to use sentinel error for early termination --- pkg/distribution/internal/bundle/parse.go | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/pkg/distribution/internal/bundle/parse.go b/pkg/distribution/internal/bundle/parse.go index 239cb8056..cd928937f 100644 --- a/pkg/distribution/internal/bundle/parse.go +++ b/pkg/distribution/internal/bundle/parse.go @@ -2,6 +2,7 @@ package bundle import ( "encoding/json" + "errors" "fmt" "os" "path/filepath" @@ -11,6 +12,10 @@ import ( "github.com/docker/model-runner/pkg/distribution/types" ) +// errFoundSafetensors is a sentinel error used to stop filepath.Walk early +// after finding the first safetensors file. +var errFoundSafetensors = fmt.Errorf("found safetensors file") + // Parse returns the Bundle at the given rootDir func Parse(rootDir string) (*Bundle, error) { if fi, err := os.Stat(rootDir); err != nil || !fi.IsDir() { @@ -129,11 +134,11 @@ func findSafetensorsFile(modelDir string) (string, error) { return relErr } firstFound = rel - return filepath.SkipAll // found one, stop walking + return errFoundSafetensors // found one, stop walking } return nil }) - if walkErr != nil { + if walkErr != nil && !errors.Is(walkErr, errFoundSafetensors) { return "", fmt.Errorf("walk for safetensors files: %w", walkErr) }