Exposes Packaging SDK (docker#74)

* Provides API for packaging models

- introduces builder package for building artifacts
- extracts registry packge for direct interaction with OCI registries

Signed-off-by: Emily Casey <[email protected]>

* cleanup

Signed-off-by: Emily Casey <[email protected]>

* slim down artifact interface

Signed-off-by: Emily Casey <[email protected]>

* cleanup

Signed-off-by: Emily Casey <[email protected]>

---------

Signed-off-by: Emily Casey <[email protected]>

Author: ekcasey

pkg/distribution/builder/builder.go

@@ -0,0 +1,49 @@
+package builder
+
+import (
+	"context"
+	"fmt"
+	"io"
+
+	"github.com/docker/model-distribution/internal/gguf"
+	"github.com/docker/model-distribution/internal/mutate"
+	"github.com/docker/model-distribution/internal/partial"
+	"github.com/docker/model-distribution/types"
+)
+
+// Builder builds a model artifact
+type Builder struct {
+	model types.ModelArtifact
+}
+
+// FromGGUF returns a *Builder that builds a model artifacts from a GGUF file
+func FromGGUF(path string) (*Builder, error) {
+	mdl, err := gguf.NewModel(path)
+	if err != nil {
+		return nil, err
+	}
+	return &Builder{
+		model: mdl,
+	}, nil
+}
+
+// WithLicense adds a license file to the artifact
+func (b *Builder) WithLicense(path string) (*Builder, error) {
+	licenseLayer, err := partial.NewLayer(path, types.MediaTypeLicense)
+	if err != nil {
+		return nil, fmt.Errorf("license layer from %q: %w", path, err)
+	}
+	return &Builder{
+		model: mutate.AppendLayers(b.model, licenseLayer),
+	}, nil
+}
+
+// Target represents a build target
+type Target interface {
+	Write(context.Context, types.ModelArtifact, io.Writer) error
+}
+
+// Build finalizes the artifact and writes it to the given target, reporting progress to the given writer
+func (b *Builder) Build(ctx context.Context, target Target, pw io.Writer) error {
+	return target.Write(ctx, b.model, pw)
+}

pkg/distribution/distribution/client.go

@@ -6,27 +6,20 @@ import (
 	"io"
 	"net/http"
 	"os"
-	"strings"
 
-	"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"
 	"github.com/sirupsen/logrus"
 
+	"github.com/docker/model-distribution/internal/progress"
 	"github.com/docker/model-distribution/internal/store"
+	"github.com/docker/model-distribution/registry"
 	"github.com/docker/model-distribution/types"
 )
 
-const (
-	defaultUserAgent = "model-distribution"
-)
-
 // Client provides model distribution functionality
 type Client struct {
-	store         *store.LocalStore
-	log           *logrus.Entry
-	remoteOptions []remote.Option
+	store    *store.LocalStore
+	log      *logrus.Entry
+	registry *registry.Client
 }
 
 // GetStorePath returns the root path where models are stored
@@ -84,8 +77,8 @@ func WithUserAgent(ua string) Option {
 func defaultOptions() *options {
 	return &options{
 		logger:    logrus.NewEntry(logrus.StandardLogger()),
-		transport: remote.DefaultTransport,
-		userAgent: defaultUserAgent,
+		transport: registry.DefaultTransport,
+		userAgent: registry.DefaultUserAgent,
 	}
 }
 
@@ -111,50 +104,29 @@ func NewClient(opts ...Option) (*Client, error) {
 	return &Client{
 		store: s,
 		log:   options.logger,
-		remoteOptions: []remote.Option{
-			remote.WithAuthFromKeychain(authn.DefaultKeychain),
-			remote.WithTransport(options.transport),
-			remote.WithUserAgent(options.userAgent),
-		},
+		registry: registry.NewClient(
+			registry.WithTransport(options.transport),
+			registry.WithUserAgent(options.userAgent),
+		),
 	}, nil
 }
 
 // PullModel pulls a model from a registry and returns the local file path
 func (c *Client) PullModel(ctx context.Context, reference string, progressWriter io.Writer) error {
 	c.log.Infoln("Starting model pull:", reference)
 
-	// Parse the reference
-	ref, err := name.ParseReference(reference)
-	if err != nil {
-		return NewReferenceError(reference, err)
-	}
-
-	// First, check the remote registry for the model's digest
-	c.log.Infoln("Checking remote registry for model:", reference)
-	opts := append([]remote.Option{remote.WithContext(ctx)}, c.remoteOptions...)
-	remoteImg, err := remote.Image(ref, opts...)
+	remoteModel, err := c.registry.Model(ctx, reference)
 	if err != nil {
-		errStr := err.Error()
-		if strings.Contains(errStr, "UNAUTHORIZED") {
-			return NewPullError(reference, "UNAUTHORIZED", "Authentication required for this model", err)
-		}
-		if strings.Contains(errStr, "MANIFEST_UNKNOWN") {
-			return NewPullError(reference, "MANIFEST_UNKNOWN", "Model not found", err)
-		}
-		if strings.Contains(errStr, "NAME_UNKNOWN") {
-			return NewPullError(reference, "NAME_UNKNOWN", "Repository not found", err)
-		}
-		c.log.Errorln("Failed to check remote image:", err, "reference:", reference)
-		return NewPullError(reference, "UNKNOWN", err.Error(), err)
+		return fmt.Errorf("reading model from registry: %w", err)
 	}
 
 	//Check for supported type
-	if err := checkCompat(remoteImg); err != nil {
+	if err := checkCompat(remoteModel); err != nil {
 		return err
 	}
 
 	// Get the remote image digest
-	remoteDigest, err := remoteImg.Digest()
+	remoteDigest, err := remoteModel.Digest()
 	if err != nil {
 		c.log.Errorln("Failed to get remote image digest:", err)
 		return fmt.Errorf("getting remote image digest: %w", err)
@@ -178,7 +150,7 @@ func (c *Client) PullModel(ctx context.Context, reference string, progressWriter
 
 		// Report progress for local model
 		size := fileInfo.Size()
-		err = writeSuccess(progressWriter, fmt.Sprintf("Using cached model: %.2f MB", float64(size)/1024/1024))
+		err = progress.WriteSuccess(progressWriter, fmt.Sprintf("Using cached model: %.2f MB", float64(size)/1024/1024))
 		if err != nil {
 			c.log.Warnf("Writing progress: %v", err)
 			// If we fail to write progress, don't try again
@@ -196,23 +168,23 @@ func (c *Client) PullModel(ctx context.Context, reference string, progressWriter
 
 	// Model doesn't exist in local store or digests don't match, pull from remote
 
-	pr := newProgressReporter(progressWriter, pullMsg)
+	pr := progress.NewProgressReporter(progressWriter, progress.PullMsg)
 	defer func() {
 		if err := pr.Wait(); err != nil {
 			c.log.Warnf("Failed to write progress: %v", err)
 		}
 	}()
 
-	if err = c.store.Write(remoteImg, []string{reference}, pr.updates()); err != nil {
-		if writeErr := writeError(progressWriter, fmt.Sprintf("Error: %s", err.Error())); writeErr != nil {
+	if err = c.store.Write(remoteModel, []string{reference}, pr.Updates()); err != nil {
+		if writeErr := progress.WriteError(progressWriter, fmt.Sprintf("Error: %s", err.Error())); writeErr != nil {
 			c.log.Warnf("Failed to write error message: %v", writeErr)
 			// If we fail to write error message, don't try again
 			progressWriter = nil
 		}
 		return fmt.Errorf("writing image to store: %w", err)
 	}
 
-	if err := writeSuccess(progressWriter, "Model pulled successfully"); err != nil {
+	if err := progress.WriteSuccess(progressWriter, "Model pulled successfully"); err != nil {
 		c.log.Warnf("Failed to write success message: %v", err)
 		// If we fail to write success message, don't try again
 		progressWriter = nil
@@ -307,9 +279,9 @@ func (c *Client) Tag(source string, target string) error {
 // PushModel pushes a tagged model from the content store to the registry.
 func (c *Client) PushModel(ctx context.Context, tag string, progressWriter io.Writer) (err error) {
 	// Parse the tag
-	ref, err := name.NewTag(tag)
+	target, err := c.registry.NewTarget(tag)
 	if err != nil {
-		return fmt.Errorf("invalid tag %q: %w", tag, err)
+		return fmt.Errorf("new tag: %w", err)
 	}
 
 	// Get the model from the store
@@ -320,36 +292,23 @@ func (c *Client) PushModel(ctx context.Context, tag string, progressWriter io.Wr
 
 	// Push the model
 	c.log.Infoln("Pushing model:", tag)
-
-	pr := newProgressReporter(progressWriter, pushMsg)
-	defer func() {
-		if err := pr.Wait(); err != nil {
-			c.log.Warnf("Failed to write progress: %v", err)
-		}
-	}()
-
-	opts := append([]remote.Option{
-		remote.WithContext(ctx),
-		remote.WithProgress(pr.updates()),
-	}, c.remoteOptions...)
-
-	if err := remote.Write(ref, mdl, opts...); err != nil {
+	if err := target.Write(ctx, mdl, progressWriter); err != nil {
 		c.log.Errorln("Failed to push image:", err, "reference:", tag)
-		if writeErr := writeError(progressWriter, fmt.Sprintf("Error: %s", err.Error())); writeErr != nil {
+		if writeErr := progress.WriteError(progressWriter, fmt.Sprintf("Error: %s", err.Error())); writeErr != nil {
 			c.log.Warnf("Failed to write error message: %v", writeErr)
 		}
 		return fmt.Errorf("pushing image: %w", err)
 	}
 
 	c.log.Infoln("Successfully pushed model:", tag)
-	if err := writeSuccess(progressWriter, "Model pushed successfully"); err != nil {
+	if err := progress.WriteSuccess(progressWriter, "Model pushed successfully"); err != nil {
 		c.log.Warnf("Failed to write success message: %v", err)
 	}
 
 	return nil
 }
 
-func checkCompat(image v1.Image) error {
+func checkCompat(image types.ModelArtifact) error {
 	manifest, err := image.Manifest()
 	if err != nil {
 		return err

pkg/distribution/distribution/client_test.go

@@ -16,13 +16,15 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/docker/model-distribution/internal/progress"
 	"github.com/google/go-containerregistry/pkg/name"
 	"github.com/google/go-containerregistry/pkg/registry"
 	"github.com/google/go-containerregistry/pkg/v1/remote"
 	"github.com/sirupsen/logrus"
 
 	"github.com/docker/model-distribution/internal/gguf"
 	"github.com/docker/model-distribution/internal/mutate"
+	mdregistry "github.com/docker/model-distribution/registry"
 )
 
 var (
@@ -158,8 +160,8 @@ func TestClientPullModel(t *testing.T) {
 			t.Fatal("Expected error for non-existent model, got nil")
 		}
 
-		// Verify it's a PullError
-		var pullErr *PullError
+		// Verify it's a registry.Error
+		var pullErr *mdregistry.Error
 		ok := errors.As(err, &pullErr)
 		if !ok {
 			t.Fatalf("Expected PullError, got %T", err)
@@ -178,7 +180,7 @@ func TestClientPullModel(t *testing.T) {
 		if pullErr.Err == nil {
 			t.Error("Expected underlying error to be non-nil")
 		}
-		if !errors.Is(pullErr, ErrModelNotFound) {
+		if !errors.Is(pullErr, mdregistry.ErrModelNotFound) {
 			t.Errorf("Expected underlying error to match ErrModelNotFound, got %v", pullErr.Err)
 		}
 	})
@@ -422,11 +424,11 @@ func TestClientPullModel(t *testing.T) {
 		}
 
 		// Parse progress output as JSON
-		var messages []ProgressMessage
+		var messages []progress.ProgressMessage
 		scanner := bufio.NewScanner(&progressBuffer)
 		for scanner.Scan() {
 			line := scanner.Text()
-			var msg ProgressMessage
+			var msg progress.ProgressMessage
 			if err := json.Unmarshal([]byte(line), &msg); err != nil {
 				t.Fatalf("Failed to parse JSON progress message: %v, line: %s", err, line)
 			}
@@ -496,10 +498,9 @@ func TestClientPullModel(t *testing.T) {
 			t.Fatal("Expected error for non-existent model, got nil")
 		}
 
-		// Verify it's a PullError
-		var pullErr *PullError
-		if !errors.As(err, &pullErr) {
-			t.Fatalf("Expected PullError, got %T", err)
+		// Verify it matches registry.ErrModelNotFound
+		if !errors.Is(err, mdregistry.ErrModelNotFound) {
+			t.Fatalf("Expected registry.ErrModelNotFound, got %T", err)
 		}
 
 		// No JSON messages should be in the buffer for this error case
@@ -815,18 +816,8 @@ func TestNewReferenceError(t *testing.T) {
 		t.Fatal("Expected error for invalid reference, got nil")
 	}
 
-	// Verify it's a ReferenceError
-	refErr, ok := err.(*ReferenceError)
-	if !ok {
-		t.Fatalf("Expected ReferenceError, got %T", err)
-	}
-
-	// Verify error fields
-	if refErr.Reference != invalidRef {
-		t.Errorf("Expected reference %q, got %q", invalidRef, refErr.Reference)
-	}
-	if refErr.Err == nil {
-		t.Error("Expected underlying error to be non-nil")
+	if !errors.Is(err, ErrInvalidReference) {
+		t.Fatalf("Expected error to match sentinel invalid reference error, got %v", err)
 	}
 }
 

pkg/distribution/distribution/errors.go

@@ -5,13 +5,13 @@ import (
 	"fmt"
 
 	"github.com/docker/model-distribution/internal/store"
+	"github.com/docker/model-distribution/registry"
 	"github.com/docker/model-distribution/types"
 )
 
 var (
-	ErrInvalidReference     = errors.New("invalid model reference")
-	ErrModelNotFound        = store.ErrModelNotFound
-	ErrUnauthorized         = errors.New("unauthorized access to model")
+	ErrInvalidReference     = registry.ErrInvalidReference
+	ErrModelNotFound        = store.ErrModelNotFound // model not found in store
 	ErrUnsupportedMediaType = errors.New(fmt.Sprintf(
 		"client supports only models of type %q and older - try upgrading",
 		types.MediaTypeModelConfigV01,
@@ -37,51 +37,3 @@ func (e *ReferenceError) Unwrap() error {
 func (e *ReferenceError) Is(target error) bool {
 	return target == ErrInvalidReference
 }
-
-// PullError represents an error that occurs when pulling a model
-type PullError struct {
-	Reference string
-	// Code should be one of error codes defined in the distribution spec
-	// (see https://github.com/opencontainers/distribution-spec/blob/583e014d15418d839d67f68152bc2c83821770e0/spec.md#error-codes)
-	Code    string
-	Message string
-	Err     error
-}
-
-func (e *PullError) Error() string {
-	return fmt.Sprintf("failed to pull model %q: %s - %s", e.Reference, e.Code, e.Message)
-}
-
-func (e *PullError) Unwrap() error {
-	return e.Err
-}
-
-// Is implements error matching for PullError
-func (e *PullError) Is(target error) bool {
-	switch target {
-	case ErrModelNotFound:
-		return e.Code == "MANIFEST_UNKNOWN" || e.Code == "NAME_UNKNOWN"
-	case ErrUnauthorized:
-		return e.Code == "UNAUTHORIZED"
-	default:
-		return false
-	}
-}
-
-// NewReferenceError creates a new ReferenceError
-func NewReferenceError(reference string, err error) error {
-	return &ReferenceError{
-		Reference: reference,
-		Err:       err,
-	}
-}
-
-// NewPullError creates a new PullError
-func NewPullError(reference, code, message string, err error) error {
-	return &PullError{
-		Reference: reference,
-		Code:      code,
-		Message:   message,
-		Err:       err,
-	}
-}