Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 40 additions & 0 deletions e2e/tests/up/gpu_detection.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
package up

import (
"os"
"os/exec"
"path/filepath"

docker "github.com/devsy-org/devsy/pkg/docker"
"github.com/onsi/ginkgo/v2"
"github.com/onsi/gomega"
)

var _ = ginkgo.Describe(
"GPU detection graceful fallback",
ginkgo.Label("up-gpu-detection"),
func() {
ginkgo.It("should not error under the test environment's container runtime", func() {
initialDir, err := os.Getwd()
gomega.Expect(err).NotTo(gomega.HaveOccurred())

dockerCmd := testDockerCommand
if _, err := exec.LookPath("podman"); err == nil {
if _, dockerErr := exec.LookPath(testDockerCommand); dockerErr != nil {
dockerCmd = "podman"
}
}

binDir := filepath.Join(initialDir, "bin")
h := &docker.DockerHelper{DockerCommand: filepath.Join(binDir, dockerCmd)}
if _, err := os.Stat(h.DockerCommand); os.IsNotExist(err) {
h.DockerCommand = dockerCmd
}

available, err := h.GPUSupportEnabled()
gomega.Expect(err).NotTo(gomega.HaveOccurred(),
"GPU detection should not error regardless of runtime")
ginkgo.GinkgoWriter.Printf("GPU available: %v (runtime: %s)\n", available, dockerCmd)
})
},
)
4 changes: 2 additions & 2 deletions pkg/agent/delivery/factory.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ type FactoryOptions struct {
HelperImage string
IsRemoteDocker bool
ContainerID string
ExecFunc inject.ExecFunc
ExecFunc inject.ExecFunc //nolint:staticcheck // legacy delivery strategies require this type
}

func NewAgentDelivery(opts FactoryOptions) AgentDelivery {
Expand Down Expand Up @@ -105,7 +105,7 @@ func CommandFunc(
stdin io.Reader, stdout io.Writer, stderr io.Writer,
) error,
workspaceID string,
) inject.ExecFunc {
) inject.ExecFunc { //nolint:staticcheck // bridges driver command signature to legacy ExecFunc
return func(
ctx context.Context,
command string,
Expand Down
2 changes: 1 addition & 1 deletion pkg/agent/delivery/kubernetes.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ import (
var _ AgentDelivery = (*KubernetesDelivery)(nil)

type KubernetesDelivery struct {
ExecFunc inject.ExecFunc
ExecFunc inject.ExecFunc //nolint:staticcheck // K8s exec routing requires this type
}

func (d *KubernetesDelivery) Phase() DeliveryPhase {
Expand Down
2 changes: 1 addition & 1 deletion pkg/agent/delivery/legacy_shell.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ import (
// Deprecated: LegacyShellDelivery is deprecated. Platform-native AgentDelivery implementations
// (LocalDockerDelivery, RemoteDockerDelivery, KubernetesDelivery) are the replacements.
type LegacyShellDelivery struct {
ExecFunc inject.ExecFunc
ExecFunc inject.ExecFunc //nolint:staticcheck // deprecated injection path
DownloadURL string
Timeout func() *agent.InjectOptions
}
Expand Down
28 changes: 23 additions & 5 deletions pkg/docker/helper.go
Original file line number Diff line number Diff line change
Expand Up @@ -59,12 +59,10 @@ type DockerHelper struct {
}

func (r *DockerHelper) GPUSupportEnabled() (bool, error) {
out, err := r.buildCmd(context.TODO(), "info", "-f", "{{.Runtimes.nvidia}}").Output()
if err != nil {
return false, command.WrapCommandError(out, err)
if r.IsPodman() {
return r.podmanGPUAvailable()
}

return strings.Contains(string(out), "nvidia-container-runtime"), nil
return r.dockerGPUAvailable()
}

func (r *DockerHelper) FindDevContainer(
Expand Down Expand Up @@ -450,6 +448,26 @@ func (r *DockerHelper) GetContainerLogs(
return cmd.Run()
}

func (r *DockerHelper) dockerGPUAvailable() (bool, error) {
out, err := r.buildCmd(context.TODO(), "info", "-f", "{{.Runtimes.nvidia}}").Output()
if err != nil {
log.Debugf("GPU detection: docker runtime query failed, assuming no GPU: %v", err)
return false, nil
}

return strings.Contains(string(out), "nvidia-container-runtime"), nil
}

func (r *DockerHelper) podmanGPUAvailable() (bool, error) {
out, err := r.buildCmd(context.TODO(), "info", "-f", "{{.Host.CDIDevices}}").Output()
if err != nil {
log.Debugf("GPU detection: podman CDI query failed, assuming no GPU: %v", err)
return false, nil
}

return strings.Contains(strings.ToLower(string(out)), "nvidia"), nil
}

func (r *DockerHelper) buildCmd(ctx context.Context, args ...string) *exec.Cmd {
cmd := exec.CommandContext(ctx, r.DockerCommand, args...)
if r.Environment != nil {
Expand Down
95 changes: 95 additions & 0 deletions pkg/docker/helper_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,95 @@
package docker

import (
"os"
"path/filepath"
"testing"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)

func writeScript(t *testing.T, dir, name, script string) string {
t.Helper()
path := filepath.Join(dir, name)
//nolint:gosec // test helper script needs exec bit
require.NoError(t, os.WriteFile(path, []byte(script), 0o755))
return path
}

func TestGPUSupportEnabled_DockerWithNvidia(t *testing.T) {
tmp := t.TempDir()
bin := writeScript(t, tmp, "docker-fake", `#!/bin/sh
case "$1" in
--version) echo "Docker version 24.0.7, build afdd53b";;
info) echo "nvidia-container-runtime";;
esac
`)

h := &DockerHelper{DockerCommand: bin}
got, err := h.GPUSupportEnabled()

assert.NoError(t, err)
assert.True(t, got, "should detect GPU when Docker nvidia runtime is present")
}

func TestGPUSupportEnabled_DockerWithoutNvidia(t *testing.T) {
tmp := t.TempDir()
bin := writeScript(t, tmp, "docker-fake", `#!/bin/sh
case "$1" in
--version) echo "Docker version 24.0.7, build afdd53b";;
info) echo "{}";;
esac
`)

h := &DockerHelper{DockerCommand: bin}
got, err := h.GPUSupportEnabled()

assert.NoError(t, err)
assert.False(t, got, "should not detect GPU when Docker nvidia runtime is absent")
}

func TestGPUSupportEnabled_PodmanWithCDINvidia(t *testing.T) {
tmp := t.TempDir()
bin := writeScript(t, tmp, "podman-fake", `#!/bin/sh
case "$1" in
--version) echo "podman version 4.9.3";;
info) echo "[nvidia.com/gpu=all]";;
esac
`)

h := &DockerHelper{DockerCommand: bin}
got, err := h.GPUSupportEnabled()

assert.NoError(t, err)
assert.True(t, got, "should detect GPU when Podman CDI has nvidia device")
}

func TestGPUSupportEnabled_PodmanWithoutCDINvidia(t *testing.T) {
tmp := t.TempDir()
bin := writeScript(t, tmp, "podman-fake", `#!/bin/sh
case "$1" in
--version) echo "podman version 4.9.3";;
info) echo "[]";;
esac
`)

h := &DockerHelper{DockerCommand: bin}
got, err := h.GPUSupportEnabled()

assert.NoError(t, err)
assert.False(t, got, "should not detect GPU when Podman CDI has no nvidia device")
}

func TestGPUSupportEnabled_CommandFailure(t *testing.T) {
tmp := t.TempDir()
bin := writeScript(t, tmp, "bad-runtime", `#!/bin/sh
exit 1
`)

h := &DockerHelper{DockerCommand: bin}
got, err := h.GPUSupportEnabled()

assert.NoError(t, err, "should not propagate error on command failure")
assert.False(t, got, "should fall back to no GPU on command failure")
}
Loading