From d155342ce7f185ba0f6f772e7dfcf5328f4e9949 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 17:27:52 +0000 Subject: [PATCH] feat: add new CDI hook to set CUDA memory limits Signed-off-by: Tariq Ibrahim get cgroup path from procfs instead Signed-off-by: Tariq Ibrahim (cherry picked from commit 2d761670b7f9d3adec465c4cf41faa86b700ba9d) --- cmd/nvidia-cdi-hook/README.md | 1 + .../apply-cuda-memory-limits.go | 218 ++++++++++++++++++ .../apply-cuda-memory-limits_test.go | 98 ++++++++ cmd/nvidia-cdi-hook/commands/commands.go | 2 + cmd/nvidia-ctk/cdi/generate/generate_test.go | 80 ++++++- internal/discover/hooks.go | 11 +- internal/info/cgroup/cgroup_path.go | 64 +++++ internal/info/cgroup/cgroup_path_test.go | 93 ++++++++ pkg/nvcdi/cuda-memory-limits.go | 71 ++++++ pkg/nvcdi/full-gpu-nvml.go | 6 + pkg/nvcdi/lib-csv_test.go | 8 + 11 files changed, 650 insertions(+), 2 deletions(-) create mode 100644 cmd/nvidia-cdi-hook/apply-cuda-memory-limits/apply-cuda-memory-limits.go create mode 100644 cmd/nvidia-cdi-hook/apply-cuda-memory-limits/apply-cuda-memory-limits_test.go create mode 100644 internal/info/cgroup/cgroup_path.go create mode 100644 internal/info/cgroup/cgroup_path_test.go create mode 100644 pkg/nvcdi/cuda-memory-limits.go diff --git a/cmd/nvidia-cdi-hook/README.md b/cmd/nvidia-cdi-hook/README.md index 79354e7d9..6352f530d 100644 --- a/cmd/nvidia-cdi-hook/README.md +++ b/cmd/nvidia-cdi-hook/README.md @@ -32,3 +32,4 @@ The `nvidia-cdi-hook` CLI provides the following functionality: * `enable-cuda-compat` - Ensure that the directory containing the CUDA compat libraries is added to the ldconfig search path if required. * `disable-device-node-modification` - Ensure that the `/proc/driver/nvidia/params` file present in the container does not allow device node modifications. * `update-application-profile` - Update driver settings through "application profiles". Currently, this hook sets `EGLVisibleDGPUDevices` to restrict EGL/Vulkan GPU visibility inside the container. +* `apply-cuda-memory-limits` - Set the soft and hard limits of CUDA memory usage on GPU device(s) in the container. diff --git a/cmd/nvidia-cdi-hook/apply-cuda-memory-limits/apply-cuda-memory-limits.go b/cmd/nvidia-cdi-hook/apply-cuda-memory-limits/apply-cuda-memory-limits.go new file mode 100644 index 000000000..ce07f6636 --- /dev/null +++ b/cmd/nvidia-cdi-hook/apply-cuda-memory-limits/apply-cuda-memory-limits.go @@ -0,0 +1,218 @@ +/** +# SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +**/ + +package cudamemorylimits + +import ( + "context" + "fmt" + "strconv" + "strings" + + "github.com/NVIDIA/go-nvml/pkg/nvml" + "github.com/urfave/cli/v3" + + "github.com/NVIDIA/nvidia-container-toolkit/internal/info/cgroup" + "github.com/NVIDIA/nvidia-container-toolkit/internal/logger" + "github.com/NVIDIA/nvidia-container-toolkit/internal/oci" + "github.com/NVIDIA/nvidia-container-toolkit/pkg/lookup" +) + +type command struct { + logger logger.Interface +} + +const ( + GPUMemoryRequestEnvName = "NVIDIA_GPU_MEMORY_REQUEST" + GPUMemoryLimitEnvName = "NVIDIA_GPU_MEMORY_LIMIT" + + mebiByteMultiplier = 1024 * 1024 + maxMebiBytes = nvml.DEVICE_MEMORY_LIMIT_MAX / mebiByteMultiplier +) + +type config struct { + driverRoot string + gpuIds []string + containerSpec string +} + +func NewCommand(logger logger.Interface) *cli.Command { + c := command{ + logger: logger, + } + return c.build() +} + +func (m command) build() *cli.Command { + cfg := config{} + + c := cli.Command{ + Name: "apply-cuda-memory-limits", + Usage: "Set the soft and hard limits of CUDA memory usage on GPU device(s) in the container. " + + "It introspects the OCI container spec and fetches the memory limits from the environment variables:\n" + + "1) NVIDIA_GPU_MEMORY_REQUEST\n2) NVIDIA_GPU_MEMORY_LIMIT\nThe values must be valid integers and are in MebiBytes.", + Before: func(ctx context.Context, cmd *cli.Command) (context.Context, error) { + return ctx, m.validateFlags(cmd, &cfg) + }, + Action: func(ctx context.Context, cmd *cli.Command) error { + return m.run(cmd, &cfg) + }, + Flags: []cli.Flag{ + &cli.StringFlag{ + Name: "driver-root", + Usage: "Specify the driver root", + Destination: &cfg.driverRoot, + }, + &cli.StringSliceFlag{ + Name: "gpu-id", + Usage: "Specify the UUID of the GPU", + Destination: &cfg.gpuIds, + }, + &cli.StringFlag{ + Name: "container-spec", + Usage: "Specify the path to the OCI container spec. If empty or '-' the spec will be read from STDIN", + Destination: &cfg.containerSpec, + }, + }, + } + + return &c +} + +func (m command) validateFlags(_ *cli.Command, cfg *config) error { + if len(cfg.gpuIds) == 0 { + return fmt.Errorf("at least one gpu-id must be specified") + } + for _, id := range cfg.gpuIds { + if strings.TrimSpace(id) == "" { + return fmt.Errorf("gpu-id must not be empty") + } + } + + return nil +} + +func (m command) run(_ *cli.Command, cfg *config) error { + s, err := oci.LoadContainerState(cfg.containerSpec) + if err != nil { + return fmt.Errorf("failed to load container state: %w", err) + } + specFilePath := oci.GetSpecFilePath(s.Bundle) + fs := oci.NewFileSpec(specFilePath, false) + _, err = fs.Load() + if err != nil { + return fmt.Errorf("failed to load OCI container spec: %w", err) + } + + memReqStr, hasRequest := fs.LookupEnv(GPUMemoryRequestEnvName) + memLimitStr, hasLimit := fs.LookupEnv(GPUMemoryLimitEnvName) + if !hasRequest && !hasLimit { + return nil + } + + if !cgroup.IsCgroupV2() { + return fmt.Errorf("setting GPU memory limits is only supported in cgroup v2") + } + + cgroupPath, err := cgroup.GetAbsolutePath(s.Pid) + if err != nil { + return fmt.Errorf("failed to resolve cgroup path: %w", err) + } + + var memoryRequestBytes, memoryLimitBytes uint64 + + if hasRequest { + memoryRequestMiB, err := parseMebiBytes(memReqStr) + if err != nil { + return fmt.Errorf("failed to parse %s: %w", GPUMemoryRequestEnvName, err) + } + memoryRequestBytes = memoryRequestMiB * mebiByteMultiplier + } + if hasLimit { + memoryLimitMiB, err := parseMebiBytes(memLimitStr) + if err != nil { + return fmt.Errorf("failed to parse %s: %w", GPUMemoryLimitEnvName, err) + } + memoryLimitBytes = memoryLimitMiB * mebiByteMultiplier + } + + if !hasLimit { + memoryLimitBytes = nvml.DEVICE_MEMORY_LIMIT_MAX + } + if !hasRequest { + memoryRequestBytes = memoryLimitBytes + } + + if memoryRequestBytes > memoryLimitBytes { + return fmt.Errorf("memory request (%d MiB) exceeds memory limit (%d MiB)", + memoryRequestBytes/mebiByteMultiplier, memoryLimitBytes/mebiByteMultiplier) + } + + return m.runApplyCudaMemoryLimits(cgroupPath, memoryRequestBytes, memoryLimitBytes, cfg.driverRoot, cfg.gpuIds) +} + +func (m command) runApplyCudaMemoryLimits(cgroupPath string, requestBytes uint64, limitBytes uint64, driverRoot string, gpuIDs []string) error { + driverLibLocator := lookup.NewLibraryLocator( + lookup.WithLogger(m.logger), + lookup.WithRoot(driverRoot), + ) + + candidates, err := driverLibLocator.Locate("libnvidia-ml.so.1") + if err != nil { + return fmt.Errorf("failed to locate libnvidia-ml.so.1: %w", err) + } + if len(candidates) == 0 { + return fmt.Errorf("no libnvidia-ml.so.1 found") + } + + m.logger.Infof("driver library found: %s", candidates[0]) + + nvmllib := nvml.New(nvml.WithLibraryPath(candidates[0])) + ret := nvmllib.Init() + if ret != nvml.SUCCESS { + return fmt.Errorf("failed to initialize nvml: %w", ret) + } + defer func() { + _ = nvmllib.Shutdown() + }() + + for _, gpuID := range gpuIDs { + device, ret := nvmllib.DeviceGetHandleByUUID(gpuID) + if ret != nvml.SUCCESS { + return fmt.Errorf("failed to get GPU device handle with uuid %s: %w", gpuID, ret) + } + if device == nil { + return fmt.Errorf("empty GPU device handle: %s", gpuID) + } + ret = device.SetMemoryLimits_v1(cgroupPath, requestBytes, limitBytes) + if ret != nvml.SUCCESS { + return fmt.Errorf("failed to set memory limits for gpu %q: %w", gpuID, ret) + } + } + return nil +} + +func parseMebiBytes(value string) (uint64, error) { + mebiBytes, err := strconv.ParseUint(value, 10, 64) + if err != nil { + return 0, err + } + if mebiBytes > maxMebiBytes { + return 0, fmt.Errorf("%d MiB exceeds the maximum supported value of %d MiB", mebiBytes, maxMebiBytes) + } + return mebiBytes, nil +} diff --git a/cmd/nvidia-cdi-hook/apply-cuda-memory-limits/apply-cuda-memory-limits_test.go b/cmd/nvidia-cdi-hook/apply-cuda-memory-limits/apply-cuda-memory-limits_test.go new file mode 100644 index 000000000..4737aac83 --- /dev/null +++ b/cmd/nvidia-cdi-hook/apply-cuda-memory-limits/apply-cuda-memory-limits_test.go @@ -0,0 +1,98 @@ +/** +# SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +**/ + +package cudamemorylimits + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestParseMebiBytes(t *testing.T) { + testCases := []struct { + description string + value string + expected uint64 + expectError bool + }{ + { + description: "valid value", + value: "1024", + expected: 1024, + }, + { + description: "zero", + value: "0", + expected: 0, + }, + { + description: "largest value convertible to bytes", + value: "17592186044415", + expected: maxMebiBytes, + }, + { + description: "one more than the largest convertible value", + value: "17592186044416", + expectError: true, + }, + { + description: "max uint64", + value: "18446744073709551615", + expectError: true, + }, + { + description: "value exceeding uint64", + value: "18446744073709551616", + expectError: true, + }, + { + description: "non-numeric value", + value: "12abc", + expectError: true, + }, + { + description: "empty value", + value: "", + expectError: true, + }, + { + description: "negative value", + value: "-1", + expectError: true, + }, + { + description: "leading whitespace", + value: " 1024", + expectError: true, + }, + } + + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + mebiBytes, err := parseMebiBytes(tc.value) + if tc.expectError { + require.Error(t, err) + return + } + require.NoError(t, err) + require.Equal(t, tc.expected, mebiBytes) + // The value must be convertible to bytes without overflowing. + require.Equal(t, tc.expected, mebiBytes*mebiByteMultiplier/mebiByteMultiplier) + }) + } +} diff --git a/cmd/nvidia-cdi-hook/commands/commands.go b/cmd/nvidia-cdi-hook/commands/commands.go index a15b36cc2..85f734204 100644 --- a/cmd/nvidia-cdi-hook/commands/commands.go +++ b/cmd/nvidia-cdi-hook/commands/commands.go @@ -22,6 +22,7 @@ import ( "github.com/urfave/cli/v3" + cudamemorylimits "github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-cdi-hook/apply-cuda-memory-limits" "github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-cdi-hook/chmod" symlinks "github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-cdi-hook/create-symlinks" "github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-cdi-hook/cudacompat" @@ -91,6 +92,7 @@ func ConfigureCDIHookCommand(logger logger.Interface, base *cli.Command) *cli.Co chmod.NewCommand(logger), cudacompat.NewCommand(logger), disabledevicenodemodification.NewCommand(logger), + cudamemorylimits.NewCommand(logger), updateapplicationprofile.NewCommand(logger), { Name: "noop", diff --git a/cmd/nvidia-ctk/cdi/generate/generate_test.go b/cmd/nvidia-ctk/cdi/generate/generate_test.go index a954d2ba7..6cac036c2 100644 --- a/cmd/nvidia-ctk/cdi/generate/generate_test.go +++ b/cmd/nvidia-ctk/cdi/generate/generate_test.go @@ -98,11 +98,35 @@ devices: deviceNodes: - path: /dev/nvidia0 hostPath: {{ .driverRoot }}/dev/nvidia0 + hooks: + - hookName: createRuntime + path: /usr/bin/nvidia-cdi-hook + args: + - nvidia-cdi-hook + - apply-cuda-memory-limits + - --driver-root + - {{ .driverRoot }} + - --gpu-id + - {{ .gpuID }} + env: + - NVIDIA_CTK_DEBUG=false - name: all containerEdits: deviceNodes: - path: /dev/nvidia0 hostPath: {{ .driverRoot }}/dev/nvidia0 + hooks: + - hookName: createRuntime + path: /usr/bin/nvidia-cdi-hook + args: + - nvidia-cdi-hook + - apply-cuda-memory-limits + - --driver-root + - {{ .driverRoot }} + - --gpu-id + - {{ .gpuID }} + env: + - NVIDIA_CTK_DEBUG=false containerEdits: env: - NVIDIA_CTK_LIBCUDA_DIR=/lib/x86_64-linux-gnu @@ -200,11 +224,35 @@ devices: deviceNodes: - path: /dev/nvidia0 hostPath: {{ .driverRoot }}/dev/nvidia0 + hooks: + - hookName: createRuntime + path: /usr/bin/nvidia-cdi-hook + args: + - nvidia-cdi-hook + - apply-cuda-memory-limits + - --driver-root + - {{ .driverRoot }} + - --gpu-id + - {{ .gpuID }} + env: + - NVIDIA_CTK_DEBUG=false - name: all containerEdits: deviceNodes: - path: /dev/nvidia0 hostPath: {{ .driverRoot }}/dev/nvidia0 + hooks: + - hookName: createRuntime + path: /usr/bin/nvidia-cdi-hook + args: + - nvidia-cdi-hook + - apply-cuda-memory-limits + - --driver-root + - {{ .driverRoot }} + - --gpu-id + - {{ .gpuID }} + env: + - NVIDIA_CTK_DEBUG=false containerEdits: env: - NVIDIA_CTK_LIBCUDA_DIR=/lib/x86_64-linux-gnu @@ -294,11 +342,35 @@ devices: deviceNodes: - path: /dev/nvidia0 hostPath: {{ .driverRoot }}/dev/nvidia0 + hooks: + - hookName: createRuntime + path: /usr/bin/nvidia-cdi-hook + args: + - nvidia-cdi-hook + - apply-cuda-memory-limits + - --driver-root + - {{ .driverRoot }} + - --gpu-id + - {{ .gpuID }} + env: + - NVIDIA_CTK_DEBUG=false - name: all containerEdits: deviceNodes: - path: /dev/nvidia0 hostPath: {{ .driverRoot }}/dev/nvidia0 + hooks: + - hookName: createRuntime + path: /usr/bin/nvidia-cdi-hook + args: + - nvidia-cdi-hook + - apply-cuda-memory-limits + - --driver-root + - {{ .driverRoot }} + - --gpu-id + - {{ .gpuID }} + env: + - NVIDIA_CTK_DEBUG=false containerEdits: env: - NVIDIA_CTK_LIBCUDA_DIR=/lib/x86_64-linux-gnu @@ -539,7 +611,13 @@ containerEdits: require.NoError(t, err) } - require.Equal(t, strings.ReplaceAll(tc.expectedSpec, "{{ .driverRoot }}", driverRoot), buf.String()) + gpuID, ret := server.Devices[0].GetUUID() + require.True(t, ret == nvml.SUCCESS, gpuID) + + expected := strings.ReplaceAll(tc.expectedSpec, "{{ .driverRoot }}", driverRoot) + expected = strings.ReplaceAll(expected, "{{ .gpuID }}", gpuID) + + require.Equal(t, expected, buf.String()) }) } } diff --git a/internal/discover/hooks.go b/internal/discover/hooks.go index 95cc35e61..173291f7f 100644 --- a/internal/discover/hooks.go +++ b/internal/discover/hooks.go @@ -63,6 +63,8 @@ const ( // An UpdateLDCacheHook is the hook used to update the ldcache in the // container. This allows injected libraries to be discoverable. UpdateLDCacheHook = HookName("update-ldcache") + // ApplyCudaMemoryLimitsHook is used to assign soft and hard limits of CUDA memory usage to a container + ApplyCudaMemoryLimitsHook = HookName("apply-cuda-memory-limits") defaultNvidiaCDIHookPath = "/usr/bin/nvidia-cdi-hook" ) @@ -222,6 +224,8 @@ func (c cdiHookCreator) getOCIHookType(name HookName) OCIHookType { switch name { case CreateSymlinksHook, ChmodHook, DisableDeviceNodeModificationHook, EnableCudaCompatHook, UpdateLDCacheHook, ApplicationProfileHook: return OCIHookTypeCreateContainer + case ApplyCudaMemoryLimitsHook: + return OCIHookTypeCreateRuntime default: return OCIHookTypeCreateContainer } @@ -238,7 +242,7 @@ func (c cdiHookCreator) isDisabled(name HookName, args ...string) bool { // still reject hooks that require args if none were provided switch name { - case CreateSymlinksHook, ChmodHook: + case CreateSymlinksHook, ChmodHook, ApplyCudaMemoryLimitsHook: return len(args) == 0 } return false @@ -267,6 +271,11 @@ func (c cdiHookCreator) transformArgs(name HookName, args ...string) []string { for _, arg := range args { transformedArgs = append(transformedArgs, "--folder", arg) } + case ApplyCudaMemoryLimitsHook: + transformedArgs = append(transformedArgs, "--driver-root", args[0]) + for _, arg := range args[1:] { + transformedArgs = append(transformedArgs, "--gpu-id", arg) + } default: return args } diff --git a/internal/info/cgroup/cgroup_path.go b/internal/info/cgroup/cgroup_path.go new file mode 100644 index 000000000..b95ed3c35 --- /dev/null +++ b/internal/info/cgroup/cgroup_path.go @@ -0,0 +1,64 @@ +/** +# SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Portions of this file are derived from github.com/opencontainers/cgroups, +# licensed under the Apache License, Version 2.0. + +**/ + +package cgroup + +import ( + "errors" + "fmt" + "os" + "path/filepath" + "strings" + + "golang.org/x/sys/unix" +) + +const ( + v2FsMagicNumber = 0x63677270 + rootDirectory = "/sys/fs/cgroup" + unifiedPrefix = "0::" +) + +func IsCgroupV2() bool { + var s unix.Statfs_t + _ = unix.Statfs(rootDirectory, &s) + return s.Type == v2FsMagicNumber +} + +func GetAbsolutePath(pid int) (string, error) { + cgroupProcFile := fmt.Sprintf("/proc/%d/cgroup", pid) + b, err := os.ReadFile(cgroupProcFile) + if err != nil { + return "", err + } + return parseCgroupProcFile(b) +} + +func parseCgroupProcFile(b []byte) (string, error) { + for line := range strings.SplitSeq(string(b), "\n") { + line = strings.TrimSpace(line) + if !strings.HasPrefix(line, unifiedPrefix) { + continue + } + return filepath.Join(rootDirectory, strings.TrimPrefix(line, unifiedPrefix)), nil + } + return "", errors.New("no cgroup v2 unified hierarchy entry found") +} diff --git a/internal/info/cgroup/cgroup_path_test.go b/internal/info/cgroup/cgroup_path_test.go new file mode 100644 index 000000000..0298aa40e --- /dev/null +++ b/internal/info/cgroup/cgroup_path_test.go @@ -0,0 +1,93 @@ +/** +# SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +**/ + +package cgroup + +import ( + "math" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestGetAbsolutePath(t *testing.T) { + _, err := GetAbsolutePath(math.MaxInt32) + require.Error(t, err) +} + +func TestParseCgroupProcFile(t *testing.T) { + testCases := []struct { + description string + contents string + expected string + expectError bool + }{ + { + description: "cgroup v2 unified hierarchy", + contents: "0::/user.slice/user-1000.slice/session-1.scope\n", + expected: filepath.Join(rootDirectory, "/user.slice/user-1000.slice/session-1.scope"), + }, + { + description: "cgroup v1 named hierarchy", + contents: "5:devices:/docker/abc123\n", + expectError: true, + }, + { + description: "hybrid hierarchy selects the unified entry", + contents: "12:devices:/docker/abc123\n0::/docker/abc123\n", + expected: filepath.Join(rootDirectory, "/docker/abc123"), + }, + { + description: "root cgroup path", + contents: "0::/\n", + expected: filepath.Join(rootDirectory, "/"), + }, + { + description: "missing trailing newline", + contents: "0::/foo", + expected: filepath.Join(rootDirectory, "/foo"), + }, + { + description: "empty contents", + contents: "", + expectError: true, + }, + { + description: "too few fields", + contents: "0:/foo", + expectError: true, + }, + { + description: "path containing a colon", + contents: "0::/foo:extra", + expected: filepath.Join(rootDirectory, "/foo:extra"), + }, + } + + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + path, err := parseCgroupProcFile([]byte(tc.contents)) + if tc.expectError { + require.Error(t, err) + return + } + require.NoError(t, err) + require.Equal(t, tc.expected, path) + }) + } +} diff --git a/pkg/nvcdi/cuda-memory-limits.go b/pkg/nvcdi/cuda-memory-limits.go new file mode 100644 index 000000000..449ec6fcb --- /dev/null +++ b/pkg/nvcdi/cuda-memory-limits.go @@ -0,0 +1,71 @@ +/** +# SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +**/ + +package nvcdi + +import ( + "fmt" + + "github.com/NVIDIA/go-nvlib/pkg/nvlib/device" + "github.com/NVIDIA/go-nvml/pkg/nvml" + + "github.com/NVIDIA/nvidia-container-toolkit/internal/discover" + "github.com/NVIDIA/nvidia-container-toolkit/internal/logger" +) + +type cudaMemoryLimits struct { + logger logger.Interface + driverRoot string + uuid string + hookCreator discover.HookCreator +} + +func (l *nvcdilib) newCudaMemoryLimits(d device.Device) (discover.Discover, error) { + uuid, nvmlRet := d.GetUUID() + if nvmlRet != nvml.SUCCESS { + return nil, fmt.Errorf("failed to get device UUID: %w", nvmlRet) + } + + cMemLimits := &cudaMemoryLimits{ + logger: l.logger, + driverRoot: l.driver.Root, + uuid: uuid, + hookCreator: l.hookCreator, + } + + return cMemLimits, nil +} + +// Devices are empty for this discoverer +func (c *cudaMemoryLimits) Devices() ([]discover.Device, error) { + return nil, nil +} + +// EnvVars are empty for this discoverer +func (c *cudaMemoryLimits) EnvVars() ([]discover.EnvVar, error) { + return nil, nil +} + +// Hooks returns a set of hooks that assigns a CUDA memory limit to the cgroup of the GPU workload container +func (c *cudaMemoryLimits) Hooks() ([]discover.Hook, error) { + return c.hookCreator.Create(discover.ApplyCudaMemoryLimitsHook, c.driverRoot, c.uuid).Hooks() +} + +// Mounts are empty for this discoverer +func (c *cudaMemoryLimits) Mounts() ([]discover.Mount, error) { + return nil, nil +} diff --git a/pkg/nvcdi/full-gpu-nvml.go b/pkg/nvcdi/full-gpu-nvml.go index 49bf0bdd1..554d8cbdb 100644 --- a/pkg/nvcdi/full-gpu-nvml.go +++ b/pkg/nvcdi/full-gpu-nvml.go @@ -173,11 +173,17 @@ func (l *fullGPUDeviceSpecGenerator) newFullGPUDiscoverer(d device.Device) (disc deviceNodes, ) + cudaMemoryLimitsHook, err := (*nvcdilib)(l.nvmllib).newCudaMemoryLimits(d) + if err != nil { + return nil, fmt.Errorf("failed to create cuda memory limits discoverer: %w", err) + } + var discoverers []discover.Discover discoverers = append(discoverers, deviceNodes, deviceFolderPermissionHooks, + cudaMemoryLimitsHook, ) discoverers = append(discoverers, l.additionalDiscoverers...) diff --git a/pkg/nvcdi/lib-csv_test.go b/pkg/nvcdi/lib-csv_test.go index c5ed3d8df..37ff0991e 100644 --- a/pkg/nvcdi/lib-csv_test.go +++ b/pkg/nvcdi/lib-csv_test.go @@ -200,6 +200,14 @@ func TestDeviceSpecGenerators(t *testing.T) { {Path: "/dev/nvidiactl", HostPath: "/dev/nvidiactl"}, {Path: "/dev/nvmap", HostPath: "/dev/nvmap", FileMode: to.Ptr(os.FileMode(0400)), Permissions: "rwm", GID: to.Ptr[uint32](44)}, }, + Hooks: []*specs.Hook{ + { + HookName: "createRuntime", + Path: "/usr/bin/nvidia-cdi-hook", + Args: []string{"nvidia-cdi-hook", "apply-cuda-memory-limits", "--driver-root", "", "--gpu-id", "GPU-1"}, + Env: []string{"NVIDIA_CTK_DEBUG=false"}, + }, + }, }, }, },