More changes

Author: vvolam

scripts/reboot

@@ -35,7 +35,6 @@ ASIC_TYPE=$(sonic-cfggen -y /etc/sonic/sonic_version.yml -v asic_type)
 SUBTYPE=$(sonic-cfggen -d -v DEVICE_METADATA.localhost.subtype)
 ASAN=$(sonic-cfggen -y /etc/sonic/sonic_version.yml -v asan)
 VERBOSE=no
-EXIT_SUCCESS=0
 EXIT_NEXT_IMAGE_NOT_EXISTS=4
 EXIT_SONIC_INSTALLER_VERIFY_REBOOT=21
 EXIT_PLATFORM_FW_AU_FAILURE=22
@@ -136,6 +135,8 @@ function show_help_and_exit()
     echo " "
     echo "    Available options:"
     echo "        -h, -? : getting this help"
+    echo "        -d : DPU module name on a smart switch, option is invalid when on DPU"
+    echo "        -p : Pre-shutdown steps on DPU, invalid on NPU"
 
     exit ${EXIT_SUCCESS}
 }
@@ -198,7 +199,7 @@ function get_dpu_ip()
 
 # Function to retrieve GNMI port from CONFIG_DB
 function get_gnmi_port() {
-    port=$(sonic-db-cli CONFIG_DB HGET "GNMI|gnmi" "port")
+    port=$(sonic-db-cli CONFIG_DB HGET "DPU_PORT|<TODO: PORT_NAME>" "gnmi")
     if [ $? -ne 0 ] || [ -z "$port" ]; then
         echo "Error: Failed to retrieve GNMI port"
         exit ${EXIT_ERROR}
@@ -293,7 +294,7 @@ function reboot_dpu_module()
 
 function parse_options()
 {
-    while getopts "h?vf" opt; do
+    while getopts "h?vfpd" opt; do
         case ${opt} in
             h|\? )
                 show_help_and_exit
@@ -348,6 +349,32 @@ function reboot_all_dpus() {
     wait
 }
 
+# Function to handle scenarios on smart switch
+function handle_smart_switch() {
+# TODO: Check if this is DPU or not, proceed only if PRE-SHUTDOWN is set on DPU, else throw error
+
+    if [ -f "$PLATFORM_JSON_PATH" ]; then
+        NUM_DPU=$(jq -r '.DPUS | length' "$PLATFORM_JSON_PATH" 2>/dev/null)
+        if [ "$NUM_DPU" -gt 0 ]; then
+            SMART_SWITCH="yes"
+        fi
+    fi
+
+    if [[ "$REBOOT_DPU" == "yes" ]]; then
+        if [[ "$SMART_SWITCH" == "yes" ]]; then
+            echo "User requested to reboot the device ${DPU_MODULE_NAME}"
+            reboot_dpu_module "$DPU_MODULE_NAME"
+        else
+            echo "Invalid '-d' option specified for a non-smart switch"
+            exit ${EXIT_ERROR}
+        fi
+    fi
+
+    if [[ "$SMART_SWITCH" == "yes" ]]; then
+        reboot_all_dpus "$NUM_DPU"
+    fi
+}
+
 parse_options $@
 
 # Exit if not superuser
@@ -358,19 +385,7 @@ fi
 
 debug "User requested rebooting device ..."
 
-if [ -f "$PLATFORM_JSON_PATH" ]; then
-    NUM_DPU=$(jq -r '.DPUS | length' "$PLATFORM_JSON_PATH" 2>/dev/null)
-    if [ "$NUM_DPU" -gt 0 ]; then
-        SMART_SWITCH="yes"
-    fi
-fi
-
-if [[ "$REBOOT_DPU" == "yes" && "$SMART_SWITCH" == "yes" ]]; then
-    echo "User requested to reboot the device ${DPU_MODULE_NAME}"
-    reboot_dpu_module "$DPU_MODULE_NAME"
-elif [ "$SMART_SWITCH" == "yes" ]; then
-    reboot_all_dpus "$NUM_DPU"
-fi
+handle_smart_switch
 
 check_conflict_boot_in_fw_update
 

scripts/reboot_helper.py

@@ -10,6 +10,7 @@
 
 chk_log_level = syslog.LOG_ERR
 
+
 def _log_msg(lvl, pfx, msg):
     if lvl <= chk_log_level:
         print("{}: {}".format(pfx, msg))
@@ -31,6 +32,7 @@ def log_debug(m):
 # Global variable for platform chassis
 platform_chassis = None
 
+
 def load_platform_chassis():
     global platform_chassis
 
@@ -47,6 +49,7 @@ def load_platform_chassis():
 
     return True
 
+
 def reboot_module(module_name):
     """Reboot the specified module by invoking the platform API"""
 
@@ -68,23 +71,24 @@ def reboot_module(module_name):
                 log_info(f"Rebooting module {module_name}...")
                 try:
                     module.reboot()
-                    log_info(f"Reboot command sent for module {module_name}")
+                    log_info("Reboot command sent for module {module_name}")
                     return True
                 except NotImplementedError:
-                    log_error(f"Reboot not implemented for module {module_name}.")
+                    log_err("Reboot not implemented for module {module_name}.")
                     return False
                 except Exception as e:
-                    log_error(f"An error occurred while rebooting module {module_name}: {e}")
+                    log_err("An error occurred while rebooting module {module_name}: {e}")
                     return False
 
         # If the module with the given name is not found
-        log_err(f"Module {module_name} not found")
+        log_err("Module {module_name} not found")
         return False
 
     except Exception as e:
-        log_err(f"Error occurred while rebooting module {module_name}: {repr(e)}")
+        log_err("Error occurred while rebooting module {module_name}: {repr(e)}")
         return False
 
+
 if __name__ == "__main__":
     if len(sys.argv) < 3:
         print("Usage: reboot_helper.py <command> <module_name>")
@@ -98,7 +102,7 @@ def reboot_module(module_name):
         if not success:
             sys.exit(1)
         else:
-            print(f"Reboot command sent for module {module_name}")
+            print("Reboot command sent for module {module_name}")
     else:
-        print(f"Unknown command: {command}")
+        print("Unknown command: {command}")
         sys.exit(1)