diff --git a/docs/environment.md b/docs/environment.md index ee9b3afc6b..f163f9c4f9 100644 --- a/docs/environment.md +++ b/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. diff --git a/docs/llms.txt b/docs/llms.txt index e8419e39a8..3a36955a55 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -925,6 +925,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. diff --git a/pkg/config/config.go b/pkg/config/config.go index 621b4bf49f..f4ba6c207b 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -191,6 +191,14 @@ func (c *Config) TorchVersion() (string, bool) { return c.pythonPackageVersion("torch") } +// TorchExactVersion returns the torch version only when it is pinned with an exact `==` +// specifier. Unlike TorchVersion, which returns the first version token regardless of the +// comparator, this reports ("", false) for ranges, inequalities, direct URLs, and unpinned +// requirements, where pip is free to resolve a different release than the token suggests. +func (c *Config) TorchExactVersion() (string, bool) { + return c.pythonPackageExactVersion("torch") +} + func (c *Config) TorchvisionVersion() (string, bool) { return c.pythonPackageVersion("torchvision") } @@ -239,6 +247,17 @@ func (c *Config) pythonPackageVersion(name string) (version string, ok bool) { return "", false } +// pythonPackageExactVersion returns the version of the named package only when it is pinned +// with an exact `==` specifier; see ExactVersion for the requirements it rejects. +func (c *Config) pythonPackageExactVersion(name string) (version string, ok bool) { + for _, pkg := range c.Build.pythonRequirementsContent { + if requirements.PackageName(pkg) == name { + return requirements.ExactVersion(pkg) + } + } + return "", false +} + func splitPythonVersion(version string) (major int, minor int, err error) { version = strings.TrimSpace(version) parts := strings.SplitN(version, ".", 3) diff --git a/pkg/doctor/gpu_compatibility.go b/pkg/doctor/gpu_compatibility.go new file mode 100644 index 0000000000..cadbe67ef6 --- /dev/null +++ b/pkg/doctor/gpu_compatibility.go @@ -0,0 +1,242 @@ +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 == nil || !ctx.Config.Build.GPU { + return nil, nil + } + + // Only an exact `==` pin names one concrete version to reason about; for a range or an + // unpinned requirement pip is free to resolve a different release, so we say nothing. + torchVersion, hasTorch := ctx.Config.TorchExactVersion() + if !hasTorch { + return nil, nil + } + + capabilities, ok := detectComputeCapabilities(ctx.ctx) + if !ok { + // No GPU on this machine (or no driver); nothing to compare against. + return nil, nil + } + + return evaluateGPUCompatAll(capabilities, torchVersion, ctx.Config.Build.CUDA), nil +} + +// evaluateGPUCompatAll evaluates every distinct compute capability the machine reports. A +// mixed-GPU host must satisfy the floor of each device, not just the numerically lowest: +// torch 2.4.1 clears the sm_90 floor but not sm_120, so a host with both still needs a +// finding. One finding per failing capability, in ascending order for stable output. +func evaluateGPUCompatAll(capabilities [][2]int, torchVersion string, cuda string) []Finding { + seen := make(map[[2]int]bool, len(capabilities)) + var findings []Finding + for _, capability := range capabilities { + if seen[capability] { + continue + } + seen[capability] = true + findings = append(findings, evaluateGPUCompat(capability, torchVersion, cuda)...) + } + return findings +} + +// 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 + } + + // Compare the torch release without its local tag: version.GreaterOrEqual folds the + // local modifier into equality, so 2.7.0+cu128 would otherwise read as neither greater + // than nor equal to the 2.7.0 floor and warn falsely at the exact boundary. The CUDA + // bound is checked separately from ctx.Config.Build.CUDA, so nothing is lost here. + torchOK := version.GreaterOrEqual(stripLocalVersion(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 +} + +// stripLocalVersion drops a PEP 440 local version tag ("+cu128") from a version string. +func stripLocalVersion(v string) string { + release, _, _ := strings.Cut(v, "+") + return release +} + +// queryComputeCaps runs nvidia-smi and returns its raw stdout. It is a package variable so +// tests can substitute a fake without a GPU or the real binary on PATH. +var queryComputeCaps = func(ctx context.Context) ([]byte, error) { + execCtx, cancel := context.WithTimeout(ctx, envCheckTimeout) + defer cancel() + return exec.CommandContext(execCtx, + "nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader").Output() +} + +// detectComputeCapabilities returns the distinct compute capabilities of the machine's GPUs, +// sorted ascending. The image must run on every device present, so all are reported rather +// than only the weakest. Returns ok=false when nvidia-smi is absent, fails, or names none. +func detectComputeCapabilities(ctx context.Context) ([][2]int, bool) { + out, err := queryComputeCaps(ctx) + if err != nil { + return nil, false + } + caps := parseComputeCapabilities(string(out)) + if len(caps) == 0 { + return nil, false + } + return caps, true +} + +// parseComputeCapabilities extracts the distinct, ascending compute capabilities from +// nvidia-smi's compute_cap output. Unparsable lines are skipped so malformed or partial +// output degrades to whatever valid rows it contains rather than failing outright. +func parseComputeCapabilities(out string) [][2]int { + seen := make(map[[2]int]bool) + var caps [][2]int + for line := range strings.SplitSeq(out, "\n") { + if c, ok := parseCapability(line); ok && !seen[c] { + seen[c] = true + caps = append(caps, c) + } + } + 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 +} + +// 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 +} diff --git a/pkg/doctor/gpu_compatibility_test.go b/pkg/doctor/gpu_compatibility_test.go new file mode 100644 index 0000000000..958ff62e84 --- /dev/null +++ b/pkg/doctor/gpu_compatibility_test.go @@ -0,0 +1,310 @@ +package doctor + +import ( + "context" + "errors" + "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() + return gpuContextRaw(t, gpu, "torch=="+torch, cuda) +} + +// gpuContextRaw builds a context from a raw torch requirement line, so tests can exercise +// ranges and unpinned requirements that gpuContext's "torch==" shorthand cannot express. +func gpuContextRaw(t *testing.T, gpu bool, torchReq string, cuda string) *CheckContext { + t.Helper() + cfg := &config.Config{ + Build: &config.Build{ + GPU: gpu, + PythonVersion: "3.11", + PythonPackages: []string{torchReq}, + CUDA: cuda, + }, + } + // Complete populates the requirements content that TorchExactVersion 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} +} + +// stubComputeCaps replaces the nvidia-smi probe with fixed output for the test's duration, +// so the production detection path runs deterministically without a GPU. +func stubComputeCaps(t *testing.T, out string, err error) { + t.Helper() + prev := queryComputeCaps + queryComputeCaps = func(context.Context) ([]byte, error) { return []byte(out), err } + t.Cleanup(func() { queryComputeCaps = prev }) +} + +func TestGPUCompatibilityCheck_FiresForIncompatibleGPU(t *testing.T) { + stubComputeCaps(t, "12.0\n", nil) + ctx := gpuContext(t, true, "2.4.1", "12.4") + + findings, err := (&GPUCompatibilityCheck{}).Check(ctx) + + require.NoError(t, err) + require.Len(t, findings, 1) + assert.Contains(t, findings[0].Message, "sm_120") +} + +func TestGPUCompatibilityCheck_PassesForCompatibleGPU(t *testing.T) { + stubComputeCaps(t, "12.0\n", nil) + ctx := gpuContext(t, true, "2.7.0", "12.8") + + findings, err := (&GPUCompatibilityCheck{}).Check(ctx) + + require.NoError(t, err) + require.Empty(t, findings) +} + +// A host with an old-but-supported GPU and a newer unsupported one must still warn: the +// numerically lowest device (sm_90) clears torch 2.4.1's floor but sm_120 does not. +func TestGPUCompatibilityCheck_MixedGPUsFireForStrictestFloor(t *testing.T) { + stubComputeCaps(t, "9.0\n12.0\n", nil) + ctx := gpuContext(t, true, "2.4.1", "12.4") + + findings, err := (&GPUCompatibilityCheck{}).Check(ctx) + + require.NoError(t, err) + require.Len(t, findings, 1) + assert.Contains(t, findings[0].Message, "sm_120") + assert.NotContains(t, findings[0].Message, "sm_90") +} + +func TestGPUCompatibilityCheck_SkipsWhenNvidiaSMIFails(t *testing.T) { + stubComputeCaps(t, "", errors.New("nvidia-smi: command not found")) + ctx := gpuContext(t, true, "2.4.1", "12.4") + + findings, err := (&GPUCompatibilityCheck{}).Check(ctx) + + require.NoError(t, err) + require.Empty(t, findings) +} + +func TestGPUCompatibilityCheck_SkipsWhenOutputMalformed(t *testing.T) { + stubComputeCaps(t, "N/A\n[Insufficient Permissions]\n", nil) + ctx := gpuContext(t, true, "2.4.1", "12.4") + + findings, err := (&GPUCompatibilityCheck{}).Check(ctx) + + require.NoError(t, err) + require.Empty(t, findings) +} + +// A non-exact pin lets pip resolve any release, so the check must not warn on the token even +// when the detected GPU would fail that token as an exact version. +func TestGPUCompatibilityCheck_SkipsNonExactPin(t *testing.T) { + stubComputeCaps(t, "12.0\n", nil) + ctx := gpuContextRaw(t, true, "torch<2.7.0", "12.4") + + findings, err := (&GPUCompatibilityCheck{}).Check(ctx) + + require.NoError(t, err) + require.Empty(t, findings) +} + +func TestGPUCompatibilityCheck_RegisteredInRunner(t *testing.T) { + var registered bool + for _, c := range AllChecks() { + if _, ok := c.(*GPUCompatibilityCheck); ok { + registered = true + break + } + } + require.True(t, registered, "GPUCompatibilityCheck should be registered in AllChecks") +} + +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_SkipsWhenNoBuildSection(t *testing.T) { + ctx := &CheckContext{ctx: context.Background(), ProjectDir: t.TempDir(), Config: &config.Config{}} + + 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}, + // Exactly at the floor with a local tag: the +cu128 modifier must not read as below + // 2.7.0. Regression for version.GreaterOrEqual folding the local tag into equality. + {"sm_120 exact floor with local tag passes", [2]int{12, 0}, "2.7.0+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 TestEvaluateGPUCompatAll(t *testing.T) { + t.Run("distinct capabilities each yield a finding", func(t *testing.T) { + findings := evaluateGPUCompatAll([][2]int{{10, 0}, {12, 0}}, "2.4.1", "12.4") + require.Len(t, findings, 2) + }) + t.Run("a passing capability suppresses its finding", func(t *testing.T) { + findings := evaluateGPUCompatAll([][2]int{{9, 0}, {12, 0}}, "2.4.1", "12.4") + require.Len(t, findings, 1) + assert.Contains(t, findings[0].Message, "sm_120") + }) + t.Run("all-compatible capabilities yield nothing", func(t *testing.T) { + findings := evaluateGPUCompatAll([][2]int{{9, 0}, {12, 0}}, "2.7.0", "12.8") + require.Empty(t, findings) + }) + t.Run("duplicate capabilities are evaluated once", func(t *testing.T) { + findings := evaluateGPUCompatAll([][2]int{{12, 0}, {12, 0}}, "2.4.1", "12.4") + require.Len(t, findings, 1) + }) +} + +func TestParseComputeCapabilities(t *testing.T) { + for _, tt := range []struct { + name string + out string + want [][2]int + }{ + {"single", "12.0\n", [][2]int{{12, 0}}}, + {"multiple sorted ascending", "12.0\n9.0\n", [][2]int{{9, 0}, {12, 0}}}, + {"duplicates collapsed", "12.0\n12.0\n", [][2]int{{12, 0}}}, + {"malformed lines skipped", "N/A\n12.0\ngarbage\n", [][2]int{{12, 0}}}, + {"all malformed yields none", "N/A\n[Insufficient Permissions]\n", nil}, + {"empty yields none", "", nil}, + } { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.want, parseComputeCapabilities(tt.out)) + }) + } +} + +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) + } + } +} diff --git a/pkg/doctor/registry.go b/pkg/doctor/registry.go index 95038f3817..a914b25dc9 100644 --- a/pkg/doctor/registry.go +++ b/pkg/doctor/registry.go @@ -19,5 +19,6 @@ func AllChecks() []Check { // Environment checks &DockerCheck{}, &PythonVersionCheck{}, + &GPUCompatibilityCheck{}, } } diff --git a/pkg/requirements/requirements.go b/pkg/requirements/requirements.go index b3b8ac59ab..d14cbf0f54 100644 --- a/pkg/requirements/requirements.go +++ b/pkg/requirements/requirements.go @@ -151,6 +151,38 @@ func PackageName(pipRequirement string) string { return "" } +// versionSpecifierRe matches a single PEP 440 version specifier: an operator followed +// by a version token. Used to distinguish an exact `==` pin from ranges and inequalities. +var versionSpecifierRe = regexp.MustCompile(`(===|==|!=|~=|<=|>=|<|>)\s*([^\s,;|]+)`) + +// directURLRe matches a PEP 508 direct reference (e.g. "torch @ https://..."). +var directURLRe = regexp.MustCompile(`@\s*\S+`) + +// ExactVersion returns the version from a pip requirement only when it is a single exact +// `==` pin, e.g. "torch==2.7.0" (inline pip options and the CUDA local tag are preserved). +// Ranges, inequalities, wildcard pins, compound specifiers, direct URLs, and unpinned +// requirements return ("", false): they do not name one concrete version to reason about. +func ExactVersion(pipRequirement string) (string, bool) { + // Drop environment markers (";" and after); a marker's own comparators are not the pin. + if idx := strings.Index(pipRequirement, ";"); idx >= 0 { + pipRequirement = pipRequirement[:idx] + } + // A direct URL reference is never an exact version pin. + if directURLRe.MatchString(pipRequirement) { + return "", false + } + matches := versionSpecifierRe.FindAllStringSubmatch(pipRequirement, -1) + if len(matches) != 1 { + // Zero specifiers (unpinned) or more than one (a range like ">=2.0,<3.0"). + return "", false + } + op, ver := matches[0][1], matches[0][2] + if op != "==" || strings.Contains(ver, "*") { + return "", false + } + return ver, true +} + func Versions(pipRequirement string) []string { var versions []string diff --git a/pkg/requirements/requirements_test.go b/pkg/requirements/requirements_test.go index dabbfe9de3..c4fb9df82a 100644 --- a/pkg/requirements/requirements_test.go +++ b/pkg/requirements/requirements_test.go @@ -429,6 +429,40 @@ func TestVersions(t *testing.T) { require.Equal(t, versions, []string{"https://some.domain/package.whl"}) } +func TestExactVersion(t *testing.T) { + for _, tt := range []struct { + req string + want string + ok bool + }{ + // Exact pins resolve to one concrete version. + {"torch==2.7.0", "2.7.0", true}, + {"torch == 2.7.0", "2.7.0", true}, + {"torch[opt]==2.7.0", "2.7.0", true}, + {"torch==2.7.0+cu128", "2.7.0+cu128", true}, + {"torch==2.7.0 --extra-index-url=https://download.pytorch.org/whl/cu128", "2.7.0", true}, + {"torch==2.7.0 ; python_version >= '3.10'", "2.7.0", true}, + // Ranges and inequalities do not pin one version. + {"torch<2.7.0", "", false}, + {"torch>=2.7.0", "", false}, + {"torch~=2.7.0", "", false}, + {"torch!=2.7.0", "", false}, + {"torch>=2.0,<3.0", "", false}, + {"torch==2.7.0,!=2.7.1", "", false}, + // Wildcards, direct URLs, and unpinned requirements are not exact. + {"torch==2.7.*", "", false}, + {"torch @ https://download.pytorch.org/whl/torch-2.7.0.whl", "", false}, + {"torch", "", false}, + {"", "", false}, + } { + t.Run(tt.req, func(t *testing.T) { + got, ok := ExactVersion(tt.req) + require.Equal(t, tt.ok, ok, "req %q", tt.req) + require.Equal(t, tt.want, got, "req %q", tt.req) + }) + } +} + func checkRequirements(t *testing.T, expected []string, actual []string) { t.Helper() for n, expectLine := range expected {