Pants can't build multi stage Docker image with named stage and image digest

[lvets]

Describe the bug

TL;DR: Trying to build a multi stage Docker build with using a <image>[@<digest>] instead of <image>[:<tag>] results in a pants error. There is a workaround though, see "Additional info" below.

I have a Dockerfile with the following FROM stages:

❯ grep FROM src/docker/u21repr/Dockerfile
FROM golang:1.23-bookworm AS builder
FROM gcr.io/distroless/python3-debian12@sha256:8e432c787b5c0697dfbfd783120351d90fd5f23ba9fff29532bbdbb87bc13160 AS runtime
FROM gcr.io/distroless/python3-debian12:debug AS upstream_debug
FROM runtime AS debug
FROM debug AS smoke_tests

❯

Trying to build this with pants results in the following error:

❯ pants --docker-build-verbose package src/docker/u21repr:runtime
09:08:08.03 [ERROR] 1 Exception encountered:

Engine traceback:
  in `package` goal

DockerBuildTargetStageError: The 'target_stage' field in `docker_image` src/docker/u21repr:runtime was set to 'runtime', but there is no such stage in `src/docker/u21repr/Dockerfile`. Available stages: builder, debug, smoke_tests, upstream_debug.

❯

The docker_image() target in the BUILD file looks like:

docker_image(
    name="runtime",
    repository=IMAGE_REPO_NAME,
    target_stage="runtime",
    image_tags=[
        "latest",
    ],
    output={
        "type": "docker", # Default output type.
        "rewrite-timestamp": "true"
    },
    dependencies=CONTEXT_INCLUDES,
    build_platform=BUILD_PLATFORMS,
)

• Using a tag in the runtime FROM line works.
• Changing runtime to foobar in both the BUILD and Dockerfile results in the same error (but it can't find foobar instead of runtime)

Pants version

2.23.1

OS

Tested on macOS.

Additional info

• Slack discussion: https://pantsbuild.slack.com/archives/C046T6T9U/p1737058948138679

Workaround

@huonw did show me a workaround that I didn't know works with Docker. Apparently, if you insert a tag and then the sha digest, Docker ignores the tag and pulls the image with sha digest. This works for docker pull as well as in the pants build.
If I change the line FROM gcr.io/distroless/python3-debian12@sha256:8e432c787b5c0697dfbfd783120351d90fd5f23ba9fff29532bbdbb87bc13160 AS runtime
to FROM gcr.io/distroless/python3-debian12:totally_made_up_not_a_real_tag@sha256:8e432c787b5c0697dfbfd783120351d90fd5f23ba9fff29532bbdbb87bc13160 AS runtime, the build works again.

[lilatomic]

Looks like this is an oversight in our parser. It only lists stages that have a tag. (I think the logic is that if it can't find a tag, it's an image reference. We just didn't consider references without a tag but with a sha).

pants/src/python/pants/backend/docker/subsystems/dockerfile_wrapper_script.py

Lines 203 to 208 in ed3f23e

	return tuple(
	f"{stage} {tag}"
	for stage, name_parts in self.from_baseimages()
	for tag in [_get_tag(name_parts[-1])]
	if tag
	)

That's also the case with the Rust/tree-sitter-based parser

pants/src/rust/engine/dep_inference/src/dockerfile/mod.rs

Lines 143 to 148 in a91e9f2

	if let Some(tag) = tag {
	let mut stage_collector = StageCollector::new(self.code);
	stage_collector.walk(&mut cursor);
	self.version_tags
	.insert(stage_collector.get_stage(self.stage_counter), tag);
	}