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

pkg/doctor/check_env_gpu_compat.go

@@ -0,0 +1,199 @@
+package doctor
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"os/exec"
+	"sort"
+	"strconv"
+	"strings"
+
+	"github.com/replicate/cog/pkg/util/version"
+)
+
+// GPUFloor is the oldest torch (and CUDA) release that ships executable kernels for a
+// given compute capability.
+//
+// CUDA cubins are binary-compatible only within a compute capability major; crossing majors
+// requires embedded PTX for the driver to JIT, which older torch wheels do not ship. Both
+// bounds matter: torch 2.7.0+cu118 is a genuine 2.7 build with no Blackwell kernels.
+type GPUFloor struct {
+	MinTorch string
+	MinCUDA  string
+}
+
+// gpuFloors maps compute capability (major, minor) to its floor. Lookup takes the highest
+// entry the device meets or exceeds; a device below every entry yields no finding.
+//
+// Measured by installing each wheel and reading torch._C._cuda_getArchFlags(), which works
+// without a GPU. Each floor is bracketed: the named version ships kernels for the
+// capability, the release below it does not.
+//
+//	torch 2.7.0+cu128   cubins: sm_75 sm_80 sm_86 sm_90 sm_100 sm_120   PTX: compute_120
+//	torch 2.6.0+cu126   cubins: sm_50 ... sm_90                         PTX: none
+//	torch 2.4.1+cu124   cubins: sm_50 ... sm_90                         PTX: none
+//	torch 2.0.1+cu118   cubins: sm_37 ... sm_90                         PTX: none
+//	torch 1.13.1+cu117  cubins: sm_37 ... sm_86                         PTX: none
+//
+// No rows below sm_90: every probed wheel already covers Turing/Ampere/Ada, so those floors
+// cannot be bracketed. When extending, bracket the new row the same way rather than guess.
+//
+// sm_89 (Ada) appears in no cubin list yet RTX 40xx works: sm_86 cubins run on sm_89 (same
+// major). sm_120 gets no such fallback from sm_90.
+var gpuFloors = map[[2]int]GPUFloor{
+	{12, 0}: {MinTorch: "2.7.0", MinCUDA: "12.8"}, // Blackwell consumer, RTX 50xx
+	{10, 0}: {MinTorch: "2.7.0", MinCUDA: "12.8"}, // Blackwell datacenter, B100/B200
+	{9, 0}:  {MinTorch: "2.0.1", MinCUDA: "11.8"}, // Hopper, H100
+}
+
+// GPUCompatibilityCheck verifies that the resolved torch/CUDA versions can execute on this
+// machine's GPU. Cog derives CUDA from the framework pin without consulting the device, so
+// a build can succeed and produce an image where every CUDA operation fails at runtime with
+// "no kernel image is available for execution on the device".
+type GPUCompatibilityCheck struct{}
+
+func (c *GPUCompatibilityCheck) Name() string        { return "env-gpu-compat" }
+func (c *GPUCompatibilityCheck) Group() Group        { return GroupEnvironment }
+func (c *GPUCompatibilityCheck) Description() string { return "GPU compatibility" }
+
+func (c *GPUCompatibilityCheck) Check(ctx *CheckContext) ([]Finding, error) {
+	if os.Getenv("COG_SKIP_GPU_CHECK") != "" {
+		return nil, nil
+	}
+	if ctx.Config == nil || !ctx.Config.Build.GPU {
+		return nil, nil
+	}
+
+	torchVersion, hasTorch := ctx.Config.TorchVersion()
+	if !hasTorch {
+		return nil, nil
+	}
+
+	capability, ok := detectComputeCapability(ctx.ctx)
+	if !ok {
+		// No GPU on this machine (or no driver); nothing to compare against.
+		return nil, nil
+	}
+
+	return evaluateGPUCompat(capability, torchVersion, ctx.Config.Build.CUDA), nil
+}
+
+// evaluateGPUCompat compares a resolved torch/CUDA pair against the floor for the given
+// compute capability. Kept free of I/O so it is testable without a GPU.
+func evaluateGPUCompat(capability [2]int, torchVersion string, cuda string) []Finding {
+	floor, ok := floorFor(capability)
+	if !ok {
+		// Older than any floor we know; say nothing rather than guess.
+		return nil
+	}
+
+	if torchVersion == "" {
+		// torch listed but unpinned; pip resolves a current release. Nothing to compare.
+		return nil
+	}
+
+	torchOK := version.GreaterOrEqual(torchVersion, floor.MinTorch)
+	cudaOK := cuda == "" || version.GreaterOrEqual(cuda, floor.MinCUDA)
+	if torchOK && cudaOK {
+		return nil
+	}
+
+	sm := fmt.Sprintf("sm_%d%d", capability[0], capability[1])
+	// Warning, not error: the image is valid and runs on other hardware; it only fails on
+	// this machine's GPU (same taxonomy as PythonVersionCheck).
+	return []Finding{{
+		Severity: SeverityWarning,
+		Message: fmt.Sprintf(
+			"torch==%s (CUDA %s) ships no kernels for %s, the compute capability of this machine's GPU. "+
+				"The image will build, but every CUDA operation in it will fail at runtime with "+
+				"\"no kernel image is available for execution on the device\".",
+			torchVersion, cudaDisplay(cuda), sm,
+		),
+		Remediation: fmt.Sprintf(
+			"%s requires torch>=%s built against CUDA>=%s. Pin a newer torch, or set "+
+				"COG_SKIP_GPU_CHECK=1 if you are building for a different GPU than the one in this machine.",
+			sm, floor.MinTorch, floor.MinCUDA,
+		),
+	}}
+}
+
+func (c *GPUCompatibilityCheck) Fix(_ *CheckContext, _ []Finding) error {
+	return ErrNoAutoFix
+}
+
+func cudaDisplay(cuda string) string {
+	if cuda == "" {
+		return "unset"
+	}
+	return cuda
+}
+
+// detectComputeCapability returns the lowest compute capability among the machine's GPUs:
+// the image has to run on the weakest device present.
+func detectComputeCapability(ctx context.Context) ([2]int, bool) {
+	execCtx, cancel := context.WithTimeout(ctx, envCheckTimeout)
+	defer cancel()
+
+	out, err := exec.CommandContext(execCtx,
+		"nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader").Output()
+	if err != nil {
+		return [2]int{}, false
+	}
+
+	var caps [][2]int
+	for line := range strings.SplitSeq(string(out), "\n") {
+		if c, ok := parseCapability(line); ok {
+			caps = append(caps, c)
+		}
+	}
+	if len(caps) == 0 {
+		return [2]int{}, false
+	}
+
+	sort.Slice(caps, func(i, j int) bool {
+		if caps[i][0] != caps[j][0] {
+			return caps[i][0] < caps[j][0]
+		}
+		return caps[i][1] < caps[j][1]
+	})
+	return caps[0], true
+}
+
+// parseCapability parses a single nvidia-smi compute_cap value, e.g. "12.0".
+func parseCapability(s string) ([2]int, bool) {
+	majorStr, minorStr, found := strings.Cut(strings.TrimSpace(s), ".")
+	if !found {
+		return [2]int{}, false
+	}
+	major, err := strconv.Atoi(majorStr)
+	if err != nil {
+		return [2]int{}, false
+	}
+	minor, err := strconv.Atoi(minorStr)
+	if err != nil {
+		return [2]int{}, false
+	}
+	return [2]int{major, minor}, true
+}
+
+// floorFor returns the floor for the highest table entry the device meets or exceeds.
+func floorFor(capability [2]int) (GPUFloor, bool) {
+	keys := make([][2]int, 0, len(gpuFloors))
+	for k := range gpuFloors {
+		keys = append(keys, k)
+	}
+	sort.Slice(keys, func(i, j int) bool {
+		if keys[i][0] != keys[j][0] {
+			return keys[i][0] > keys[j][0]
+		}
+		return keys[i][1] > keys[j][1]
+	})
+
+	for _, k := range keys {
+		if capability[0] > k[0] || (capability[0] == k[0] && capability[1] >= k[1]) {
+			return gpuFloors[k], true
+		}
+	}
+	return GPUFloor{}, false
+}

pkg/doctor/check_env_gpu_compat_test.go

@@ -0,0 +1,171 @@
+package doctor
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/replicate/cog/pkg/config"
+)
+
+func gpuContext(t *testing.T, gpu bool, torch string, cuda string) *CheckContext {
+	t.Helper()
+	cfg := &config.Config{
+		Build: &config.Build{
+			GPU:            gpu,
+			PythonVersion:  "3.11",
+			PythonPackages: []string{"torch==" + torch},
+			CUDA:           cuda,
+		},
+	}
+	// Complete populates the requirements content that TorchVersion reads; without it the
+	// check skips before reaching the comparison.
+	require.NoError(t, cfg.Complete(t.TempDir()))
+	return &CheckContext{ctx: context.Background(), ProjectDir: t.TempDir(), Config: cfg}
+}
+
+func TestGPUCompatibilityCheck_RunsWithoutError(t *testing.T) {
+	// A GPU may or may not be present in the test environment; just ensure no panic or error.
+	ctx := gpuContext(t, true, "2.4.1", "12.4")
+	_, err := (&GPUCompatibilityCheck{}).Check(ctx)
+	require.NoError(t, err)
+}
+
+func TestGPUCompatibilityCheck_SkipsWhenDisabled(t *testing.T) {
+	t.Setenv("COG_SKIP_GPU_CHECK", "1")
+	ctx := gpuContext(t, true, "2.4.1", "12.4")
+
+	findings, err := (&GPUCompatibilityCheck{}).Check(ctx)
+
+	require.NoError(t, err)
+	require.Empty(t, findings)
+}
+
+func TestGPUCompatibilityCheck_SkipsWhenGPUNotRequested(t *testing.T) {
+	ctx := gpuContext(t, false, "2.4.1", "12.4")
+
+	findings, err := (&GPUCompatibilityCheck{}).Check(ctx)
+
+	require.NoError(t, err)
+	require.Empty(t, findings)
+}
+
+func TestGPUCompatibilityCheck_SkipsWhenNoConfig(t *testing.T) {
+	ctx := &CheckContext{ctx: context.Background(), ProjectDir: t.TempDir()}
+
+	findings, err := (&GPUCompatibilityCheck{}).Check(ctx)
+
+	require.NoError(t, err)
+	require.Empty(t, findings)
+}
+
+func TestGPUCompatibilityCheck_FixReturnsNoAutoFix(t *testing.T) {
+	require.ErrorIs(t, (&GPUCompatibilityCheck{}).Fix(nil, nil), ErrNoAutoFix)
+}
+
+// The comparison logic is pure; these exercise it without a GPU.
+
+func TestEvaluateGPUCompat(t *testing.T) {
+	for _, tt := range []struct {
+		name       string
+		capability [2]int
+		torch      string
+		cuda       string
+		fires      bool
+	}{
+		// sm_120 (Blackwell consumer): needs torch>=2.7.0 AND CUDA>=12.8.
+		{"sm_120 old torch fires", [2]int{12, 0}, "2.4.1", "12.4", true},
+		{"sm_120 new torch passes", [2]int{12, 0}, "2.7.0", "12.8", false},
+		{"sm_120 local version tag passes", [2]int{12, 0}, "2.7.1+cu128", "12.8", false},
+		// Both bounds are load-bearing: torch 2.7.0+cu118 is a genuine 2.7 build with no
+		// Blackwell kernels in it.
+		{"sm_120 new torch but old CUDA fires", [2]int{12, 0}, "2.7.0", "11.8", true},
+		{"sm_120 unset CUDA defers to torch bound", [2]int{12, 0}, "2.7.0", "", false},
+		// torch listed but unpinned: pip resolves a current release; nothing to compare.
+		{"unpinned torch is silent", [2]int{12, 0}, "", "", false},
+		// sm_90 (Hopper): needs torch>=2.0.1.
+		{"sm_90 old torch fires", [2]int{9, 0}, "1.13.1", "11.7", true},
+		{"sm_90 new torch passes", [2]int{9, 0}, "2.0.1", "11.8", false},
+		// An unknown newer capability falls back to the highest floor it exceeds.
+		{"unknown sm_130 old torch fires", [2]int{13, 0}, "2.4.1", "12.4", true},
+		{"unknown sm_130 new torch passes", [2]int{13, 0}, "2.7.0", "12.8", false},
+		// Below every floor: no row, no finding.
+		{"sm_86 below all floors is silent", [2]int{8, 6}, "1.13.1", "11.7", false},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			findings := evaluateGPUCompat(tt.capability, tt.torch, tt.cuda)
+			if !tt.fires {
+				require.Empty(t, findings)
+				return
+			}
+			require.Len(t, findings, 1)
+			assert.Equal(t, SeverityWarning, findings[0].Severity)
+			assert.Contains(t, findings[0].Message, "torch=="+tt.torch)
+			assert.Contains(t, findings[0].Message, "no kernels for")
+			assert.Contains(t, findings[0].Remediation, "COG_SKIP_GPU_CHECK=1")
+		})
+	}
+}
+
+func TestEvaluateGPUCompat_Message(t *testing.T) {
+	findings := evaluateGPUCompat([2]int{12, 0}, "2.4.1", "12.4")
+
+	require.Len(t, findings, 1)
+	assert.Contains(t, findings[0].Message, "torch==2.4.1 (CUDA 12.4) ships no kernels for sm_120")
+	assert.Contains(t, findings[0].Message, "no kernel image is available for execution on the device")
+	assert.Contains(t, findings[0].Remediation, "sm_120 requires torch>=2.7.0 built against CUDA>=12.8")
+}
+
+func TestFloorFor(t *testing.T) {
+	for _, tt := range []struct {
+		name  string
+		cc    [2]int
+		torch string
+		cuda  string
+		found bool
+	}{
+		{"blackwell consumer sm_120", [2]int{12, 0}, "2.7.0", "12.8", true},
+		{"blackwell datacenter sm_100", [2]int{10, 0}, "2.7.0", "12.8", true},
+		{"hopper sm_90", [2]int{9, 0}, "2.0.1", "11.8", true},
+		// An unknown newer capability falls back to the highest floor it exceeds.
+		{"unknown newer sm_130", [2]int{13, 0}, "2.7.0", "12.8", true},
+		// No rows below sm_90: every probed wheel already covers Ada/Ampere/Turing, so those
+		// floors cannot be bracketed. Say nothing rather than guess.
+		{"ada sm_89 has no floor", [2]int{8, 9}, "", "", false},
+		{"ampere sm_86 has no floor", [2]int{8, 6}, "", "", false},
+		{"turing sm_75 has no floor", [2]int{7, 5}, "", "", false},
+		{"pascal sm_61 has no floor", [2]int{6, 1}, "", "", false},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			floor, ok := floorFor(tt.cc)
+			require.Equal(t, tt.found, ok)
+			if tt.found {
+				require.Equal(t, tt.torch, floor.MinTorch)
+				require.Equal(t, tt.cuda, floor.MinCUDA)
+			}
+		})
+	}
+}
+
+func TestParseCapability(t *testing.T) {
+	for _, tt := range []struct {
+		in    string
+		want  [2]int
+		valid bool
+	}{
+		{"12.0", [2]int{12, 0}, true},
+		{"8.6\n", [2]int{8, 6}, true},
+		{"  9.0  ", [2]int{9, 0}, true},
+		{"", [2]int{}, false},
+		{"not a version", [2]int{}, false},
+		{"12", [2]int{}, false},
+	} {
+		got, ok := parseCapability(tt.in)
+		require.Equal(t, tt.valid, ok, "input %q", tt.in)
+		if tt.valid {
+			require.Equal(t, tt.want, got, "input %q", tt.in)
+		}
+	}
+}

pkg/doctor/registry.go

@@ -19,5 +19,6 @@ func AllChecks() []Check {
 		// Environment checks
 		&DockerCheck{},
 		&PythonVersionCheck{},
+		&GPUCompatibilityCheck{},
 	}
 }