diff --git a/architecture/06-cli.md b/architecture/06-cli.md index e1048a3f89..945b8726eb 100644 --- a/architecture/06-cli.md +++ b/architecture/06-cli.md @@ -14,7 +14,7 @@ The Cog CLI is a Go binary that provides commands for the full model lifecycle: | `cog exec` | Run arbitrary commands in a container | | `cog serve` | Start HTTP server in a container | | `cog playground` | Open a local UI for a model API | -| `cog push` | Deploy to Replicate | +| `cog push` | Publish to an OCI registry | | `cog login` | Authenticate with Replicate | ## Development Commands @@ -146,21 +146,19 @@ Key flags: ### cog push -**Job**: Build and push to Replicate. +**Job**: Build and publish an image or model bundle to an OCI registry. ```bash -cog push r8.im/username/model-name +cog push registry.example.com/username/model-name:v1 ``` -What happens: +The `image` or `model` field in `cog.yaml` selects the artifact format. A positional target changes the destination without changing that format. Image projects publish one container image. Bundle projects publish an OCI image index containing the model image and any managed weight manifests. -1. Builds image (like `cog build`) -2. Pushes to Replicate's registry -3. Registers model with Replicate API +Pushes attempt to resolve the published artifacts to digest-pinned references, which the normal terminal output renders as a tree. `--json` makes digest resolution mandatory and emits one versioned document to stdout after the push and registry-provider post-processing succeed; progress and diagnostics remain on stderr. This lets automation consume immutable references without hiding a long-running push's live output. -The image tag must be a Replicate model reference (`r8.im/owner/name`). +Registry-specific authentication, diagnostics, and post-push work sit behind the provider boundary. Replicate is one provider, not a requirement of the push path. -**Code**: `pkg/cli/push.go`, `pkg/web/` +**Code**: `pkg/cli/` owns the command and output contract; `pkg/model/` owns image and bundle publication; `pkg/registry/` and `pkg/provider/` own registry operations and provider-specific behavior. --- diff --git a/crates/Cargo.lock b/crates/Cargo.lock index 2d173f3b8f..9ae6b005cb 100644 --- a/crates/Cargo.lock +++ b/crates/Cargo.lock @@ -509,7 +509,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -2389,14 +2389,14 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] name = "rustls" -version = "0.23.42" +version = "0.23.45" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c54fcab019b409d04215d3a17cb438fd7fbf192ee61461f20f4fe18704bc138" +checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634" dependencies = [ "log", "once_cell", @@ -2446,7 +2446,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -2457,9 +2457,9 @@ checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f" [[package]] name = "rustls-webpki" -version = "0.103.13" +version = "0.103.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" dependencies = [ "ring", "rustls-pki-types", @@ -2882,7 +2882,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3676,7 +3676,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/docs/cli.md b/docs/cli.md index a15bed8be0..398c35e201 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -220,25 +220,36 @@ cog playground [flags] ## `cog push` -Build a Docker image from cog.yaml and push it to a container registry. +Build from cog.yaml and push to an OCI-compliant registry. Run 'cog login' +first when pushing to Replicate's registry (r8.im). -Cog can push to any OCI-compliant registry. When pushing to Replicate's -registry (r8.im), run 'cog login' first to authenticate. +TARGET overrides the configured destination and all COG_MODEL\* environment +variables. It doesn't change the project format: projects configured with +'image' push an image, while projects configured with 'model' push an OCI +bundle. Push targets must use tags, not digests. Untagged image targets use +Docker's default 'latest' tag; untagged bundle targets get a timestamp tag. + +With --json, Cog writes one versioned JSON result to stdout after the entire +push succeeds. Progress, warnings, and diagnostics continue on stderr. Every +reference in the result is digest-pinned. ``` -cog push [IMAGE] [flags] +cog push [TARGET] [flags] ``` **Examples** ``` - # Push to Replicate + # Push an image to Replicate cog push r8.im/your-username/my-model - # Push to any OCI registry + # Push an image to any OCI registry cog push registry.example.com/your-username/model-name - # Push with model weights in a separate layer (Replicate only) + # Push a bundle project and print its immutable references as JSON + cog push registry.example.com/your-username/model-name:v1 --json + + # Push with model weights in a separate image layer (Replicate only) cog push r8.im/your-username/my-model --separate-weights ``` @@ -247,6 +258,7 @@ cog push [IMAGE] [flags] ``` -f, --file string The name of the config file. (default "cog.yaml") -h, --help help for push + --json Output the pushed references as JSON --no-cache Do not use cache when building the image --openapi-schema string Load OpenAPI schema from a file --progress string Set type of build progress output, 'auto' (default), 'tty', 'plain', or 'quiet' (default "auto") diff --git a/docs/deploy.md b/docs/deploy.md index f028b95e28..3dba245455 100644 --- a/docs/deploy.md +++ b/docs/deploy.md @@ -22,6 +22,83 @@ cog build -t my-model The image contains your model code, dependencies, the Cog runtime, and everything in between. It serves an HTTP server on port 5000 when run. +## Push to a registry + +`cog push` builds the project and publishes it to an OCI-compliant registry: + +```console +cog push registry.example.com/acme/my-model:v1 +``` + +The project configuration decides what Cog publishes. A project with `image` in +`cog.yaml` publishes a container image. A project with `model` publishes an OCI +bundle containing the image and any managed weights. Passing a target doesn't +switch between these formats. + +The positional target is optional when `image` or `model` already supplies a +destination. When present, it overrides that configured destination and every +`COG_MODEL*` environment variable. The target must be tag-addressable -- Cog +rejects digest-pinned targets. An untagged image target uses Docker's `latest` +tag, while an untagged bundle target gets a timestamp tag generated by Cog. + +For a bundle with managed weights, use the same repository where `cog weights +import` published the weight manifests. The model tag may differ. If the +repository doesn't contain every required weight manifest, Cog fails before +pushing the image and tells you to run `cog weights import` for that repository. + +Without `--json`, a successful push prints the published references as a +human-readable tree on stderr. + +### Machine-readable output + +Use `--json` when another program needs the immutable references produced by a +push. For example, this pushes a bundle and prints its model, image, and managed +weight references: + +```console +cog push registry.example.com/acme/my-model:v1 --json +``` + +Version 1 output for a bundle with managed weights has this shape: + +```json +{ + "version": 1, + "model": "registry.example.com/acme/my-model@sha256:1111111111111111111111111111111111111111111111111111111111111111", + "image": "registry.example.com/acme/my-model@sha256:2222222222222222222222222222222222222222222222222222222222222222", + "weights": [ + { + "name": "transformer", + "reference": "registry.example.com/acme/my-model@sha256:3333333333333333333333333333333333333333333333333333333333333333" + } + ] +} +``` + +The version 1 fields are: + +- `version`: Always `1`. +- `image`: Always present. The value is the published image's complete, + digest-pinned reference. +- `model`: Present only for bundle pushes. The value is the complete, + digest-pinned bundle reference. +- `weights`: Present only when the bundle has managed weights. Each entry has a + `name` and a complete, digest-pinned `reference`. Entries keep their + `cog.yaml` order. + +JSON whitespace and object key order aren't part of the contract. + +Stdout stays empty while the command runs. Build and upload progress, warnings, +diagnostics, and registry-provider messages go to stderr. After the push and +provider post-processing both succeed, stdout receives exactly one JSON +document. JSON mode suppresses the human-readable reference tree. + +Any build, push, digest-resolution, provider, or serialization failure exits +nonzero without writing a JSON document to stdout. JSON mode requires Cog to +resolve every published reference to an immutable digest, so it can report +failure even when an image was pushed successfully. Non-JSON image pushes keep +their existing fallback when a registry can't resolve the pushed image's digest. + ## Run the model You have several options for running a built image. diff --git a/docs/environment.md b/docs/environment.md index 5eb0677f4c..fa00688c8e 100644 --- a/docs/environment.md +++ b/docs/environment.md @@ -137,7 +137,7 @@ Overrides the full model reference used by commands that need a model destinatio The value is parsed as a complete model reference (`registry/repo`, `registry/repo:tag`, or `registry/repo@digest`). If no tag is supplied, Cog generates a timestamp tag. -When `COG_MODEL` is set, it takes precedence over `COG_MODEL_REGISTRY`, `COG_MODEL_REPO`, and `COG_MODEL_TAG`. +When `COG_MODEL` is set, it takes precedence over `COG_MODEL_REGISTRY`, `COG_MODEL_REPO`, and `COG_MODEL_TAG`. A positional target passed to `cog push` takes precedence over all four variables. Although commands that inspect existing artifacts accept a digest, `cog push` rejects digest-pinned destinations because a push requires a tag. ```console $ COG_MODEL=r8.im/acme/my-model:v1 cog push diff --git a/docs/llms.txt b/docs/llms.txt index ced96c765c..8c20345eb8 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -541,25 +541,36 @@ cog playground [flags] ## `cog push` -Build a Docker image from cog.yaml and push it to a container registry. +Build from cog.yaml and push to an OCI-compliant registry. Run 'cog login' +first when pushing to Replicate's registry (r8.im). -Cog can push to any OCI-compliant registry. When pushing to Replicate's -registry (r8.im), run 'cog login' first to authenticate. +TARGET overrides the configured destination and all COG_MODEL\* environment +variables. It doesn't change the project format: projects configured with +'image' push an image, while projects configured with 'model' push an OCI +bundle. Push targets must use tags, not digests. Untagged image targets use +Docker's default 'latest' tag; untagged bundle targets get a timestamp tag. + +With --json, Cog writes one versioned JSON result to stdout after the entire +push succeeds. Progress, warnings, and diagnostics continue on stderr. Every +reference in the result is digest-pinned. ``` -cog push [IMAGE] [flags] +cog push [TARGET] [flags] ``` **Examples** ``` - # Push to Replicate + # Push an image to Replicate cog push r8.im/your-username/my-model - # Push to any OCI registry + # Push an image to any OCI registry cog push registry.example.com/your-username/model-name - # Push with model weights in a separate layer (Replicate only) + # Push a bundle project and print its immutable references as JSON + cog push registry.example.com/your-username/model-name:v1 --json + + # Push with model weights in a separate image layer (Replicate only) cog push r8.im/your-username/my-model --separate-weights ``` @@ -568,6 +579,7 @@ cog push [IMAGE] [flags] ``` -f, --file string The name of the config file. (default "cog.yaml") -h, --help help for push + --json Output the pushed references as JSON --no-cache Do not use cache when building the image --openapi-schema string Load OpenAPI schema from a file --progress string Set type of build progress output, 'auto' (default), 'tty', 'plain', or 'quiet' (default "auto") @@ -711,6 +723,83 @@ cog build -t my-model The image contains your model code, dependencies, the Cog runtime, and everything in between. It serves an HTTP server on port 5000 when run. +## Push to a registry + +`cog push` builds the project and publishes it to an OCI-compliant registry: + +```console +cog push registry.example.com/acme/my-model:v1 +``` + +The project configuration decides what Cog publishes. A project with `image` in +`cog.yaml` publishes a container image. A project with `model` publishes an OCI +bundle containing the image and any managed weights. Passing a target doesn't +switch between these formats. + +The positional target is optional when `image` or `model` already supplies a +destination. When present, it overrides that configured destination and every +`COG_MODEL*` environment variable. The target must be tag-addressable -- Cog +rejects digest-pinned targets. An untagged image target uses Docker's `latest` +tag, while an untagged bundle target gets a timestamp tag generated by Cog. + +For a bundle with managed weights, use the same repository where `cog weights +import` published the weight manifests. The model tag may differ. If the +repository doesn't contain every required weight manifest, Cog fails before +pushing the image and tells you to run `cog weights import` for that repository. + +Without `--json`, a successful push prints the published references as a +human-readable tree on stderr. + +### Machine-readable output + +Use `--json` when another program needs the immutable references produced by a +push. For example, this pushes a bundle and prints its model, image, and managed +weight references: + +```console +cog push registry.example.com/acme/my-model:v1 --json +``` + +Version 1 output for a bundle with managed weights has this shape: + +```json +{ + "version": 1, + "model": "registry.example.com/acme/my-model@sha256:1111111111111111111111111111111111111111111111111111111111111111", + "image": "registry.example.com/acme/my-model@sha256:2222222222222222222222222222222222222222222222222222222222222222", + "weights": [ + { + "name": "transformer", + "reference": "registry.example.com/acme/my-model@sha256:3333333333333333333333333333333333333333333333333333333333333333" + } + ] +} +``` + +The version 1 fields are: + +- `version`: Always `1`. +- `image`: Always present. The value is the published image's complete, + digest-pinned reference. +- `model`: Present only for bundle pushes. The value is the complete, + digest-pinned bundle reference. +- `weights`: Present only when the bundle has managed weights. Each entry has a + `name` and a complete, digest-pinned `reference`. Entries keep their + `cog.yaml` order. + +JSON whitespace and object key order aren't part of the contract. + +Stdout stays empty while the command runs. Build and upload progress, warnings, +diagnostics, and registry-provider messages go to stderr. After the push and +provider post-processing both succeed, stdout receives exactly one JSON +document. JSON mode suppresses the human-readable reference tree. + +Any build, push, digest-resolution, provider, or serialization failure exits +nonzero without writing a JSON document to stdout. JSON mode requires Cog to +resolve every published reference to an immutable digest, so it can report +failure even when an image was pushed successfully. Non-JSON image pushes keep +their existing fallback when a registry can't resolve the pushed image's digest. + ## Run the model You have several options for running a built image. @@ -1078,7 +1167,7 @@ Overrides the full model reference used by commands that need a model destinatio The value is parsed as a complete model reference (`registry/repo`, `registry/repo:tag`, or `registry/repo@digest`). If no tag is supplied, Cog generates a timestamp tag. -When `COG_MODEL` is set, it takes precedence over `COG_MODEL_REGISTRY`, `COG_MODEL_REPO`, and `COG_MODEL_TAG`. +When `COG_MODEL` is set, it takes precedence over `COG_MODEL_REGISTRY`, `COG_MODEL_REPO`, and `COG_MODEL_TAG`. A positional target passed to `cog push` takes precedence over all four variables. Although commands that inspect existing artifacts accept a digest, `cog push` rejects digest-pinned destinations because a push requires a tag. ```console $ COG_MODEL=r8.im/acme/my-model:v1 cog push @@ -4187,7 +4276,21 @@ If you don't set this, then a name will be generated from the directory name. If you set this, then you can run `cog push` without specifying the model name. -If you specify an image name argument when pushing (like `cog push your-username/custom-model-name`), the argument will be used and the value of `image` in cog.yaml will be ignored. +If you specify a target when pushing (like `cog push your-username/custom-model-name`), the target will be used and the value of `image` in cog.yaml will be ignored. See [Push to a registry](deploy.md#push-to-a-registry) for target precedence, tagging, and machine-readable output. + +## `model` + +The repository where Cog publishes an OCI bundle containing the model image and any managed weights. `model` must be a bare repository without a tag or digest, and it can't be set together with `image`. + +For example: + +```yaml +model: "r8.im/your-username/your-model" +``` + +`cog push` generates a timestamp tag unless you select another tag with a positional target or a [`COG_MODEL*` environment variable](environment.md#model-reference-and-registry-variables). A positional target changes the destination but still publishes a bundle. + +Managed weights require `model`. Run `cog weights import` before pushing so the target repository contains the required weight manifests. See [Push to a registry](deploy.md#push-to-a-registry) for the repository requirement and output behavior. ## `run` diff --git a/docs/yaml.md b/docs/yaml.md index 772725c302..fa6406801f 100644 --- a/docs/yaml.md +++ b/docs/yaml.md @@ -265,7 +265,21 @@ If you don't set this, then a name will be generated from the directory name. If you set this, then you can run `cog push` without specifying the model name. -If you specify an image name argument when pushing (like `cog push your-username/custom-model-name`), the argument will be used and the value of `image` in cog.yaml will be ignored. +If you specify a target when pushing (like `cog push your-username/custom-model-name`), the target will be used and the value of `image` in cog.yaml will be ignored. See [Push to a registry](deploy.md#push-to-a-registry) for target precedence, tagging, and machine-readable output. + +## `model` + +The repository where Cog publishes an OCI bundle containing the model image and any managed weights. `model` must be a bare repository without a tag or digest, and it can't be set together with `image`. + +For example: + +```yaml +model: "r8.im/your-username/your-model" +``` + +`cog push` generates a timestamp tag unless you select another tag with a positional target or a [`COG_MODEL*` environment variable](environment.md#model-reference-and-registry-variables). A positional target changes the destination but still publishes a bundle. + +Managed weights require `model`. Run `cog weights import` before pushing so the target repository contains the required weight manifests. See [Push to a registry](deploy.md#push-to-a-registry) for the repository requirement and output behavior. ## `run` diff --git a/integration-tests/tests/oci_bundle_push.txtar b/integration-tests/tests/oci_bundle_push.txtar index f23a3e7437..c295811a55 100644 --- a/integration-tests/tests/oci_bundle_push.txtar +++ b/integration-tests/tests/oci_bundle_push.txtar @@ -8,11 +8,9 @@ registry-start -# Drive both weight import and the final push to the same model ref. -# COG_MODEL pins the full ref (registry + repo + tag) so the bundle -# lands at $TEST_REGISTRY/test/bundle-model:v1 deterministically; -# without an explicit tag, ResolveModelRef would append a fresh -# timestamp and break registry-inspect below. +# Import the weights into the repository that the positional push target +# will use. The tag doesn't affect weight manifests, but keeping it explicit +# makes the setup deterministic. env COG_MODEL=$TEST_REGISTRY/test/bundle-model:v1 # Two weight directories. Each holds a single small file; the packer @@ -24,22 +22,22 @@ exec sh -c 'printf AAAA > weights-alpha/data.bin' exec sh -c 'printf BBBB > weights-beta/data.bin' # Step 1: import generates weights.lock and pushes weight manifests. -# cog push later requires the lockfile to already exist. +# cog push later requires the lockfile to already exist. Reverse the lockfile +# entries afterward to prove output follows cog.yaml rather than lockfile order. cog weights import +exec python3 -c 'import json; p="weights.lock"; d=json.load(open(p)); d["weights"].reverse(); open(p, "w").write(json.dumps(d, indent=2) + "\n")' -# Step 2: cog push assembles the OCI index (image + weight manifests) -# and uploads it to the same model ref. Positional [IMAGE] args are -# rejected in model mode, so the ref comes entirely from COG_MODEL. -cog push +# Step 2: use the positional target to override a deliberately wrong +# COG_MODEL while keeping the configured bundle format. The positional +# target's repository contains the manifests imported above. +env COG_MODEL=$TEST_REGISTRY/test/ignored:v9 +cog push $TEST_REGISTRY/test/bundle-model:v1 --json -# Push prints a structured ref tree on success: the bundle (model) -# digest plus each contained artifact. Weights were already pushed -# during `cog weights import`, so they only show up in the tree. +# Progress stays on stderr. Stdout is one versioned document containing +# only complete digest references, with weights in cog.yaml order. stderr 'Pushing to localhost' -stderr 'model localhost' -stderr 'image localhost' -stderr 'weight alpha localhost' -stderr 'weight beta localhost' +! stderr 'model localhost' +stdout '^\{"version":1,"model":"localhost:[0-9]+/test/bundle-model@sha256:[0-9a-f]{64}","image":"localhost:[0-9]+/test/bundle-model@sha256:[0-9a-f]{64}","weights":\[\{"name":"alpha","reference":"localhost:[0-9]+/test/bundle-model@sha256:[0-9a-f]{64}"\},\{"name":"beta","reference":"localhost:[0-9]+/test/bundle-model@sha256:[0-9a-f]{64}"\}\]\}$' # Resolve the pushed ref and verify the top-level index shape. registry-inspect $TEST_REGISTRY/test/bundle-model:v1 @@ -49,6 +47,52 @@ stdout 'run.cog.weight.name.*alpha' stdout 'run.cog.weight.name.*beta' stdout 'run.cog.weight.set-digest' +# Positional bundle targets also work without --json. The existing reference +# tree stays on stderr and stdout stays empty. +cog push $TEST_REGISTRY/test/bundle-model:v2 +! stdout . +stderr 'model localhost:[0-9]+/test/bundle-model@sha256:[0-9a-f]{64}' +stderr 'weight alpha.*localhost:[0-9]+/test/bundle-model@sha256:[0-9a-f]{64}' +stderr 'weight beta.*localhost:[0-9]+/test/bundle-model@sha256:[0-9a-f]{64}' + +# Untagged bundle targets receive a generated timestamp tag while JSON still +# returns immutable digest references. +cog push $TEST_REGISTRY/test/bundle-model --json +stderr 'Pushing to localhost:[0-9]+/test/bundle-model:[0-9]{8}T[0-9]{6}Z' +stdout '^\{"version":1,"model":"localhost:[0-9]+/test/bundle-model@sha256:[0-9a-f]{64}","image":"localhost:[0-9]+/test/bundle-model@sha256:[0-9a-f]{64}","weights":\[\{"name":"alpha","reference":"localhost:[0-9]+/test/bundle-model@sha256:[0-9a-f]{64}"\},\{"name":"beta","reference":"localhost:[0-9]+/test/bundle-model@sha256:[0-9a-f]{64}"\}\]\}$' + +# A different repository can't reuse managed weight manifests. Verification +# fails before the model image is pushed and emits no JSON document. +! cog push $TEST_REGISTRY/test/wrong-repository:v1 --json +! stdout . +stderr 'not found in target repository' +stderr 'cog weights import' +! registry-inspect $TEST_REGISTRY/test/wrong-repository:cog-image.v1 + +# Digest-pinned, reserved, and OCI-invalid positional bundle tags all fail +# before a build starts and leave stdout empty. +! cog push $TEST_REGISTRY/test/bundle-model@sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa --json +! stdout . +stderr 'digest-pinned' + +! cog push $TEST_REGISTRY/test/bundle-model:cog-image.user --json +! stdout . +stderr 'reserved prefix' + +! cog push $TEST_REGISTRY/test/bundle-model:-invalid --json +! stdout . +stderr 'OCI tag regex' + +# Command-level output covers one managed weight and no managed weights too. +# Re-importing with the one-weight config updates weights.lock to match it. +env COG_MODEL=$TEST_REGISTRY/test/bundle-model:v1 +cog weights import -f cog-one.yaml +cog push $TEST_REGISTRY/test/bundle-model:one-weight --json -f cog-one.yaml +stdout '^\{"version":1,"model":"localhost:[0-9]+/test/bundle-model@sha256:[0-9a-f]{64}","image":"localhost:[0-9]+/test/bundle-model@sha256:[0-9a-f]{64}","weights":\[\{"name":"alpha","reference":"localhost:[0-9]+/test/bundle-model@sha256:[0-9a-f]{64}"\}\]\}$' + +cog push $TEST_REGISTRY/test/bundle-model:no-weights --json -f cog-zero.yaml +stdout '^\{"version":1,"model":"localhost:[0-9]+/test/bundle-model@sha256:[0-9a-f]{64}","image":"localhost:[0-9]+/test/bundle-model@sha256:[0-9a-f]{64}"\}$' + -- cog.yaml -- # 'model:' satisfies the schema requirement that weight-bearing configs # declare a model ref. The full registry/repo:tag target is overridden @@ -67,6 +111,23 @@ weights: source: uri: file://./weights-beta +-- cog-one.yaml -- +build: + python_version: "3.12" +predict: "predict.py:Predictor" +model: test/bundle-model +weights: + - name: alpha + target: /src/weights-alpha + source: + uri: file://./weights-alpha + +-- cog-zero.yaml -- +build: + python_version: "3.12" +predict: "predict.py:Predictor" +model: test/bundle-model + -- predict.py -- from cog import BasePredictor diff --git a/integration-tests/tests/oci_image_push.txtar b/integration-tests/tests/oci_image_push.txtar new file mode 100644 index 0000000000..ce0b8da096 --- /dev/null +++ b/integration-tests/tests/oci_image_push.txtar @@ -0,0 +1,64 @@ +# Exercise JSON and human `cog push` output for legacy image projects against +# the local registry. Repeated pushes reuse the Docker build cache. + +[short] skip 'requires local registry' + +registry-start + +# A positional target keeps the configured image format and ignores every +# COG_MODEL* override, even one that would otherwise conflict with image:. +env COG_MODEL=invalid-target +cog push $TEST_REGISTRY/test/image-model:v1 --json + +# JSON is the only stdout output. Progress remains on stderr and the result is +# the digest reported by the upload, without bundle-only fields. +stderr 'Pushing to localhost' +stdout '^\{"version":1,"image":"localhost:[0-9]+/test/image-model@sha256:[0-9a-f]{64}"\}$' +registry-inspect $TEST_REGISTRY/test/image-model:v1 +stdout 'application/vnd.(docker.distribution.manifest.v2|oci.image.manifest.v1)\+json' + +# Untagged image targets retain Docker's :latest behavior. +cog push $TEST_REGISTRY/test/image-untagged --json +stdout '^\{"version":1,"image":"localhost:[0-9]+/test/image-untagged@sha256:[0-9a-f]{64}"\}$' +registry-inspect $TEST_REGISTRY/test/image-untagged:latest +stdout 'application/vnd.(docker.distribution.manifest.v2|oci.image.manifest.v1)\+json' + +# Existing non-JSON output stays on stderr and stdout remains empty. +cog push $TEST_REGISTRY/test/image-model:v2 +! stdout . +stderr 'image localhost:[0-9]+/test/image-model@sha256:[0-9a-f]{64}' + +# Invalid destinations fail without a success document. +! cog push $TEST_REGISTRY/test/image-model@sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa --json +! stdout . +stderr 'digest-pinned' + +# Build failures are command failures too -- no partial result reaches stdout. +! cog push $TEST_REGISTRY/test/image-model:build-failure --json -f cog-invalid.yaml +! stdout . +stderr 'missing.py' + +# A registry push failure also leaves stdout empty. The preceding builds warm +# the cache, so this reaches the push without rebuilding every layer. +! cog push localhost:1/test/image-model:v1 --json +! stdout . +stderr 'failed to push image|connection refused|server misbehaving' + +-- cog.yaml -- +build: + python_version: "3.12" +predict: "predict.py:Predictor" +image: registry.invalid/configured/image + +-- cog-invalid.yaml -- +build: + python_version: "3.12" +predict: "missing.py:Predictor" +image: registry.invalid/configured/image + +-- predict.py -- +from cog import BasePredictor + +class Predictor(BasePredictor): + def predict(self, text: str) -> str: + return f"processed: {text}" diff --git a/pkg/cli/push.go b/pkg/cli/push.go index 6246c01772..23fdb99ed0 100644 --- a/pkg/cli/push.go +++ b/pkg/cli/push.go @@ -1,6 +1,8 @@ package cli import ( + "context" + "encoding/json" "errors" "fmt" "os" @@ -20,22 +22,37 @@ import ( ) func newPushCommand() *cobra.Command { + var jsonOutput bool + cmd := &cobra.Command{ - Use: "push [IMAGE]", + Use: "push [TARGET]", Short: "Build and push model in current directory to a Docker registry", - Long: `Build a Docker image from cog.yaml and push it to a container registry. - -Cog can push to any OCI-compliant registry. When pushing to Replicate's -registry (r8.im), run 'cog login' first to authenticate.`, - Example: ` # Push to Replicate + Long: `Build from cog.yaml and push to an OCI-compliant registry. Run 'cog login' +first when pushing to Replicate's registry (r8.im). + +TARGET overrides the configured destination and all COG_MODEL* environment +variables. It doesn't change the project format: projects configured with +'image' push an image, while projects configured with 'model' push an OCI +bundle. Push targets must use tags, not digests. Untagged image targets use +Docker's default 'latest' tag; untagged bundle targets get a timestamp tag. + +With --json, Cog writes one versioned JSON result to stdout after the entire +push succeeds. Progress, warnings, and diagnostics continue on stderr. Every +reference in the result is digest-pinned.`, + Example: ` # Push an image to Replicate cog push r8.im/your-username/my-model - # Push to any OCI registry + # Push an image to any OCI registry cog push registry.example.com/your-username/model-name - # Push with model weights in a separate layer (Replicate only) + # Push a bundle project and print its immutable references as JSON + cog push registry.example.com/your-username/model-name:v1 --json + + # Push with model weights in a separate image layer (Replicate only) cog push r8.im/your-username/my-model --separate-weights`, - RunE: push, + RunE: func(cmd *cobra.Command, args []string) error { + return push(cmd, args, jsonOutput) + }, Args: cobra.MaximumNArgs(1), } addSecretsFlag(cmd) @@ -49,11 +66,12 @@ registry (r8.im), run 'cog login' first to authenticate.`, addStripFlag(cmd) addPrecompileFlag(cmd) addConfigFlag(cmd) + cmd.Flags().BoolVar(&jsonOutput, "json", false, "Output the pushed references as JSON") return cmd } -func push(cmd *cobra.Command, args []string) error { +func push(cmd *cobra.Command, args []string, jsonOutput bool) error { ctx := cmd.Context() // Initialize the provider registry @@ -74,28 +92,14 @@ func push(cmd *cobra.Command, args []string) error { return err } - // Resolve up front so user-input errors (malformed COG_MODEL_TAG, - // image:+env mode mix-up, bad positional arg) fail in seconds - // rather than after a multi-minute Docker build. We also need the - // resolved ref before Resolver.Build so the provider lookup below - // can drive credential selection from the correct host even when - // Build fails (the build-error path calls p.PostPush(...) for - // Replicate-specific guidance). - modelRef, err := validatePushArgs(src.Config.Image, src.Config.Model, args) + // Resolve up front so malformed destinations fail before a + // multi-minute Docker build. The resolved destination is also passed + // into Build so it can't disagree on a generated timestamp tag. + destination, err := resolvePushDestination(src.Config.Image, src.Config.Model, args) if err != nil { return err } - var pushTarget string - switch { - case modelRef != nil: - pushTarget = modelRef.String() - case len(args) > 0: - pushTarget = args[0] - case src.Config.Image != "": - pushTarget = src.Config.Image - default: - return errors.New("To push images, you must either set the 'image' option in cog.yaml or pass an image name as an argument. For example, 'cog push registry.example.com/your-username/model-name'") - } + pushTarget := destination.target // Look up the provider for the target registry p := provider.DefaultRegistry().ForImage(pushTarget) @@ -123,6 +127,8 @@ func push(cmd *cobra.Command, args []string) error { console.Infof("Building Docker image from environment in cog.yaml as %s...", console.Bold(pushTarget)) console.Info("") buildOpts := buildOptionsFromFlags(cmd, pushTarget, annotations) + buildOpts.Format = destination.format + buildOpts.ModelRef = destination.modelRef m, err := resolver.Build(ctx, src, buildOpts) if err != nil { // Call PostPush to handle error logging/analytics @@ -150,6 +156,7 @@ func push(cmd *cobra.Command, args []string) error { defer pw.Close() pushed, pushErr := resolver.Push(ctx, m, model.PushOptions{ + RequireDigest: jsonOutput, ImageProgressFn: func(prog model.PushProgress) { if prog.Phase != "" { switch prog.Phase { @@ -173,52 +180,195 @@ func push(cmd *cobra.Command, args []string) error { pw.Close() - // Bypass console.InfoUnformatted: it wraps at terminal width and - // would hard-break the digest refs we want to be copy-pasteable. - if pushErr == nil && pushed != nil { - if tree := formatPushResult(pushed); tree != "" { - _, _ = fmt.Fprintln(os.Stderr) - _, _ = fmt.Fprintln(os.Stderr, tree) + return completePush(ctx, p, pushOpts, pushed, pushErr, jsonOutput, console.Output) +} + +func completePush( + ctx context.Context, + p provider.Provider, + pushOpts provider.PushOptions, + pushed *model.Model, + pushErr error, + jsonOutput bool, + output func(string), +) error { + var jsonResult []byte + resultErr := pushErr + if pushErr == nil { + if jsonOutput { + jsonResult, resultErr = marshalPushOutput(pushed) + } else if pushed != nil { + // Bypass console.InfoUnformatted: it wraps at terminal width and + // would hard-break the digest refs we want to be copy-pasteable. + if tree := formatPushResult(pushed); tree != "" { + _, _ = fmt.Fprintln(os.Stderr) + _, _ = fmt.Fprintln(os.Stderr, tree) + } } } - // PostPush: the provider handles formatting errors and any - // provider-specific success output (e.g. the Replicate model URL). - if err := p.PostPush(ctx, pushOpts, pushErr); err != nil { + // PostPush may add provider-specific diagnostics and success output. + // JSON stays buffered until it returns successfully. + if err := p.PostPush(ctx, pushOpts, resultErr); err != nil { return err } - // If there was a push error but PostPush didn't return one, - // return a generic error - if pushErr != nil { - return fmt.Errorf("failed to push image: %w", pushErr) + if resultErr != nil { + return fmt.Errorf("failed to push image: %w", resultErr) } + if jsonOutput { + output(string(jsonResult)) + } return nil } -// validatePushArgs resolves the model ref and runs the push-specific -// user-input checks before any Docker work. Returns the resolved ref -// (or nil for FormatImage paths) so the caller can drive the provider -// lookup and push target from the same resolution — avoiding a second -// call to ResolveModelRef that could disagree on the timestamp tag. -// -// The positional [IMAGE] arg is rejected in FormatBundle mode where -// its meaning is ambiguous; it's still valid for FormatImage models, -// so the rejection is conditional on a resolvable model ref. -func validatePushArgs(configImage, configModel string, args []string) (*model.ResolvedRef, error) { +type pushDestination struct { + target string + format model.Format + modelRef *model.ResolvedRef +} + +// PushOutput is the versioned machine-readable result of a successful push. +type PushOutput struct { + Version int `json:"version"` + Model string `json:"model,omitempty"` + Image string `json:"image"` + Weights []PushWeightOutput `json:"weights,omitempty"` +} + +// PushWeightOutput identifies one managed weight manifest. +type PushWeightOutput struct { + Name string `json:"name"` + Reference string `json:"reference"` +} + +// resolvePushDestination resolves and validates the destination before any +// Docker work. A positional target replaces all COG_MODEL* destinations, but +// the configured image/model field still decides the artifact format. +func resolvePushDestination(configImage, configModel string, args []string) (*pushDestination, error) { + if len(args) > 0 { + target := args[0] + if configModel != "" { + ref, err := resolvedBundleTarget(target) + if err != nil { + return nil, err + } + return &pushDestination{target: ref.String(), format: model.FormatBundle, modelRef: ref}, nil + } + if err := validatePushTarget(target); err != nil { + return nil, err + } + return &pushDestination{target: target, format: model.FormatImage}, nil + } + ref, err := model.ResolveModelRef(configImage, configModel) if err != nil && !errors.Is(err, model.ErrNoModelRef) { return nil, err } - if ref != nil && len(args) > 0 { - return nil, errors.New( - "positional image argument not supported with 'model' config\n" + - " use COG_MODEL to override the full reference\n" + - " use COG_MODEL_TAG to override just the tag", - ) + if ref != nil { + if ref.Digest != "" { + return nil, fmt.Errorf("cannot push to digest-pinned target %q: use a tag instead", ref.String()) + } + return &pushDestination{target: ref.String(), format: model.FormatBundle, modelRef: ref}, nil + } + if configImage == "" { + return nil, errors.New("To push images, you must either set the 'image' option in cog.yaml or pass an image name as an argument. For example, 'cog push registry.example.com/your-username/model-name'") + } + if err := validatePushTarget(configImage); err != nil { + return nil, err + } + return &pushDestination{target: configImage, format: model.FormatImage}, nil +} + +func resolvedBundleTarget(target string) (*model.ResolvedRef, error) { + parsed, err := model.ParseRef(target, model.Insecure(), model.WithDefaultTag(model.GenerateTimestampTag())) + if err != nil { + return nil, err + } + if parsed.IsDigest() { + return nil, fmt.Errorf("cannot push to digest-pinned target %q: use a tag instead", target) + } + if err := model.ValidateTag(parsed.Tag()); err != nil { + return nil, fmt.Errorf("invalid bundle push target %q: tag %q: %w", target, parsed.Tag(), err) + } + return &model.ResolvedRef{ + Registry: parsed.Registry(), + Repo: parsed.Repository(), + Tag: parsed.Tag(), + }, nil +} + +func validatePushTarget(target string) error { + parsed, err := model.ParseRef(target, model.Insecure()) + if err != nil { + return err + } + if parsed.IsDigest() { + return fmt.Errorf("cannot push to digest-pinned target %q: use a tag instead", target) + } + return nil +} + +func marshalPushOutput(m *model.Model) ([]byte, error) { + out, err := newPushOutput(m) + if err != nil { + return nil, err + } + data, err := json.Marshal(out) + if err != nil { + return nil, fmt.Errorf("serialize push result: %w", err) } - return ref, nil + return data, nil +} + +func newPushOutput(m *model.Model) (*PushOutput, error) { + if m == nil { + return nil, errors.New("push completed without a model result") + } + img := m.GetImageArtifact() + if img == nil { + return nil, errors.New("push result has no image artifact") + } + if err := validateDigestReference("image", img.Reference); err != nil { + return nil, err + } + + out := &PushOutput{Version: 1, Image: img.Reference} + if m.Format == model.FormatBundle { + if m.Ref == nil { + return nil, errors.New("bundle push result has no model reference") + } + out.Model = m.Ref.String() + if err := validateDigestReference("model", out.Model); err != nil { + return nil, err + } + } + + if len(m.Weights) > 0 { + out.Weights = make([]PushWeightOutput, len(m.Weights)) + for i, weight := range m.Weights { + if err := validateDigestReference(fmt.Sprintf("weight %q", weight.Name), weight.Reference); err != nil { + return nil, err + } + out.Weights[i] = PushWeightOutput{Name: weight.Name, Reference: weight.Reference} + } + } + return out, nil +} + +func validateDigestReference(kind, ref string) error { + if ref == "" { + return fmt.Errorf("%s reference is empty", kind) + } + parsed, err := model.ParseRef(ref, model.Insecure()) + if err != nil { + return fmt.Errorf("invalid %s reference %q: %w", kind, ref, err) + } + if !parsed.IsDigest() { + return fmt.Errorf("%s reference %q is not digest-pinned", kind, ref) + } + return nil } // formatPushResult renders a tree of the digest-pinned refs published diff --git a/pkg/cli/push_test.go b/pkg/cli/push_test.go index 1e567f36a2..1b53267fe0 100644 --- a/pkg/cli/push_test.go +++ b/pkg/cli/push_test.go @@ -1,6 +1,10 @@ package cli import ( + "context" + "encoding/json" + "errors" + "regexp" "strings" "testing" @@ -9,6 +13,7 @@ import ( "github.com/replicate/cog/pkg/model" "github.com/replicate/cog/pkg/model/modeltest" + "github.com/replicate/cog/pkg/provider" ) // Refs are full 64-char digests so the assertions exercise the @@ -21,94 +26,383 @@ const ( testWeightRef2 = testRepo + "@sha256:e4f5a60000000000000000000000000000000000000000000000000000000000" ) -func TestValidatePushArgs(t *testing.T) { - const bundleModel = "registry.example.com/user/model" +func TestResolvePushDestination(t *testing.T) { + const digestTarget = "registry.example.com/user/model@sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" tests := []struct { - name string - configImage string - configModel string - envRepo string // COG_MODEL_REPO override - envTag string // COG_MODEL_TAG override - args []string - errContains []string // empty = expect no error - wantRef bool // expect a non-nil ResolvedRef (FormatBundle path) + name string + configImage string + configModel string + env map[string]string + args []string + wantFormat model.Format + wantTarget string + wantTargetRE string + wantRef bool + wantErr string }{ { - // cog.yaml `model:` means FormatBundle. The legacy - // positional IMAGE arg is ambiguous here — error message - // must direct the user to the env var overrides instead. - name: "FormatBundle with positional arg rejected with helpful message", - configModel: bundleModel, - args: []string{"some/other:tag"}, - errContains: []string{"positional image argument not supported", model.EnvModel, model.EnvModelTag}, + name: "tagged positional bundle target replaces full environment override", + configModel: "registry.example.com/config/model", + env: map[string]string{model.EnvModel: "invalid target"}, + args: []string{"registry.example.com/target/model:v2"}, + wantFormat: model.FormatBundle, + wantTarget: "registry.example.com/target/model:v2", + wantRef: true, }, { - name: "FormatBundle with no args proceeds and returns the ref", - configModel: bundleModel, - wantRef: true, + name: "positional bundle target replaces partial environment overrides", + configModel: "registry.example.com/config/model", + env: map[string]string{ + model.EnvModelRegistry: "invalid/registry", + model.EnvModelRepo: "invalid:repo", + model.EnvModelTag: "invalid tag", + }, + args: []string{"registry.example.com/target/model:v3"}, + wantFormat: model.FormatBundle, + wantTarget: "registry.example.com/target/model:v3", + wantRef: true, + }, + { + name: "untagged positional bundle target gets timestamp tag", + configModel: "registry.example.com/config/model", + args: []string{"registry.example.com/target/model"}, + wantFormat: model.FormatBundle, + wantTargetRE: `^registry\.example\.com/target/model:[0-9]{8}T[0-9]{6}Z$`, + wantRef: true, + }, + { + name: "tagged positional image target replaces model environment", + configImage: "registry.example.com/config/image", + env: map[string]string{model.EnvModel: "registry.example.com/env/model:v1"}, + args: []string{"registry.example.com/target/image:v2"}, + wantFormat: model.FormatImage, + wantTarget: "registry.example.com/target/image:v2", + }, + { + name: "untagged positional image target stays untagged", + configImage: "registry.example.com/config/image", + args: []string{"registry.example.com/target/image"}, + wantFormat: model.FormatImage, + wantTarget: "registry.example.com/target/image", + }, + { + name: "existing configured image behavior", + configImage: "registry.example.com/config/image", + wantFormat: model.FormatImage, + wantTarget: "registry.example.com/config/image", + }, + { + name: "existing configured bundle behavior", + configModel: "registry.example.com/config/model", + wantFormat: model.FormatBundle, + wantRef: true, + wantTargetRE: `^registry\.example\.com/config/model:[0-9]{8}T[0-9]{6}Z$`, }, { - // FormatImage path: positional arg is the legacy way to - // specify the image ref and must keep working. - name: "FormatImage with positional arg proceeds", - args: []string{"r8.im/user/model"}, + name: "positional image digest rejected", + configImage: "registry.example.com/config/image", + args: []string{digestTarget}, + wantErr: "digest-pinned", }, { - // validatePushArgs is not responsible for the "no image: - // and no arg" error — that's the downstream caller's job. - name: "FormatImage with no args proceeds", + name: "positional bundle digest rejected", + configModel: "registry.example.com/config/model", + args: []string{digestTarget}, + wantErr: "digest-pinned", }, { - // COG_MODEL_REPO alone is enough to flip to FormatBundle — - // the CI override path. Same rejection applies. - name: "env var promotion to FormatBundle rejects positional arg", - envRepo: "user/model", - args: []string{"r8.im/user/model"}, - errContains: []string{"positional image argument not supported"}, + name: "positional bundle reserved tag rejected", + configModel: "registry.example.com/config/model", + args: []string{"registry.example.com/target/model:cog-image.v2"}, + wantErr: "reserved prefix", }, { - // Validation errors surface fast, even with no positional - // arg — the whole point of the pre-flight check. - name: "invalid env var surfaces before positional check", - configModel: bundleModel, - envTag: "cog-reserved", - errContains: []string{"reserved prefix"}, + name: "positional bundle tag outside OCI grammar rejected", + configModel: "registry.example.com/config/model", + args: []string{"registry.example.com/target/model:-invalid"}, + wantErr: "OCI tag regex", }, { - // Smoke test that the image:+env mode-mix rejection - // propagates through the CLI boundary. The full matrix - // lives in TestResolveModelRef_ImageModelEnvConflict. - name: "image: in cog.yaml + COG_MODEL_REPO is rejected", - configImage: "ghcr.io/owner/repo", - envRepo: "acct/model", - errContains: []string{"'image' in cog.yaml cannot be combined with COG_MODEL"}, + name: "environment bundle digest rejected", + configModel: "registry.example.com/config/model", + env: map[string]string{model.EnvModel: digestTarget}, + wantErr: "digest-pinned", + }, + { + name: "invalid environment still fails without positional target", + configModel: "registry.example.com/config/model", + env: map[string]string{model.EnvModelTag: "cog-reserved"}, + wantErr: "reserved prefix", + }, + { + name: "image and model environment conflict remains without positional target", + configImage: "registry.example.com/config/image", + env: map[string]string{model.EnvModelRepo: "acct/model"}, + wantErr: "'image' in cog.yaml cannot be combined with COG_MODEL", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { modeltest.ClearEnv(t) - if tt.envRepo != "" { - t.Setenv(model.EnvModelRepo, tt.envRepo) + for key, value := range tt.env { + t.Setenv(key, value) + } + + destination, err := resolvePushDestination(tt.configImage, tt.configModel, tt.args) + if tt.wantErr != "" { + require.ErrorContains(t, err, tt.wantErr) + return } - if tt.envTag != "" { - t.Setenv(model.EnvModelTag, tt.envTag) + require.NoError(t, err) + require.NotNil(t, destination) + assert.Equal(t, tt.wantFormat, destination.format) + if tt.wantTarget != "" { + assert.Equal(t, tt.wantTarget, destination.target) } + if tt.wantTargetRE != "" { + assert.Regexp(t, regexp.MustCompile(tt.wantTargetRE), destination.target) + } + if tt.wantRef { + assert.NotNil(t, destination.modelRef) + } else { + assert.Nil(t, destination.modelRef) + } + }) + } +} - ref, err := validatePushArgs(tt.configImage, tt.configModel, tt.args) - if len(tt.errContains) == 0 { - require.NoError(t, err) - if tt.wantRef { - require.NotNil(t, ref, "FormatBundle path should return a resolved ref") - } else { - require.Nil(t, ref, "FormatImage path should return nil ref") +func TestMarshalPushOutput(t *testing.T) { + bundle := func(weights []model.Weight) *model.Model { + img := &model.ImageArtifact{Reference: testImageRef} + return &model.Model{ + Format: model.FormatBundle, + Ref: &model.ResolvedRef{ + Registry: "registry.example.com", + Repo: "acct/resnet-50", + Digest: testModelDigest, + }, + Image: img, + Artifacts: []model.Artifact{img}, + Weights: weights, + } + } + + tests := []struct { + name string + model *model.Model + wantJSON string + wantWeights []PushWeightOutput + }{ + { + name: "ordinary image", + model: func() *model.Model { + img := &model.ImageArtifact{Reference: testImageRef} + return &model.Model{Format: model.FormatImage, Image: img, Artifacts: []model.Artifact{img}} + }(), + wantJSON: `{"version":1,"image":"` + testImageRef + `"}`, + }, + { + name: "bundle without weights omits weights", + model: bundle(nil), + wantJSON: `{"version":1,"model":"` + testRepo + `@` + testModelDigest + `","image":"` + testImageRef + `"}`, + }, + { + name: "bundle with one weight", + model: bundle([]model.Weight{ + {Name: "transformer", Reference: testWeightRef1}, + }), + wantWeights: []PushWeightOutput{ + {Name: "transformer", Reference: testWeightRef1}, + }, + }, + { + name: "bundle preserves multiple weight order", + model: bundle([]model.Weight{ + {Name: "transformer", Reference: testWeightRef1}, + {Name: "text-encoder", Reference: testWeightRef2}, + }), + wantWeights: []PushWeightOutput{ + {Name: "transformer", Reference: testWeightRef1}, + {Name: "text-encoder", Reference: testWeightRef2}, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + data, err := marshalPushOutput(tt.model) + require.NoError(t, err) + + if tt.wantJSON != "" { + assert.JSONEq(t, tt.wantJSON, string(data)) + } + var got PushOutput + require.NoError(t, json.Unmarshal(data, &got)) + assert.Equal(t, 1, got.Version) + assert.Equal(t, tt.wantWeights, got.Weights) + assert.NotContains(t, string(data), "\n", "output must be one JSON document on one line") + }) + } +} + +func TestMarshalPushOutput_RejectsIncompleteReferences(t *testing.T) { + tests := []struct { + name string + model *model.Model + wantErr string + }{ + {name: "nil model", wantErr: "without a model result"}, + { + name: "missing image", + model: &model.Model{Format: model.FormatImage}, + wantErr: "no image artifact", + }, + { + name: "tagged image", + model: func() *model.Model { + img := &model.ImageArtifact{Reference: testRepo + ":latest"} + return &model.Model{Format: model.FormatImage, Image: img, Artifacts: []model.Artifact{img}} + }(), + wantErr: "not digest-pinned", + }, + { + name: "bundle without model ref", + model: func() *model.Model { + img := &model.ImageArtifact{Reference: testImageRef} + return &model.Model{Format: model.FormatBundle, Image: img, Artifacts: []model.Artifact{img}} + }(), + wantErr: "no model reference", + }, + { + name: "weight without digest ref", + model: func() *model.Model { + img := &model.ImageArtifact{Reference: testImageRef} + return &model.Model{ + Format: model.FormatBundle, + Ref: &model.ResolvedRef{Registry: "registry.example.com", Repo: "acct/resnet-50", Digest: testModelDigest}, + Image: img, + Artifacts: []model.Artifact{img}, + Weights: []model.Weight{{Name: "transformer", Reference: testRepo + ":cog-weight.transformer"}}, } - return + }(), + wantErr: `weight "transformer" reference`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + data, err := marshalPushOutput(tt.model) + require.ErrorContains(t, err, tt.wantErr) + assert.Nil(t, data) + }) + } +} + +func TestNewPushCommand_JSONFlag(t *testing.T) { + cmd := newPushCommand() + flag := cmd.Flags().Lookup("json") + require.NotNil(t, flag) + assert.Equal(t, "false", flag.DefValue) +} + +type pushTestProvider struct { + postPush func(pushErr error) error +} + +func (p *pushTestProvider) Name() string { return "test" } + +func (p *pushTestProvider) MatchesRegistry(string) bool { return true } + +func (p *pushTestProvider) Login(context.Context, provider.LoginOptions) error { return nil } + +func (p *pushTestProvider) PostPush(_ context.Context, _ provider.PushOptions, pushErr error) error { + return p.postPush(pushErr) +} + +func TestCompletePush_JSONOutputWaitsForFullSuccess(t *testing.T) { + validImage := func() *model.Model { + img := &model.ImageArtifact{Reference: testImageRef} + return &model.Model{Format: model.FormatImage, Image: img, Artifacts: []model.Artifact{img}} + } + + tests := []struct { + name string + pushed *model.Model + pushErr error + providerErr error + wantErr string + wantProviderErr string + wantOutput bool + }{ + { + name: "success", + pushed: validImage(), + wantOutput: true, + }, + { + name: "push failure", + pushErr: errors.New("registry rejected push"), + wantErr: "registry rejected push", + wantProviderErr: "registry rejected push", + }, + { + name: "provider failure", + pushed: validImage(), + providerErr: errors.New("provider post-processing failed"), + wantErr: "provider post-processing failed", + }, + { + name: "digest validation failure", + pushed: &model.Model{Format: model.FormatImage}, + wantErr: "no image artifact", + wantProviderErr: "no image artifact", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + providerCalled := false + providerDone := false + var providerPushErr error + p := &pushTestProvider{postPush: func(pushErr error) error { + providerCalled = true + providerPushErr = pushErr + providerDone = true + return tt.providerErr + }} + var outputs []string + + err := completePush( + context.Background(), + p, + provider.PushOptions{}, + tt.pushed, + tt.pushErr, + true, + func(output string) { + assert.True(t, providerDone, "JSON must be emitted after provider post-processing") + outputs = append(outputs, output) + }, + ) + + assert.True(t, providerCalled) + if tt.wantErr != "" { + require.ErrorContains(t, err, tt.wantErr) + } else { + require.NoError(t, err) + } + if tt.wantProviderErr != "" { + require.ErrorContains(t, providerPushErr, tt.wantProviderErr) + } else { + assert.NoError(t, providerPushErr) } - require.Error(t, err) - for _, s := range tt.errContains { - assert.Contains(t, err.Error(), s) + if tt.wantOutput { + require.Len(t, outputs, 1) + assert.JSONEq(t, `{"version":1,"image":"`+testImageRef+`"}`, outputs[0]) + } else { + assert.Empty(t, outputs) } }) } diff --git a/pkg/docker/command/command.go b/pkg/docker/command/command.go index 4751a843cf..d716b18c75 100644 --- a/pkg/docker/command/command.go +++ b/pkg/docker/command/command.go @@ -8,6 +8,18 @@ import ( "github.com/docker/docker/api/types/image" ) +// PushResult identifies the manifest uploaded by a Docker push. +type PushResult struct { + Digest string + Size int64 +} + +// PushResultCommand is implemented by Docker commands that can return the +// manifest digest reported by the daemon's push stream. +type PushResultCommand interface { + PushWithResult(ctx context.Context, ref string) (PushResult, error) +} + type Command interface { // Pull pulls an image from a remote registry and returns the inspect response for the local image. // If the image already exists, it will return the inspect response for the local image without pulling. diff --git a/pkg/docker/docker.go b/pkg/docker/docker.go index 29803173e5..3933ae9249 100644 --- a/pkg/docker/docker.go +++ b/pkg/docker/docker.go @@ -1,12 +1,15 @@ package docker import ( + "bytes" "context" + "encoding/json" "errors" "fmt" "io" "net" "os" + "regexp" "strconv" "strings" @@ -195,11 +198,16 @@ func (c *apiClient) ContainerLogs(ctx context.Context, containerID string, w io. } func (c *apiClient) Push(ctx context.Context, imageRef string) error { + _, err := c.PushWithResult(ctx, imageRef) + return err +} + +func (c *apiClient) PushWithResult(ctx context.Context, imageRef string) (command.PushResult, error) { console.Debugf("=== APIClient.Push %s", imageRef) parsedName, err := name.ParseReference(imageRef) if err != nil { - return fmt.Errorf("failed to parse image reference: %w", err) + return command.PushResult{}, fmt.Errorf("failed to parse image reference: %w", err) } console.Debugf("fully qualified image ref: %s", parsedName) @@ -224,35 +232,89 @@ func (c *apiClient) Push(ctx context.Context, imageRef string) error { var opts image.PushOptions encodedAuth, err := registry.EncodeAuthConfig(authConfig) if err != nil { - return fmt.Errorf("failed to encode auth config: %w", err) + return command.PushResult{}, fmt.Errorf("failed to encode auth config: %w", err) } opts.RegistryAuth = encodedAuth output, err := c.client.ImagePush(ctx, imageRef, opts) if err != nil { - return fmt.Errorf("failed to push image: %w", err) + return command.PushResult{}, fmt.Errorf("failed to push image: %w", err) } defer output.Close() - // output is a json stream, so we need to parse it, handle errors, and write progress to stderr - isTTY := console.IsTTY(os.Stderr) - if err := jsonmessage.DisplayJSONMessagesStream(output, os.Stderr, os.Stderr.Fd(), isTTY, nil); err != nil { + return displayPushOutput(output, os.Stderr, os.Stderr.Fd(), console.IsTTY(os.Stderr), imageRef) +} + +var pushStatusPattern = regexp.MustCompile(`(?:^|: )digest: (sha256:[0-9a-f]{64}) size: ([0-9]+)$`) + +type pushOutputCapture struct { + pending []byte + result command.PushResult +} + +func (c *pushOutputCapture) Write(p []byte) (int, error) { + c.pending = append(c.pending, p...) + for { + newline := bytes.IndexByte(c.pending, '\n') + if newline < 0 { + break + } + c.capture(c.pending[:newline]) + c.pending = c.pending[newline+1:] + } + return len(p), nil +} + +func (c *pushOutputCapture) finish() { + if len(c.pending) > 0 { + c.capture(c.pending) + c.pending = nil + } +} + +func (c *pushOutputCapture) capture(line []byte) { + var message jsonmessage.JSONMessage + if err := json.Unmarshal(line, &message); err != nil { + return + } + if message.Aux != nil { + var pushed types.PushResult + if err := json.Unmarshal(*message.Aux, &pushed); err == nil && pushed.Digest != "" { + c.result = command.PushResult{Digest: pushed.Digest, Size: int64(pushed.Size)} + } + return + } + matches := pushStatusPattern.FindStringSubmatch(message.Status) + if len(matches) != 3 { + return + } + size, err := strconv.ParseInt(matches[2], 10, 64) + if err != nil { + return + } + c.result = command.PushResult{Digest: matches[1], Size: size} +} + +func displayPushOutput(output io.Reader, stderr io.Writer, terminalFD uintptr, isTTY bool, imageRef string) (command.PushResult, error) { + var capture pushOutputCapture + err := jsonmessage.DisplayJSONMessagesStream(io.TeeReader(output, &capture), stderr, terminalFD, isTTY, nil) + capture.finish() + if err != nil { var streamErr *jsonmessage.JSONError if errors.As(err, &streamErr) { if isTagNotFoundError(err) { - return &command.NotFoundError{Ref: imageRef, Object: "tag"} + return command.PushResult{}, &command.NotFoundError{Ref: imageRef, Object: "tag"} } if isRepositoryNotFoundError(err) { - return &command.NotFoundError{Ref: imageRef, Object: "repository"} + return command.PushResult{}, &command.NotFoundError{Ref: imageRef, Object: "repository"} } if isAuthorizationFailedError(err) { - return command.ErrAuthorizationFailed + return command.PushResult{}, command.ErrAuthorizationFailed } } - return fmt.Errorf("error during image push: %w", err) + return command.PushResult{}, fmt.Errorf("error during image push: %w", err) } - - return nil + return capture.result, nil } func (c *apiClient) ImageSave(ctx context.Context, imageRef string) (io.ReadCloser, error) { diff --git a/pkg/docker/dockertest/command_mocks.go b/pkg/docker/dockertest/command_mocks.go index 035e18962e..483d128eea 100644 --- a/pkg/docker/dockertest/command_mocks.go +++ b/pkg/docker/dockertest/command_mocks.go @@ -14,6 +14,99 @@ import ( mock "github.com/stretchr/testify/mock" ) +// NewMockPushResultCommand2 creates a new instance of MockPushResultCommand2. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewMockPushResultCommand2(t interface { + mock.TestingT + Cleanup(func()) +}) *MockPushResultCommand2 { + mock := &MockPushResultCommand2{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} + +// MockPushResultCommand2 is an autogenerated mock type for the PushResultCommand type +type MockPushResultCommand2 struct { + mock.Mock +} + +type MockPushResultCommand2_Expecter struct { + mock *mock.Mock +} + +func (_m *MockPushResultCommand2) EXPECT() *MockPushResultCommand2_Expecter { + return &MockPushResultCommand2_Expecter{mock: &_m.Mock} +} + +// PushWithResult provides a mock function for the type MockPushResultCommand2 +func (_mock *MockPushResultCommand2) PushWithResult(ctx context.Context, ref string) (command.PushResult, error) { + ret := _mock.Called(ctx, ref) + + if len(ret) == 0 { + panic("no return value specified for PushWithResult") + } + + var r0 command.PushResult + var r1 error + if returnFunc, ok := ret.Get(0).(func(context.Context, string) (command.PushResult, error)); ok { + return returnFunc(ctx, ref) + } + if returnFunc, ok := ret.Get(0).(func(context.Context, string) command.PushResult); ok { + r0 = returnFunc(ctx, ref) + } else { + r0 = ret.Get(0).(command.PushResult) + } + if returnFunc, ok := ret.Get(1).(func(context.Context, string) error); ok { + r1 = returnFunc(ctx, ref) + } else { + r1 = ret.Error(1) + } + return r0, r1 +} + +// MockPushResultCommand2_PushWithResult_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'PushWithResult' +type MockPushResultCommand2_PushWithResult_Call struct { + *mock.Call +} + +// PushWithResult is a helper method to define mock.On call +// - ctx context.Context +// - ref string +func (_e *MockPushResultCommand2_Expecter) PushWithResult(ctx interface{}, ref interface{}) *MockPushResultCommand2_PushWithResult_Call { + return &MockPushResultCommand2_PushWithResult_Call{Call: _e.mock.On("PushWithResult", ctx, ref)} +} + +func (_c *MockPushResultCommand2_PushWithResult_Call) Run(run func(ctx context.Context, ref string)) *MockPushResultCommand2_PushWithResult_Call { + _c.Call.Run(func(args mock.Arguments) { + var arg0 context.Context + if args[0] != nil { + arg0 = args[0].(context.Context) + } + var arg1 string + if args[1] != nil { + arg1 = args[1].(string) + } + run( + arg0, + arg1, + ) + }) + return _c +} + +func (_c *MockPushResultCommand2_PushWithResult_Call) Return(pushResult command.PushResult, err error) *MockPushResultCommand2_PushWithResult_Call { + _c.Call.Return(pushResult, err) + return _c +} + +func (_c *MockPushResultCommand2_PushWithResult_Call) RunAndReturn(run func(ctx context.Context, ref string) (command.PushResult, error)) *MockPushResultCommand2_PushWithResult_Call { + _c.Call.Return(run) + return _c +} + // NewMockCommand2 creates a new instance of MockCommand2. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. // The first argument is typically a *testing.T value. func NewMockCommand2(t interface { diff --git a/pkg/docker/push_test.go b/pkg/docker/push_test.go new file mode 100644 index 0000000000..6cee892cf2 --- /dev/null +++ b/pkg/docker/push_test.go @@ -0,0 +1,38 @@ +package docker + +import ( + "bytes" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestDisplayPushOutput(t *testing.T) { + const digest = "sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + output := strings.NewReader( + `{"status":"Pushing layers"}` + "\n" + + `{"aux":{"Tag":"v1","Digest":"` + digest + `","Size":1234}}` + "\n", + ) + var stderr bytes.Buffer + + result, err := displayPushOutput(output, &stderr, 0, false, "registry.example.com/user/model:v1") + + require.NoError(t, err) + assert.Equal(t, digest, result.Digest) + assert.Equal(t, int64(1234), result.Size) + assert.Contains(t, stderr.String(), "Pushing layers") + assert.NotContains(t, stderr.String(), digest, "the machine-readable push result shouldn't be rendered as progress") +} + +func TestDisplayPushOutput_ContainerdStatus(t *testing.T) { + const digest = "sha256:bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" + output := strings.NewReader(`{"status":"v1: digest: ` + digest + ` size: 5678"}` + "\n") + + result, err := displayPushOutput(output, &bytes.Buffer{}, 0, false, "registry.example.com/user/model:v1") + + require.NoError(t, err) + assert.Equal(t, digest, result.Digest) + assert.Equal(t, int64(5678), result.Size) +} diff --git a/pkg/model/image_pusher.go b/pkg/model/image_pusher.go index ad43db4589..7084848cef 100644 --- a/pkg/model/image_pusher.go +++ b/pkg/model/image_pusher.go @@ -65,17 +65,31 @@ func WithOnFallback(fn func()) ImagePushOption { } } +// ImagePushResult identifies the manifest uploaded by ImagePusher. +type ImagePushResult struct { + Digest string + Size int64 +} + // Push pushes a container image to the registry. // // Tries the OCI chunked push path first (if enabled and registry client is // available), then falls back to Docker push on any non-fatal error. // The artifact must have a valid Reference. func (p *ImagePusher) Push(ctx context.Context, artifact *ImageArtifact, opts ...ImagePushOption) error { + _, err := p.PushWithResult(ctx, artifact, opts...) + return err +} + +// PushWithResult pushes an image and returns the digest reported by the path +// that uploaded its manifest. The digest is empty only when the Docker command +// doesn't expose push results. +func (p *ImagePusher) PushWithResult(ctx context.Context, artifact *ImageArtifact, opts ...ImagePushOption) (*ImagePushResult, error) { if artifact == nil { - return fmt.Errorf("image artifact is nil") + return nil, fmt.Errorf("image artifact is nil") } if artifact.Reference == "" { - return fmt.Errorf("image artifact has no reference") + return nil, fmt.Errorf("image artifact has no reference") } var opt imagePushOptions @@ -86,12 +100,12 @@ func (p *ImagePusher) Push(ctx context.Context, artifact *ImageArtifact, opts .. imageRef := artifact.Reference if p.canOCIPush() { - err := p.ociPush(ctx, imageRef, opt) + result, err := p.ociPush(ctx, imageRef, opt) if err == nil { - return nil + return result, nil } if !shouldFallbackToDocker(err) { - return fmt.Errorf("OCI chunked push: %w", err) + return nil, fmt.Errorf("OCI chunked push: %w", err) } if opt.onFallback != nil { opt.onFallback() @@ -99,7 +113,17 @@ func (p *ImagePusher) Push(ctx context.Context, artifact *ImageArtifact, opts .. console.Warnf("OCI chunked push failed, falling back to Docker push: %v", sanitizeError(err)) } - return p.docker.Push(ctx, imageRef) + if dockerWithResult, ok := p.docker.(command.PushResultCommand); ok { + result, err := dockerWithResult.PushWithResult(ctx, imageRef) + if err != nil { + return nil, err + } + return &ImagePushResult{Digest: result.Digest, Size: result.Size}, nil + } + if err := p.docker.Push(ctx, imageRef); err != nil { + return nil, err + } + return &ImagePushResult{}, nil } // canOCIPush returns true if OCI chunked push is enabled. @@ -109,7 +133,7 @@ func (p *ImagePusher) canOCIPush() bool { // ociPush exports the image from Docker daemon as a tar, then pushes all layers, // config, and manifest to the registry using chunked uploads. -func (p *ImagePusher) ociPush(ctx context.Context, imageRef string, opt imagePushOptions) error { +func (p *ImagePusher) ociPush(ctx context.Context, imageRef string, opt imagePushOptions) (*ImagePushResult, error) { console.Debugf("Exporting image %s from Docker daemon...", imageRef) if opt.progressFn != nil { @@ -118,20 +142,20 @@ func (p *ImagePusher) ociPush(ctx context.Context, imageRef string, opt imagePus ref, err := name.ParseReference(imageRef, name.Insecure) if err != nil { - return fmt.Errorf("parse image reference %q: %w", imageRef, err) + return nil, fmt.Errorf("parse image reference %q: %w", imageRef, err) } // Get the Docker tar stream directly from the docker command rc, err := p.docker.ImageSave(ctx, imageRef) if err != nil { - return fmt.Errorf("export image from daemon: %w", err) + return nil, fmt.Errorf("export image from daemon: %w", err) } defer rc.Close() //nolint:errcheck // Write the tar to a temp file so we can seek on it tmpTar, err := os.CreateTemp("", "cog-image-*.tar") if err != nil { - return fmt.Errorf("create temp tar file: %w", err) + return nil, fmt.Errorf("create temp tar file: %w", err) } defer func() { @@ -140,7 +164,7 @@ func (p *ImagePusher) ociPush(ctx context.Context, imageRef string, opt imagePus }() if _, err := io.Copy(tmpTar, rc); err != nil { - return fmt.Errorf("write image tar: %w", err) + return nil, fmt.Errorf("write image tar: %w", err) } _ = rc.Close() @@ -155,7 +179,7 @@ func (p *ImagePusher) ociPush(ctx context.Context, imageRef string, opt imagePus img, err := tarball.ImageFromPath(tmpTar.Name(), &tag) if err != nil { - return fmt.Errorf("load image from tar: %w", err) + return nil, fmt.Errorf("load image from tar: %w", err) } if opt.progressFn != nil { @@ -166,23 +190,31 @@ func (p *ImagePusher) ociPush(ctx context.Context, imageRef string, opt imagePus } // pushImage pushes a v1.Image (layers, config, manifest) to the registry. -func (p *ImagePusher) pushImage(ctx context.Context, imageRef string, img v1.Image, opt imagePushOptions) error { +func (p *ImagePusher) pushImage(ctx context.Context, imageRef string, img v1.Image, opt imagePushOptions) (*ImagePushResult, error) { repo := repoFromReference(imageRef) if err := p.pushLayers(ctx, repo, img, opt); err != nil { - return fmt.Errorf("push layers: %w", err) + return nil, fmt.Errorf("push layers: %w", err) } if err := p.pushConfig(ctx, repo, img); err != nil { - return fmt.Errorf("push config: %w", err) + return nil, fmt.Errorf("push config: %w", err) } console.Debugf("Pushing image manifest for %s", imageRef) if err := p.registry.PushImage(ctx, imageRef, img); err != nil { - return fmt.Errorf("push manifest: %w", err) + return nil, fmt.Errorf("push manifest: %w", err) } - return nil + digest, err := img.Digest() + if err != nil { + return nil, fmt.Errorf("compute pushed image digest: %w", err) + } + rawManifest, err := img.RawManifest() + if err != nil { + return nil, fmt.Errorf("read pushed image manifest: %w", err) + } + return &ImagePushResult{Digest: digest.String(), Size: int64(len(rawManifest))}, nil } // pushLayers pushes all image layers concurrently using the registry client's diff --git a/pkg/model/image_pusher_test.go b/pkg/model/image_pusher_test.go index 300cbb4fb9..2ba3845b31 100644 --- a/pkg/model/image_pusher_test.go +++ b/pkg/model/image_pusher_test.go @@ -119,8 +119,13 @@ func TestImagePusher_Push(t *testing.T) { docker := &mockDocker{imageSaveFunc: fakeImageSaveFunc(img, tag)} pusher := newImagePusher(docker, mock) - err = pusher.Push(context.Background(), testArtifact(tag)) + result, err := pusher.PushWithResult(context.Background(), testArtifact(tag)) + require.NoError(t, err) + expectedDigest, err := img.Digest() require.NoError(t, err) + require.NotNil(t, result) + assert.Equal(t, expectedDigest.String(), result.Digest) + assert.Positive(t, result.Size) // Should have pushed 2 layers + 1 config blob = 3 WriteLayer calls assert.Equal(t, 3, mock.writeLayerCount) diff --git a/pkg/model/options.go b/pkg/model/options.go index bd96ea80f5..20ec84ccf2 100644 --- a/pkg/model/options.go +++ b/pkg/model/options.go @@ -8,6 +8,16 @@ type BuildOptions struct { // ImageName is the output image name (required). ImageName string + // Format overrides the artifact format resolved from cog.yaml and + // COG_MODEL*. A zero value keeps the normal resolution behavior. + // Push uses this when a positional target must replace environment + // destinations without changing the project's configured format. + Format Format + + // ModelRef is the resolved bundle destination when Format is + // FormatBundle. It is ignored for FormatImage. + ModelRef *ResolvedRef + // NoCache disables build cache. NoCache bool diff --git a/pkg/model/pusher.go b/pkg/model/pusher.go index 9d331f1f4e..adeddafeea 100644 --- a/pkg/model/pusher.go +++ b/pkg/model/pusher.go @@ -19,6 +19,10 @@ type PushOptions struct { // Default: linux/amd64 Platform *Platform + // RequireDigest rejects tag-based fallback when an image push doesn't + // report its manifest digest. + RequireDigest bool + // ImageProgressFn is an optional callback for reporting push progress. // It receives both phase transitions (Phase set, byte fields zero) and // per-layer byte progress (Phase empty, Complete/Total set). @@ -66,7 +70,8 @@ func NewBundlePusher(docker command.Command, reg registry.Client) *BundlePusher // the image manifest lands at a stable, namespaced tag in the // registry, independent of whatever tag the model index will // carry. Push from that tag. -// 3. HEAD the image to capture its registry-side digest. +// 3. Capture the manifest digest reported by the upload and HEAD that +// digest-pinned reference to obtain the descriptor without racing the tag. // 4. Build and push the OCI index at the model ref; capture its // descriptor locally from the index bytes (content-addressed — // no second registry round-trip needed). @@ -140,13 +145,34 @@ func (p *BundlePusher) Push(ctx context.Context, m *Model, opts PushOptions) (*M if opts.OnFallback != nil { imagePushOpts = append(imagePushOpts, WithOnFallback(opts.OnFallback)) } - if err := p.imagePusher.Push(ctx, pushedImage, imagePushOpts...); err != nil { + imageResult, err := p.imagePusher.PushWithResult(ctx, pushedImage, imagePushOpts...) + if err != nil { return nil, fmt.Errorf("push image %q: %w", imageRef, err) } + var imgDesc v1.Descriptor + switch { + case imageResult.Digest != "": + if _, err := v1.NewHash(imageResult.Digest); err != nil { + return nil, fmt.Errorf("push image %q: the upload reported invalid manifest digest %q: %w", imageRef, imageResult.Digest, err) + } - imgDesc, err := p.registry.GetDescriptor(ctx, imageRef) - if err != nil { - return nil, fmt.Errorf("get image descriptor: %w", err) + imageDigestRef := repo + "@" + imageResult.Digest + imgDesc, err = p.registry.GetDescriptor(ctx, imageDigestRef) + if err != nil { + return nil, fmt.Errorf("get uploaded image descriptor %q: %w", imageDigestRef, err) + } + if imgDesc.Digest.String() != imageResult.Digest { + return nil, fmt.Errorf("registry returned digest %s for uploaded image %s", imgDesc.Digest.String(), imageResult.Digest) + } + case opts.RequireDigest: + return nil, fmt.Errorf("push image %q: the upload did not report the manifest digest", imageRef) + default: + // Older or alternate Docker commands may not expose the upload digest. + // Preserve the previous non-JSON lookup by tag. + imgDesc, err = p.registry.GetDescriptor(ctx, imageRef) + if err != nil { + return nil, fmt.Errorf("get image descriptor %q: %w", imageRef, err) + } } platform := opts.Platform @@ -251,8 +277,8 @@ func (p *BundlePusher) verifyWeights( desc, err := p.registry.GetDescriptor(ctx, ref) if err != nil { return fmt.Errorf( - "weight %q not found in registry (%s); run 'cog weights import' to push weights first: %w", - w.Name, ref, err, + "weight %q from weights.lock was not found in target repository %q (%s); run 'cog weights import' for that repository before pushing: %w", + w.Name, repo, ref, err, ) } if desc.Digest.String() != w.Digest { diff --git a/pkg/model/pusher_test.go b/pkg/model/pusher_test.go index 61898eb3ff..2fbf7a98d9 100644 --- a/pkg/model/pusher_test.go +++ b/pkg/model/pusher_test.go @@ -12,6 +12,8 @@ import ( v1 "github.com/google/go-containerregistry/pkg/v1" "github.com/google/go-containerregistry/pkg/v1/types" "github.com/stretchr/testify/require" + + "github.com/replicate/cog/pkg/docker/command" ) const ( @@ -29,6 +31,10 @@ var ( testModelRef = testRepo + ":" + testModelTag ) +type commandWithoutPushResult struct { + command.Command +} + // Valid 64-char hex digests for use in test fixtures. v1.NewHash // rejects non-hex strings, so the verifyWeights digest-equality check // requires real-looking digests. @@ -118,10 +124,12 @@ func TestBundlePusher_Push(t *testing.T) { pushFunc: func(ctx context.Context, ref string) error { return nil }, } + imageHash, err := v1.NewHash(testDockerPushDigest) + require.NoError(t, err) imgDesc := v1.Descriptor{ MediaType: types.OCIManifestSchema1, Size: 1234, - Digest: v1.Hash{Algorithm: "sha256", Hex: "imgonly"}, + Digest: imageHash, } reg := &mockRegistry{ @@ -141,10 +149,82 @@ func TestBundlePusher_Push(t *testing.T) { pusher := NewBundlePusher(docker, reg) - _, err := pusher.Push(context.Background(), testBundleModel(), PushOptions{}) + _, err = pusher.Push(context.Background(), testBundleModel(), PushOptions{}) require.NoError(t, err) }) + for _, test := range []struct { + name string + docker command.Command + }{ + { + name: "empty upload digest", + docker: &mockDocker{ + pushFunc: func(ctx context.Context, ref string) error { return nil }, + pushResult: &command.PushResult{}, + }, + }, + { + name: "command without push result support", + docker: commandWithoutPushResult{Command: &mockDocker{ + pushFunc: func(ctx context.Context, ref string) error { return nil }, + }}, + }, + } { + t.Run("falls back to tag lookup for "+test.name, func(t *testing.T) { + imageHash, err := v1.NewHash(testDockerPushDigest) + require.NoError(t, err) + imgDesc := v1.Descriptor{ + MediaType: types.OCIManifestSchema1, + Size: 1234, + Digest: imageHash, + } + + reg := &mockRegistry{ + getDescriptorFunc: func(ctx context.Context, ref string) (v1.Descriptor, error) { + require.Equal(t, testCogImageRef, ref) + return imgDesc, nil + }, + pushIndexFunc: func(ctx context.Context, ref string, idx v1.ImageIndex) error { + idxManifest, err := idx.IndexManifest() + require.NoError(t, err) + require.Equal(t, imgDesc.Digest, idxManifest.Manifests[0].Digest) + return nil + }, + } + + pushed, err := NewBundlePusher(test.docker, reg).Push(context.Background(), testBundleModel(), PushOptions{}) + + require.NoError(t, err) + require.Equal(t, testRepo+"@"+testDockerPushDigest, pushed.Image.Reference) + }) + } + + t.Run("requires upload digest when requested", func(t *testing.T) { + docker := &mockDocker{ + pushFunc: func(ctx context.Context, ref string) error { return nil }, + pushResult: &command.PushResult{}, + } + reg := &mockRegistry{ + getDescriptorFunc: func(ctx context.Context, ref string) (v1.Descriptor, error) { + require.FailNow(t, "descriptor lookup must not fall back to a mutable tag when a digest is required") + return v1.Descriptor{}, nil + }, + pushIndexFunc: func(ctx context.Context, ref string, idx v1.ImageIndex) error { + require.FailNow(t, "index must not be pushed without an upload digest") + return nil + }, + } + + _, err := NewBundlePusher(docker, reg).Push( + context.Background(), + testBundleModel(), + PushOptions{RequireDigest: true}, + ) + + require.ErrorContains(t, err, "upload did not report the manifest digest") + }) + t.Run("full push flow succeeds with single weight", func(t *testing.T) { w := Weight{ Name: "model-v1", @@ -178,10 +258,12 @@ func TestBundlePusher_Push(t *testing.T) { }, } + imageHash, err := v1.NewHash(testDockerPushDigest) + require.NoError(t, err) imgDesc := v1.Descriptor{ MediaType: types.OCIManifestSchema1, Size: 1234, - Digest: v1.Hash{Algorithm: "sha256", Hex: "imgdigestabc1234567"}, + Digest: imageHash, } weightDesc := v1.Descriptor{ @@ -203,7 +285,7 @@ func TestBundlePusher_Push(t *testing.T) { switch ref { case weightRef: return weightDesc, nil - case testCogImageRef: + case testRepo + "@" + testDockerPushDigest: return imgDesc, nil } return v1.Descriptor{}, fmt.Errorf("unexpected descriptor lookup: %s", ref) @@ -242,7 +324,7 @@ func TestBundlePusher_Push(t *testing.T) { // Verify call sequence: weight verified first (HEAD by digest, // before anything mutates the registry), then local re-tag, - // then docker push, then image HEAD, then index push, then + // then docker push, then digest-pinned image HEAD, then index push, then // the deferred local-tag cleanup. The index descriptor is // computed locally from the v1.ImageIndex bytes — no HEAD. require.Equal(t, @@ -250,7 +332,7 @@ func TestBundlePusher_Push(t *testing.T) { "registry:getDescriptor:" + weightRef, "docker:tag:" + testImageRef + "->" + testCogImageRef, "docker:push:" + testCogImageRef, - "registry:getDescriptor:" + testCogImageRef, + "registry:getDescriptor:" + testRepo + "@" + testDockerPushDigest, "registry:pushIndex:" + testModelRef, "docker:remove:" + testCogImageRef, }, @@ -387,6 +469,29 @@ func TestBundlePusher_Push(t *testing.T) { require.Contains(t, err.Error(), "manifest not found") }) + t.Run("rejects a descriptor that doesn't match the uploaded image digest", func(t *testing.T) { + docker := &mockDocker{ + pushFunc: func(ctx context.Context, ref string) error { return nil }, + } + mismatchedDigest, err := v1.NewHash("sha256:bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb") + require.NoError(t, err) + reg := &mockRegistry{ + getDescriptorFunc: func(ctx context.Context, ref string) (v1.Descriptor, error) { + require.Equal(t, testRepo+"@"+testDockerPushDigest, ref) + return v1.Descriptor{Digest: mismatchedDigest}, nil + }, + pushIndexFunc: func(ctx context.Context, ref string, idx v1.ImageIndex) error { + require.FailNow(t, "index must not be pushed with a mismatched image descriptor") + return nil + }, + } + + _, err = NewBundlePusher(docker, reg).Push(context.Background(), testBundleModel(), PushOptions{}) + + require.ErrorContains(t, err, "registry returned digest") + require.ErrorContains(t, err, testDockerPushDigest) + }) + t.Run("returns error when weight manifest not in registry", func(t *testing.T) { docker := &mockDocker{ pushFunc: func(ctx context.Context, ref string) error { @@ -408,7 +513,8 @@ func TestBundlePusher_Push(t *testing.T) { require.Error(t, err) require.Contains(t, err.Error(), "w1") - require.Contains(t, err.Error(), "not found in registry") + require.Contains(t, err.Error(), "not found in target repository") + require.Contains(t, err.Error(), testRepo) require.Contains(t, err.Error(), "cog weights import") require.Contains(t, err.Error(), "manifest unknown") }) @@ -494,8 +600,7 @@ func TestBundlePusher_Push(t *testing.T) { var headCheckCount atomic.Int32 reg := &mockRegistry{ getDescriptorFunc: func(ctx context.Context, ref string) (v1.Descriptor, error) { - // Count weight HEADs only (those are by repo@digest). - if strings.Contains(ref, "@") { + if strings.HasSuffix(ref, testW1Digest) || strings.HasSuffix(ref, testW2Digest) { headCheckCount.Add(1) } return descriptorFromRef(ref), nil @@ -552,14 +657,17 @@ func TestResolver_Push(t *testing.T) { return nil }, } + imageHash, err := v1.NewHash(testDockerPushDigest) + require.NoError(t, err) imgDesc := v1.Descriptor{ MediaType: types.OCIManifestSchema1, Size: 1234, - Digest: v1.Hash{Algorithm: "sha256", Hex: "imagedigestformatimage"}, + Digest: imageHash, } reg := &mockRegistry{ getDescriptorFunc: func(ctx context.Context, ref string) (v1.Descriptor, error) { - return imgDesc, nil + require.Failf(t, "unexpected descriptor lookup", "image result must use the digest reported by push, not HEAD mutable tag %q", ref) + return v1.Descriptor{}, nil }, } resolver := NewResolver(docker, reg) @@ -600,6 +708,7 @@ func TestResolver_Push(t *testing.T) { dockerPushed = true return nil }, + pushResult: &command.PushResult{}, } reg := &mockRegistry{ getDescriptorFunc: func(ctx context.Context, ref string) (v1.Descriptor, error) { @@ -622,6 +731,27 @@ func TestResolver_Push(t *testing.T) { "on HEAD failure, Push should return the input Model unchanged") }) + t.Run("FormatImage returns digest lookup error when required", func(t *testing.T) { + var dockerPushed bool + docker := &mockDocker{ + pushFunc: func(ctx context.Context, ref string) error { + dockerPushed = true + return nil + }, + pushResult: &command.PushResult{}, + } + resolver := NewResolver(docker, &mockRegistry{}) + + img := &ImageArtifact{name: "model", Reference: testImageRef} + m := &Model{Format: FormatImage, Image: img, Artifacts: []Artifact{img}} + + pushed, err := resolver.Push(context.Background(), m, PushOptions{RequireDigest: true}) + require.ErrorContains(t, err, "resolve pushed image digest") + require.ErrorContains(t, err, "did not report") + require.Nil(t, pushed) + require.True(t, dockerPushed) + }) + t.Run("FormatBundle with no weights produces a single-entry index", func(t *testing.T) { // Behavioral change: a FormatBundle model with zero weights is // still pushed as an OCI index (containing only the image diff --git a/pkg/model/resolve.go b/pkg/model/resolve.go index 7bb72eb8bb..19566cb5fb 100644 --- a/pkg/model/resolve.go +++ b/pkg/model/resolve.go @@ -244,17 +244,25 @@ func validateRepoEnv(v string) error { return nil } -// validateTagEnv checks that COG_MODEL_TAG is a valid OCI tag and -// does not use the reserved "cog-" prefix. -func validateTagEnv(v string) error { - if v == "" { +// ValidateTag checks that a user-supplied bundle tag follows the OCI tag +// grammar and doesn't use Cog's reserved "cog-" prefix. +func ValidateTag(tag string) error { + if tag == "" { return nil } - if IsReservedTag(v) { - return fmt.Errorf("invalid %s %q: %q is a reserved prefix — choose a different tag", EnvModelTag, v, ReservedTagPrefix) + if IsReservedTag(tag) { + return fmt.Errorf("uses reserved prefix %q — choose a different tag", ReservedTagPrefix) } - if !tagRegex.MatchString(v) { - return fmt.Errorf("invalid %s %q: must match OCI tag regex %s", EnvModelTag, v, tagRegex.String()) + if !tagRegex.MatchString(tag) { + return fmt.Errorf("doesn't match OCI tag regex %s", tagRegex.String()) + } + return nil +} + +// validateTagEnv adds the environment variable name and value to tag errors. +func validateTagEnv(v string) error { + if err := ValidateTag(v); err != nil { + return fmt.Errorf("invalid %s %q: %w", EnvModelTag, v, err) } return nil } diff --git a/pkg/model/resolver.go b/pkg/model/resolver.go index 2e6a72f0e8..7cd9086777 100644 --- a/pkg/model/resolver.go +++ b/pkg/model/resolver.go @@ -7,6 +7,7 @@ import ( "strings" "github.com/docker/docker/api/types/image" + v1 "github.com/google/go-containerregistry/pkg/v1" "github.com/replicate/cog/pkg/config" "github.com/replicate/cog/pkg/docker/command" @@ -226,14 +227,23 @@ func (r *Resolver) Build(ctx context.Context, src *Source, opts BuildOptions) (* // errors (malformed COG_MODEL_TAG, reserved prefix, image:+env // mode mix-up, etc.) surface before kicking off a Docker build. // ErrNoModelRef just means there's no model field — fall back to - // FormatImage. - format := FormatImage - ref, err := ResolveModelRef(src.Config.Image, src.Config.Model) - if err != nil && !errors.Is(err, ErrNoModelRef) { - return nil, err + // FormatImage. Push can supply an already-resolved format and ref + // so a positional target overrides every COG_MODEL* variable. + format := opts.Format + ref := opts.ModelRef + if format == 0 { + format = FormatImage + var err error + ref, err = ResolveModelRef(src.Config.Image, src.Config.Model) + if err != nil && !errors.Is(err, ErrNoModelRef) { + return nil, err + } + if ref != nil { + format = FormatBundle + } } - if ref != nil { - format = FormatBundle + if format == FormatBundle && ref == nil { + return nil, fmt.Errorf("bundle build requires a model ref") } ib := NewImageBuilder(r.factory, r.docker, src, opts) @@ -264,12 +274,32 @@ func (r *Resolver) Build(ctx context.Context, src *Source, opts BuildOptions) (* if weightErr != nil { return nil, weightErr } - m.Weights = weights + m.Weights, weightErr = orderWeightsByConfig(weights, src.Config.Weights) + if weightErr != nil { + return nil, weightErr + } } return m, nil } +func orderWeightsByConfig(weights []Weight, configured []config.WeightSource) ([]Weight, error) { + byName := make(map[string]Weight, len(weights)) + for _, weight := range weights { + byName[weight.Name] = weight + } + + ordered := make([]Weight, len(configured)) + for i, configuredWeight := range configured { + weight, ok := byName[configuredWeight.Name] + if !ok { + return nil, fmt.Errorf("weight %q is missing from weights.lock; re-run 'cog weights import'", configuredWeight.Name) + } + ordered[i] = weight + } + return ordered, nil +} + // Push pushes a Model to a container registry and returns an // enriched copy of the Model with all registry references populated // (Model.Ref pinned to the resulting digest, image artifact and @@ -296,18 +326,32 @@ func (r *Resolver) Push(ctx context.Context, m *Model, opts PushOptions) (*Model if opts.OnFallback != nil { imagePushOpts = append(imagePushOpts, WithOnFallback(opts.OnFallback)) } - if err := r.imagePusher.Push(ctx, imgArtifact, imagePushOpts...); err != nil { + imageResult, err := r.imagePusher.PushWithResult(ctx, imgArtifact, imagePushOpts...) + if err != nil { return nil, err } - // Enrich the image artifact with the pushed digest. On registries - // that don't support HEAD on tags, fall back to the original Model - // rather than failing a successful push. FormatImage is the legacy - // single-image path; bundle pushes surface HEAD failures as errors. - desc, err := r.registry.GetDescriptor(ctx, imgArtifact.Reference) - if err != nil { - console.Debugf("post-push HEAD on %q failed; returning Model without digest enrichment: %v", imgArtifact.Reference, err) - return m, nil + // Prefer the digest reported by the upload itself. This pins the result to + // the manifest Cog pushed even if another process immediately moves the tag. + var desc v1.Descriptor + if imageResult.Digest != "" { + digest, digestErr := v1.NewHash(imageResult.Digest) + if digestErr != nil { + return nil, fmt.Errorf("image upload reported invalid manifest digest %q: %w", imageResult.Digest, digestErr) + } + desc = v1.Descriptor{Digest: digest, Size: imageResult.Size} + } else { + // Older or alternate command implementations may not expose Docker's + // push result. Preserve the legacy non-JSON tag lookup fallback, but + // callers requiring immutable output fail rather than trusting it. + if opts.RequireDigest { + return nil, fmt.Errorf("resolve pushed image digest: the upload did not report the manifest digest") + } + desc, err = r.registry.GetDescriptor(ctx, imgArtifact.Reference) + if err != nil { + console.Debugf("post-push HEAD on %q failed; returning Model without digest enrichment: %v", imgArtifact.Reference, err) + return m, nil + } } enrichedImage := imgArtifact.WithDigest(repoFromReference(imgArtifact.Reference), desc) diff --git a/pkg/model/resolver_test.go b/pkg/model/resolver_test.go index 62ff410503..11cf2ffbd8 100644 --- a/pkg/model/resolver_test.go +++ b/pkg/model/resolver_test.go @@ -20,11 +20,14 @@ import ( "github.com/replicate/cog/pkg/weights/lockfile" ) +const testDockerPushDigest = "sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + // mockDocker implements command.Command for testing. type mockDocker struct { inspectFunc func(ctx context.Context, ref string) (*image.InspectResponse, error) pullFunc func(ctx context.Context, ref string, force bool) (*image.InspectResponse, error) pushFunc func(ctx context.Context, ref string) error + pushResult *command.PushResult tagFunc func(ctx context.Context, source, target string) error removeFunc func(ctx context.Context, ref string) error imageSaveFunc func(ctx context.Context, imageRef string) (io.ReadCloser, error) @@ -51,6 +54,16 @@ func (m *mockDocker) Push(ctx context.Context, ref string) error { return errors.New("mockDocker.Push not implemented") } +func (m *mockDocker) PushWithResult(ctx context.Context, ref string) (command.PushResult, error) { + if err := m.Push(ctx, ref); err != nil { + return command.PushResult{}, err + } + if m.pushResult != nil { + return *m.pushResult, nil + } + return command.PushResult{Digest: testDockerPushDigest}, nil +} + func (m *mockDocker) LoadUserInformation(ctx context.Context, registryHost string) (*command.UserInfo, error) { return nil, errors.New("mockDocker.LoadUserInformation not implemented") } @@ -1176,6 +1189,22 @@ func TestResolver_Build_PopulatesWeights(t *testing.T) { require.True(t, m.IsBundle()) } +func TestOrderWeightsByConfig(t *testing.T) { + weights := []Weight{ + {Name: "beta", Digest: "sha256:bbbb"}, + {Name: "alpha", Digest: "sha256:aaaa"}, + } + configured := []config.WeightSource{ + {Name: "alpha"}, + {Name: "beta"}, + } + + ordered, err := orderWeightsByConfig(weights, configured) + require.NoError(t, err) + require.Equal(t, []string{"alpha", "beta"}, []string{ordered[0].Name, ordered[1].Name}) + require.Equal(t, []string{"sha256:aaaa", "sha256:bbbb"}, []string{ordered[0].Digest, ordered[1].Digest}) +} + func TestResolver_Build_WithWeightsIsBundle(t *testing.T) { clearModelEnv(t) imageDigest := "sha256:a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2" @@ -1324,6 +1353,38 @@ func TestResolver_Build_InvalidEnvVarSurfaces(t *testing.T) { require.Contains(t, err.Error(), "reserved prefix") } +func TestResolver_Build_ExplicitFormatIgnoresModelEnvironment(t *testing.T) { + clearModelEnv(t) + t.Setenv(EnvModel, "invalid target") + resolver, src := newFormatTestResolver(t) + src.Config.Image = "registry.example.com/user/image" + + m, err := resolver.Build(context.Background(), src, BuildOptions{ + ImageName: "registry.example.com/target/image", + Format: FormatImage, + }) + require.NoError(t, err) + require.Equal(t, FormatImage, m.Format) + require.Nil(t, m.Ref) +} + +func TestResolver_Build_ExplicitBundleRefIgnoresModelEnvironment(t *testing.T) { + clearModelEnv(t) + t.Setenv(EnvModelTag, "cog-reserved") + resolver, src := newFormatTestResolver(t) + src.Config.Model = "registry.example.com/user/model" + ref := &ResolvedRef{Registry: "registry.example.com", Repo: "target/model", Tag: "v2"} + + m, err := resolver.Build(context.Background(), src, BuildOptions{ + ImageName: "registry.example.com/target/model:v2", + Format: FormatBundle, + ModelRef: ref, + }) + require.NoError(t, err) + require.Equal(t, FormatBundle, m.Format) + require.Same(t, ref, m.Ref) +} + func TestIndexDetectionHelpers(t *testing.T) { t.Run("findImageManifest", func(t *testing.T) { manifests := []registry.PlatformManifest{