/
spec_nvidia_hook.go
48 lines (39 loc) · 1.08 KB
/
spec_nvidia_hook.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
package mgr
import (
"os/exec"
"path"
"github.com/alibaba/pouch/pkg/utils"
"github.com/opencontainers/runtime-spec/specs-go"
)
var (
// nvidiaHookName is a custom OCI prestart hook binary to runc in order to enable GPU containers.
nvidiaHookName = "nvidia-container-runtime-hook"
)
func setNvidiaHook(c *Container, spec *SpecWrapper) error {
n := c.HostConfig.NvidiaConfig
// to make compatible for k8s.
// if user set environments of NVIDIA, then set prestart hook
kv := utils.ConvertKVStrToMapWithNoErr(c.Config.Env)
_, hasEnvCapabilities := kv["NVIDIA_DRIVER_CAPABILITIES"]
_, hasEnvDevices := kv["NVIDIA_VISIBLE_DEVICES"]
if n == nil && !hasEnvCapabilities && !hasEnvDevices {
return nil
}
hookPath := ""
if !path.IsAbs(nvidiaHookName) {
var err error
hookPath, err = exec.LookPath(nvidiaHookName)
if err != nil {
return err
}
} else {
hookPath = nvidiaHookName
}
args := []string{hookPath}
nvidiaPrestart := specs.Hook{
Path: hookPath,
Args: append(args, "prestart"),
}
spec.s.Hooks.Prestart = append(spec.s.Hooks.Prestart, nvidiaPrestart)
return nil
}