feat: detect when the resolved torch/CUDA ships no kernels for the local GPU

docs/environment.md

@@ -107,6 +107,18 @@ Skips the `cog doctor` Docker environment check. Set it to any non-empty value.
 $ COG_SKIP_DOCKER_CHECK=1 cog doctor
 ```
 
+### `COG_SKIP_GPU_CHECK`
+
+Skips the `cog doctor` GPU compatibility check, which warns when the resolved torch/CUDA
+versions ship no kernels for the GPU in the current machine. Set it to any non-empty value.
+
+Use this when you are deliberately building for a different GPU than the one in this
+machine, so the check should not warn about the local device.
+
+```console
+$ COG_SKIP_GPU_CHECK=1 cog doctor
+```
+
 ### `COG_CACHE_DIR`
 
 Overrides Cog's local cache root.

pkg/config/compatibility.go

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"regexp"
 	"sort"
 	"strings"
 
@@ -119,6 +120,20 @@ func init() {
 	TorchCompatibilityMatrix = filteredTorchCompatibilityMatrix
 }
 
+var cudaIndexURLRe = regexp.MustCompile(`cu(\d{2,})`)
+
+// cudaVersionFromIndexURL derives the CUDA version from a pytorch package index URL, e.g.
+// "https://download.pytorch.org/whl/cu128/" -> "12.8". Returns ok=false when the URL has no
+// cuXYZ segment, as with the cpu and rocm indexes.
+func cudaVersionFromIndexURL(indexURL string) (string, bool) {
+	m := cudaIndexURLRe.FindStringSubmatch(indexURL)
+	if m == nil {
+		return "", false
+	}
+	digits := m[1]
+	return digits[:len(digits)-1] + "." + digits[len(digits)-1:], true
+}
+
 func cudaVersionFromTorchPlusVersion(ver string) (string, string) {
 	const cudaVersionPrefix = "cu"
 

pkg/config/compatibility_test.go

@@ -18,3 +18,24 @@ func TestCudasFromTorchWithCUVersionModifier(t *testing.T) {
 	require.Equal(t, cudas[0], "11.8")
 	require.Nil(t, err)
 }
+
+func TestCUDAVersionFromIndexURL(t *testing.T) {
+	for _, tt := range []struct {
+		url  string
+		want string
+		ok   bool
+	}{
+		{"https://download.pytorch.org/whl/cu128/", "12.8", true},
+		{"https://download.pytorch.org/whl/cu118", "11.8", true},
+		{"https://download.pytorch.org/whl/cu92/", "9.2", true},
+		{"https://download.pytorch.org/whl/cpu/", "", false},
+		{"https://download.pytorch.org/whl/rocm6.2/", "", false},
+		{"", "", false},
+	} {
+		t.Run(tt.url, func(t *testing.T) {
+			got, ok := cudaVersionFromIndexURL(tt.url)
+			require.Equal(t, tt.ok, ok)
+			require.Equal(t, tt.want, got)
+		})
+	}
+}

pkg/config/config.go

@@ -191,6 +191,36 @@ func (c *Config) TorchVersion() (string, bool) {
 	return c.pythonPackageVersion("torch")
 }
 
+// ResolvedTorchWheel reports the torch wheel a GPU build will actually install: the exact
+// pinned version and the CUDA it was built against, read from the resolved package index
+// (e.g. .../whl/cu128 -> "12.8"). It mirrors PythonRequirementsForArch, so it reflects the
+// installed wheel rather than the raw requirement line, and falls back to Build.CUDA when the
+// resolved index carries no CUDA. Returns ok=false unless torch is pinned to one exact version.
+func (c *Config) ResolvedTorchWheel(goos string, goarch string) (torchVersion string, cuda string, ok bool) {
+	for _, pkg := range c.Build.pythonRequirementsContent {
+		if requirements.NormalizePackageName(requirements.PackageName(pkg)) != "torch" {
+			continue
+		}
+		resolved, _, extraIndexURLs, err := c.pythonPackageForArch(pkg, goos, goarch)
+		if err != nil {
+			return "", "", false
+		}
+		version, hasExact := requirements.ExactVersion(resolved)
+		if !hasExact {
+			return "", "", false
+		}
+		resolvedCUDA := c.Build.CUDA
+		for _, indexURL := range extraIndexURLs {
+			if derived, found := cudaVersionFromIndexURL(indexURL); found {
+				resolvedCUDA = derived
+				break
+			}
+		}
+		return version, resolvedCUDA, true
+	}
+	return "", "", false
+}
+
 func (c *Config) TorchvisionVersion() (string, bool) {
 	return c.pythonPackageVersion("torchvision")
 }

pkg/config/config_test.go

@@ -133,6 +133,78 @@ foo==1.0.0`
 	require.Equal(t, expected, requirements)
 }
 
+func TestResolvedTorchWheel(t *testing.T) {
+	for _, tt := range []struct {
+		name         string
+		requirements string
+		cuda         string
+		wantVersion  string
+		wantCUDA     string
+		wantOK       bool
+	}{
+		{
+			name:         "plain pin resolves cuda from the matrix index",
+			requirements: "torch==2.4.1",
+			cuda:         "12.4",
+			wantVersion:  "2.4.1",
+			wantCUDA:     "12.4",
+			wantOK:       true,
+		},
+		{
+			// The +cu118 tag is discarded during resolution: cuda: "12.8" selects the cu128
+			// wheel, so the reported CUDA is 12.8, not the requirement's 11.8.
+			name:         "local tag yields to build cuda via the resolved index",
+			requirements: "torch==2.7.0+cu118",
+			cuda:         "12.8",
+			wantVersion:  "2.7.0",
+			wantCUDA:     "12.8",
+			wantOK:       true,
+		},
+		{
+			// An explicit index is passed through verbatim and wins over cuda: "12.8".
+			name:         "explicit extra-index-url wins over build cuda",
+			requirements: "torch==2.7.0 --extra-index-url=https://download.pytorch.org/whl/cu118",
+			cuda:         "12.8",
+			wantVersion:  "2.7.0",
+			wantCUDA:     "11.8",
+			wantOK:       true,
+		},
+		{
+			name:         "non-exact pin is not resolvable",
+			requirements: "torch<2.7.0",
+			cuda:         "12.4",
+			wantOK:       false,
+		},
+		{
+			name:         "no torch present",
+			requirements: "numpy==1.26.0",
+			cuda:         "12.4",
+			wantOK:       false,
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			tmpDir := t.TempDir()
+			require.NoError(t, os.WriteFile(path.Join(tmpDir, "requirements.txt"), []byte(tt.requirements), 0o644))
+			cfg := &Config{
+				Build: &Build{
+					GPU:                true,
+					PythonVersion:      "3.11",
+					PythonRequirements: "requirements.txt",
+					CUDA:               tt.cuda,
+				},
+			}
+			require.NoError(t, cfg.Complete(tmpDir))
+
+			version, cuda, ok := cfg.ResolvedTorchWheel("linux", "amd64")
+			require.Equal(t, tt.wantOK, ok)
+			if tt.wantOK {
+				require.Equal(t, tt.wantVersion, version)
+				require.Equal(t, tt.wantCUDA, cuda)
+			}
+		})
+	}
+}
+
 func TestPythonRequirementsResolvesPythonPackagesAndCudaVersionsWithExtraIndexURL(t *testing.T) {
 	tmpDir := t.TempDir()
 	err := os.WriteFile(path.Join(tmpDir, "requirements.txt"), []byte(`torch==1.12.1