From ebe18fbb7f540838c8c8eca11fd827d6691524cf Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Mon, 6 Mar 2023 14:34:40 +0200
Subject: [PATCH] Add gds mode to nvcdi API

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 pkg/nvcdi/api.go |  2 ++
 pkg/nvcdi/gds.go | 82 ++++++++++++++++++++++++++++++++++++++++++++++++
 pkg/nvcdi/lib.go |  5 +++
 3 files changed, 89 insertions(+)
 create mode 100644 pkg/nvcdi/gds.go

diff --git a/pkg/nvcdi/api.go b/pkg/nvcdi/api.go
index a1a26e2d..df97528b 100644
--- a/pkg/nvcdi/api.go
+++ b/pkg/nvcdi/api.go
@@ -32,6 +32,8 @@ const (
 	ModeWsl = "wsl"
 	// ModeManagement configures the CDI spec generator to generate a management spec.
 	ModeManagement = "management"
+	// ModeGds configures the CDI spec generator to generate a GDS spec.
+	ModeGds = "gds"
 )
 
 // Interface defines the API for the nvcdi package
diff --git a/pkg/nvcdi/gds.go b/pkg/nvcdi/gds.go
new file mode 100644
index 00000000..2aba8d06
--- /dev/null
+++ b/pkg/nvcdi/gds.go
@@ -0,0 +1,82 @@
+/**
+# Copyright (c) NVIDIA CORPORATION.  All rights reserved.
+#
+# 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/nvidia-container-toolkit/internal/discover"
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/edits"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/spec"
+	"github.com/container-orchestrated-devices/container-device-interface/pkg/cdi"
+	"github.com/container-orchestrated-devices/container-device-interface/specs-go"
+	"gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvlib/device"
+)
+
+type gdslib nvcdilib
+
+var _ Interface = (*gdslib)(nil)
+
+// GetAllDeviceSpecs returns the device specs for all available devices.
+func (l *gdslib) GetAllDeviceSpecs() ([]specs.Device, error) {
+	discoverer, err := discover.NewGDSDiscoverer(l.logger, l.driverRoot)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create GPUDirect Storage discoverer: %v", err)
+	}
+	edits, err := edits.FromDiscoverer(discoverer)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create container edits for GPUDirect Storage: %v", err)
+	}
+
+	deviceSpec := specs.Device{
+		Name:           "all",
+		ContainerEdits: *edits.ContainerEdits,
+	}
+
+	return []specs.Device{deviceSpec}, nil
+}
+
+// GetCommonEdits generates a CDI specification that can be used for ANY devices
+func (l *gdslib) GetCommonEdits() (*cdi.ContainerEdits, error) {
+	return edits.FromDiscoverer(discover.None{})
+}
+
+// GetSpec is unsppported for the gdslib specs.
+// gdslib is typically wrapped by a spec that implements GetSpec.
+func (l *gdslib) GetSpec() (spec.Interface, error) {
+	return nil, fmt.Errorf("GetSpec is not supported")
+}
+
+// GetGPUDeviceEdits is unsupported for the gdslib specs
+func (l *gdslib) GetGPUDeviceEdits(device.Device) (*cdi.ContainerEdits, error) {
+	return nil, fmt.Errorf("GetGPUDeviceEdits is not supported")
+}
+
+// GetGPUDeviceSpecs is unsupported for the gdslib specs
+func (l *gdslib) GetGPUDeviceSpecs(int, device.Device) (*specs.Device, error) {
+	return nil, fmt.Errorf("GetGPUDeviceSpecs is not supported")
+}
+
+// GetMIGDeviceEdits is unsupported for the gdslib specs
+func (l *gdslib) GetMIGDeviceEdits(device.Device, device.MigDevice) (*cdi.ContainerEdits, error) {
+	return nil, fmt.Errorf("GetMIGDeviceEdits is not supported")
+}
+
+// GetMIGDeviceSpecs is unsupported for the gdslib specs
+func (l *gdslib) GetMIGDeviceSpecs(int, device.Device, int, device.MigDevice) (*specs.Device, error) {
+	return nil, fmt.Errorf("GetMIGDeviceSpecs is not supported")
+}
diff --git a/pkg/nvcdi/lib.go b/pkg/nvcdi/lib.go
index ba94822e..95f47fb4 100644
--- a/pkg/nvcdi/lib.go
+++ b/pkg/nvcdi/lib.go
@@ -89,6 +89,11 @@ func New(opts ...Option) Interface {
 		lib = (*nvmllib)(l)
 	case ModeWsl:
 		lib = (*wsllib)(l)
+	case ModeGds:
+		if l.class == "" {
+			l.class = "gds"
+		}
+		lib = (*gdslib)(l)
 	default:
 		// TODO: We would like to return an error here instead of panicking
 		panic("Unknown mode")