Merge pull request #791 from docker/modelpack-support

Modelpack support

Author: ericcurtin

pkg/distribution/builder/builder_test.go

@@ -10,7 +10,7 @@ import (
 	"time"
 
 	"github.com/docker/model-runner/pkg/distribution/builder"
-	"github.com/docker/model-runner/pkg/distribution/oci"
+	"github.com/docker/model-runner/pkg/distribution/internal/testutil"
 	"github.com/docker/model-runner/pkg/distribution/types"
 )
 
@@ -398,8 +398,7 @@ func TestFromModelWithAdditionalLayers(t *testing.T) {
 
 // TestFromModelErrorHandling tests that FromModel properly handles and surfaces errors from mdl.Layers()
 func TestFromModelErrorHandling(t *testing.T) {
-	// Create a mock model that fails when Layers() is called
-	mockModel := &mockFailingModel{}
+	mockModel := testutil.WithLayersError(testutil.NewGGUFArtifact(t, filepath.Join("..", "assets", "dummy.gguf")), fmt.Errorf("simulated layers error"))
 
 	// Attempt to create a builder from the failing model
 	_, err := builder.FromModel(mockModel)
@@ -424,12 +423,3 @@ func (ft *fakeTarget) Write(ctx context.Context, artifact types.ModelArtifact, w
 	ft.artifact = artifact
 	return nil
 }
-
-// mockFailingModel is a mock that fails when Layers() is called
-type mockFailingModel struct {
-	types.ModelArtifact
-}
-
-func (m *mockFailingModel) Layers() ([]oci.Layer, error) {
-	return nil, fmt.Errorf("simulated layers error")
-}

pkg/distribution/distribution/bundle_test.go

@@ -6,9 +6,7 @@ import (
 	"path/filepath"
 	"testing"
 
-	"github.com/docker/model-runner/pkg/distribution/builder"
-	"github.com/docker/model-runner/pkg/distribution/internal/mutate"
-	"github.com/docker/model-runner/pkg/distribution/internal/partial"
+	"github.com/docker/model-runner/pkg/distribution/internal/testutil"
 	"github.com/docker/model-runner/pkg/distribution/types"
 )
 
@@ -22,12 +20,7 @@ func TestBundle(t *testing.T) {
 		t.Fatalf("Failed to create client: %v", err)
 	}
 
-	// Load dummy model from assets directory
-	b, err := builder.FromPath(filepath.Join("..", "assets", "dummy.gguf"))
-	if err != nil {
-		t.Fatalf("Failed to create model: %v", err)
-	}
-	mdl := b.Model()
+	mdl := testutil.NewGGUFArtifact(t, filepath.Join("..", "assets", "dummy.gguf"))
 	singleGGUFID, err := mdl.ID()
 	if err != nil {
 		t.Fatalf("Failed to get model ID: %v", err)
@@ -36,12 +29,11 @@ func TestBundle(t *testing.T) {
 		t.Fatalf("Failed to write model to store: %v", err)
 	}
 
-	// Load model with multi-modal projector file
-	mmprojLayer, err := partial.NewLayer(filepath.Join("..", "assets", "dummy.mmproj"), types.MediaTypeMultimodalProjector)
-	if err != nil {
-		t.Fatalf("Failed to create mmproj layer: %v", err)
-	}
-	mmprojMdl := mutate.AppendLayers(mdl, mmprojLayer)
+	mmprojMdl := testutil.NewGGUFArtifact(
+		t,
+		filepath.Join("..", "assets", "dummy.gguf"),
+		testutil.Layer(filepath.Join("..", "assets", "dummy.mmproj"), types.MediaTypeMultimodalProjector),
+	)
 	mmprojMdlID, err := mmprojMdl.ID()
 	if err != nil {
 		t.Fatalf("Failed to get model ID: %v", err)
@@ -50,12 +42,11 @@ func TestBundle(t *testing.T) {
 		t.Fatalf("Failed to write model to store: %v", err)
 	}
 
-	// Load model with template file
-	templateLayer, err := partial.NewLayer(filepath.Join("..", "assets", "template.jinja"), types.MediaTypeChatTemplate)
-	if err != nil {
-		t.Fatalf("Failed to create chat template layer: %v", err)
-	}
-	templateMdl := mutate.AppendLayers(mdl, templateLayer)
+	templateMdl := testutil.NewGGUFArtifact(
+		t,
+		filepath.Join("..", "assets", "dummy.gguf"),
+		testutil.Layer(filepath.Join("..", "assets", "template.jinja"), types.MediaTypeChatTemplate),
+	)
 	templateMdlID, err := templateMdl.ID()
 	if err != nil {
 		t.Fatalf("Failed to get model ID: %v", err)
@@ -64,12 +55,12 @@ func TestBundle(t *testing.T) {
 		t.Fatalf("Failed to write model to store: %v", err)
 	}
 
-	// Load sharded dummy model from asset directory
-	shardedB, err := builder.FromPath(filepath.Join("..", "assets", "dummy-00001-of-00002.gguf"))
-	if err != nil {
-		t.Fatalf("Failed to create model: %v", err)
-	}
-	shardedMdl := shardedB.Model()
+	shardedMdl := testutil.NewDockerArtifact(
+		t,
+		types.Config{Format: types.FormatGGUF},
+		testutil.Layer(filepath.Join("..", "assets", "dummy-00001-of-00002.gguf"), types.MediaTypeGGUF),
+		testutil.Layer(filepath.Join("..", "assets", "dummy-00002-of-00002.gguf"), types.MediaTypeGGUF),
+	)
 	shardedGGUFID, err := shardedMdl.ID()
 	if err != nil {
 		t.Fatalf("Failed to get model ID: %v", err)

pkg/distribution/distribution/client.go

@@ -15,6 +15,7 @@ import (
 	"github.com/docker/model-runner/pkg/distribution/internal/mutate"
 	"github.com/docker/model-runner/pkg/distribution/internal/progress"
 	"github.com/docker/model-runner/pkg/distribution/internal/store"
+	"github.com/docker/model-runner/pkg/distribution/modelpack"
 	"github.com/docker/model-runner/pkg/distribution/oci"
 	"github.com/docker/model-runner/pkg/distribution/oci/authn"
 	"github.com/docker/model-runner/pkg/distribution/oci/remote"
@@ -777,7 +778,9 @@ func checkCompat(image types.ModelArtifact, log *slog.Logger, reference string,
 	if err != nil {
 		return err
 	}
-	if manifest.Config.MediaType != types.MediaTypeModelConfigV01 && manifest.Config.MediaType != types.MediaTypeModelConfigV02 {
+	if manifest.Config.MediaType != types.MediaTypeModelConfigV01 &&
+		manifest.Config.MediaType != types.MediaTypeModelConfigV02 &&
+		manifest.Config.MediaType != modelpack.MediaTypeModelConfigV1 {
 		return fmt.Errorf("config type %q is unsupported: %w", manifest.Config.MediaType, ErrUnsupportedMediaType)
 	}