dilim/moby
1
0
Fork 0
mirror of https://github.com/moby/moby.git synced 2026-08-08 17:11:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #49952 from sgopinath1/49824-amd-gpu

Browse Source 
Added support for AMD GPUs in "docker run --gpus".
This commit is contained in:
Sebastiaan van Stijn
2025-06-11 14:48:58 +02:00
committed by GitHub
parent 349a2d0050 e32715ec03
commit 3b1d2f746d
2 changed files with 51 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
daemon/devices_amd_linux.go Normal file
View File

@@ -0,0 +1,27 @@
package daemon
import (
"strings"
"github.com/opencontainers/runtime-spec/specs-go"
)
func setAMDGPUs(s *specs.Spec, dev *deviceInstance) error {
req := dev.req
if req.Count != 0 && len(req.DeviceIDs) > 0 {
return errConflictCountDeviceIDs
}
switch {
case len(req.DeviceIDs) > 0:
s.Process.Env = append(s.Process.Env, "AMD_VISIBLE_DEVICES="+strings.Join(req.DeviceIDs, ","))
case req.Count > 0:
s.Process.Env = append(s.Process.Env, "AMD_VISIBLE_DEVICES="+countToDevices(req.Count))
case req.Count < 0:
s.Process.Env = append(s.Process.Env, "AMD_VISIBLE_DEVICES=all")
case req.Count == 0:
s.Process.Env = append(s.Process.Env, "AMD_VISIBLE_DEVICES=void")
}
return nil
}

35
daemon/nvidia_linux.go → daemon/devices_nvidia_linux.go
View File

@@ -17,7 +17,10 @@ import (
var errConflictCountDeviceIDs = errors.New("cannot set both Count and DeviceIDs on device request")
const nvidiaHook = "nvidia-container-runtime-hook"
const (
nvidiaHook = "nvidia-container-runtime-hook"
amdContainerRuntimeExecutableName = "amd-container-runtime"
)
// These are NVIDIA-specific capabilities stolen from github.com/containerd/containerd/contrib/nvidia.allCaps
var allNvidiaCaps = map[nvidia.Capability]struct{}{
@@ -30,19 +33,29 @@ var allNvidiaCaps = map[nvidia.Capability]struct{}{
}
func init() {
if _, err := exec.LookPath(nvidiaHook); err != nil {
// do not register Nvidia driver if helper binary is not present.
// Register Nvidia driver if Nvidia helper binary is present.
if _, err := exec.LookPath(nvidiaHook); err == nil {
capset := capabilities.Set{"gpu": struct{}{}, "nvidia": struct{}{}}
for c := range allNvidiaCaps {
capset[string(c)] = struct{}{}
}
registerDeviceDriver("nvidia", &deviceDriver{
capset: capset,
updateSpec: setNvidiaGPUs,
})
return
}
capset := capabilities.Set{"gpu": struct{}{}, "nvidia": struct{}{}}
nvidiaDriver := &deviceDriver{
capset: capset,
updateSpec: setNvidiaGPUs,
// Register AMD driver if AMD helper binary is present.
if _, err := exec.LookPath(amdContainerRuntimeExecutableName); err == nil {
registerDeviceDriver("amd", &deviceDriver{
capset: capabilities.Set{"gpu": struct{}{}, "amd": struct{}{}},
updateSpec: setAMDGPUs,
})
return
}
for c := range allNvidiaCaps {
nvidiaDriver.capset[string(c)] = struct{}{}
}
registerDeviceDriver("nvidia", nvidiaDriver)
// No "gpu" capability
}
func setNvidiaGPUs(s *specs.Spec, dev *deviceInstance) error {