[FABRIC MANAGER] Using common library for getting NVSwitch count

Author: Himani Anil Deshpande

cookbooks/aws-parallelcluster-platform/libraries/nvidia.rb

@@ -18,3 +18,21 @@ def is_process_running(process_name)
 
   !ps.stdout.strip.empty?
 end
+
+#
+# Get Count of GPUs in instance
+#
+def get_nvswitch_count(device_id)
+  shell_out("lspci -d #{device_id} | wc -l").stdout.strip.to_i
+end
+
+def get_device_ids
+  #  A100 (P4), H100(P5), B200(P6) and GB200()p6e) systems have NVSwitches
+  # NVSwitch device id is 10de:1af1 for P4 instance
+  # NVSwitch device id is 10de:22a3 for P5 instance
+  # NVSwitch device id is 10de:2901 for P6 instance
+  # NVSwitch device id is 10de:2941 for P6e instance
+  # We sum the count for all these deviceIds as output of lscpi command will be >0
+  # for only one device ID based on the instance type
+  { 'a100' => '10de:1af1', 'h100' => '10de:22a3', 'b200' => '10de:2901', 'gb200' => '10de:2941' }
+end

cookbooks/aws-parallelcluster-platform/resources/fabric_manager/partial/_fabric_manager_common.rb

@@ -54,12 +54,6 @@ def _nvidia_driver_version
 
 # Get number of nv switches
 def get_nvswitches
-  #  A100 (P4), H100(P5) and B200(P6) systems have NVSwitches
-  # NVSwitch device id is 10de:1af1 for P4 instance
-  # NVSwitch device id is 10de:22a3 for P5 instance
-  # NVSwitch device id is 10de:2901 for P6 instance
-  # We sum the count for all these deviceIds as output of lscpi command will be >0
-  # for only one device ID based on the instance type
-  nvswitch_device_ids = ['10de:1af1', '10de:22a3', '10de:2901']
-  nvswitch_device_ids.sum { |id| shell_out("lspci -d #{id} | wc -l").stdout.strip.to_i }
+  nvswitch_device_ids = get_device_ids.values
+  nvswitch_device_ids.sum { |id| get_nvswitch_count(id) }
 end