chore(docker): reduce size between docker builds

[keturn]

by adding a layer with all the pytorch dependencies that don't change most of the time.

Summary

Every time the main docker images rebuild and I pull main-cuda, it gets another 3+ GB, which seems like about a zillion times too much since most things don't change from one commit on main to the next.

This is an attempt to follow the guidance in Using uv in Docker: Intermediate Layers so there's one layer that installs all the dependencies—including PyTorch with its bundled nvidia libraries—before the project's own frequently-changing files are copied in to the image.

Related Issues / Discussions

• Improved docker layer cache with uv
• astral: Can uv pip install torch, but not uv sync it

QA Instructions

Hopefully the CI system building the docker images is sufficient.

But there is one change to pyproject.toml related to xformers, so it'd be worth checking that python -m xformers.info still says it has triton on the platforms that expect it.

Merge Plan

I don't expect this to be a disruptive merge.

(An earlier revision of this PR moved the venv, but I've reverted that change at ebr's recommendation.)

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

[keturn]

This Dockerfile is also quirky in that it separates builder and runtime stages, but then it puts all the build-deps in the runtime stage anyway (to build patchmatch?), which kinda defeats the purpose. But I think we can leave that alone for now as an independent concern.

There is one thing I haven't confirmed for the space savings: my test builds have been with podman (buildah), not buildkit. buildah doesn't support COPY --link, so the interaction between the stages and layers isn't exactly the same…

[ebr]

hey @keturn , thanks for this! we tried some time ago with pip, but it caused more headaches than was worth the trouble. Now with uv this seems like a very sound approach.
This will work well for someone building the images locally. But just to set expectations: we're not currently doing any caching of Docker layers in GHA runners (we might revisit that at some point), so the intermediate layers will be rebuilt anyway. So IF the expectation is to not have to pull the 2.5GB pytorch layer every time - that is unfortunately still going to continue happening, at least for now. But again, if you're building the image locally, it should help quite a bit on rebuilds.
That said, the GHA docker builds are failing right now, so once I fix that perhaps we can rebase this PR and I'll come back to re-reviewing it. Will keep you posted.

[ebr]

This Dockerfile is also quirky in that it separates builder and runtime stages, but then it puts all the build-deps in the runtime stage anyway (to build patchmatch?), which kinda defeats the purpose. But I think we can leave that alone for now as an independent concern.

Indeed, this is needed to build patchmatch. Hope at some point we no longer have to do this, but for now it's ok. There's still a small benefit to using the runner image because we just don't need to worry about any other cruft that may be left in the builder image, so i'd like to keep it this way for the time being.

[keturn]

But just to set expectations: we're not currently doing any caching of Docker layers in GHA runners (we might revisit that at some point), so the intermediate layers will be rebuilt anyway. So IF the expectation is to not have to pull the 2.5GB pytorch layer every time - that is unfortunately still going to continue happening, at least for now.

This got me to re-visit the docs on docker cache invalidation. From what I gather, the upshot is

• no build-cache means those layers will be re-built every time, and
• bit-for-bit reproducible docker builds are still an esoteric subject, so
• a re-built layer gets published as a new thing, even if its contents are semantically the same?

well, that's a bit disappointing, but I guess this PR is still a step in the right direction if we are going all-in with uv.

[ebr]

OK, looks good to me. We'll merge it and I will work on some ideas around re-introducing caching.