Safetensors as OCI Artifact #186
Conversation
Reviewer's GuideIntroduce safetensors as a first-class OCI artifact by extending the CLI to package directory-based or file-based safetensors models with accompanying configs, adding builder and partial helpers for safetensors layers and config archives, augmenting distribution types and interfaces, updating unpack logic and bundle representation, and implementing a new safetensors package for model creation and metadata. Sequence diagram for packaging a safetensors model with config archivesequenceDiagram
actor User
participant CLI as "mdltool CLI"
participant Safetensors as "safetensors package"
participant Builder as "Builder"
participant Registry as "Registry"
User->>CLI: Run 'mdltool package ./model-dir --tag ...'
CLI->>Safetensors: packageFromDirectory(model-dir)
Safetensors->>CLI: Return safetensorsPaths, configArchive
CLI->>Builder: FromSafetensorsWithConfig(safetensorsPaths, configArchive)
Builder->>Registry: Build(ctx, target, os.Stdout)
Registry-->>User: Model artifact pushed
Class diagram for updated Model interfaces and BundleclassDiagram
class Model {
+ID() string
+GGUFPaths() []string
+SafetensorsPaths() []string
+ConfigArchivePath() string
+MMPROJPath() string
+Config() Config
+Tags() []string
+ChatTemplatePath() string
}
class ModelBundle {
+RootDir() string
+GGUFPath() string
+SafetensorsPath() string
+ConfigDir() string
+ChatTemplatePath() string
+MMPROJPath() string
+RuntimeConfig() Config
}
class Bundle {
-dir string
-mmprojPath string
-ggufFile string
-safetensorsFile string
-configDir string
-runtimeConfig Config
-chatTemplatePath string
+SafetensorsPath() string
+ConfigDir() string
+RuntimeConfig() Config
}
Model <|.. ModelBundle
ModelBundle <|.. Bundle
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
There was a problem hiding this comment.
Pull Request Overview
This PR introduces support for Safetensors model format as OCI artifacts, expanding the system beyond GGUF models. The implementation adds a parallel set of interfaces and handlers for Safetensors models while maintaining backward compatibility with existing GGUF functionality.
Key Changes:
- Add Safetensors model format support with new media types and configuration
- Implement auto-discovery of sharded Safetensors files and configuration archives
- Extend CLI tooling to detect and package both GGUF and Safetensors models from directories
Reviewed Changes
Copilot reviewed 13 out of 13 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| pkg/distribution/types/model.go | Add interface methods for Safetensors paths and config archive handling |
| pkg/distribution/types/config.go | Define new media types and format constants for Safetensors |
| pkg/distribution/internal/safetensors/model.go | Core Safetensors model implementation with OCI artifact interface |
| pkg/distribution/internal/safetensors/create.go | Model creation logic with shard discovery and config extraction |
| pkg/distribution/internal/bundle/unpack.go | Enhanced unpacking logic to handle both GGUF and Safetensors formats |
| cmd/mdltool/main.go | CLI enhancements for directory-based packaging and format detection |
| pkg/distribution/builder/builder.go | Builder pattern extension for Safetensors models with config archives |
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 13 out of 13 changed files in this pull request and generated 3 comments.
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 13 out of 13 changed files in this pull request and generated 3 comments.
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 13 out of 13 changed files in this pull request and generated 8 comments.
Comments suppressed due to low confidence (1)
pkg/distribution/types/model.go:1
- Adding required methods (SafetensorsPaths, ConfigArchivePath, SafetensorsPath, ConfigDir) to existing interfaces is a breaking change for all external implementations. Consider introducing a new extended interface (e.g., ModelWithSafetensors) or providing adapter shims to preserve backward compatibility.
package types
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 13 out of 13 changed files in this pull request and generated 3 comments.
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 12 out of 12 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (2)
pkg/distribution/internal/bundle/unpack.go:1
- Files from the config archive can overwrite previously unpacked model assets if the archive includes conflicting filenames (e.g., model.safetensors, model.gguf). Consider rejecting or renaming entries when a destination file already exists (using os.Lstat + fail) to prevent malicious or accidental overwrite of critical bundle contents.
package bundle
pkg/distribution/internal/bundle/unpack.go:1
- Files from the config archive can overwrite previously unpacked model assets if the archive includes conflicting filenames (e.g., model.safetensors, model.gguf). Consider rejecting or renaming entries when a destination file already exists (using os.Lstat + fail) to prevent malicious or accidental overwrite of critical bundle contents.
package bundle
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 12 out of 12 changed files in this pull request and generated 3 comments.
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
| rel, err := filepath.Rel(absBaseDir, absTarget) | ||
| if err != nil { | ||
| return fmt.Errorf("compute relative path: %w", err) | ||
| } | ||
|
|
||
| // Use filepath.IsLocal() to check if the relative path is local (doesn't escape baseDir) | ||
| // This handles all edge cases including empty strings, ".", "..", symlinks, etc. | ||
| if !filepath.IsLocal(rel) { | ||
| return fmt.Errorf("invalid entry %q: path attempts to escape destination directory", targetPath) | ||
| } |
There was a problem hiding this comment.
The directory traversal check is ineffective: filepath.IsLocal(rel) returns true for paths beginning with "../", so entries like ../../etc/passwd would not be rejected. Replace the check with a robust containment test (e.g., ensure strings.HasPrefix(absTarget+string(os.PathSeparator), absBaseDir+string(os.PathSeparator)) or verify that rel does not start with ".."+path separator or equal "..") to properly prevent escaping the destination directory.
| // Add config archive layer | ||
| if configArchivePath != "" { | ||
| // Check if a config archive layer already exists | ||
| for _, layer := range model.layers { | ||
| mediaType, err := layer.MediaType() | ||
| if err == nil && mediaType == types.MediaTypeVLLMConfigArchive { | ||
| return nil, fmt.Errorf("model already has a config archive layer") | ||
| } | ||
| } |
There was a problem hiding this comment.
[nitpick] Duplicate logic for detecting an existing config archive layer also appears in builder.WithConfigArchive; consider extracting a shared helper (e.g., hasConfigArchiveLayer(layers)) to reduce repetition and keep validation consistent.
| // Check if config archive already exists | ||
| layers, err := b.model.Layers() | ||
| if err != nil { | ||
| return nil, fmt.Errorf("get model layers: %w", err) | ||
| } | ||
|
|
||
| for _, layer := range layers { | ||
| mediaType, err := layer.MediaType() | ||
| if err == nil && mediaType == types.MediaTypeVLLMConfigArchive { | ||
| return nil, fmt.Errorf("model already has a config archive layer") | ||
| } | ||
| } |
There was a problem hiding this comment.
[nitpick] Logic duplicating the config archive presence check in safetensors.NewModelWithConfigArchive; factor into a shared utility to ensure consistent validation and simplify future changes.
ilopezluna
changed the title
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
We encountered an error and are unable to review this PR. We have been notified and are working to fix it.
You can try again by commenting this pull request with @sourcery-ai review, or contact us for help.
There was a problem hiding this comment.
We encountered an error and are unable to review this PR. We have been notified and are working to fix it.
You can try again by commenting this pull request with @sourcery-ai review, or contact us for help.
There was a problem hiding this comment.
We encountered an error and are unable to review this PR. We have been notified and are working to fix it.
You can try again by commenting this pull request with @sourcery-ai review, or contact us for help.
Adds support for packaging Safetensors models as OCI Artifacts. It supports multiple Safetensors files (by naming convention).
Additional *.json files and
merge.txtare included in a separate tar-packaged layer.When the bundle is created, this layer is unpacked to restore the exact same folder structure.
Summary by Sourcery
Add comprehensive support for packaging and unpacking safetensors model artifacts alongside existing GGUF format.
New Features:
Enhancements:
Tests:
Chores: