OCI bundle format for model weights #2676

michaeldwan · 2026-02-02T22:03:03Z

Summary

Adds OCI bundle format support for separating model weights from container images. This enables weight deduplication across model versions and faster cold starts by allowing weights to be cached and shared.

What's Changed

New Types & Formats

ModelFormat (pkg/model/format.go) - Enum for image formats: FormatStandalone (traditional) and FormatBundle (OCI index with weights)
WeightsManifest / WeightFile (pkg/model/weights.go) - Types for weight file metadata (digests, sizes, destinations)
WeightsLock (pkg/model/weights_lock.go) - Lockfile format mapping weight identifiers to deployment blob metadata
Index / IndexManifest (pkg/model/index.go) - OCI Image Index types with manifest type detection

OCI Index Building

IndexFactory (pkg/model/index_factory.go) - Builds OCI Image Indexes containing:
- Model image manifest (platform-specific)
- Weights artifact manifest (platform: unknown/unknown, with vnd.cog.reference.type: weights annotation)
WeightsArtifactBuilder - Creates OCI artifacts from weight files with proper media types and annotations
IndexBuilder - Combines model image + weights artifact into an OCI index

Push Flow

Pusher interface (pkg/model/pusher.go) - Strategy pattern for push operations
- ImagePusher - Standard docker push for standalone images
- BundlePusher - OCI index push with weights artifact
Resolver.Push() - Selects appropriate pusher based on Model.ImageFormat

Build Integration

BuildOptions.ImageFormat - Specify bundle vs standalone format
BuildOptions.WeightsLockPath - Custom path for weights.lock file
Resolver loads weights.lock and attaches WeightsManifest to built models

Testing

weights-lock-gen (tools/weights-lock-gen/) - CLI tool to generate random weight files for testing
Mock registry (integration-tests/harness/) - Commands for testing OCI bundle push/inspect:
- registry-start / registry-stop - Ephemeral Docker registry
- registry-inspect - Inspect pushed manifests
- mock-weights - Generate test weight files and lockfile
Integration tests for OCI bundle build and push flow

Cleanup

Removed deprecated fast/monobase build system (pkg/docker/fast_push.go, pkg/monobeam/, etc.)
Removed Source field from WeightFile (lockfile maps name→blob, source comes from cog.yaml)

Design Notes

Weight names are identifiers, not filenames (e.g., "llama-3.1-8b" not "model.safetensors")
File paths passed separately via filePaths map[string]string parameter
Platform unknown/unknown used for non-platform-specific artifacts (weights)
Annotation vnd.cog.reference.type: weights identifies weights manifests in indexes

Testing

# Unit tests
go test ./pkg/model/... -v

# Integration tests (requires Docker)
make test-integration

Next Steps

Wire cog build --oci-index flag to use bundle format
Implement declarative weights in cog.yaml
Add weight deduplication logic
Chunked uploads for large weight files

markphelps

one minor nit but lgtm!

pkg/model/resolver.go

markphelps

constantly approve

…istry Implements WeightsArtifactBuilder for creating OCI artifacts containing model weights. The builder: - Creates OCI-format manifests (application/vnd.oci.image.manifest.v1+json) - Sets artifactType to application/vnd.cog.weights.v1 - Adds layers with custom media types and annotations for weight files - Uses mutate.Addendum to properly attach layer annotations Layer annotations include: - vnd.cog.weights.name: original filename - vnd.cog.weights.dest: container path - vnd.cog.weights.source: origin URL (optional) - vnd.cog.weights.digest.original: uncompressed digest - vnd.cog.weights.size.uncompressed: uncompressed size Test coverage: 93.8%

Add helper functions to detect and handle OCI indexes: - isOCIIndex: checks if ManifestResult is an OCI Image Index - findWeightsManifest: finds weights manifest by annotation - findImageManifest: finds model image manifest with platform filtering Also: - Add Annotations field to PlatformManifest struct - Populate annotations when reading OCI index manifests - Add PushImage/PushIndex to registry Client interface for OCI index push support

- Move WeightsArtifactBuilder, IndexBuilder, IndexFactory from pkg/ociartifact to pkg/model/index_factory.go - Delete pkg/ociartifact/ package (code duplication eliminated) - Add OCI index push flow in pkg/cli/push.go (COG_OCI_INDEX=1 env var) - Add index detection helpers in resolver for loading OCI indexes - Simplify BuildOptions: remove WeightsLockPath (hardcoded to weights.lock) - Add Platform.Variant field to pkg/model/image.go - Add push_test.go placeholder for registry push tests

Replace the --files flag with random file generation: - Add --count (-n) flag for number of files (default: 3) - Add --min-size and --max-size flags (default: 25mb-50mb) - Add --output-dir to optionally keep generated files - Files are named weights-001.bin, weights-002.bin, etc. - Generate random data in 1MB chunks for memory efficiency This simplifies testing the OCI bundle format without needing real weight files.

- Add ModelImageFormat enum (FormatStandalone, FormatBundle) in format.go - Update Model struct: replace Format with ImageFormat field - Update BuildOptions: replace OCIIndex bool with ImageFormat - Create Pusher interface with ImagePusher and BundlePusher implementations - Add Resolver.Push() that delegates to appropriate pusher based on format - Add BuildWeightsArtifactFromManifest to IndexFactory - Simplify pkg/cli/push.go to use resolver.Push() for bundles - Move ImageFormatFromEnv() to model package The push logic is now cleanly separated: - ImagePusher: Simple docker push for standalone images - BundlePusher: Full OCI Index workflow (push image, build weights artifact, build index, push index) Closes cog-cz1

- Add inline mock weights generator to avoid pkg/wheels init panic - Add StartTestRegistryWithCleanup for harness (no *testing.T required) - Add registry-start, registry-inspect, docker-push, mock-weights commands - Add oci_bundle_build.txtar and oci_bundle_push.txtar integration tests - Update push_test.go to use testcontainers instead of TEST_REGISTRY env

…ht resolution - Remove Source field from WeightFile - lockfile maps name→blob, not source - Change Name semantics to identifier/handle (e.g., 'model-v1') not filename - Add filePaths map[string]string parameter to AddLayersFromLock for file location - Add context cancellation support to AddLayersFromLock for long operations - Add AnnotationValueWeights constant for 'weights' annotation value - Improve registry push error messages to include reference - Remove binary test files (~109MB) that were accidentally committed - Update all tests to use new API and identifier-style names

- Add PlatformUnknown constant for non-platform-specific artifacts - Use AnnotationValueWeights constant instead of "weights" string literal - Update all usages in index_factory.go, resolver.go, and tests

markphelps

i will not approve again

michaeldwan force-pushed the md/oci-bundle branch 3 times, most recently from a8a88f6 to 0fe2ec4 Compare February 3, 2026 21:00

michaeldwan marked this pull request as ready for review February 3, 2026 21:41

michaeldwan requested a review from a team as a code owner February 3, 2026 21:41

markphelps approved these changes Feb 4, 2026

View reviewed changes

pkg/model/resolver.go Outdated Show resolved Hide resolved

tempusfrangit approved these changes Feb 4, 2026

View reviewed changes

markphelps approved these changes Feb 4, 2026

View reviewed changes

michaeldwan changed the title ~~WIP: OCI bundle format for model weights~~ OCI bundle format for model weights Feb 4, 2026

tempusfrangit approved these changes Feb 4, 2026

View reviewed changes

michaeldwan added 16 commits February 4, 2026 13:35

feat(model): add ModelFormat and Index types for OCI bundle support

5d057f9

feat(model): add WeightsManifest and WeightFile types

54a883e

feat(model): add WeightsLock file format for weight metadata

5865f34

feat(tools): add weights-lock-gen for testing OCI bundle format

65b34de

feat(ociartifact): add AddLayersFromLock for file-based weight loading

335768e

feat(model): add OCIIndex and WeightsLockPath build options

f5de0ad

chore: add test weight files to gitignore

1dbc353

fix(weights-lock-gen): suppress gosec warning for test data RNG

7521e40

refactor: replace magic strings with constants for platform and weights

26be175

- Add PlatformUnknown constant for non-platform-specific artifacts - Use AnnotationValueWeights constant instead of "weights" string literal - Update all usages in index_factory.go, resolver.go, and tests

michaeldwan force-pushed the md/oci-bundle branch from 36064a1 to 26be175 Compare February 4, 2026 20:36

markphelps approved these changes Feb 4, 2026

View reviewed changes

michaeldwan merged commit d6cb965 into main Feb 4, 2026
16 checks passed

michaeldwan deleted the md/oci-bundle branch February 4, 2026 21:04

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

OCI bundle format for model weights #2676

OCI bundle format for model weights #2676

Uh oh!

michaeldwan commented Feb 2, 2026 •

edited

Loading

Uh oh!

markphelps left a comment

Uh oh!

Uh oh!

markphelps left a comment

Uh oh!

markphelps left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

OCI bundle format for model weights #2676

OCI bundle format for model weights #2676

Uh oh!

Conversation

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

Uh oh!

Summary

What's Changed

New Types & Formats

OCI Index Building

Push Flow

Build Integration

Testing

Cleanup

Design Notes

Testing

Next Steps

Uh oh!

markphelps left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

markphelps left a comment

Choose a reason for hiding this comment

Uh oh!

markphelps left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

michaeldwan commented Feb 2, 2026 •

edited

Loading