added new gpu name got emr nodes

Signed-off-by: cindyyuanjiang <[email protected]>
NVIDIA · Dec 28, 2023 · b3fd46f · b3fd46f
1 parent 5585c42
commit b3fd46f
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/user_tools/src/spark_rapids_pytools/cloud_api/sp_types.py b/user_tools/src/spark_rapids_pytools/cloud_api/sp_types.py
@@ -50,6 +50,7 @@ class GpuDevice(EnumeratedType):
     P4 = 'P4'
     L4 = 'l4'
     A10 = 'a10'
+    A10G = 'a10g'
 
     @classmethod
     def get_default_gpu(cls):
@@ -64,7 +65,8 @@ def get_gpu_mem(self) -> list:
             self.K80: [12288],
             self.V100: [16384],
             self.P100: [16384],
-            self.A10: [24576]
+            self.A10: [24576],
+            self.A10G: [24576]
         }
         return memory_hash.get(self)
 
@@ -115,8 +117,6 @@ class GpuHWInfo:
     gpu_device: GpuDevice = GpuDevice.get_default_gpu()
 
     def get_gpu_device_name(self) -> str:
-        print("self.gpu_device =", self.gpu_device)
-        print("GpuDevice.get_default_gpu()=", GpuDevice.get_default_gpu())
         return GpuDevice.tostring(self.gpu_device)
 
 

diff --git a/user_tools/src/spark_rapids_tools/enums.py b/user_tools/src/spark_rapids_tools/enums.py
@@ -37,9 +37,7 @@ def tostring(cls, value: Union[Enum, str]) -> str:
         :return: the uppercase string that represents the state object
         :rtype: str
         """
-        print("value =", value)
         value = cast(Enum, value)
-        print("value =", value)
         return str(value._value_).upper()  # pylint: disable=protected-access
 
     @classmethod