Add pkg/model Resolver API for build encapsulation #2663

michaeldwan · 2026-01-29T21:39:35Z

Summary

This PR introduces a new pkg/model package that provides a unified API for building, inspecting, and pulling Cog models. This encapsulates the build process behind a clean interface.

Key Changes

New pkg/model package with:
- ParsedRef for validated image reference parsing
- Image struct with Cog label accessors and convenience methods
- Model struct representing a validated Cog model with config and schema
- Source struct for model source (directory + config)
- BuildOptions for build configuration
- Factory interface for pluggable build backends
- Resolver for orchestrating builds and image loading
API:
- resolver.Build(ctx, source, opts) - Build a model from source
- resolver.Inspect(ctx, ref) - Inspect a local or remote image
- resolver.InspectByID(ctx, id) - Inspect by image ID
- resolver.Pull(ctx, ref) - Ensure image is locally available
Sentinel errors ErrNotCogModel and ErrNotFound for proper error handling
CLI migration - All CLI commands migrated to use the new Resolver:
- build.go
- push.go
- run.go
- predict.go
- train.go
- serve.go

Testing

go test ./pkg/model/... -v
go test ./pkg/cli/... -v

Introduces the model package as the foundation for the Build API Encapsulation epic (cog-fuo). This first phase adds: - ParsedRef: wrapper around go-containerregistry/pkg/name.Reference - ParseRef(): validates and parses image references with options - ParseOption functions: Insecure(), WithDefaultRegistry(), WithDefaultTag() - Methods: Registry(), Repository(), Tag(), Digest(), IsTag(), IsDigest(), IsReplicate() The lazy method design queries the underlying reference on-demand, ensuring IsReplicate() always reflects the current registry host setting.

Source wraps config + projectDir for build inputs. BuildOptions consolidates all build flags with WithDefaults().

Factory interface enables pluggable build backends. DockerfileFactory wraps existing image.Build() function.

- Resolver orchestrates building and loading Models - Load() with functional options: LocalOnly, RemoteOnly, PreferRemote, WithPlatform - LoadByID() for stable image ID-based lookups (short or full SHA) - Build() delegates to Factory and inspects result - Default PreferLocal behavior with smart fallback (only on not-found errors) - Comprehensive tests with mocked docker/registry

Implements the Ref pattern for declarative model resolution: - Ref interface with internal resolve() method - TagRef/FromTag: resolves via Load (local-first, fallback to remote) - LocalRef/FromLocal: resolves from local docker daemon only - RemoteRef/FromRemote: resolves from remote registry only - BuildRef/FromBuild: resolves by building from Source - Resolver.Resolve(ctx, ref) dispatches to appropriate resolution

- ImageBuild now returns (string, error) with manifest digest - image.Build() returns the final image ID from label-adding step - DockerfileFactory.Build() populates Image.Digest - Resolver.Build() uses digest for inspection, falls back to tag This fixes instability where tag resolution could pull from remote registry instead of using the locally-built image.

Refactor build.go to use the new model.Resolver.Build() API instead of directly calling image.Build() with 19 parameters. This encapsulates build complexity behind the Resolver abstraction. Changes: - Replace config.GetConfig() with model.NewSource() - Replace image.Build() call with resolver.Build() using BuildOptions - Use m.ImageRef() for output message instead of input imageName - Add defensive nil checks for src.Config.Build

…dels - Rename Load() to Inspect() for clearer semantics (metadata only, no pull) - Add Pull() method that ensures image is locally available - Add sentinel errors: ErrNotCogModel, ErrNotFound - Validate Cog model labels in both local and remote inspection - Update registry.Inspect() to fetch config blob for labels - Support multi-platform indexes by picking default image for labels - Migrate CLI commands (predict, train) to use resolver.Pull() - Add BuildBase support for dev mode base image builds

- Add Image.ParsedConfig() - parses cog config JSON from labels - Add Image.ParsedOpenAPISchema() - parses OpenAPI schema from labels - Add Image.ToModel() - converts Image to Model by parsing labels - Refactor resolver to use new methods, eliminating duplicate parsing code - Fix: schema parsing errors now propagate instead of being swallowed

- Remove redundant inspectLocal wrapper in Pull() - Drop unused Model fields (Weights, Runtime, GitCommit, GitTag) - Fix pickDefaultImage to log errors instead of silently swallowing - Deprecate Fast field in BuildOptions (read from config at build time) - Fix test mocks to return errors instead of panic

Copilot

Pull request overview

This PR adds a new pkg/model package that centralizes Cog model build/inspect/pull behavior behind a Resolver API, and migrates CLI commands to use it. It also extends registry inspection to expose image labels (including for multi-platform images) and updates the Docker build interface to return an image ID.

Changes:

Introduce pkg/model with ParsedRef, Source, Image, Model, Factory, and Resolver (plus tests).
Migrate CLI commands to build/inspect/pull via the new model.Resolver.
Plumb image IDs/labels through Docker build and registry inspect to support model metadata extraction.

Reviewed changes

Copilot reviewed 32 out of 32 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
pkg/registry/registry_client.go	Adds label extraction (including default-manifest selection for indexes) to support model inspection.
pkg/registry/manifest_result.go	Extends manifest result to include labels.
pkg/model/source.go	Adds `Source` abstraction for project dir + parsed config.
pkg/model/source_test.go	Tests `Source` constructors.
pkg/model/ref.go	Adds validated/parsed image reference type and parsing options.
pkg/model/ref_test.go	Tests reference parsing behavior.
pkg/model/ref_types.go	Adds declarative `Ref` types and `Resolver.Resolve` dispatcher.
pkg/model/ref_types_test.go	Tests `Ref` resolution behavior across tag/local/remote/build refs.
pkg/model/image.go	Adds `Image` with Cog label accessors and parsing helpers.
pkg/model/image_test.go	Tests image label accessors and parsing/to-model behavior.
pkg/model/model.go	Adds `Model` wrapper with convenience methods (GPU/fast/schema/image ref).
pkg/model/model_test.go	Tests `Model` convenience methods.
pkg/model/options.go	Adds build/base-build option structs and defaulting.
pkg/model/options_test.go	Tests build option defaulting behavior.
pkg/model/errors.go	Adds sentinel errors for “not found” and “not a Cog model”.
pkg/model/errors_test.go	Tests sentinel error wrapping semantics.
pkg/model/factory.go	Adds build backend interface + Dockerfile-backed default factory.
pkg/model/factory_test.go	Tests factory wiring and interface compliance.
pkg/model/resolver.go	Adds `Resolver` orchestration for local/remote inspect, pull, build, build-base.
pkg/model/resolver_test.go	Comprehensive tests for resolver behavior and fallback rules.
pkg/image/build.go	Changes build functions to return image IDs and propagates through label-adding step.
pkg/docker/command/command.go	Updates `ImageBuild` interface to return an image ID.
pkg/docker/docker.go	Implements `ImageBuild` returning the build-produced image digest.
pkg/docker/dockertest/mock_command.go	Updates mock to match new `ImageBuild` signature.
pkg/docker/dockertest/command_mocks.go	Updates generated mocks for new `ImageBuild` signature/return values.
pkg/cli/build.go	Migrates build command to `model.Resolver` and factors flag→BuildOptions helpers.
pkg/cli/push.go	Migrates push workflow to `model.Resolver` build output.
pkg/cli/run.go	Migrates run build/base-build selection to `model.Resolver`.
pkg/cli/predict.go	Migrates predict build/pull/inspect flow to `model.Resolver`.
pkg/cli/train.go	Migrates train build/pull/inspect flow to `model.Resolver`.
pkg/cli/serve.go	Migrates serve base-build flow to `model.Resolver`.
pkg/cli/baseimage.go	Adapts to new `ImageBuild` return signature.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-01-29T22:31:31Z

pkg/registry/registry_client.go

+	desc, err := remote.Get(digestRef, remote.WithAuthFromKeychain(authn.DefaultKeychain))
+	if err != nil {
+		console.Debugf("pickDefaultImage: failed to fetch image %s: %v", digestRef.String(), err)
+		return nil
+	}


pickDefaultImage performs a remote.Get without using the caller's context, so cancellations/timeouts on Inspect(ctx, ...) won’t abort this request. Consider adding a ctx parameter to pickDefaultImage and passing remote.WithContext(ctx) (and matching options) to remote.Get.

Copilot · 2026-01-29T22:31:31Z

pkg/docker/docker.go

 	}

-	console.Debugf("image digest %s", res.ExporterResponse[exptypes.ExporterImageDigestKey])
+	imageID := res.ExporterResponse[exptypes.ExporterImageDigestKey]


ImageBuild now returns the image ID, but this implementation returns res.ExporterResponse[ExporterImageDigestKey] without checking that the key is present/non-empty. Consider validating the value and returning an error if it’s missing to avoid propagating an empty image ID to callers.

Suggested change

imageID := res.ExporterResponse[exptypes.ExporterImageDigestKey]

imageID := res.ExporterResponse[exptypes.ExporterImageDigestKey]

if imageID == "" {

return "", fmt.Errorf("buildkit did not return an image digest")

}

Copilot · 2026-01-29T22:31:32Z

pkg/registry/registry_client.go

+			if defaultImg := pickDefaultImage(ref, indexManifest); defaultImg != nil {
+				if configFile, err := defaultImg.ConfigFile(); err == nil {
+					result.Labels = configFile.Config.Labels
+				}
+			}


For multi-platform indexes, failing to fetch the default manifest’s config currently just leaves result.Labels nil and returns success. Callers that rely on labels (e.g., Cog model inspection) will then misclassify the image as not a Cog model. Consider propagating an error when the default manifest/config fetch fails (or at least distinguishing "couldn’t fetch labels" from "not a Cog model"), and optionally populating result.Config based on the chosen manifest too.

Copilot · 2026-01-29T22:31:32Z

pkg/model/resolver.go

+// Build creates a Model by building from source.
+func (r *Resolver) Build(ctx context.Context, src *Source, opts BuildOptions) (*Model, error) {
+	opts = opts.WithDefaults(src)
+
+	img, err := r.factory.Build(ctx, src, opts)
+	if err != nil {
+		return nil, err


Resolver.Build/BuildBase assume src is non-nil and (for Build) that src.Config is non-nil; BuildRef explicitly allows a nil Source, which will currently panic at opts.WithDefaults(src) / downstream build calls. Add explicit validation (e.g., return a wrapped error when src == nil, src.ProjectDir == "", or src.Config == nil) so callers get a proper error instead of a panic.

Copilot · 2026-01-29T22:31:32Z

pkg/model/errors.go

+// Sentinel errors for Resolver operations.
+var (
+	// ErrNotCogModel indicates the image exists but is not a valid Cog model.
+	// This occurs when the image lacks the required org.cogmodel.config label.


This comment mentions the "org.cogmodel.config" label, but the codebase uses the "run.cog.config" label (global.LabelNamespace + "config"). Please update the comment to match the actual required label key to avoid confusing API consumers.

Suggested change

// This occurs when the image lacks the required org.cogmodel.config label.

// This occurs when the image lacks the required run.cog.config label.

Copilot · 2026-01-29T22:31:32Z

pkg/registry/registry_client.go

+	digestRef, err := name.NewDigest(ref.Context().Name() + "@" + targetDigest)
+	if err != nil {
+		console.Debugf("pickDefaultImage: failed to create digest ref: %v", err)
+		return nil
+	}


pickDefaultImage constructs a digest reference without carrying over the same name options used for the original reference (this client parses with name.Insecure). Consider passing name.Insecure (and any other relevant options) to name.NewDigest so HTTP/insecure registries behave consistently.

- Pass context to pickDefaultImage for cancellation support - Add name.Insecure option to NewDigest for HTTP registry compatibility - Validate ImageBuild returns non-empty image ID from buildkit - Fix comment: org.cogmodel.config -> run.cog.config - Add nil validation to Resolver.Build/BuildBase methods

michaeldwan added 12 commits January 29, 2026 14:43

Add pkg/model/image.go with Image struct and Cog label helpers

987f989

Add pkg/model/model.go with Model struct and related types

c59309d

Add pkg/model Source and BuildOptions types

6b53b2f

Source wraps config + projectDir for build inputs. BuildOptions consolidates all build flags with WithDefaults().

Add pkg/model Factory interface and DockerfileFactory

d679149

Factory interface enables pluggable build backends. DockerfileFactory wraps existing image.Build() function.

michaeldwan force-pushed the md/model-resolver branch from bbb432c to 0c88da8 Compare January 29, 2026 21:58

michaeldwan marked this pull request as ready for review January 29, 2026 22:18

michaeldwan requested a review from a team as a code owner January 29, 2026 22:18

michaeldwan requested review from markphelps and tempusfrangit January 29, 2026 22:18

michaeldwan changed the title ~~WIP: Add pkg/model Resolver API for build encapsulation~~ Add pkg/model Resolver API for build encapsulation Jan 29, 2026

michaeldwan requested a review from Copilot January 29, 2026 22:20

Copilot started reviewing on behalf of michaeldwan January 29, 2026 22:20 View session

Copilot AI reviewed Jan 29, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add pkg/model Resolver API for build encapsulation #2663

Add pkg/model Resolver API for build encapsulation #2663

Uh oh!

michaeldwan commented Jan 29, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Jan 29, 2026

Uh oh!

Copilot AI Jan 29, 2026

Uh oh!

Copilot AI Jan 29, 2026

Uh oh!

Copilot AI Jan 29, 2026

Uh oh!

Copilot AI Jan 29, 2026

Uh oh!

Copilot AI Jan 29, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	// This occurs when the image lacks the required org.cogmodel.config label.
	// This occurs when the image lacks the required run.cog.config label.

Add pkg/model Resolver API for build encapsulation #2663

Are you sure you want to change the base?

Add pkg/model Resolver API for build encapsulation #2663

Uh oh!

Conversation

michaeldwan commented Jan 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Key Changes

Testing

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Jan 29, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 29, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 29, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 29, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 29, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 29, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

michaeldwan commented Jan 29, 2026 •

edited

Loading