feat: add safe tensor model type to SBOM output (#4844)

---------
Signed-off-by: Christopher Phillips <32073428+spiffcs@users.noreply.github.com>
Signed-off-by: Alex Goodman <wagoodman@users.noreply.github.com>
Co-authored-by: Alex Goodman <wagoodman@users.noreply.github.com>
This commit is contained in:
Christopher Angelo Phillips 2026-07-15 17:14:03 -04:00 committed by GitHub
parent ed499fd2d0
commit aff2c247db
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
37 changed files with 7834 additions and 94 deletions

View File

@ -3,7 +3,7 @@ package internal
const ( const (
// JSONSchemaVersion is the current schema version output by the JSON encoder // JSONSchemaVersion is the current schema version output by the JSON encoder
// This is roughly following the "SchemaVer" guidelines for versioning the JSON schema. Please see schema/json/README.md for details on how to increment. // This is roughly following the "SchemaVer" guidelines for versioning the JSON schema. Please see schema/json/README.md for details on how to increment.
JSONSchemaVersion = "16.1.9" JSONSchemaVersion = "16.1.10"
// Changelog // Changelog
// 16.1.0 - reformulated the python pdm fields (added "URL" and removed the unused "path" field). // 16.1.0 - reformulated the python pdm fields (added "URL" and removed the unused "path" field).
@ -16,4 +16,5 @@ const (
// 16.1.7 - add AppleAppBundleEntry metadata type for the apple app bundle cataloger // 16.1.7 - add AppleAppBundleEntry metadata type for the apple app bundle cataloger
// 16.1.8 - add VcpkgManifest metadata type for vcpkg manifest support // 16.1.8 - add VcpkgManifest metadata type for vcpkg manifest support
// 16.1.9 - add Symbols (grouped by owning package import path) to GolangBinaryBuildinfoEntry metadata // 16.1.9 - add Symbols (grouped by owning package import path) to GolangBinaryBuildinfoEntry metadata
// 16.1.10 - add SafeTensorsModelInfo metadata type for the safetensors AI model cataloger
) )

View File

@ -27,6 +27,7 @@ var knownNonMetadataTypeNames = strset.New(
var knownMetadataTypeNames = strset.New( var knownMetadataTypeNames = strset.New(
"DotnetPortableExecutableEntry", "DotnetPortableExecutableEntry",
"GGUFFileHeader", "GGUFFileHeader",
"SafeTensorsModelInfo",
) )
func DiscoverTypeNames() ([]string, error) { func DiscoverTypeNames() ([]string, error) {

View File

@ -68,6 +68,7 @@ func AllTypes() []any {
pkg.RubyGemspec{}, pkg.RubyGemspec{},
pkg.RustBinaryAuditEntry{}, pkg.RustBinaryAuditEntry{},
pkg.RustCargoLockEntry{}, pkg.RustCargoLockEntry{},
pkg.SafeTensorsModelInfo{},
pkg.SnapEntry{}, pkg.SnapEntry{},
pkg.SwiftPackageManagerResolvedEntry{}, pkg.SwiftPackageManagerResolvedEntry{},
pkg.SwiplPackEntry{}, pkg.SwiplPackEntry{},

View File

@ -130,6 +130,7 @@ var jsonTypes = makeJSONTypes(
jsonNames(pkg.DotnetPackagesLockEntry{}, "dotnet-packages-lock-entry"), jsonNames(pkg.DotnetPackagesLockEntry{}, "dotnet-packages-lock-entry"),
jsonNames(pkg.CondaMetaPackage{}, "conda-metadata-entry", "CondaPackageMetadata"), jsonNames(pkg.CondaMetaPackage{}, "conda-metadata-entry", "CondaPackageMetadata"),
jsonNames(pkg.GGUFFileHeader{}, "gguf-file-header"), jsonNames(pkg.GGUFFileHeader{}, "gguf-file-header"),
jsonNames(pkg.SafeTensorsModelInfo{}, "safetensors-model-info"),
jsonNames(pkg.VcpkgManifest{}, "vcpkg-manifest"), jsonNames(pkg.VcpkgManifest{}, "vcpkg-manifest"),
) )

View File

@ -191,6 +191,7 @@ func DefaultPackageTaskFactories() Factories {
newSimplePackageTaskFactory(conda.NewCondaMetaCataloger, pkgcataloging.DirectoryTag, pkgcataloging.InstalledTag, pkgcataloging.PackageTag, "conda"), newSimplePackageTaskFactory(conda.NewCondaMetaCataloger, pkgcataloging.DirectoryTag, pkgcataloging.InstalledTag, pkgcataloging.PackageTag, "conda"),
newSimplePackageTaskFactory(snap.NewCataloger, pkgcataloging.DirectoryTag, pkgcataloging.InstalledTag, pkgcataloging.ImageTag, "snap"), newSimplePackageTaskFactory(snap.NewCataloger, pkgcataloging.DirectoryTag, pkgcataloging.InstalledTag, pkgcataloging.ImageTag, "snap"),
newSimplePackageTaskFactory(ai.NewGGUFCataloger, pkgcataloging.DirectoryTag, pkgcataloging.ImageTag, "ai", "model", "gguf", "ml"), newSimplePackageTaskFactory(ai.NewGGUFCataloger, pkgcataloging.DirectoryTag, pkgcataloging.ImageTag, "ai", "model", "gguf", "ml"),
newSimplePackageTaskFactory(ai.NewSafeTensorsCataloger, pkgcataloging.DirectoryTag, pkgcataloging.ImageTag, "ai", "model", "safetensors", "ml"),
// deprecated catalogers //////////////////////////////////////// // deprecated catalogers ////////////////////////////////////////
// these are catalogers that should not be selectable other than specific inclusion via name or "deprecated" tag (to remain backwards compatible) // these are catalogers that should not be selectable other than specific inclusion via name or "deprecated" tag (to remain backwards compatible)

File diff suppressed because it is too large Load Diff

View File

@ -1,6 +1,6 @@
{ {
"$schema": "https://json-schema.org/draft/2020-12/schema", "$schema": "https://json-schema.org/draft/2020-12/schema",
"$id": "anchore.io/schema/syft/json/16.1.9/document", "$id": "anchore.io/schema/syft/json/16.1.10/document",
"$ref": "#/$defs/Document", "$ref": "#/$defs/Document",
"$defs": { "$defs": {
"AlpmDbEntry": { "AlpmDbEntry": {
@ -2925,6 +2925,9 @@
{ {
"$ref": "#/$defs/RustCargoLockEntry" "$ref": "#/$defs/RustCargoLockEntry"
}, },
{
"$ref": "#/$defs/SafetensorsModelInfo"
},
{ {
"$ref": "#/$defs/SnapEntry" "$ref": "#/$defs/SnapEntry"
}, },
@ -4216,6 +4219,55 @@
], ],
"description": "RustCargoLockEntry represents a locked dependency from a Cargo.lock file with precise version and checksum information." "description": "RustCargoLockEntry represents a locked dependency from a Cargo.lock file with precise version and checksum information."
}, },
"SafetensorsModelInfo": {
"properties": {
"format": {
"type": "string",
"description": "Format is the source format label (always \"safetensors\" for this metadata type).\nPresent because the Docker AI model config blob carries an explicit format field"
},
"architecture": {
"type": "string",
"description": "Architecture is the model architecture (e.g., \"LlamaForCausalLM\",\n\"Qwen3MoeForConditionalGeneration\"). It is not present in the SafeTensors\nheader itself; it is enriched from the companion config.json\n\"architectures\" array when one is found alongside the model."
},
"quantization": {
"type": "string",
"description": "Quantization describes tensor precision (e.g., \"BF16\", \"F16\", \"F32\", \"INT8\")."
},
"parameters": {
"type": "integer",
"description": "Parameters is the total number of model parameters, computed from the tensor\nshapes in the SafeTensors header(s). For a sharded model it is the sum across\nevery shard."
},
"tensorCount": {
"type": "integer",
"description": "TensorCount is the number of tensor entries in the file header."
},
"totalSize": {
"type": "string",
"description": "TotalSize is the total byte size of tensor data across all shards when known\n(from the Docker AI model config \"size\" field)."
},
"shardCount": {
"type": "integer",
"description": "ShardCount is the number of .safetensors shards for a sharded model (1 for a\nsingle-file model)."
},
"userMetadata": {
"$ref": "#/$defs/KeyValues",
"description": "UserMetadata is the optional \"__metadata__\" map from a .safetensors file header\n(string-to-string key/values set by the producer)."
},
"metadataHash": {
"type": "string",
"description": "MetadataHash is an xxhash over the on-disk SafeTensors header (sorted tensor\nentries + __metadata__). It is derived ONLY from the safetensors file bytes."
},
"parts": {
"items": {
"$ref": "#/$defs/SafetensorsModelInfo"
},
"type": "array",
"description": "Parts contains metadata from additional SafeTensors shards or OCI layers that\nwere merged into this package during post-processing."
}
},
"type": "object",
"description": "SafeTensorsModelInfo holds the model details extracted from SafeTensors content."
},
"Schema": { "Schema": {
"properties": { "properties": {
"version": { "version": {

View File

@ -25,6 +25,7 @@ type MockResolver struct {
mediaTypeIndex map[string][]Location mediaTypeIndex map[string][]Location
extension map[string][]Location extension map[string][]Location
basename map[string][]Location basename map[string][]Location
ociRef string
} }
// NewMockResolverForPaths creates a new MockResolver, where the only resolvable // NewMockResolverForPaths creates a new MockResolver, where the only resolvable
@ -74,6 +75,16 @@ func NewMockResolverForPathsWithMetadata(metadata map[Coordinates]Metadata) *Moc
} }
} }
// NewMockResolverForOCIArtifact creates a MockResolver that can resolve files
// by media type AND surfaces the given OCI ref via the ImageReference method.
// Intended for tests that exercise the catalogers' OCI-artifact-aware naming
// code paths.
func NewMockResolverForOCIArtifact(ref string, mediaTypes map[string][]Location) *MockResolver {
r := NewMockResolverForMediaTypes(mediaTypes)
r.ociRef = ref
return r
}
// NewMockResolverForMediaTypes creates a MockResolver that can resolve files by media type. // NewMockResolverForMediaTypes creates a MockResolver that can resolve files by media type.
// The mediaTypes map specifies which locations should be returned for each media type. // The mediaTypes map specifies which locations should be returned for each media type.
func NewMockResolverForMediaTypes(mediaTypes map[string][]Location) *MockResolver { func NewMockResolverForMediaTypes(mediaTypes map[string][]Location) *MockResolver {
@ -102,6 +113,11 @@ func NewMockResolverForMediaTypes(mediaTypes map[string][]Location) *MockResolve
} }
} }
// ImageReference returns the image reference associated with this mock, if any.
func (r MockResolver) ImageReference() string {
return r.ociRef
}
// HasPath indicates if the given path exists in the underlying source. // HasPath indicates if the given path exists in the underlying source.
func (r MockResolver) HasPath(path string) bool { func (r MockResolver) HasPath(path string) bool {
for _, l := range r.locations { for _, l := range r.locations {

View File

@ -58,6 +58,7 @@ func Test_OriginatorSupplier(t *testing.T) {
pkg.YarnLockEntry{}, pkg.YarnLockEntry{},
pkg.TerraformLockProviderEntry{}, pkg.TerraformLockProviderEntry{},
pkg.GGUFFileHeader{}, pkg.GGUFFileHeader{},
pkg.SafeTensorsModelInfo{},
pkg.DenoLockEntry{}, pkg.DenoLockEntry{},
pkg.DenoRemoteLockEntry{}, pkg.DenoRemoteLockEntry{},
pkg.VcpkgManifest{}, pkg.VcpkgManifest{},

View File

@ -85,7 +85,7 @@ func SourceInfo(p pkg.Package) string {
case pkg.TerraformPkg: case pkg.TerraformPkg:
answer = "acquired package info from Terraform dependency lock file" answer = "acquired package info from Terraform dependency lock file"
case pkg.ModelPkg: case pkg.ModelPkg:
answer = "acquired package info from AI artifact (e.g. GGUF File)" answer = "acquired package info from AI model artifact"
case pkg.VcpkgPkg: case pkg.VcpkgPkg:
answer = "acquired package info from vcpkg manifest file" answer = "acquired package info from vcpkg manifest file"
default: default:

View File

@ -20,16 +20,22 @@ type LayerInfo struct {
} }
// ContainerImageModel is a file.Resolver implementation that provides access to // ContainerImageModel is a file.Resolver implementation that provides access to
// GGUF header data fetched from OCI model artifacts via range-GET requests. // model header and metadata data (GGUF and SafeTensors headers, the model config
// This does not fetch the entire model from the registry, only a sliver of it. // blob, and companion layers) fetched from OCI model artifacts via range-GET
// requests. This does not fetch the entire model from the registry, only a
// sliver of it.
type ContainerImageModel struct { type ContainerImageModel struct {
tempDir string // temp directory containing all layer files tempDir string // temp directory containing all layer files
layerFiles map[string]LayerInfo // digest -> layer info (temp path + media type) layerFiles map[string]LayerInfo // digest -> layer info (temp path + media type)
locations map[string]file.Location // digest -> location locations map[string]file.Location // digest -> location
ref string // image reference the artifact was fetched with
} }
// NewContainerImageModel creates a new resolver with the given temp directory and layer files. // NewContainerImageModel creates a new resolver with the given temp directory
func NewContainerImageModel(tempDir string, layerFiles map[string]LayerInfo) *ContainerImageModel { // and layer files. The ref is surfaced through the ImageReference method so
// catalogers can derive context-level naming hints from the artifact reference
// when the layer contents don't carry one.
func NewContainerImageModel(tempDir string, layerFiles map[string]LayerInfo, ref string) *ContainerImageModel {
// Create locations for all layer files // Create locations for all layer files
// Each location has RealPath="/", FileSystemID=digest, AccessPath="/" // Each location has RealPath="/", FileSystemID=digest, AccessPath="/"
locations := make(map[string]file.Location, len(layerFiles)) locations := make(map[string]file.Location, len(layerFiles))
@ -43,9 +49,15 @@ func NewContainerImageModel(tempDir string, layerFiles map[string]LayerInfo) *Co
tempDir: tempDir, tempDir: tempDir,
layerFiles: layerFiles, layerFiles: layerFiles,
locations: locations, locations: locations,
ref: ref,
} }
} }
// ImageReference returns the image reference the artifact was fetched with.
func (r *ContainerImageModel) ImageReference() string {
return r.ref
}
// FilesByMediaType returns locations for layers matching the given media type patterns. // FilesByMediaType returns locations for layers matching the given media type patterns.
// Patterns support glob-style matching (e.g., "application/vnd.docker.ai*"). // Patterns support glob-style matching (e.g., "application/vnd.docker.ai*").
func (r *ContainerImageModel) FilesByMediaType(types ...string) ([]file.Location, error) { func (r *ContainerImageModel) FilesByMediaType(types ...string) ([]file.Location, error) {

View File

@ -60,7 +60,8 @@ func TestOCIModelResolver_FilesByMediaType(t *testing.T) {
for _, test := range tests { for _, test := range tests {
t.Run(test.name, func(t *testing.T) { t.Run(test.name, func(t *testing.T) {
resolver := NewContainerImageModel(tempDir, test.layerFiles) ref := ""
resolver := NewContainerImageModel(tempDir, test.layerFiles, ref)
locations, err := resolver.FilesByMediaType(test.patterns...) locations, err := resolver.FilesByMediaType(test.patterns...)
require.NoError(t, err) require.NoError(t, err)
@ -81,7 +82,8 @@ func TestOCIModelResolver_FileContentsByLocation(t *testing.T) {
digest: {TempPath: tempFile, MediaType: ggufLayerMediaType}, digest: {TempPath: tempFile, MediaType: ggufLayerMediaType},
} }
resolver := NewContainerImageModel(tempDir, layerFiles) ref := ""
resolver := NewContainerImageModel(tempDir, layerFiles, ref)
tests := []struct { tests := []struct {
name string name string

View File

@ -44,3 +44,46 @@ catalogers:
default: true default: true
evidence: evidence:
- GGUFFileHeader.MetadataKeyValuesHash - GGUFFileHeader.MetadataKeyValuesHash
- ecosystem: ai # MANUAL
name: safetensors-cataloger # AUTO-GENERATED
type: generic # AUTO-GENERATED
source: # AUTO-GENERATED
file: syft/pkg/cataloger/ai/cataloger.go
function: NewSafeTensorsCataloger
selectors: # AUTO-GENERATED
- ai
- directory
- image
- ml
- model
- package
- safetensors
parsers: # AUTO-GENERATED structure
- function: parseSafeTensorsFile
detector: # AUTO-GENERATED
method: glob # AUTO-GENERATED
criteria: # AUTO-GENERATED
- '**/*.safetensors'
metadata_types: # AUTO-GENERATED
- pkg.SafeTensorsModelInfo
package_types: # AUTO-GENERATED
- model
json_schema_types: # AUTO-GENERATED
- SafetensorsModelInfo
capabilities: # MANUAL - preserved across regeneration
- name: license
default: true
- name: dependency.depth
default: []
- name: dependency.edges
default: ""
- name: dependency.kinds
default: []
- name: package_manager.files.listing
default: false
- name: package_manager.files.digests
default: false
- name: package_manager.package_integrity_hash
default: true
evidence:
- SafeTensorsModelInfo.MetadataHash

View File

@ -1,6 +1,6 @@
/* /*
Package ai provides concrete Cataloger implementations for AI artifacts and machine learning models, Package ai provides concrete Cataloger implementations for AI artifacts and machine learning models,
including support for GGUF (GPT-Generated Unified Format) model files. including support for GGUF (GPT-Generated Unified Format) and SafeTensors model files.
*/ */
package ai package ai
@ -12,6 +12,7 @@ import (
const ( const (
catalogerName = "gguf-cataloger" catalogerName = "gguf-cataloger"
ggufLayerMediaType = "application/vnd.docker.ai*" ggufLayerMediaType = "application/vnd.docker.ai*"
safeTensorsCatalogerName = "safetensors-cataloger"
) )
// NewGGUFCataloger returns a new cataloger instance for GGUF model files. // NewGGUFCataloger returns a new cataloger instance for GGUF model files.
@ -23,3 +24,26 @@ func NewGGUFCataloger() pkg.Cataloger {
WithParserByMediaType(parseGGUFModel, ggufLayerMediaType). WithParserByMediaType(parseGGUFModel, ggufLayerMediaType).
WithProcessors(ggufMergeProcessor) WithProcessors(ggufMergeProcessor)
} }
// NewSafeTensorsCataloger returns a cataloger for SafeTensors model files,
// covering three discovery paths:
// - **/*.safetensors files (single-file models and individual shards;
// header-only parse)
// - application/vnd.docker.ai.model.config.v0.1+json / v0.2+json OCI layers
// (Docker Model Runner artifacts whose config advertises format=="safetensors")
// - application/vnd.docker.ai.safetensors OCI layers (per-shard JSON headers,
// fetched as a prefix by the OCI model source; emitted as nameless
// packages and merged into the config-derived package as Parts)
//
// model.safetensors.index.json files are intentionally not parsed today: the
// index describes how tensors map to shards but contributes no metadata the
// cataloger can't derive from the shard headers themselves. If a model is
// distributed as just an index.json with no accompanying shard files, the
// cataloger emits nothing for that directory.
func NewSafeTensorsCataloger() pkg.Cataloger {
return generic.NewCataloger(safeTensorsCatalogerName).
WithParserByGlobs(parseSafeTensorsFile, "**/*.safetensors").
WithParserByMediaType(parseSafeTensorsOCIConfig, dockerAIModelConfigMediaTypes...).
WithParserByMediaType(parseSafeTensorsOCILayer, dockerAISafeTensorsMediaType).
WithResolvingProcessors(safeTensorsMergeProcessor)
}

View File

@ -0,0 +1,56 @@
package ai
import (
"github.com/anchore/syft/syft/artifact"
"github.com/anchore/syft/syft/pkg"
)
// ggufMergeProcessor consolidates multiple GGUF packages into a single package
// representing the AI model. When scanning OCI images with multiple layers,
// each layer may produce a separate package. This processor finds the package
// with a name and merges metadata from nameless packages into its GGUFFileParts field.
// Only packages with a non-empty name are returned in the final result.
func ggufMergeProcessor(pkgs []pkg.Package, rels []artifact.Relationship, err error) ([]pkg.Package, []artifact.Relationship, error) {
if err != nil {
return pkgs, rels, err
}
if len(pkgs) == 0 {
return pkgs, rels, err
}
// Separate packages with names from those without
var namedPkgs []pkg.Package
var namelessHeaders []pkg.GGUFFileHeader
for _, p := range pkgs {
if p.Name != "" {
namedPkgs = append(namedPkgs, p)
} else {
if header, ok := p.Metadata.(pkg.GGUFFileHeader); ok {
// We do not want a kv hash for nameless headers
header.MetadataKeyValuesHash = ""
namelessHeaders = append(namelessHeaders, header)
}
}
}
// If there are no named packages, return nothing
if len(namedPkgs) == 0 {
return nil, rels, err
}
// merge nameless headers into a single named package;
// if there are multiple named packages, return them without trying to merge headers.
// we cannot determine which nameless headers belong to which package
// this is because the order we receive the gguf headers in is not guaranteed
// to match the layer order in the original oci image
if len(namedPkgs) == 1 && len(namelessHeaders) > 0 {
winner := &namedPkgs[0]
if header, ok := winner.Metadata.(pkg.GGUFFileHeader); ok {
header.Parts = namelessHeaders
winner.Metadata = header
}
}
return namedPkgs, rels, err
}

View File

@ -0,0 +1,93 @@
package ai
import (
"bytes"
"gopkg.in/yaml.v3"
"github.com/anchore/syft/internal/log"
"github.com/anchore/syft/syft/pkg"
)
// hfConfig is a minimal projection of Hugging Face config.json fields.
type hfConfig struct {
Architectures []string `json:"architectures"`
NameOrPath string `json:"_name_or_path"`
}
func applyHFConfig(md *pkg.SafeTensorsModelInfo, cfg *hfConfig) {
if md.Architecture == "" && len(cfg.Architectures) > 0 {
md.Architecture = cfg.Architectures[0]
}
}
// readmeFrontmatter holds the subset of YAML frontmatter fields we extract.
type readmeFrontmatter struct {
License string `yaml:"license"`
BaseModel []string `yaml:"base_model"`
}
type licenseFrontmatter struct {
SPDXID string `yaml:"spdx-id"`
}
// extractFrontmatterBlock returns the YAML bytes between the first and second
// "---" delimiters of a file
func extractFrontmatterBlock(buf []byte) []byte {
trimmed := bytes.TrimLeft(buf, "\xef\xbb\xbf \t\r\n")
if !bytes.HasPrefix(trimmed, []byte("---")) {
return nil
}
rest := trimmed[3:]
if i := bytes.IndexByte(rest, '\n'); i >= 0 {
rest = rest[i+1:]
}
block, _, found := bytes.Cut(rest, []byte("\n---"))
if !found {
return nil
}
return block
}
// parseFrontmatter decodes a Hugging Face model card YAML frontmatter block
// and returns the license and base_model fields.
func parseFrontmatter(buf []byte) *readmeFrontmatter {
block := extractFrontmatterBlock(buf)
if block == nil {
return nil
}
var raw struct {
License string `yaml:"license"`
BaseModel yaml.Node `yaml:"base_model"`
}
if err := yaml.Unmarshal(block, &raw); err != nil {
log.Debugf("failed to parse README frontmatter: %v", err)
return nil
}
fm := readmeFrontmatter{License: raw.License}
switch raw.BaseModel.Kind {
case yaml.ScalarNode:
if raw.BaseModel.Value != "" {
fm.BaseModel = []string{raw.BaseModel.Value}
}
case yaml.SequenceNode:
_ = raw.BaseModel.Decode(&fm.BaseModel)
}
return &fm
}
// parseLicenseFrontmatter returns the producer-declared SPDX identifier
func parseLicenseFrontmatter(buf []byte) string {
block := extractFrontmatterBlock(buf)
if block == nil {
return ""
}
var fm licenseFrontmatter
if err := yaml.Unmarshal(block, &fm); err != nil {
log.Debugf("failed to parse license frontmatter: %v", err)
return ""
}
return fm.SPDXID
}

View File

@ -0,0 +1,93 @@
package ai
import (
"context"
"encoding/json"
"io"
"path"
"github.com/anchore/syft/internal"
"github.com/anchore/syft/internal/log"
"github.com/anchore/syft/syft/file"
"github.com/anchore/syft/syft/pkg"
)
// resolveSafeTensorsDirIdentity handles the directory-scan case for safe tensors
// find config.json beside the model files (walking up parent directories to the
// scanned source root if no sibling exists) and a sibling README.md. It returns
// the group's name candidates, resolved licenses, and supporting evidence.
func resolveSafeTensorsDirIdentity(ctx context.Context, resolver file.Resolver, dir string, md *pkg.SafeTensorsModelInfo) safeTensorsIdentity {
id := safeTensorsIdentity{fallbackName: safeTensorsDirName(dir)}
if loc, cfg := findDirHFConfig(resolver, dir); cfg != nil {
applyHFConfig(md, cfg)
id.nameOrPath = cfg.NameOrPath
id.supporting = append(id.supporting, *loc)
}
if loc, fm := readDirReadmeFrontmatter(resolver, path.Join(dir, "README.md")); fm != nil {
if fm.License != "" {
id.licenses = pkg.NewLicensesFromValuesWithContext(ctx, fm.License)
}
if id.nameOrPath == "" && len(fm.BaseModel) > 0 {
id.nameOrPath = fm.BaseModel[0]
}
id.supporting = append(id.supporting, *loc)
}
return id
}
// findDirHFConfig looks for a config.json beside the model files
func findDirHFConfig(resolver file.Resolver, dir string) (*file.Location, *hfConfig) {
for {
if loc, cfg := readDirHFConfig(resolver, path.Join(dir, "config.json")); cfg != nil {
return loc, cfg
}
parent := path.Dir(dir)
if parent == dir {
return nil, nil // reached the source root
}
dir = parent
}
}
func readDirHFConfig(resolver file.Resolver, p string) (*file.Location, *hfConfig) {
locations, err := resolver.FilesByPath(p)
if err != nil || len(locations) == 0 {
return nil, nil
}
rc, err := resolver.FileContentsByLocation(locations[0])
if err != nil {
return nil, nil
}
defer internal.CloseAndLogError(rc, p)
var cfg hfConfig
if err := json.NewDecoder(rc).Decode(&cfg); err != nil {
log.Debugf("failed to decode %s: %v", p, err)
return nil, nil
}
return &locations[0], &cfg
}
func readDirReadmeFrontmatter(resolver file.Resolver, p string) (*file.Location, *readmeFrontmatter) {
locations, err := resolver.FilesByPath(p)
if err != nil || len(locations) == 0 {
return nil, nil
}
rc, err := resolver.FileContentsByLocation(locations[0])
if err != nil {
return nil, nil
}
defer internal.CloseAndLogError(rc, p)
buf, err := io.ReadAll(io.LimitReader(rc, 1024*1024))
if err != nil {
return nil, nil
}
fm := parseFrontmatter(buf)
if fm == nil {
return nil, nil
}
return &locations[0], fm
}

View File

@ -0,0 +1,173 @@
package ai
import (
"bytes"
"context"
"encoding/json"
"io"
"path"
gcrname "github.com/google/go-containerregistry/pkg/name"
"github.com/anchore/syft/internal"
"github.com/anchore/syft/internal/log"
"github.com/anchore/syft/syft/file"
"github.com/anchore/syft/syft/pkg"
"github.com/anchore/syft/syft/pkg/cataloger/internal/licenses"
)
// resolveSafeTensorsOCIIdentity handles the OCI-artifact case
func resolveSafeTensorsOCIIdentity(ctx context.Context, resolver file.Resolver, md *pkg.SafeTensorsModelInfo) safeTensorsIdentity {
ociResolver, ok := resolver.(file.OCIMediaTypeResolver)
if !ok {
return safeTensorsIdentity{}
}
modelFileLocs, err := ociResolver.FilesByMediaType(dockerAIModelFileMediaType)
if err != nil {
log.Debugf("failed to list docker AI model-file layers: %v", err)
}
// Collect config / readme candidates separately so the layer-iteration order
// returned by the resolver doesn't decide the precedence.
var configName, readmeName, readmeLicense string
var supporting []file.Location
for _, loc := range modelFileLocs {
cfg, fm := classifyOCIModelFileLayer(resolver, loc)
switch {
case cfg != nil:
applyHFConfig(md, cfg)
if configName == "" {
configName = cfg.NameOrPath
}
supporting = append(supporting, loc)
case fm != nil:
if readmeLicense == "" {
readmeLicense = fm.License
}
if readmeName == "" && len(fm.BaseModel) > 0 {
readmeName = fm.BaseModel[0]
}
supporting = append(supporting, loc)
}
}
// Precedence: config.json _name_or_path > README base_model.
nameOrPath := configName
if nameOrPath == "" {
nameOrPath = readmeName
}
id := safeTensorsIdentity{
nameOrPath: nameOrPath,
fallbackName: ociImageRefBasename(resolver),
supporting: supporting,
}
// License precedence: a dedicated vnd.docker.ai.license layer is a
// outranks the free-text license field in a model card's README frontmatter.
licLocs, err := ociResolver.FilesByMediaType(dockerAILicenseMediaType)
if err != nil {
log.Debugf("failed to list docker AI license layers: %v", err)
}
switch {
case len(licLocs) > 0:
id.licenses = identifyLicenseLayers(ctx, resolver, licLocs)
id.supporting = append(id.supporting, licLocs...)
case readmeLicense != "":
id.licenses = pkg.NewLicensesFromValuesWithContext(ctx, readmeLicense)
}
return id
}
// ociImageReferencer is the minimal capability ociImageRefBasename needs: a
// resolver that can surface the OCI image reference it was built from. It is
// kept local to this package (rather than exported from the file package) so the
// assertion stays with its only consumer.
type ociImageReferencer interface {
ImageReference() string
}
func ociImageRefBasename(resolver file.Resolver) string {
// TODO: we don't think this approach is generalizable quite yet, but we really do need this information.
// (Ideally we should be NOT be type asserting on the file resolver directly).
info, ok := resolver.(ociImageReferencer)
if !ok {
return ""
}
ref := info.ImageReference()
if ref == "" {
return ""
}
parsed, err := gcrname.ParseReference(ref)
if err != nil {
log.Debugf("failed to parse OCI ref %q: %v", ref, err)
return ""
}
return path.Base(parsed.Context().RepositoryStr())
}
// classifyOCIModelFileLayer reads up to 4 MiB of a model.file layer and decodes
// it as either an HF config.json or a README model card's YAML frontmatter,
// based on its leading bytes. It returns whichever it recognized; both are nil
// when the layer is neither (or fails to decode). The caller owns precedence and
// metadata enrichment.
func classifyOCIModelFileLayer(resolver file.Resolver, loc file.Location) (*hfConfig, *readmeFrontmatter) {
rc, err := resolver.FileContentsByLocation(loc)
if err != nil {
return nil, nil
}
defer internal.CloseAndLogError(rc, loc.RealPath)
buf, err := io.ReadAll(io.LimitReader(rc, 4*1024*1024))
if err != nil {
return nil, nil
}
trimmed := bytes.TrimLeft(buf, "\xef\xbb\xbf \t\r\n")
switch {
case bytes.HasPrefix(trimmed, []byte("---")):
return nil, parseFrontmatter(buf)
case bytes.HasPrefix(trimmed, []byte("{")):
var cfg hfConfig
if err := json.Unmarshal(buf, &cfg); err != nil {
return nil, nil
}
return &cfg, nil
}
return nil, nil
}
// identifyLicenseLayers turns Docker AI license-layer locations into pkg.License values.
func identifyLicenseLayers(ctx context.Context, resolver file.Resolver, locs []file.Location) []pkg.License {
var out []pkg.License
var scanFallback []file.Location
for i := range locs {
loc := locs[i]
if spdx := readLicenseSPDXIDFromFrontmatter(resolver, loc); spdx != "" {
out = append(out, pkg.NewLicenseFromFieldsWithContext(ctx, spdx, "", &loc))
continue
}
scanFallback = append(scanFallback, loc)
}
if len(scanFallback) > 0 {
out = append(out, licenses.FindAtLocations(ctx, resolver, scanFallback...)...)
}
return out
}
// readLicenseSPDXIDFromFrontmatter reads a bounded prefix of a license-layer
// blob and returns the spdx-id declared in its YAML frontmatter
func readLicenseSPDXIDFromFrontmatter(resolver file.Resolver, loc file.Location) string {
rc, err := resolver.FileContentsByLocation(loc)
if err != nil {
return ""
}
defer internal.CloseAndLogError(rc, loc.RealPath)
buf, err := io.ReadAll(io.LimitReader(rc, 64*1024))
if err != nil {
return ""
}
return parseLicenseFrontmatter(buf)
}

View File

@ -0,0 +1,139 @@
package ai
import (
"fmt"
"sort"
"strings"
"github.com/cespare/xxhash/v2"
"github.com/anchore/syft/syft/file"
"github.com/anchore/syft/syft/pkg"
)
// mergeSafeTensorsGroup folds a group's per-member metadata into a single package.
func mergeSafeTensorsGroup(members []pkg.Package) pkg.Package {
locSet := unionLocations(members)
aggregates, shards := bucketSafeTensorsMembers(members)
merged := pkg.SafeTensorsModelInfo{Format: "safetensors"}
mergeAggregatesInto(&merged, aggregates)
shardTensorTotal, hashes := mergeShardsInto(&merged, shards)
// Keep merged UserMetadata globally key-sorted so the SBOM is stable
sort.Slice(merged.UserMetadata, func(i, j int) bool {
return merged.UserMetadata[i].Key < merged.UserMetadata[j].Key
})
if merged.TensorCount == 0 {
merged.TensorCount = shardTensorTotal
}
if merged.ShardCount == 0 {
if len(shards) > 0 {
merged.ShardCount = len(shards)
} else {
merged.ShardCount = 1
}
}
merged.MetadataHash = rollupHash(hashes)
// Parts only carry value for multi-shard models; for a single shard the
// outer view already exposes every per-shard field.
if len(shards) > 1 {
parts := append([]pkg.SafeTensorsModelInfo(nil), shards...)
sort.Slice(parts, func(i, j int) bool {
return parts[i].MetadataHash < parts[j].MetadataHash
})
merged.Parts = parts
}
return pkg.Package{
Locations: locSet,
Type: pkg.ModelPkg,
Metadata: merged,
}
}
func mergeAggregatesInto(merged *pkg.SafeTensorsModelInfo, aggregates []pkg.SafeTensorsModelInfo) {
for _, a := range aggregates {
if merged.TensorCount == 0 {
merged.TensorCount = a.TensorCount
}
firstNonEmpty(&merged.TotalSize, a.TotalSize)
firstNonEmpty(&merged.Quantization, a.Quantization)
}
}
// mergeShardsInto folds the per-shard header metadata into merged, returning
// the summed shard TensorCount and the list of non-empty per-shard hashes for
// the rollup. Shards carry only the content-derived fields (Quantization,
// Parameters, UserMetadata), so those are the only fields folded in here.
// TensorCount and Parameters are summed because each shard holds a distinct
// slice of the model; Quantization takes the first value since all shards share
// one precision.
func mergeShardsInto(merged *pkg.SafeTensorsModelInfo, shards []pkg.SafeTensorsModelInfo) (shardTensorTotal uint64, hashes []string) {
seenKV := map[string]bool{}
for _, s := range shards {
shardTensorTotal += s.TensorCount
merged.Parameters += s.Parameters
firstNonEmpty(&merged.Quantization, s.Quantization)
for _, kv := range s.UserMetadata {
if seenKV[kv.Key] {
continue
}
seenKV[kv.Key] = true
merged.UserMetadata = append(merged.UserMetadata, kv)
}
if s.MetadataHash != "" {
hashes = append(hashes, s.MetadataHash)
}
}
return shardTensorTotal, hashes
}
func firstNonEmpty(dst *string, v string) {
if *dst == "" {
*dst = v
}
}
// unionLocations gathers every location from every member into a single set.
func unionLocations(members []pkg.Package) file.LocationSet {
out := file.NewLocationSet()
for _, m := range members {
for _, l := range m.Locations.ToSlice() {
out.Add(l)
}
}
return out
}
// bucketSafeTensorsMembers splits group members into aggregate-flavored entries
// (no MetadataHash — Docker AI config blob or sharded index) and shard-flavored
// entries (carry a content-derived MetadataHash from a header parser).
func bucketSafeTensorsMembers(members []pkg.Package) (aggregates, shards []pkg.SafeTensorsModelInfo) {
for _, m := range members {
md, ok := m.Metadata.(pkg.SafeTensorsModelInfo)
if !ok {
continue
}
if md.MetadataHash != "" {
shards = append(shards, md)
continue
}
aggregates = append(aggregates, md)
}
return aggregates, shards
}
func rollupHash(hashes []string) string {
if len(hashes) == 0 {
return ""
}
if len(hashes) == 1 {
return hashes[0]
}
sorted := append([]string(nil), hashes...)
sort.Strings(sorted)
return fmt.Sprintf("%016x", xxhash.Sum64String(strings.Join(sorted, "|")))
}

View File

@ -0,0 +1,166 @@
package ai
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/anchore/syft/syft/file"
"github.com/anchore/syft/syft/pkg"
)
// stPkg builds a model package carrying the given metadata, with each path
// recorded as a primary-evidence location.
func stPkg(md pkg.SafeTensorsModelInfo, paths ...string) pkg.Package {
locs := make([]file.Location, 0, len(paths))
for _, p := range paths {
locs = append(locs, file.NewLocation(p).WithAnnotation(pkg.EvidenceAnnotationKey, pkg.PrimaryEvidenceAnnotation))
}
return pkg.Package{Type: pkg.ModelPkg, Metadata: md, Locations: file.NewLocationSet(locs...)}
}
// shardMeta is a content-derived shard entry: it carries a MetadataHash, which is
// what marks a group member as a shard (vs. a hash-less aggregate config blob).
func shardMeta(hash string, tensorCount uint64) pkg.SafeTensorsModelInfo {
return pkg.SafeTensorsModelInfo{
Format: "safetensors",
TensorCount: tensorCount,
Quantization: "BF16",
Parameters: 1000,
MetadataHash: hash,
}
}
// TestMergeSafeTensorsGroup exercises the rollup contract directly (the cataloger
// tests cover it only as a side effect of the merge processor). It locks how a
// group's per-member metadata folds into one package: tensor-count summing,
// aggregate-over-shard field precedence, UserMetadata dedup + sorting, Parts
// rollup, ShardCount derivation, and the content MetadataHash rollup.
func TestMergeSafeTensorsGroup(t *testing.T) {
t.Run("single shard: hash passes through, ShardCount 1, no Parts", func(t *testing.T) {
out := mergeSafeTensorsGroup([]pkg.Package{stPkg(shardMeta("aaaa", 5), "/m/a.safetensors")})
md := out.Metadata.(pkg.SafeTensorsModelInfo)
assert.Equal(t, pkg.ModelPkg, out.Type)
assert.Equal(t, 1, md.ShardCount)
assert.Equal(t, uint64(5), md.TensorCount)
assert.Equal(t, "aaaa", md.MetadataHash, "a single shard's hash passes through unchanged")
assert.Nil(t, md.Parts, "single-shard models do not populate Parts")
})
t.Run("multi-shard: tensors summed, Parts sorted by hash, rollup is order-independent", func(t *testing.T) {
in := []pkg.Package{
stPkg(shardMeta("cccc", 3), "/m/c.safetensors"),
stPkg(shardMeta("aaaa", 3), "/m/a.safetensors"),
stPkg(shardMeta("bbbb", 3), "/m/b.safetensors"),
}
out := mergeSafeTensorsGroup(in)
md := out.Metadata.(pkg.SafeTensorsModelInfo)
assert.Equal(t, 3, md.ShardCount)
assert.Equal(t, uint64(9), md.TensorCount, "tensor counts are summed across shards")
assert.Equal(t, uint64(3000), md.Parameters, "parameter counts are summed across shards")
require.Len(t, md.Parts, 3)
assert.Equal(t,
[]string{"aaaa", "bbbb", "cccc"},
[]string{md.Parts[0].MetadataHash, md.Parts[1].MetadataHash, md.Parts[2].MetadataHash},
"Parts are sorted by metadata hash",
)
assert.Equal(t, rollupHash([]string{"aaaa", "bbbb", "cccc"}), md.MetadataHash)
// the rollup hash must not depend on the order members arrive in
shuffled := []pkg.Package{
stPkg(shardMeta("bbbb", 3), "/m/b.safetensors"),
stPkg(shardMeta("cccc", 3), "/m/c.safetensors"),
stPkg(shardMeta("aaaa", 3), "/m/a.safetensors"),
}
out2 := mergeSafeTensorsGroup(shuffled)
assert.Equal(t, md.MetadataHash, out2.Metadata.(pkg.SafeTensorsModelInfo).MetadataHash)
})
t.Run("aggregate fields win over shard-derived fields", func(t *testing.T) {
// an aggregate (OCI config blob) carries no MetadataHash but declares the
// authoritative totals.
aggregate := pkg.SafeTensorsModelInfo{
Format: "safetensors",
TensorCount: 999,
TotalSize: "5.00GB",
Quantization: "Q4_K_M",
}
in := []pkg.Package{
stPkg(aggregate, "/"),
stPkg(shardMeta("aaaa", 3), "/"),
stPkg(shardMeta("bbbb", 3), "/"),
}
out := mergeSafeTensorsGroup(in)
md := out.Metadata.(pkg.SafeTensorsModelInfo)
assert.Equal(t, uint64(999), md.TensorCount, "aggregate TensorCount is authoritative; shard counts are not summed in")
assert.Equal(t, "5.00GB", md.TotalSize)
assert.Equal(t, uint64(2000), md.Parameters, "parameters are always measured from the shards (summed), not taken from the aggregate")
assert.Equal(t, "Q4_K_M", md.Quantization, "aggregate quantization wins over the shard dtype")
assert.Equal(t, 2, md.ShardCount, "ShardCount comes from the number of shards, not the aggregate")
assert.Equal(t, rollupHash([]string{"aaaa", "bbbb"}), md.MetadataHash, "the content hash still rolls up the shard hashes")
})
t.Run("aggregate-only group: ShardCount 1, empty hash, no Parts", func(t *testing.T) {
aggregate := pkg.SafeTensorsModelInfo{Format: "safetensors", TensorCount: 42, TotalSize: "1.00GB"}
out := mergeSafeTensorsGroup([]pkg.Package{stPkg(aggregate, "/")})
md := out.Metadata.(pkg.SafeTensorsModelInfo)
assert.Equal(t, uint64(42), md.TensorCount)
assert.Equal(t, 1, md.ShardCount, "a group with no shards still reports a single shard")
assert.Equal(t, "", md.MetadataHash, "there are no shard hashes to roll up")
assert.Nil(t, md.Parts)
})
t.Run("UserMetadata: keys merged and sorted, first value wins on conflict", func(t *testing.T) {
// keys are intentionally unsorted within each shard so the assertion proves
// the merge re-sorts globally; "format" appears in both shards so dedup
// precedence (first wins) is exercised too.
s1 := shardMeta("aaaa", 1)
s1.UserMetadata = pkg.KeyValues{{Key: "format", Value: "pt"}, {Key: "author", Value: "alice"}}
s2 := shardMeta("bbbb", 1)
s2.UserMetadata = pkg.KeyValues{{Key: "format", Value: "gguf"}, {Key: "license", Value: "mit"}}
out := mergeSafeTensorsGroup([]pkg.Package{stPkg(s1, "/m/a.safetensors"), stPkg(s2, "/m/b.safetensors")})
md := out.Metadata.(pkg.SafeTensorsModelInfo)
assert.Equal(t, pkg.KeyValues{
{Key: "author", Value: "alice"},
{Key: "format", Value: "pt"}, // first shard's value wins over s2's "gguf"
{Key: "license", Value: "mit"},
}, md.UserMetadata)
})
t.Run("members without safetensors metadata are ignored in the rollup", func(t *testing.T) {
notST := pkg.Package{
Type: pkg.ModelPkg,
Metadata: pkg.GGUFFileHeader{},
Locations: file.NewLocationSet(file.NewLocation("/m/x.gguf")),
}
out := mergeSafeTensorsGroup([]pkg.Package{stPkg(shardMeta("aaaa", 2), "/m/a.safetensors"), notST})
md := out.Metadata.(pkg.SafeTensorsModelInfo)
assert.Equal(t, uint64(2), md.TensorCount, "only the safetensors shard contributes")
assert.Equal(t, 1, md.ShardCount)
assert.Equal(t, "aaaa", md.MetadataHash)
})
}
// TestRollupHash locks the cross-source content-fingerprint rollup: empty input
// yields no hash, a lone shard's hash passes through unchanged (so a single-shard
// model fingerprints identically across directory and OCI sources), and multiple
// shards fold into one order-independent digest.
func TestRollupHash(t *testing.T) {
assert.Equal(t, "", rollupHash(nil), "no hashes → empty")
assert.Equal(t, "solo", rollupHash([]string{"solo"}), "a single hash passes through unchanged")
ab := rollupHash([]string{"a", "b"})
ba := rollupHash([]string{"b", "a"})
assert.Equal(t, ab, ba, "the rollup is independent of input order")
assert.Len(t, ab, 16, "a multi-hash rollup is a 16-char xxhash")
assert.NotEqual(t, "a", ab)
assert.NotEqual(t, "b", ab)
}

View File

@ -0,0 +1,22 @@
package ai
import "path"
// pickSafeTensorsName implements the documented naming precedence chain:
func pickSafeTensorsName(nameOrPath, fallbackName string) string {
if nameOrPath != "" {
return path.Base(nameOrPath)
}
return fallbackName
}
// safeTensorsDirName returns the directory-scan naming fallback: the base name
// of the group's parent directory (the group key is already that directory).
func safeTensorsDirName(directory string) string {
base := path.Base(directory)
switch base {
case "/", ".", "":
return ""
}
return base
}

View File

@ -20,3 +20,17 @@ func newGGUFPackage(metadata *pkg.GGUFFileHeader, modelName, version, license st
return p return p
} }
// newSafeTensorsPackage creates a SafeTensors package with the given metadata
// and locations. Name and Licenses are intentionally not set here and done at the processor level
func newSafeTensorsPackage(metadata *pkg.SafeTensorsModelInfo, locations ...file.Location) pkg.Package {
p := pkg.Package{
Locations: file.NewLocationSet(locations...),
Type: pkg.ModelPkg,
Metadata: *metadata,
// PURL is intentionally not set: package-url has not yet finalized ML model support.
}
p.SetID()
return p
}

View File

@ -0,0 +1,202 @@
package ai
import (
"encoding/binary"
"encoding/json"
"fmt"
"io"
"sort"
"strings"
"github.com/cespare/xxhash/v2"
"github.com/anchore/syft/syft/pkg"
)
// SafeTensors file format: [8 bytes u64 LE header size] [N bytes JSON header] [tensor data].
// Reference: https://github.com/huggingface/safetensors#format
const (
maxSafeTensorsHeaderSize = 100 * 1024 * 1024 // 100MB ceiling on header JSON to prevent OOM
)
// safeTensorsHeader is the decoded JSON header. Tensor entries live alongside a
// reserved "__metadata__" key holding a string-to-string producer map. We decode
// tensor entries into a generic map so we can iterate and count without a fixed
// schema for every field.
type safeTensorsHeader struct {
metadata map[string]string
tensors map[string]safeTensorsEntry
}
// safeTensorsEntry describes a single tensor within the header JSON.
type safeTensorsEntry struct {
DType string `json:"dtype"`
Shape []int64 `json:"shape"`
DataOffsets []int64 `json:"data_offsets"`
}
// readSafeTensorsHeader reads and parses the JSON header from a .safetensors
// file (the leading `[8-byte LE length] [length bytes of JSON]` block) and
// returns the decoded header.
func readSafeTensorsHeader(r io.Reader) (*safeTensorsHeader, error) {
var lenBuf [8]byte
if _, err := io.ReadFull(r, lenBuf[:]); err != nil {
return nil, fmt.Errorf("failed to read header length: %w", err)
}
headerLen := binary.LittleEndian.Uint64(lenBuf[:])
if headerLen == 0 {
return nil, fmt.Errorf("safetensors header length is zero")
}
if headerLen > maxSafeTensorsHeaderSize {
return nil, fmt.Errorf("safetensors header size %d exceeds maximum %d", headerLen, maxSafeTensorsHeaderSize)
}
// Read incrementally rather than pre-allocating headerLen up front
body, err := io.ReadAll(io.LimitReader(r, int64(headerLen)))
if err != nil {
return nil, fmt.Errorf("failed to read header body: %w", err)
}
if uint64(len(body)) != headerLen {
return nil, fmt.Errorf("safetensors header truncated: read %d of %d bytes", len(body), headerLen)
}
var raw map[string]json.RawMessage
if err := json.Unmarshal(body, &raw); err != nil {
return nil, fmt.Errorf("failed to decode safetensors header JSON: %w", err)
}
h := &safeTensorsHeader{tensors: make(map[string]safeTensorsEntry, len(raw))}
for key, val := range raw {
if key == "__metadata__" {
if err := json.Unmarshal(val, &h.metadata); err != nil {
return nil, fmt.Errorf("failed to decode __metadata__: %w", err)
}
continue
}
var entry safeTensorsEntry
if err := json.Unmarshal(val, &entry); err != nil {
// Not all entries must conform; skip anything we cannot decode rather than fail.
continue
}
h.tensors[key] = entry
}
return h, nil
}
// parameterCount sums the element counts across all tensors in the header.
func (h *safeTensorsHeader) parameterCount() uint64 {
var total uint64
for _, t := range h.tensors {
count := uint64(1)
for _, dim := range t.Shape {
if dim <= 0 {
count = 0
break
}
count *= uint64(dim)
}
total += count
}
return total
}
// dominantDType returns the dtype that accounts for the largest fraction of parameters.
// For mixed-precision models the "dominant" dtype is still a useful summary.
func (h *safeTensorsHeader) dominantDType() string {
sizeByDType := make(map[string]uint64)
for _, t := range h.tensors {
count := uint64(1)
for _, dim := range t.Shape {
if dim <= 0 {
count = 0
break
}
count *= uint64(dim)
}
sizeByDType[t.DType] += count
}
var best string
var bestSize uint64
for dtype, size := range sizeByDType {
if size > bestSize || (size == bestSize && dtype < best) {
best = dtype
bestSize = size
}
}
return best
}
// metadataHash returns a stable xxhash64 over the logical tensor content
// (name + dtype + shape) plus the __metadata__ map. Tensor keys are sorted to
// keep the hash deterministic across producers.
func (h *safeTensorsHeader) metadataHash() string {
type logicalEntry struct {
Name string `json:"name"`
DType string `json:"dtype"`
Shape []int64 `json:"shape"`
}
entries := make([]logicalEntry, 0, len(h.tensors))
for name, t := range h.tensors {
entries = append(entries, logicalEntry{Name: name, DType: t.DType, Shape: t.Shape})
}
sort.Slice(entries, func(i, j int) bool { return entries[i].Name < entries[j].Name })
type hashInput struct {
Tensors []logicalEntry `json:"tensors"`
Metadata map[string]string `json:"metadata,omitempty"`
}
b, err := json.Marshal(hashInput{Tensors: entries, Metadata: h.metadata})
if err != nil {
return ""
}
return fmt.Sprintf("%016x", xxhash.Sum64(b))
}
// userMetadataKeyValues converts the safetensors __metadata__ map into a
// KeyValues slice sorted by key, so SBOM output is stable across runs. Returns
// nil for empty input (omitempty then drops the field).
func userMetadataKeyValues(m map[string]string) pkg.KeyValues {
if len(m) == 0 {
return nil
}
keys := make([]string, 0, len(m))
for k := range m {
keys = append(keys, k)
}
sort.Strings(keys)
out := make(pkg.KeyValues, 0, len(keys))
for _, k := range keys {
out = append(out, pkg.KeyValue{Key: k, Value: m[k]})
}
return out
}
// normalizeDType maps a safetensors/torch dtype label to an uppercase quantization
// shorthand matching conventions used elsewhere in syft (e.g., BF16, F16, I8).
func normalizeDType(dtype string) string {
switch strings.ToUpper(dtype) {
case "BF16":
return "BF16"
case "F16", "FP16", "FLOAT16", "HALF":
return "F16"
case "F32", "FP32", "FLOAT32", "FLOAT":
return "F32"
case "F64", "FP64", "FLOAT64", "DOUBLE":
return "F64"
case "I8", "INT8":
return "I8"
case "U8", "UINT8":
return "U8"
case "I16", "INT16":
return "I16"
case "I32", "INT32":
return "I32"
case "I64", "INT64":
return "I64"
case "BOOL":
return "BOOL"
default:
return strings.ToUpper(dtype)
}
}

View File

@ -0,0 +1,48 @@
package ai
import (
"context"
"fmt"
"io"
"github.com/anchore/syft/internal"
"github.com/anchore/syft/internal/unknown"
"github.com/anchore/syft/syft/artifact"
"github.com/anchore/syft/syft/file"
"github.com/anchore/syft/syft/pkg"
"github.com/anchore/syft/syft/pkg/cataloger/generic"
)
// parseSafeTensorsFile decodes the JSON header of a single .safetensors file
// (also called once per shard for sharded models) and emits a nameless package
// whose metadata is derived purely from the header bytes. Naming, license
// resolution, sibling enrichment, and cross-shard rollup are all handled by
// safeTensorsMergeProcessor.
func parseSafeTensorsFile(_ context.Context, _ file.Resolver, _ *generic.Environment, reader file.LocationReadCloser) ([]pkg.Package, []artifact.Relationship, error) {
defer internal.CloseAndLogError(reader, reader.Path())
header, err := readSafeTensorsHeader(&io.LimitedReader{R: reader, N: maxSafeTensorsHeaderSize + 8})
if err != nil {
return nil, nil, fmt.Errorf("failed to read safetensors header: %w", err)
}
// ShardCount is intentionally not set here: the merge processor is the single
// owner of ShardCount and derives it from the number of shards in the group.
md := pkg.SafeTensorsModelInfo{
Format: "safetensors",
TensorCount: uint64(len(header.tensors)),
Parameters: header.parameterCount(),
Quantization: normalizeDType(header.dominantDType()),
UserMetadata: userMetadataKeyValues(header.metadata),
MetadataHash: header.metadataHash(),
}
p := newSafeTensorsPackage(
&md,
reader.WithAnnotation(pkg.EvidenceAnnotationKey, pkg.PrimaryEvidenceAnnotation),
)
return []pkg.Package{p}, nil, unknown.IfEmptyf([]pkg.Package{p}, "unable to parse safetensors file")
}
// integrity check
var _ generic.Parser = parseSafeTensorsFile

View File

@ -0,0 +1,117 @@
package ai
import (
"context"
"encoding/json"
"fmt"
"io"
"strings"
"github.com/anchore/syft/internal"
"github.com/anchore/syft/internal/unknown"
"github.com/anchore/syft/syft/artifact"
"github.com/anchore/syft/syft/file"
"github.com/anchore/syft/syft/pkg"
"github.com/anchore/syft/syft/pkg/cataloger/generic"
)
// Docker AI OCI media types used by Docker Model Runner artifacts.
const (
dockerAIModelFileMediaType = "application/vnd.docker.ai.model.file"
dockerAILicenseMediaType = "application/vnd.docker.ai.license"
dockerAISafeTensorsMediaType = "application/vnd.docker.ai.safetensors"
)
// dockerAIModelConfigMediaTypes are the model-config schema versions this
// cataloger understands. Versions are enumerated explicitly rather than matched
// with a wildcard so that a future, potentially breaking, config schema is not
// silently consumed; add a new version here only after confirming the fields we
// parse still apply.
var dockerAIModelConfigMediaTypes = []string{
"application/vnd.docker.ai.model.config.v0.1+json",
"application/vnd.docker.ai.model.config.v0.2+json",
}
// dockerAIModelConfig mirrors the JSON shape of the vnd.docker.ai.model.config
// blob written by Docker Model Runner for AI artifacts. Only fields we use are
// declared; unknown fields are ignored.
type dockerAIModelConfig struct {
Config struct {
Format string `json:"format"`
Quantization string `json:"quantization"`
Size string `json:"size"`
SafeTensors struct {
TensorCount json.Number `json:"tensor_count"`
} `json:"safetensors"`
} `json:"config"`
}
// parseSafeTensorsOCIConfig decodes the Docker AI model-config blob
func parseSafeTensorsOCIConfig(_ context.Context, _ file.Resolver, _ *generic.Environment, reader file.LocationReadCloser) ([]pkg.Package, []artifact.Relationship, error) {
defer internal.CloseAndLogError(reader, reader.Path())
body, err := io.ReadAll(io.LimitReader(reader, 1024*1024))
if err != nil {
return nil, nil, fmt.Errorf("failed to read docker AI model config: %w", err)
}
var cfg dockerAIModelConfig
if err := json.Unmarshal(body, &cfg); err != nil {
return nil, nil, fmt.Errorf("failed to decode docker AI model config: %w", err)
}
if !strings.EqualFold(cfg.Config.Format, "safetensors") {
return nil, nil, nil
}
// Parameters is intentionally not read from the config blob: we measure the
// true parameter count from the SafeTensors layer headers (parseSafeTensorsOCILayer)
// so OCI and directory scans of the same model agree, rather than trusting the
// producer-supplied label here.
md := pkg.SafeTensorsModelInfo{
Format: "safetensors",
Quantization: cfg.Config.Quantization,
TotalSize: cfg.Config.Size,
}
if n, err := cfg.Config.SafeTensors.TensorCount.Int64(); err == nil && n > 0 {
md.TensorCount = uint64(n)
}
p := newSafeTensorsPackage(
&md,
reader.WithAnnotation(pkg.EvidenceAnnotationKey, pkg.PrimaryEvidenceAnnotation),
)
return []pkg.Package{p}, nil, unknown.IfEmptyf([]pkg.Package{p}, "unable to parse docker AI safetensors config")
}
// parseSafeTensorsOCILayer decodes the JSON header of a SafeTensors weight
// layer fetched from an OCI model artifact
func parseSafeTensorsOCILayer(_ context.Context, _ file.Resolver, _ *generic.Environment, reader file.LocationReadCloser) ([]pkg.Package, []artifact.Relationship, error) {
defer internal.CloseAndLogError(reader, reader.Path())
header, err := readSafeTensorsHeader(&io.LimitedReader{R: reader, N: maxSafeTensorsHeaderSize + 8})
if err != nil {
return nil, nil, fmt.Errorf("failed to read safetensors layer header: %w", err)
}
md := pkg.SafeTensorsModelInfo{
Format: "safetensors",
TensorCount: uint64(len(header.tensors)),
Parameters: header.parameterCount(),
Quantization: normalizeDType(header.dominantDType()),
UserMetadata: userMetadataKeyValues(header.metadata),
MetadataHash: header.metadataHash(),
}
p := newSafeTensorsPackage(
&md,
reader.WithAnnotation(pkg.EvidenceAnnotationKey, pkg.PrimaryEvidenceAnnotation),
)
return []pkg.Package{p}, nil, nil
}
// integrity checks
var (
_ generic.Parser = parseSafeTensorsOCIConfig
_ generic.Parser = parseSafeTensorsOCILayer
)

File diff suppressed because it is too large Load Diff

View File

@ -1,59 +1,172 @@
package ai package ai
import ( import (
"context"
"path"
"sort"
"github.com/anchore/syft/internal/log"
"github.com/anchore/syft/syft/artifact" "github.com/anchore/syft/syft/artifact"
"github.com/anchore/syft/syft/file"
"github.com/anchore/syft/syft/pkg" "github.com/anchore/syft/syft/pkg"
) )
// ggufMergeProcessor consolidates multiple GGUF packages into a single package // safeTensorsMergeProcessor owns naming, license resolution, and final package
// representing the AI model. When scanning OCI images with multiple layers, // assembly. SafeTensors packages reach it nameless from the parsers; it groups
// each layer may produce a separate package. This processor finds the package // them per model, merges the per-shard metadata, resolves a name + licenses, and
// with a name and merges metadata from nameless packages into its GGUFFileParts field. // drops any model it cannot name.
// Only packages with a non-empty name are returned in the final result. func safeTensorsMergeProcessor(ctx context.Context, resolver file.Resolver, pkgs []pkg.Package, rels []artifact.Relationship, err error) ([]pkg.Package, []artifact.Relationship, error) {
func ggufMergeProcessor(pkgs []pkg.Package, rels []artifact.Relationship, err error) ([]pkg.Package, []artifact.Relationship, error) {
if err != nil {
return pkgs, rels, err
}
if len(pkgs) == 0 { if len(pkgs) == 0 {
return pkgs, rels, err return pkgs, rels, err
} }
// Separate packages with names from those without // Note: we do NOT early-return when err != nil. A non-nil err here means some
var namedPkgs []pkg.Package // file in the run failed to parse, but the successfully-parsed packages are
var namelessHeaders []pkg.GGUFFileHeader // still nameless until this processor names or drops them. Skipping that work
// would let those nameless packages flow downstream, where they are silently
for _, p := range pkgs { // dropped by missing-name compliance — turning one bad file into the loss of
if p.Name != "" { // every otherwise-valid model. So we always name/drop and propagate err.
namedPkgs = append(namedPkgs, p) stPkgs, other := partitionSafeTensorsPackages(pkgs)
} else { if len(stPkgs) == 0 {
if header, ok := p.Metadata.(pkg.GGUFFileHeader); ok { return pkgs, rels, err
// We do not want a kv hash for nameless headers
header.MetadataKeyValuesHash = ""
namelessHeaders = append(namelessHeaders, header)
}
}
} }
// If there are no named packages, return nothing if fromOCIArtifact(stPkgs) {
if len(namedPkgs) == 0 { return append(other, mergeOCIModel(ctx, resolver, stPkgs)...), rels, err
return nil, rels, err
} }
return append(other, mergeDirModels(ctx, resolver, stPkgs)...), rels, err
// merge nameless headers into a single named package; }
// if there are multiple named packages, return them without trying to merge headers.
// we cannot determine which nameless headers belong to which package // partitionSafeTensorsPackages separates safetensors packages from anything else
// this is because the order we receive the gguf headers in is not guaranteed // flowing through the processor.
// to match the layer order in the original oci image func partitionSafeTensorsPackages(pkgs []pkg.Package) (safeTensors, other []pkg.Package) {
if len(namedPkgs) == 1 && len(namelessHeaders) > 0 { for _, p := range pkgs {
winner := &namedPkgs[0] if _, ok := p.Metadata.(pkg.SafeTensorsModelInfo); ok {
if header, ok := winner.Metadata.(pkg.GGUFFileHeader); ok { safeTensors = append(safeTensors, p)
header.Parts = namelessHeaders continue
winner.Metadata = header }
} other = append(other, p)
} }
return safeTensors, other
// Largest number of key value }
return namedPkgs, rels, err // fromOCIArtifact reports whether the packages came from an OCI model artifact.
// That source (the ContainerImageModel resolver) presents every layer at the
// virtual path "/", whereas a filesystem scan always carries a real file path. A
// single scan is one source, so the first package is representative of the rest.
func fromOCIArtifact(pkgs []pkg.Package) bool {
loc := primaryEvidenceLocation(pkgs[0])
return loc != nil && loc.RealPath == "/"
}
// mergeOCIModel treats the whole OCI artifact as a single model: every layer
// merges into one package, named from the artifact's config.json/README or its
// image reference.
func mergeOCIModel(ctx context.Context, resolver file.Resolver, pkgs []pkg.Package) []pkg.Package {
merged := mergeSafeTensorsGroup(pkgs)
md := merged.Metadata.(pkg.SafeTensorsModelInfo)
id := resolveSafeTensorsOCIIdentity(ctx, resolver, &md)
merged.Metadata = md // write architecture enrichment back before assembly
if p, ok := assembleSafeTensorsPackage(merged, id); ok {
return []pkg.Package{p}
}
return nil
}
// mergeDirModels groups filesystem-scanned files by their parent directory and
// emits one model per directory
func mergeDirModels(ctx context.Context, resolver file.Resolver, pkgs []pkg.Package) []pkg.Package {
groups := groupByParentDir(pkgs)
// deterministic iteration order so the SBOM doesn't depend on map order
dirs := make([]string, 0, len(groups))
for dir := range groups {
dirs = append(dirs, dir)
}
sort.Strings(dirs)
var out []pkg.Package
for _, dir := range dirs {
merged := mergeSafeTensorsGroup(groups[dir])
md := merged.Metadata.(pkg.SafeTensorsModelInfo)
id := resolveSafeTensorsDirIdentity(ctx, resolver, dir, &md)
merged.Metadata = md // write architecture enrichment back before assembly
if p, ok := assembleSafeTensorsPackage(merged, id); ok {
out = append(out, p)
}
}
return out
}
// groupByParentDir buckets filesystem-scanned models by the directory their
// primary-evidence file lives in.
// This encodes a deliberate assumption: a directory holds one logical
// model, so every .safetensors file in a directory is treated as a shard of the
// same modeland merged into one package.
// The trade-off is that if a directory happens to contain several unrelated models,
// they are merged into one package rather than reported separately.
// We accept that because the conventional on-disk layout gives each model
// (with all of its shards) its own directory.
// We have no reliable per-file signal to tell co-located-but-independent models apart.
func groupByParentDir(pkgs []pkg.Package) map[string][]pkg.Package {
out := make(map[string][]pkg.Package)
for _, p := range pkgs {
loc := primaryEvidenceLocation(p)
if loc == nil {
continue
}
dir := path.Dir(loc.RealPath)
out[dir] = append(out[dir], p)
}
return out
}
func primaryEvidenceLocation(p pkg.Package) *file.Location {
locs := p.Locations.ToSlice()
for i, l := range locs {
if l.Annotations != nil && l.Annotations[pkg.EvidenceAnnotationKey] == pkg.PrimaryEvidenceAnnotation {
return &locs[i]
}
}
if len(locs) > 0 {
return &locs[0]
}
return nil
}
// safeTensorsIdentity is the fully-resolved naming/license result for a model.
// Each source resolver (dir or OCI) populates it so assembly stays source-agnostic.
type safeTensorsIdentity struct {
nameOrPath string
fallbackName string
licenses []pkg.License
supporting []file.Location
}
// assembleSafeTensorsPackage finalizes a merged model from its resolved identity:
// it picks the name, attaches licenses and supporting evidence, and sets the ID.
// A model with no name source is dropped (ok=false).
func assembleSafeTensorsPackage(merged pkg.Package, id safeTensorsIdentity) (pkg.Package, bool) {
name := pickSafeTensorsName(id.nameOrPath, id.fallbackName)
if name == "" {
log.Debugf("dropped safetensors model package (metadata hash %q): no name source",
merged.Metadata.(pkg.SafeTensorsModelInfo).MetadataHash)
return pkg.Package{}, false
}
if len(id.licenses) > 0 {
merged.Licenses = pkg.NewLicenseSet(id.licenses...)
}
for _, loc := range id.supporting {
merged.Locations.Add(loc.WithAnnotation(pkg.EvidenceAnnotationKey, pkg.SupportingEvidenceAnnotation))
}
merged.Name = name
merged.SetID()
return merged, true
} }

View File

@ -0,0 +1,28 @@
# SafeTensors header fixtures
These fixtures are `[8-byte length prefix + JSON header]` captures from
public Docker AI model artifacts on the registry.
`extract_header.go` does a range-GET of the first several MB of the layer,
slices off just `[prefix + JSON header]`, and writes
that to disk.
## Refreshing a fixture
```sh
# from the package root
go run ./testdata/safetensors/extract_header.go \
docker.io/ai/nomic-embed-text-v2-moe-safetensors:475M \
./testdata/safetensors/nomic-embed-475M.header.safetensors
```
The tool prints the layer digest it selected and the number of top-level keys
in the captured header. If you see `header length N does not fit in M fetched
bytes`, raise `fetchBytes` in `extract_header.go` and rerun.
## Notes
- Pick one shard, not the full sharded set. The fixture is meant to exercise
the per-shard parser; merging across shards has its own tests.
- Don't commit anything larger than ~1 MB. If a model has an unusually large
header, capture a smaller model instead.

View File

@ -0,0 +1,149 @@
// extract_header is a manual fixture tool that captures the real on-disk
// safetensors header from a Docker AI OCI model artifact (a vnd.docker.ai.safetensors
// layer) and writes just [8-byte length prefix + JSON header] to a destination
// file. Tensor data following the header is never downloaded, so the resulting
// fixture is a few KB to a few MB even for multi-GB models.
//
// This file lives under testdata/ so the Go build system ignores it. Run it
// manually when refreshing fixtures:
//
// go run ./testdata/safetensors/extract_header.go \
// docker.io/ai/nomic-embed-text-v2-moe-safetensors:475M \
// ./testdata/safetensors/nomic-embed-475M.header.safetensors
package main
import (
"context"
"encoding/binary"
"encoding/json"
"fmt"
"io"
"os"
"github.com/google/go-containerregistry/pkg/authn"
"github.com/google/go-containerregistry/pkg/name"
v1 "github.com/google/go-containerregistry/pkg/v1"
"github.com/google/go-containerregistry/pkg/v1/remote"
)
const (
safetensorsLayerMediaType = "application/vnd.docker.ai.safetensors"
// 8 MB matches maxHeaderBytes in the OCI model source. Real model headers
// are well under 1 MB; the extra slack covers outliers.
fetchBytes = 8 * 1024 * 1024
)
func main() {
if len(os.Args) != 3 {
fmt.Fprintf(os.Stderr, "usage: %s <registry-ref> <output-path>\n", os.Args[0])
os.Exit(2)
}
if err := run(os.Args[1], os.Args[2]); err != nil {
fmt.Fprintln(os.Stderr, err)
os.Exit(1)
}
}
func run(refStr, outPath string) error {
ctx := context.Background()
ref, err := name.ParseReference(refStr)
if err != nil {
return fmt.Errorf("parse reference: %w", err)
}
opts := []remote.Option{
remote.WithAuthFromKeychain(authn.DefaultKeychain),
remote.WithContext(ctx),
}
desc, err := remote.Get(ref, opts...)
if err != nil {
return fmt.Errorf("fetch descriptor: %w", err)
}
manifest := &v1.Manifest{}
if err := json.Unmarshal(desc.Manifest, manifest); err != nil {
return fmt.Errorf("decode manifest: %w", err)
}
weightLayer := pickWeightLayer(manifest)
if weightLayer == nil {
return fmt.Errorf("no %q layer found in %s", safetensorsLayerMediaType, ref)
}
fmt.Fprintf(os.Stderr, "selected layer %s (%d bytes on-disk)\n", weightLayer.Digest, weightLayer.Size)
prefix, err := fetchPrefix(ctx, ref, weightLayer.Digest, opts)
if err != nil {
return fmt.Errorf("fetch layer prefix: %w", err)
}
header, err := sliceHeader(prefix)
if err != nil {
return fmt.Errorf("extract header: %w", err)
}
if err := os.WriteFile(outPath, header, 0o644); err != nil {
return fmt.Errorf("write fixture: %w", err)
}
fmt.Fprintf(os.Stderr, "wrote %d bytes to %s\n", len(header), outPath)
return nil
}
// pickWeightLayer returns the first vnd.docker.ai.safetensors layer in the
// manifest, or nil if none exists. For sharded models we deliberately only
// capture one shard: the fixture is meant to exercise the parser, not the
// merge step.
func pickWeightLayer(manifest *v1.Manifest) *v1.Descriptor {
for i := range manifest.Layers {
if string(manifest.Layers[i].MediaType) == safetensorsLayerMediaType {
return &manifest.Layers[i]
}
}
return nil
}
// fetchPrefix range-reads the first fetchBytes of a layer. Closing the reader
// terminates the underlying HTTP body, so we never download the tensor data
// that follows the header.
func fetchPrefix(_ context.Context, ref name.Reference, digest v1.Hash, opts []remote.Option) ([]byte, error) {
layer, err := remote.Layer(ref.Context().Digest(digest.String()), opts...)
if err != nil {
return nil, err
}
reader, err := layer.Compressed()
if err != nil {
return nil, err
}
defer reader.Close()
buf := make([]byte, fetchBytes)
n, err := io.ReadFull(reader, buf)
if err != nil && err != io.ErrUnexpectedEOF {
return nil, err
}
return buf[:n], nil
}
// sliceHeader reads the 8-byte little-endian length prefix and returns just
// [prefix + JSON header]. It also probes the JSON to make sure the captured
// fixture is well-formed, so we never commit a half-truncated header.
func sliceHeader(buf []byte) ([]byte, error) {
if len(buf) < 8 {
return nil, fmt.Errorf("short read: only %d bytes", len(buf))
}
headerLen := binary.LittleEndian.Uint64(buf[:8])
if headerLen == 0 {
return nil, fmt.Errorf("header length is zero")
}
if headerLen > uint64(len(buf)-8) {
return nil, fmt.Errorf("header length %d does not fit in %d fetched bytes; increase fetchBytes", headerLen, len(buf))
}
out := buf[:8+int(headerLen)]
var probe map[string]json.RawMessage
if err := json.Unmarshal(out[8:], &probe); err != nil {
return nil, fmt.Errorf("captured JSON does not parse: %w", err)
}
fmt.Fprintf(os.Stderr, "header parses cleanly: %d top-level keys\n", len(probe))
return out, nil
}

48
syft/pkg/safetensors.go Normal file
View File

@ -0,0 +1,48 @@
package pkg
// SafeTensorsModelInfo holds the model details extracted from SafeTensors content.
// SafeTensors is a simple, safe serialization format for storing tensors, used
// as the default weight format for Hugging Face transformer models.
// Model name, license, and version live on the syft Package
type SafeTensorsModelInfo struct {
// Format is the source format label (always "safetensors" for this metadata type).
// Present because the Docker AI model config blob carries an explicit format field
Format string `json:"format,omitempty" cyclonedx:"format"`
// Architecture is the model architecture (e.g., "LlamaForCausalLM",
// "Qwen3MoeForConditionalGeneration"). It is not present in the SafeTensors
// header itself; it is enriched from the companion config.json
// "architectures" array when one is found alongside the model.
Architecture string `json:"architecture,omitempty" cyclonedx:"architecture"`
// Quantization describes tensor precision (e.g., "BF16", "F16", "F32", "INT8").
Quantization string `json:"quantization,omitempty" cyclonedx:"quantization"`
// Parameters is the total number of model parameters, computed from the tensor
// shapes in the SafeTensors header(s). For a sharded model it is the sum across
// every shard.
Parameters uint64 `json:"parameters,omitempty" cyclonedx:"parameters"`
// TensorCount is the number of tensor entries in the file header.
TensorCount uint64 `json:"tensorCount,omitempty" cyclonedx:"tensorCount"`
// TotalSize is the total byte size of tensor data across all shards when known
// (from the Docker AI model config "size" field).
TotalSize string `json:"totalSize,omitempty" cyclonedx:"totalSize"`
// ShardCount is the number of .safetensors shards for a sharded model (1 for a
// single-file model).
ShardCount int `json:"shardCount,omitempty" cyclonedx:"shardCount"`
// UserMetadata is the optional "__metadata__" map from a .safetensors file header
// (string-to-string key/values set by the producer).
UserMetadata KeyValues `json:"userMetadata,omitempty" cyclonedx:"userMetadata"`
// MetadataHash is an xxhash over the on-disk SafeTensors header (sorted tensor
// entries + __metadata__). It is derived ONLY from the safetensors file bytes.
MetadataHash string `json:"metadataHash,omitempty" cyclonedx:"metadataHash"`
// Parts contains metadata from additional SafeTensors shards or OCI layers that
// were merged into this package during post-processing.
Parts []SafeTensorsModelInfo `json:"parts,omitempty" cyclonedx:"parts"`
}

View File

@ -3,9 +3,10 @@ package pkg
import ( import (
"testing" "testing"
"github.com/anchore/packageurl-go"
"github.com/scylladb/go-set/strset" "github.com/scylladb/go-set/strset"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/assert"
"github.com/anchore/packageurl-go"
) )
func TestTypeFromPURL(t *testing.T) { func TestTypeFromPURL(t *testing.T) {

View File

@ -52,7 +52,7 @@ func NewFromRegistry(ctx context.Context, cfg Config) (source.Source, error) {
} }
metadata := buildMetadata(art) metadata := buildMetadata(art)
tempDir, resolver, err := fetchAndStoreGGUFHeaders(ctx, client, art) tempDir, resolver, err := fetchAndStoreModelHeaders(ctx, client, art)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -77,41 +77,121 @@ func validateAndFetchArtifact(ctx context.Context, client *registryClient, refer
return nil, err return nil, err
} }
if len(art.GGUFLayers) == 0 { if art.Format == "" {
return nil, fmt.Errorf("model artifact has no GGUF layers") return nil, fmt.Errorf("model artifact has no GGUF or SafeTensors weight layers")
} }
return art, nil return art, nil
} }
// fetchAndStoreGGUFHeaders fetches GGUF layer headers and stores them in temp files. // fetchAndStoreModelHeaders fetches the blobs needed to catalog a Docker AI
func fetchAndStoreGGUFHeaders(ctx context.Context, client *registryClient, artifact *modelArtifact) (string, *fileresolver.ContainerImageModel, error) { // model artifact and stores them on disk so the ContainerImageModel resolver
tempDir, err := os.MkdirTemp("", "syft-oci-gguf") // can serve them by media type
func fetchAndStoreModelHeaders(ctx context.Context, client *registryClient, artifact *modelArtifact) (string, *fileresolver.ContainerImageModel, error) {
tempDir, err := os.MkdirTemp("", "syft-oci-model")
if err != nil { if err != nil {
return "", nil, fmt.Errorf("failed to create temp directory: %w", err) return "", nil, fmt.Errorf("failed to create temp directory: %w", err)
} }
cleanup := func() {
if osErr := os.RemoveAll(tempDir); osErr != nil {
log.Errorf("unable to remove temp directory (%s): %v", tempDir, osErr)
}
}
layerFiles := make(map[string]fileresolver.LayerInfo) layerFiles := make(map[string]fileresolver.LayerInfo)
// GGUF weight-layer headers.
for _, layer := range artifact.GGUFLayers { for _, layer := range artifact.GGUFLayers {
li, err := fetchSingleGGUFHeader(ctx, client, artifact.Reference, layer, tempDir) li, err := fetchSingleGGUFHeader(ctx, client, artifact.Reference, layer, tempDir)
if err != nil { if err != nil {
osErr := os.RemoveAll(tempDir) cleanup()
if osErr != nil {
log.Errorf("unable to remove temp directory (%s): %v", tempDir, err)
}
return "", nil, err return "", nil, err
} }
layerFiles[layer.Digest.String()] = li layerFiles[layer.Digest.String()] = li
} }
resolver := fileresolver.NewContainerImageModel(tempDir, layerFiles) // For SafeTensors artifacts, expose the model-config blob to the resolver
// so parseSafeTensorsOCIConfig can match it by media type.
if artifact.Format == modelFormatSafeTensors && len(artifact.RawConfig) > 0 {
li, err := storeConfigBlobAsLayer(artifact, tempDir)
if err != nil {
cleanup()
return "", nil, err
}
layerFiles[artifact.Manifest.Config.Digest.String()] = li
}
// Companion layers (README, config.json, tokenizer.json, LICENSE). Small by
// convention; fetched in full up to maxCompanionBytes.
if artifact.Format == modelFormatSafeTensors {
for _, layer := range artifact.CompanionLayers {
li, err := fetchCompanionLayer(ctx, client, artifact.Reference, layer, tempDir)
if err != nil {
cleanup()
return "", nil, err
}
layerFiles[layer.Digest.String()] = li
}
}
// SafeTensors weight-layer headers. We only pull the leading prefix (same
// budget as a GGUF header)
if artifact.Format == modelFormatSafeTensors {
for _, layer := range artifact.SafeTensorsLayers {
li, err := fetchSafeTensorsLayerHeader(ctx, client, artifact.Reference, layer, tempDir)
if err != nil {
cleanup()
return "", nil, err
}
layerFiles[layer.Digest.String()] = li
}
}
resolver := fileresolver.NewContainerImageModel(
tempDir,
layerFiles,
artifact.Reference.String(),
)
return tempDir, resolver, nil return tempDir, resolver, nil
} }
// fetchSingleGGUFHeader fetches a single GGUF layer header and writes it to a temp file. // storeConfigBlobAsLayer writes the already-fetched raw config bytes to a temp
// file so the resolver can serve them via media type
func storeConfigBlobAsLayer(artifact *modelArtifact, tempDir string) (fileresolver.LayerInfo, error) {
digest := artifact.Manifest.Config.Digest.String()
safeDigest := strings.ReplaceAll(digest, ":", "-")
tempPath := filepath.Join(tempDir, safeDigest+".config.json")
if err := os.WriteFile(tempPath, artifact.RawConfig, 0600); err != nil {
return fileresolver.LayerInfo{}, fmt.Errorf("failed to write config blob: %w", err)
}
return fileresolver.LayerInfo{
TempPath: tempPath,
MediaType: string(artifact.Manifest.Config.MediaType),
}, nil
}
// fetchCompanionLayer downloads a companion (non-weight) layer to a temp file
func fetchCompanionLayer(ctx context.Context, client *registryClient, ref name.Reference, layer v1.Descriptor, tempDir string) (fileresolver.LayerInfo, error) {
data, err := client.fetchBlobRange(ctx, ref, layer.Digest, maxCompanionBytes)
if err != nil {
return fileresolver.LayerInfo{}, fmt.Errorf("failed to fetch companion layer: %w", err)
}
safeDigest := strings.ReplaceAll(layer.Digest.String(), ":", "-")
tempPath := filepath.Join(tempDir, safeDigest+".blob")
if err := os.WriteFile(tempPath, data, 0600); err != nil {
return fileresolver.LayerInfo{}, fmt.Errorf("failed to write companion temp file: %w", err)
}
return fileresolver.LayerInfo{
TempPath: tempPath,
MediaType: string(layer.MediaType),
}, nil
}
// fetchSingleGGUFHeader fetches a single GGUF layer header and writes it to a temp file
func fetchSingleGGUFHeader(ctx context.Context, client *registryClient, ref name.Reference, layer v1.Descriptor, tempDir string) (fileresolver.LayerInfo, error) { func fetchSingleGGUFHeader(ctx context.Context, client *registryClient, ref name.Reference, layer v1.Descriptor, tempDir string) (fileresolver.LayerInfo, error) {
headerData, err := client.fetchBlobRange(ctx, ref, layer.Digest, maxHeaderBytes) headerData, err := client.fetchBlobRange(ctx, ref, layer.Digest, maxWeightHeaderBytes)
if err != nil { if err != nil {
return fileresolver.LayerInfo{}, fmt.Errorf("failed to fetch GGUF layer header: %w", err) return fileresolver.LayerInfo{}, fmt.Errorf("failed to fetch GGUF layer header: %w", err)
} }
@ -129,7 +209,27 @@ func fetchSingleGGUFHeader(ctx context.Context, client *registryClient, ref name
}, nil }, nil
} }
// buildMetadata constructs OCIModelMetadata from a modelArtifact. // fetchSafeTensorsLayerHeader fetches the leading bytes of a SafeTensors weight
// layer (enough to cover the JSON header) and writes them to a temp file
func fetchSafeTensorsLayerHeader(ctx context.Context, client *registryClient, ref name.Reference, layer v1.Descriptor, tempDir string) (fileresolver.LayerInfo, error) {
headerData, err := client.fetchBlobRange(ctx, ref, layer.Digest, maxWeightHeaderBytes)
if err != nil {
return fileresolver.LayerInfo{}, fmt.Errorf("failed to fetch safetensors layer header: %w", err)
}
safeDigest := strings.ReplaceAll(layer.Digest.String(), ":", "-")
tempPath := filepath.Join(tempDir, safeDigest+".safetensors")
if err := os.WriteFile(tempPath, headerData, 0600); err != nil {
return fileresolver.LayerInfo{}, fmt.Errorf("failed to write temp file: %w", err)
}
return fileresolver.LayerInfo{
TempPath: tempPath,
MediaType: string(layer.MediaType),
}, nil
}
// buildMetadata constructs OCIModelMetadata from a modelArtifact
func buildMetadata(artifact *modelArtifact) source.OCIModelMetadata { func buildMetadata(artifact *modelArtifact) source.OCIModelMetadata {
// layers // layers
layers := make([]source.LayerMetadata, len(artifact.Manifest.Layers)) layers := make([]source.LayerMetadata, len(artifact.Manifest.Layers))
@ -224,7 +324,8 @@ func (s *ociModelSource) Describe() source.Description {
} }
} }
// FileResolver returns a file resolver for accessing header of GGUF files. // FileResolver returns a file resolver for accessing model headers and companion
// metadata (GGUF/SafeTensors headers, the model config blob, and companion layers).
func (s *ociModelSource) FileResolver(_ source.Scope) (file.Resolver, error) { func (s *ociModelSource) FileResolver(_ source.Scope) (file.Resolver, error) {
return s.resolver, nil return s.resolver, nil
} }

View File

@ -26,9 +26,29 @@ const (
// Reference: https://www.docker.com/blog/oci-artifacts-for-ai-model-packaging/ // Reference: https://www.docker.com/blog/oci-artifacts-for-ai-model-packaging/
modelConfigMediaTypePrefix = "application/vnd.docker.ai.model.config." modelConfigMediaTypePrefix = "application/vnd.docker.ai.model.config."
ggufLayerMediaType = "application/vnd.docker.ai.gguf.v3" ggufLayerMediaType = "application/vnd.docker.ai.gguf.v3"
safetensorsLayerMediaType = "application/vnd.docker.ai.safetensors"
// Maximum bytes to read/return for GGUF headers // Companion metadata layers packaged alongside the weight tensors.
maxHeaderBytes = 8 * 1024 * 1024 // 8 MB // model.file covers README.md / config.json / tokenizer.json / generation_config.json.
modelFileMediaType = "application/vnd.docker.ai.model.file"
licenseMediaType = "application/vnd.docker.ai.license"
// Weight format labels surfaced on modelArtifact.Format.
modelFormatGGUF = "gguf"
modelFormatSafeTensors = "safetensors"
// maxWeightHeaderBytes is the leading slice we range-GET from a (multi-GB)
// weight layer — enough to cover the GGUF/safetensors header. Note this is
// smaller than the ai cataloger's own maxSafeTensorsHeaderSize (100 MB) parse
// ceiling: a safetensors header between the two is parseable from a directory
// scan but truncated here, so its shard would go uncounted on an OCI scan.
// Keep this comfortably above real-world header sizes.
maxWeightHeaderBytes = 8 * 1024 * 1024 // 8 MB
// maxCompanionBytes caps a whole companion blob (README, config.json,
// license); these are small by convention. Matches the 4 MB read cap in
// classifyOCIModelFileLayer.
maxCompanionBytes = 4 * 1024 * 1024 // 4 MB
) )
// registryClient handles OCI registry interactions for model artifacts. // registryClient handles OCI registry interactions for model artifacts.
@ -110,7 +130,20 @@ type modelArtifact struct {
RawManifest []byte RawManifest []byte
RawConfig []byte RawConfig []byte
ManifestDigest string ManifestDigest string
// Format identifies the weight storage format advertised by the manifest's
// layer media types. Empty means no recognized weight layers were found.
Format string
// GGUFLayers are descriptors for layers carrying GGUF-format weights.
// We fetch the first few MB of each to read the header data
GGUFLayers []v1.Descriptor GGUFLayers []v1.Descriptor
// SafeTensorsLayers are descriptors for layers carrying SafeTensors-format weights.
SafeTensorsLayers []v1.Descriptor
// CompanionLayers are non-weight layers (README, config.json, license)
CompanionLayers []v1.Descriptor
} }
func (c *registryClient) fetchModelArtifact(ctx context.Context, refStr string) (*modelArtifact, error) { func (c *registryClient) fetchModelArtifact(ctx context.Context, refStr string) (*modelArtifact, error) {
@ -151,6 +184,8 @@ func (c *registryClient) fetchModelArtifact(ctx context.Context, refStr string)
} }
ggufLayers := extractGGUFLayers(manifest) ggufLayers := extractGGUFLayers(manifest)
safetensorsLayers := extractSafeTensorsLayers(manifest)
companionLayers := extractCompanionLayers(manifest)
return &modelArtifact{ return &modelArtifact{
Reference: ref, Reference: ref,
@ -159,10 +194,28 @@ func (c *registryClient) fetchModelArtifact(ctx context.Context, refStr string)
RawManifest: desc.Manifest, RawManifest: desc.Manifest,
RawConfig: rawConfig, RawConfig: rawConfig,
ManifestDigest: desc.Digest.String(), ManifestDigest: desc.Digest.String(),
Format: detectModelFormat(len(ggufLayers), len(safetensorsLayers)),
GGUFLayers: ggufLayers, GGUFLayers: ggufLayers,
SafeTensorsLayers: safetensorsLayers,
CompanionLayers: companionLayers,
}, nil }, nil
} }
// detectModelFormat returns a single format string when either GGUF or
// SafeTensors weight layers are present. GGUF wins if both somehow appear in one
// artifact; Docker AI artifacts carry a single weight format in practice, so the
// mixed case is not expected and the safetensors side would go uncataloged.
func detectModelFormat(ggufCount, safetensorsCount int) string {
switch {
case ggufCount > 0:
return modelFormatGGUF
case safetensorsCount > 0:
return modelFormatSafeTensors
default:
return ""
}
}
// isModelArtifact checks if the manifest represents a model artifact. // isModelArtifact checks if the manifest represents a model artifact.
func isModelArtifact(manifest *v1.Manifest) bool { func isModelArtifact(manifest *v1.Manifest) bool {
return strings.HasPrefix(string(manifest.Config.MediaType), modelConfigMediaTypePrefix) return strings.HasPrefix(string(manifest.Config.MediaType), modelConfigMediaTypePrefix)
@ -179,6 +232,33 @@ func extractGGUFLayers(manifest *v1.Manifest) []v1.Descriptor {
return ggufLayers return ggufLayers
} }
// extractSafeTensorsLayers extracts SafeTensors weight-layer descriptors from
// the manifest.
func extractSafeTensorsLayers(manifest *v1.Manifest) []v1.Descriptor {
var out []v1.Descriptor
for _, layer := range manifest.Layers {
if string(layer.MediaType) == safetensorsLayerMediaType {
out = append(out, layer)
}
}
return out
}
// extractCompanionLayers extracts small, non-weight layers that carry
// cataloger-relevant metadata: README.md / config.json / tokenizer.json /
// generation_config.json under vnd.docker.ai.model.file, and the LICENSE under
// vnd.docker.ai.license.
func extractCompanionLayers(manifest *v1.Manifest) []v1.Descriptor {
var out []v1.Descriptor
for _, layer := range manifest.Layers {
switch string(layer.MediaType) {
case modelFileMediaType, licenseMediaType:
out = append(out, layer)
}
}
return out
}
func (c *registryClient) fetchBlobRange(ctx context.Context, ref name.Reference, digest v1.Hash, maxBytes int64) ([]byte, error) { func (c *registryClient) fetchBlobRange(ctx context.Context, ref name.Reference, digest v1.Hash, maxBytes int64) ([]byte, error) {
repo := ref.Context() repo := ref.Context()
@ -189,18 +269,14 @@ func (c *registryClient) fetchBlobRange(ctx context.Context, ref name.Reference,
return nil, fmt.Errorf("failed to fetch layer: %w", err) return nil, fmt.Errorf("failed to fetch layer: %w", err)
} }
// Compressed() returns the raw stored blob. Docker AI weight/companion layers
// are stored uncompressed (their media types carry no +gzip suffix), so this
// is the header bytes as-is. A gzip-compressed layer would parse as garbage
// downstream rather than being transparently decompressed here.
reader, err := layer.Compressed() reader, err := layer.Compressed()
if err != nil { if err != nil {
return nil, fmt.Errorf("failed to get layer reader: %w", err) return nil, fmt.Errorf("failed to get layer reader: %w", err)
} }
// this defer is what causes the download to stop
// 1. io.ReadFull(reader, data) reads exactly 8MB into the buffer
// 2. The function returns with data[:n]
// 3. defer reader.Close() executes, closing the HTTP response body
// 4. Closing the response body closes the underlying TCP connection
// 5. The server receives TCP FIN/RST and stops sending
// note: some data is already in flight when we close so we will see > 8mb over the wire
// the full image will not download given we terminate the reader early here
defer reader.Close() defer reader.Close()
// Note: this is not some arbitrary number picked out of the blue. // Note: this is not some arbitrary number picked out of the blue.
@ -208,8 +284,10 @@ func (c *registryClient) fetchBlobRange(ctx context.Context, ref name.Reference,
// https://github.com/ggml-org/ggml/blob/master/docs/gguf.md#file-structure // https://github.com/ggml-org/ggml/blob/master/docs/gguf.md#file-structure
data := make([]byte, maxBytes) data := make([]byte, maxBytes)
n, err := io.ReadFull(reader, data) n, err := io.ReadFull(reader, data)
if err != nil && err != io.ErrUnexpectedEOF {
// ErrUnexpectedEOF is okay - it means the file is smaller than maxBytes // ErrUnexpectedEOF means the layer is smaller than maxBytes; EOF means it is
// empty. Both mean we read everything there was, not a failure.
if err != nil && !errors.Is(err, io.ErrUnexpectedEOF) && !errors.Is(err, io.EOF) {
return nil, fmt.Errorf("failed to read layer data: %w", err) return nil, fmt.Errorf("failed to read layer data: %w", err)
} }

View File

@ -0,0 +1,112 @@
package ocimodelsource
import (
"context"
"io"
stdlog "log"
"net/http/httptest"
"net/url"
"os"
"testing"
"github.com/google/go-containerregistry/pkg/name"
"github.com/google/go-containerregistry/pkg/registry"
"github.com/google/go-containerregistry/pkg/v1/empty"
"github.com/google/go-containerregistry/pkg/v1/mutate"
"github.com/google/go-containerregistry/pkg/v1/remote"
"github.com/google/go-containerregistry/pkg/v1/static"
"github.com/google/go-containerregistry/pkg/v1/types"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/anchore/stereoscope/pkg/image"
"github.com/anchore/syft/syft/internal/fileresolver"
)
// TestFetchAndStore_safetensors_inMemoryRegistry drives the full OCI fetch +
// on-disk staging pipeline against a real (but in-process) registry, so nothing
// touches the network. It pushes a synthetic Docker AI safetensors artifact —
// config blob + one safetensors weight layer + a model.file companion + a
// license layer — then asserts that fetchModelArtifact classifies the layers and
// that fetchAndStoreModelHeaders stages every blob under the correct media type
// with its bytes intact. This is the seam the mock-resolver merge tests assume is
// correct; the staged content here is deliberately format-agnostic (header
// parsing is covered in the ai package).
func TestFetchAndStore_safetensors_inMemoryRegistry(t *testing.T) {
ctx := context.Background()
configMediaType := modelConfigMediaTypePrefix + "v0.2+json"
weightHeader := []byte("safetensors-header-bytes")
modelFile := []byte(`{"architectures":["LlamaForCausalLM"],"_name_or_path":"org/model"}`)
license := []byte("MIT license text")
img := mutate.ConfigMediaType(empty.Image, types.MediaType(configMediaType))
img = mutate.MediaType(img, types.OCIManifestSchema1)
img, err := mutate.Append(img,
layer(weightHeader, safetensorsLayerMediaType),
layer(modelFile, modelFileMediaType),
layer(license, licenseMediaType),
)
require.NoError(t, err)
// in-memory registry; localhost so go-containerregistry selects the http scheme
server := httptest.NewServer(registry.New(registry.Logger(stdlog.New(io.Discard, "", 0))))
defer server.Close()
u, err := url.Parse(server.URL)
require.NoError(t, err)
refStr := "localhost:" + u.Port() + "/testmodel:latest"
ref, err := name.ParseReference(refStr)
require.NoError(t, err)
require.NoError(t, remote.Write(ref, img, remote.WithContext(ctx)))
client := newRegistryClient(&image.RegistryOptions{InsecureUseHTTP: true})
art, err := client.fetchModelArtifact(ctx, refStr)
require.NoError(t, err)
assert.Equal(t, modelFormatSafeTensors, art.Format)
assert.Len(t, art.SafeTensorsLayers, 1)
assert.Len(t, art.CompanionLayers, 2) // model.file + license
assert.Empty(t, art.GGUFLayers)
assert.NotEmpty(t, art.RawConfig)
assert.NotEmpty(t, art.ManifestDigest)
tempDir, resolver, err := fetchAndStoreModelHeaders(ctx, client, art)
require.NoError(t, err)
defer os.RemoveAll(tempDir)
assert.Equal(t, refStr, resolver.ImageReference())
// every blob is staged under its own media type with bytes intact — this
// exercises fetchBlobRange's short-read (ErrUnexpectedEOF) branch, since each
// blob is far smaller than the fetch cap, plus all four staging helpers.
assertServesBytes(t, resolver, safetensorsLayerMediaType, weightHeader)
assertServesBytes(t, resolver, modelFileMediaType, modelFile)
assertServesBytes(t, resolver, licenseMediaType, license)
cfgLocs, err := resolver.FilesByMediaType(configMediaType)
require.NoError(t, err)
require.Len(t, cfgLocs, 1)
}
func layer(content []byte, mediaType string) mutate.Addendum {
return mutate.Addendum{
Layer: static.NewLayer(content, types.MediaType(mediaType)),
MediaType: types.MediaType(mediaType),
}
}
func assertServesBytes(t *testing.T, resolver *fileresolver.ContainerImageModel, mediaType string, want []byte) {
t.Helper()
locs, err := resolver.FilesByMediaType(mediaType)
require.NoError(t, err)
require.Len(t, locs, 1)
rc, err := resolver.FileContentsByLocation(locs[0])
require.NoError(t, err)
defer rc.Close()
got, err := io.ReadAll(rc)
require.NoError(t, err)
assert.Equal(t, want, got)
}

View File

@ -0,0 +1,55 @@
package ocimodelsource
import (
"testing"
v1 "github.com/google/go-containerregistry/pkg/v1"
"github.com/google/go-containerregistry/pkg/v1/types"
"github.com/stretchr/testify/assert"
"github.com/anchore/syft/syft/source"
)
func TestDetectModelFormat(t *testing.T) {
tests := []struct {
name string
gguf int
safetensors int
expected string
}{
{name: "gguf only", gguf: 2, safetensors: 0, expected: modelFormatGGUF},
{name: "safetensors only", gguf: 0, safetensors: 3, expected: modelFormatSafeTensors},
{name: "both prefers gguf", gguf: 1, safetensors: 1, expected: modelFormatGGUF},
{name: "neither", gguf: 0, safetensors: 0, expected: ""},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
assert.Equal(t, test.expected, detectModelFormat(test.gguf, test.safetensors))
})
}
}
func TestExtractSafeTensorsLayers(t *testing.T) {
manifest := &v1.Manifest{Layers: []v1.Descriptor{
{MediaType: types.MediaType(safetensorsLayerMediaType), Digest: v1.Hash{Algorithm: "sha256", Hex: "a"}},
{MediaType: types.MediaType(ggufLayerMediaType), Digest: v1.Hash{Algorithm: "sha256", Hex: "b"}},
{MediaType: types.MediaType(safetensorsLayerMediaType), Digest: v1.Hash{Algorithm: "sha256", Hex: "c"}},
}}
assert.Len(t, extractSafeTensorsLayers(manifest), 2)
}
func TestExtractCompanionLayers(t *testing.T) {
manifest := &v1.Manifest{Layers: []v1.Descriptor{
{MediaType: types.MediaType(modelFileMediaType), Digest: v1.Hash{Algorithm: "sha256", Hex: "readme"}},
{MediaType: types.MediaType(licenseMediaType), Digest: v1.Hash{Algorithm: "sha256", Hex: "license"}},
{MediaType: types.MediaType(safetensorsLayerMediaType), Digest: v1.Hash{Algorithm: "sha256", Hex: "weights"}},
{MediaType: types.DockerLayer, Digest: v1.Hash{Algorithm: "sha256", Hex: "other"}},
}}
// only the model.file and license layers should be selected (not weights or arbitrary layers)
assert.Len(t, extractCompanionLayers(manifest), 2)
}
func TestCalculateTotalSize(t *testing.T) {
layers := []source.LayerMetadata{{Size: 100}, {Size: 250}, {Size: 0}}
assert.Equal(t, int64(350), calculateTotalSize(layers))
}