From 2671a25dc9724a6273cc651ea1218ab0205ea2c8 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Fri, 21 Aug 2026 11:40:00 -0500 Subject: [PATCH 1/3] fix: support uv-managed CUDA 13 Cog bases --- pkg/dockerfile/standard_generator.go | 13 ++++----- pkg/dockerfile/standard_generator_test.go | 35 +++++++++++++++++++++++ 2 files changed, 41 insertions(+), 7 deletions(-) diff --git a/pkg/dockerfile/standard_generator.go b/pkg/dockerfile/standard_generator.go index e1af7fe7c4..36472fe11b 100644 --- a/pkg/dockerfile/standard_generator.go +++ b/pkg/dockerfile/standard_generator.go @@ -150,15 +150,14 @@ func (g *StandardGenerator) SetBreakSystemPackages(breakSystemPackages bool) { } // needsBreakSystemPackages reports whether pip invocations need -// --break-system-packages. True when either the caller opted in explicitly -// (SetBreakSystemPackages) or the generated Dockerfile installs Python via -// `uv python install` (inside installPythonCUDA). The latter happens when the -// base image is nvidia/cuda — which has no Python — as opposed to python:X-slim -// or r8.im/cog-base, which ship their own. uv marks its installed Pythons as -// externally managed (PEP 668). +// --break-system-packages. True when the caller opts in explicitly, the +// generated Dockerfile installs Python via uv, or a CUDA 13+ Cog base supplies +// uv-managed Python. uv marks its installed Pythons as externally managed +// (PEP 668). func (g *StandardGenerator) needsBreakSystemPackages() bool { return g.breakSystemPackages || - (g.Config.Build.GPU && g.useCudaBaseImage && !g.IsUsingCogBaseImage()) + (g.Config.Build.GPU && g.useCudaBaseImage && !g.IsUsingCogBaseImage()) || + (g.IsUsingCogBaseImage() && version.GreaterOrEqual(g.Config.Build.CUDA, "13.0")) } func (g *StandardGenerator) uvPipInstallFlags(flags string) string { diff --git a/pkg/dockerfile/standard_generator_test.go b/pkg/dockerfile/standard_generator_test.go index ca94fdb883..5f13a0ea49 100644 --- a/pkg/dockerfile/standard_generator_test.go +++ b/pkg/dockerfile/standard_generator_test.go @@ -586,6 +586,41 @@ predict: predict.py:Predictor require.Contains(t, actual, `uv pip install --break-system-packages --no-cache cog`) } +// CUDA 13+ Cog base images install Python with uv, so their Python environment +// is externally managed regardless of whether the config also sets gpu: true. +func TestGPUCogBasePathIncludesBreakSystemPackages(t *testing.T) { + for _, gpu := range []bool{true, false} { + t.Run(fmt.Sprintf("gpu=%t", gpu), func(t *testing.T) { + tmpDir := t.TempDir() + yaml := fmt.Sprintf(` +build: + gpu: %t + cuda: "13.0" + python_version: "3.13" + python_packages: + - torch==2.11.0 + - pandas==2.0.3 +predict: predict.py:Predictor +`, gpu) + conf, err := config.FromYAML([]byte(yaml)) + require.NoError(t, err) + require.NoError(t, conf.Complete("")) + command := dockertest.NewMockCommand() + client := registrytest.NewMockRegistryClient() + client.AddMockImage(BaseImageName("13.0", "3.13", "2.11.0")) + gen, err := NewStandardGenerator(conf, tmpDir, t.TempDir(), "", command, client, true) + require.NoError(t, err) + gen.SetUseCogBaseImage(true) + pypiWheels(gen) + _, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), "r8.im/replicate/cog-test") + require.NoError(t, err) + + require.Contains(t, actual, `uv run pip install --break-system-packages --cache-dir /root/.cache/pip -r /tmp/requirements.txt`) + require.Contains(t, actual, `uv pip install --break-system-packages --no-cache cog`) + }) + } +} + // CPU builds use python:X-slim where Python is not uv-managed, // so --break-system-packages must NOT appear. func TestCPUPathOmitsBreakSystemPackages(t *testing.T) { From 73e890ed7d94102462ac3613b2d9eba3f69b642d Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Fri, 21 Aug 2026 11:53:11 -0500 Subject: [PATCH 2/3] fix: add python alias for uv-managed CUDA images --- pkg/dockerfile/standard_generator.go | 13 +++++++++++-- pkg/dockerfile/standard_generator_test.go | 6 ++++-- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/pkg/dockerfile/standard_generator.go b/pkg/dockerfile/standard_generator.go index 36472fe11b..5ef494670c 100644 --- a/pkg/dockerfile/standard_generator.go +++ b/pkg/dockerfile/standard_generator.go @@ -214,6 +214,7 @@ func (g *StandardGenerator) GenerateInitialSteps(ctx context.Context) (string, e envs, aptInstalls, g.installUV(), + g.installPythonAlias(), } // Install user packages before the SDK so that changing the SDK // wheel (e.g. via --cog-ref or COG_SDK_WHEEL) does not invalidate @@ -505,6 +506,13 @@ func (g *StandardGenerator) installPython() (string, error) { return "", nil } +func (g *StandardGenerator) installPythonAlias() string { + if g.IsUsingCogBaseImage() && version.GreaterOrEqual(g.Config.Build.CUDA, "13.0") { + return `RUN ln -sf /usr/bin/python3 /usr/local/bin/python` + } + return "" +} + func (g *StandardGenerator) installUV() string { return `COPY --from=ghcr.io/astral-sh/uv:` + UVVersion + ` /uv /uvx /usr/local/bin/ ENV UV_SYSTEM_PYTHON=true` @@ -522,9 +530,10 @@ func (g *StandardGenerator) installPythonCUDA() (string, error) { ca-certificates \ && rm -rf /var/lib/apt/lists/* ` + g.installUV() + "\n" + fmt.Sprintf(`RUN uv python install %s && \ - ln -sf $(uv python find %s) /usr/bin/python3 + ln -sf $(uv python find %s) /usr/bin/python3 && \ + ln -sf $(uv python find %s) /usr/local/bin/python ENV UV_PYTHON=%s -ENV PATH="/usr/local/bin:$PATH"`, py, py, py), nil +ENV PATH="/usr/local/bin:$PATH"`, py, py, py, py), nil } // resolveCogWheelConfigs resolves and caches the cog and coglet wheel configs. diff --git a/pkg/dockerfile/standard_generator_test.go b/pkg/dockerfile/standard_generator_test.go index 5f13a0ea49..3b9f19d77d 100644 --- a/pkg/dockerfile/standard_generator_test.go +++ b/pkg/dockerfile/standard_generator_test.go @@ -75,10 +75,11 @@ func testInstallPython(version string) string { COPY --from=ghcr.io/astral-sh/uv:`+UVVersion+` /uv /uvx /usr/local/bin/ ENV UV_SYSTEM_PYTHON=true RUN uv python install %s && \ - ln -sf $(uv python find %s) /usr/bin/python3 + ln -sf $(uv python find %s) /usr/bin/python3 && \ + ln -sf $(uv python find %s) /usr/local/bin/python ENV UV_PYTHON=%s ENV PATH="/usr/local/bin:$PATH" -`, version, version, version) +`, version, version, version, version) } func TestGenerateEmptyCPU(t *testing.T) { @@ -617,6 +618,7 @@ predict: predict.py:Predictor require.Contains(t, actual, `uv run pip install --break-system-packages --cache-dir /root/.cache/pip -r /tmp/requirements.txt`) require.Contains(t, actual, `uv pip install --break-system-packages --no-cache cog`) + require.Contains(t, actual, `RUN ln -sf /usr/bin/python3 /usr/local/bin/python`) }) } } From 2f9327608c57d6bae93e97e4c8f7b63375734469 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Fri, 21 Aug 2026 12:09:30 -0500 Subject: [PATCH 3/3] chore: update h2 to 0.4.16 --- crates/Cargo.lock | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/crates/Cargo.lock b/crates/Cargo.lock index db63e70db6..f0a6cf6075 100644 --- a/crates/Cargo.lock +++ b/crates/Cargo.lock @@ -503,7 +503,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -712,9 +712,9 @@ checksum = "e629b9b98ef3dd8afe6ca2bd0f89306cec16d43d907889945bc5d6687f2f13c7" [[package]] name = "h2" -version = "0.4.13" +version = "0.4.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f44da3a8150a6703ed5d34e164b875fd14c2cdab9af1252a9a1020bde2bdc54" +checksum = "a9f37a958b41b3b19ee2707c06439c0e9e547e847223eb791ecb0cb821c65e27" dependencies = [ "atomic-waker", "bytes", @@ -2234,7 +2234,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2291,7 +2291,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2727,7 +2727,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -3433,7 +3433,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]]