feat(unpack): improve runtime config handling and streamline layer unpacking logic

Author: ilopezluna

pkg/distribution/builder/from_directory.go

@@ -14,6 +14,8 @@ import (
 	"github.com/docker/model-runner/pkg/distribution/types"
 )
 
+const rootFSType = "rootfs"
+
 // DirectoryOptions configures the behavior of FromDirectory.
 type DirectoryOptions struct {
 	// Exclusions is a list of patterns to exclude from packaging.
@@ -130,7 +132,7 @@ func FromDirectory(dirPath string, opts ...DirectoryOption) (*Builder, error) {
 		fileType := files.Classify(path)
 		mediaType := fileTypeToMediaType(fileType)
 
-		// Track format from weight files
+		// Track format from weight files (only weight file types affect format detection)
 		switch fileType {
 		case files.FileTypeSafetensors:
 			if detectedFormat == "" {
@@ -147,10 +149,8 @@ func FromDirectory(dirPath string, opts ...DirectoryOption) (*Builder, error) {
 				detectedFormat = types.FormatDiffusers
 			}
 			weightFiles = append(weightFiles, path)
-		case files.FileTypeUnknown:
-		case files.FileTypeConfig:
-		case files.FileTypeLicense:
-		case files.FileTypeChatTemplate:
+		case files.FileTypeUnknown, files.FileTypeConfig, files.FileTypeLicense, files.FileTypeChatTemplate:
+			// Non-weight files don't affect format detection
 		}
 
 		// Create layer with relative path annotation
@@ -182,7 +182,7 @@ func FromDirectory(dirPath string, opts ...DirectoryOption) (*Builder, error) {
 		return nil, fmt.Errorf("no weight files (safetensors, GGUF, or DDUF) found in directory: %s", dirPath)
 	}
 
-	// Build config - use the first weight file for metadata extraction
+	// Build config
 	config := types.Config{
 		Format: detectedFormat,
 	}
@@ -199,12 +199,12 @@ func FromDirectory(dirPath string, opts ...DirectoryOption) (*Builder, error) {
 				Created: &created,
 			},
 			RootFS: oci.RootFS{
-				Type:    "rootfs",
+				Type:    rootFSType,
 				DiffIDs: diffIDs,
 			},
 		},
 		LayerList:       layers,
-		ConfigMediaType: types.MediaTypeModelConfigV02, // V0.2: layer-per-file with filepath annotations
+		ConfigMediaType: types.MediaTypeModelConfigV02,
 	}
 
 	return &Builder{
@@ -290,10 +290,11 @@ func fileTypeToMediaType(ft files.FileType) oci.MediaType {
 		return types.MediaTypeLicense
 	case files.FileTypeChatTemplate:
 		return types.MediaTypeChatTemplate
+	case files.FileTypeUnknown:
+		return types.MediaTypeModelFile
 	case files.FileTypeConfig:
 		return types.MediaTypeModelFile
 	default:
-		// For unknown files, use the generic model file type
 		return types.MediaTypeModelFile
 	}
 }

pkg/distribution/internal/bundle/unpack.go

@@ -175,7 +175,13 @@ func hasLayerWithMediaType(model types.Model, targetMediaType oci.MediaType) boo
 	}
 }
 
-func unpackRuntimeConfig(bundle *Bundle, mdl types.Model) error {
+// configProvider is an interface for types that provide a Config() method.
+// Both types.Model and types.ModelArtifact satisfy this interface.
+type configProvider interface {
+	Config() (types.ModelConfig, error)
+}
+
+func unpackRuntimeConfig(bundle *Bundle, mdl configProvider) error {
 	cfg, err := mdl.Config()
 	if err != nil {
 		return err
@@ -317,11 +323,8 @@ func unpackSafetensorsWithAnnotations(bundle *Bundle, modelDir string, safetenso
 		digestHex := filepath.Base(srcPath)
 
 		var destRelPath string
-		if annotatedPath, ok := blobToFilepath[digestHex]; ok && strings.Contains(annotatedPath, "/") {
-			// Use the annotated path (contains subdirectory)
-			destRelPath = annotatedPath
-		} else if annotatedPath, ok := blobToFilepath[digestHex]; ok {
-			// Annotation exists but is just a filename - use it
+		if annotatedPath, ok := blobToFilepath[digestHex]; ok {
+			// Use the annotated path
 			destRelPath = annotatedPath
 		} else {
 			// No annotation found - use legacy naming
@@ -422,7 +425,6 @@ func unpackDirTarArchives(bundle *Bundle, mdl types.Model) error {
 	for _, layer := range layers {
 		mediaType, err := layer.MediaType()
 		if err != nil {
-			fmt.Printf("Warning: failed to get media type for layer: %v", err)
 			continue
 		}
 
@@ -641,21 +643,18 @@ func UnpackFromLayers(dir string, model types.ModelArtifact) (*Bundle, error) {
 	for _, layer := range layers {
 		mediaType, err := layer.MediaType()
 		if err != nil {
-			fmt.Printf("Warning: error getting media type: %v\n", err)
 			continue
 		}
 
 		// Get the filepath annotation
 		dp, ok := layer.(descriptorProvider)
 		if !ok {
-			fmt.Printf("Warning: layer is not a descriptorProvider\n")
 			continue
 		}
 
 		desc := dp.GetDescriptor()
 		relPath, exists := desc.Annotations[types.AnnotationFilePath]
 		if !exists || relPath == "" {
-			fmt.Printf("Warning: layer missing filepath annotation\n")
 			continue
 		}
 
@@ -687,28 +686,16 @@ func UnpackFromLayers(dir string, model types.ModelArtifact) (*Bundle, error) {
 		updateBundleFieldsFromLayer(bundle, mediaType, relPath)
 	}
 
-	// Create the runtime config from the model
-	cfg, err := model.Config()
-	if err != nil {
-		return nil, fmt.Errorf("get model config: %w", err)
-	}
-
-	// Write runtime config to bundle root
-	f, err := os.Create(filepath.Join(bundle.dir, "config.json"))
-	if err != nil {
-		return nil, fmt.Errorf("create runtime config file: %w", err)
-	}
-	defer f.Close()
-	if err := json.NewEncoder(f).Encode(cfg); err != nil {
-		return nil, fmt.Errorf("encode runtime config: %w", err)
+	// Create the runtime config
+	if err := unpackRuntimeConfig(bundle, model); err != nil {
+		return nil, fmt.Errorf("add config.json to runtime bundle: %w", err)
 	}
-	bundle.runtimeConfig = cfg
 
 	return bundle, nil
 }
 
-// unpackLayerToFile unpacks a single layer to the destination path.
-// It tries to use hard linking for local layers, falling back to copying for remote layers.
+// unpackLayerToFile unpacks a single layer to the destination path using hard linking.
+// It requires the layer to be a local layer with a file path (pathProvider interface).
 func unpackLayerToFile(destPath string, layer oci.Layer) error {
 	// Try to get the layer's local path for hard linking
 	type pathProvider interface {
@@ -724,6 +711,7 @@ func unpackLayerToFile(destPath string, layer oci.Layer) error {
 
 // updateBundleFieldsFromLayer updates the bundle tracking fields based on the unpacked layer.
 func updateBundleFieldsFromLayer(bundle *Bundle, mediaType oci.MediaType, relPath string) {
+	//nolint:exhaustive // only tracking specific model-related media types
 	switch mediaType {
 	case types.MediaTypeGGUF:
 		if bundle.ggufFile == "" {

pkg/distribution/internal/bundle/unpack_test.go

@@ -8,11 +8,7 @@ import (
 
 func TestValidatePathWithinDirectory(t *testing.T) {
 	// Create a temporary directory for testing
-	baseDir, err := os.MkdirTemp("", "unpack-test-*")
-	if err != nil {
-		t.Fatalf("Failed to create temp dir: %v", err)
-	}
-	defer os.RemoveAll(baseDir)
+	baseDir := t.TempDir()
 
 	tests := []struct {
 		name        string
@@ -98,7 +94,7 @@ func TestValidatePathWithinDirectory(t *testing.T) {
 
 		// Tricky paths that might bypass naive checks
 		{
-			name:        "dot dot in middle",
+			name:        ".. in middle",
 			targetPath:  "foo/../bar/model.safetensors",
 			expectError: false,
 			description: "Path with .. that stays within directory should be valid",
@@ -127,11 +123,7 @@ func TestValidatePathWithinDirectory(t *testing.T) {
 
 func TestValidatePathWithinDirectory_RealFilesystem(t *testing.T) {
 	// Create a temporary directory structure
-	baseDir, err := os.MkdirTemp("", "unpack-realfs-test-*")
-	if err != nil {
-		t.Fatalf("Failed to create temp dir: %v", err)
-	}
-	defer os.RemoveAll(baseDir)
+	baseDir := t.TempDir()
 
 	// Create a sibling directory that attacker might try to access
 	siblingDir := filepath.Join(filepath.Dir(baseDir), "sibling-secret")
@@ -148,7 +140,7 @@ func TestValidatePathWithinDirectory_RealFilesystem(t *testing.T) {
 
 	// Try to escape to the sibling directory
 	escapePath := "../sibling-secret/secret.txt"
-	err = validatePathWithinDirectory(baseDir, escapePath)
+	err := validatePathWithinDirectory(baseDir, escapePath)
 	if err == nil {
 		t.Errorf("Expected error when attempting to escape to sibling directory, but validation passed")
 	}