From 42c739296dde6be790e2e82480a36f7235ea5f2b Mon Sep 17 00:00:00 2001
From: Vishesh Tanksale <vtanksale@nvidia.com>
Date: Mon, 26 Aug 2024 19:03:08 +0000
Subject: [PATCH] Adding alias to gpu resource name

---
 cmd/main.go                        | 75 ++++++++++++++++--------------
 pkg/device_plugin/device_plugin.go | 26 ++++++++---
 2 files changed, 59 insertions(+), 42 deletions(-)

diff --git a/cmd/main.go b/cmd/main.go
index f81c1692..31354390 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -1,35 +1,40 @@
-/*
- * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
- *
- * Redistribution and use in source and binary forms, with or without
- * modification, are permitted provided that the following conditions
- * are met:
- *  * Redistributions of source code must retain the above copyright
- *    notice, this list of conditions and the following disclaimer.
- *  * Redistributions in binary form must reproduce the above copyright
- *    notice, this list of conditions and the following disclaimer in the
- *    documentation and/or other materials provided with the distribution.
- *  * Neither the name of NVIDIA CORPORATION nor the names of its
- *    contributors may be used to endorse or promote products derived
- *    from this software without specific prior written permission.
- *
- * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
- * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
- * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
- * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
- * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
- * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
- * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
- * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
- * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
- */
-
-package main
-
-import "kubevirt-gpu-device-plugin/pkg/device_plugin"
-
-func main() {
-	device_plugin.InitiateDevicePlugin()
-}
+/*
+ * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *  * Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ *  * Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *  * Neither the name of NVIDIA CORPORATION nor the names of its
+ *    contributors may be used to endorse or promote products derived
+ *    from this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+ * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+ * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+ * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+ * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+ * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+ * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+ * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+package main
+
+import (
+	"kubevirt-gpu-device-plugin/pkg/device_plugin"
+	"os"
+)
+
+func main() {
+	device_plugin.PGPUAlias = os.Getenv("P_GPU_ALIAS")
+	device_plugin.VGPUAlias = os.Getenv("V_GPU_ALIAS")
+	device_plugin.InitiateDevicePlugin()
+}
diff --git a/pkg/device_plugin/device_plugin.go b/pkg/device_plugin/device_plugin.go
index c5ba462c..014ec76a 100644
--- a/pkg/device_plugin/device_plugin.go
+++ b/pkg/device_plugin/device_plugin.go
@@ -72,6 +72,8 @@ var readVgpuIDFromFile = readVgpuIDFromFileFunc
 var readGpuIDForVgpu = readGpuIDForVgpuFunc
 var startVgpuDevicePlugin = startVgpuDevicePluginFunc
 var stop = make(chan struct{})
+var PGPUAlias string
+var VGPUAlias string
 
 func InitiateDevicePlugin() {
 	//Identifies GPUs and represents it in appropriate structures
@@ -101,10 +103,15 @@ func createDevicePlugins() {
 				Health: pluginapi.Healthy,
 			})
 		}
-		deviceName := getDeviceName(k)
-		if deviceName == "" {
-			log.Printf("Error: Could not find device name for device id: %s", k)
-			deviceName = k
+		deviceName := ""
+		if PGPUAlias != "" {
+			deviceName = PGPUAlias
+		} else {
+			deviceName = getDeviceName(k)
+			if deviceName == "" {
+				log.Printf("Error: Could not find device name for device id: %s", k)
+				deviceName = k
+			}
 		}
 		log.Printf("DP Name %s", deviceName)
 		dp := NewGenericDevicePlugin(deviceName, "/dev/vfio/", devs)
@@ -124,9 +131,14 @@ func createDevicePlugins() {
 				Health: pluginapi.Healthy,
 			})
 		}
-		deviceName := getDeviceName(k)
-		if deviceName == "" {
-			deviceName = k
+		deviceName := ""
+		if VGPUAlias != "" {
+			deviceName = VGPUAlias
+		} else {
+			deviceName = getDeviceName(k)
+			if deviceName == "" {
+				deviceName = k
+			}
 		}
 		log.Printf("DP Name %s", deviceName)
 		dp := NewGenericVGpuDevicePlugin(deviceName, vGpuBasePath, devs)