Fix uv-managed CUDA base compatibility

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 {
@@ -215,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
@@ -506,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`
@@ -523,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.

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) {
@@ -586,6 +587,42 @@ 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`)
+			require.Contains(t, actual, `RUN ln -sf /usr/bin/python3 /usr/local/bin/python`)
+		})
+	}
+}
+
 // CPU builds use python:X-slim where Python is not uv-managed,
 // so --break-system-packages must NOT appear.
 func TestCPUPathOmitsBreakSystemPackages(t *testing.T) {