From 6e413d84456af0c31290ca3a1260b43ce9074fd4 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Fri, 6 Dec 2024 13:54:35 +0100
Subject: [PATCH] [no-relnote] Move nvcdi wrapper to separate file

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 pkg/nvcdi/lib.go     | 44 ---------------------------
 pkg/nvcdi/wrapper.go | 72 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+), 44 deletions(-)
 create mode 100644 pkg/nvcdi/wrapper.go

diff --git a/pkg/nvcdi/lib.go b/pkg/nvcdi/lib.go
index 409cea76b..7afa5096b 100644
--- a/pkg/nvcdi/lib.go
+++ b/pkg/nvcdi/lib.go
@@ -22,26 +22,14 @@ import (
 	"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
 	"github.com/NVIDIA/go-nvlib/pkg/nvlib/info"
 	"github.com/NVIDIA/go-nvml/pkg/nvml"
-	"tags.cncf.io/container-device-interface/pkg/cdi"
 
-	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/image"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup/root"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/nvsandboxutils"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/platform-support/tegra/csv"
-	"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/spec"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/transform"
 )
 
-type wrapper struct {
-	Interface
-
-	vendor string
-	class  string
-
-	mergedDeviceOptions []transform.MergedDeviceOption
-}
-
 type nvcdilib struct {
 	logger             logger.Interface
 	nvmllib            nvml.Interface
@@ -180,38 +168,6 @@ func New(opts ...Option) (Interface, error) {
 	return &w, nil
 }
 
-// GetSpec combines the device specs and common edits from the wrapped Interface to a single spec.Interface.
-func (l *wrapper) GetSpec() (spec.Interface, error) {
-	deviceSpecs, err := l.GetAllDeviceSpecs()
-	if err != nil {
-		return nil, err
-	}
-
-	edits, err := l.GetCommonEdits()
-	if err != nil {
-		return nil, err
-	}
-
-	return spec.New(
-		spec.WithDeviceSpecs(deviceSpecs),
-		spec.WithEdits(*edits.ContainerEdits),
-		spec.WithVendor(l.vendor),
-		spec.WithClass(l.class),
-		spec.WithMergedDeviceOptions(l.mergedDeviceOptions...),
-	)
-}
-
-// GetCommonEdits returns the wrapped edits and adds additional edits on top.
-func (m *wrapper) GetCommonEdits() (*cdi.ContainerEdits, error) {
-	edits, err := m.Interface.GetCommonEdits()
-	if err != nil {
-		return nil, err
-	}
-	edits.Env = append(edits.Env, image.EnvVarNvidiaVisibleDevices+"=void")
-
-	return edits, nil
-}
-
 // getCudaVersion returns the CUDA version of the current system.
 func (l *nvcdilib) getCudaVersion() (string, error) {
 	version, err := l.getCudaVersionNvsandboxutils()
diff --git a/pkg/nvcdi/wrapper.go b/pkg/nvcdi/wrapper.go
new file mode 100644
index 000000000..f45b34e70
--- /dev/null
+++ b/pkg/nvcdi/wrapper.go
@@ -0,0 +1,72 @@
+/**
+# Copyright 2025 NVIDIA CORPORATION
+#
+# 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 (
+	"tags.cncf.io/container-device-interface/pkg/cdi"
+	"tags.cncf.io/container-device-interface/specs-go"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/image"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/spec"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/transform"
+)
+
+type wrapper struct {
+	Interface
+
+	vendor string
+	class  string
+
+	mergedDeviceOptions []transform.MergedDeviceOption
+}
+
+// GetSpec combines the device specs and common edits from the wrapped Interface to a single spec.Interface.
+func (l *wrapper) GetSpec() (spec.Interface, error) {
+	deviceSpecs, err := l.GetAllDeviceSpecs()
+	if err != nil {
+		return nil, err
+	}
+
+	edits, err := l.GetCommonEdits()
+	if err != nil {
+		return nil, err
+	}
+
+	return spec.New(
+		spec.WithDeviceSpecs(deviceSpecs),
+		spec.WithEdits(*edits.ContainerEdits),
+		spec.WithVendor(l.vendor),
+		spec.WithClass(l.class),
+		spec.WithMergedDeviceOptions(l.mergedDeviceOptions...),
+	)
+}
+
+// GetAllDeviceSpecs returns the device specs for all available devices.
+func (l *wrapper) GetAllDeviceSpecs() ([]specs.Device, error) {
+	return l.Interface.GetAllDeviceSpecs()
+}
+
+// GetCommonEdits returns the wrapped edits and adds additional edits on top.
+func (m *wrapper) GetCommonEdits() (*cdi.ContainerEdits, error) {
+	edits, err := m.Interface.GetCommonEdits()
+	if err != nil {
+		return nil, err
+	}
+	edits.Env = append(edits.Env, image.EnvVarNvidiaVisibleDevices+"=void")
+
+	return edits, nil
+}