[SmartSwitch] Extend reboot script for rebooting SmartSwitch

scripts/reboot

@@ -5,6 +5,8 @@ declare -r EXIT_ERROR=1
 declare -r WATCHDOG_UTIL="/usr/local/bin/watchdogutil"
 declare -r PRE_REBOOT_HOOK="pre_reboot_hook"
 
+source reboot_smartswitch_helper
+
 DEVPATH="/usr/share/sonic/device"
 PLAT_REBOOT="platform_reboot"
 PLATFORM_UPDATE_REBOOT_CAUSE="platform_update_reboot_cause"
@@ -37,10 +39,17 @@ EXIT_NEXT_IMAGE_NOT_EXISTS=4
 EXIT_SONIC_INSTALLER_VERIFY_REBOOT=21
 EXIT_PLATFORM_FW_AU_FAILURE=22
 PLATFORM_FWUTIL_AU_REBOOT_HANDLE="platform_fw_au_reboot_handle"
+PLATFORM_JSON_FILE="platform.json"
+PLATFORM_JSON_PATH="${DEVPATH}/${PLATFORM}/${PLATFORM_JSON_FILE}"
 REBOOT_SCRIPT_NAME=$(basename $0)
 REBOOT_TYPE="${REBOOT_SCRIPT_NAME}"
 TAG_LATEST=no
 REBOOT_FLAGS=""
+FORCE_REBOOT="no"
+SMART_SWITCH="no"
+DPU_MODULE_NAME=""
+REBOOT_DPU="no"
+PRE_SHUTDOWN="no"
 
 function debug()
 {
@@ -128,6 +137,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}
 }
@@ -154,7 +165,7 @@ function reboot_pre_check()
         ${DEVPATH}/${PLATFORM}/${PLATFORM_REBOOT_PRE_CHECK}
         [[ $? -ne 0 ]] && exit $?
     fi
-    
+
     # Verify the next image by sonic-installer
     local message=$(sonic-installer verify-next-image 2>&1)
     if [ $? -ne 0 ]; then
@@ -178,7 +189,7 @@ function check_conflict_boot_in_fw_update()
 
 function parse_options()
 {
-    while getopts "h?vf" opt; do
+    while getopts "h?vfpd:" opt; do
         case ${opt} in
             h|\? )
                 show_help_and_exit
@@ -192,6 +203,13 @@ function parse_options()
             f )
                 REBOOT_FLAGS+=" -f"
                 ;;
+            d )
+                REBOOT_DPU="yes"
+                DPU_MODULE_NAME="$OPTARG"
+                ;;
+            p )
+                PRE_SHUTDOWN="yes"
+                ;;
         esac
     done
 }
@@ -225,6 +243,12 @@ fi
 
 debug "User requested rebooting device ..."
 
+handle_smart_switch "$REBOOT_DPU" "$PRE_SHUTDOWN" "$DPU_MODULE_NAME"
+smart_switch_result=$?
+if [[ $smart_switch_result -ne 0 ]]; then
+    exit $smart_switch_result
+fi
+
 check_conflict_boot_in_fw_update
 
 setup_reboot_variables
@@ -287,6 +311,11 @@ if [ -x ${WATCHDOG_UTIL} ]; then
     ${WATCHDOG_UTIL} arm
 fi
 
+if [[ "${PRE_SHUTDOWN}" == "yes" ]]; then
+    echo "${DPU_MODULE_NAME} pre-shutdown steps are completed"
+    exit ${EXIT_SUCCESS}
+fi
+
 if [ -x ${DEVPATH}/${PLATFORM}/${PLAT_REBOOT} ]; then
     VERBOSE=yes debug "Rebooting with platform ${PLATFORM} specific tool ..."
     ${DEVPATH}/${PLATFORM}/${PLAT_REBOOT} $@

scripts/reboot_helper.py

@@ -0,0 +1,161 @@
+#!/usr/bin/env python3
+#
+# reboot_helper.py
+#
+# Utility helper for reboot within SONiC
+
+import sonic_platform
+from sonic_py_common import logger
+from utilities_common.chassis import is_smartswitch, get_dpu_list
+
+SYSLOG_IDENTIFIER = "reboot_helper"
+
+EXIT_FAIL = -1
+EXIT_SUCCESS = 0
+
+# Global logger instance
+log = logger.Logger(SYSLOG_IDENTIFIER)
+
+# Global variable for platform chassis
+platform_chassis = None
+
+
+def load_platform_chassis():
+    """
+    Load the platform chassis using the SONiC platform API.
+
+    This function attempts to instantiate the platform chassis object.
+    If successful, it sets the global variable `platform_chassis` to the instantiated object.
+
+    Returns:
+        bool: True if the platform chassis is successfully loaded, False otherwise.
+    """
+    global platform_chassis
+
+    # Load new platform API class
+    try:
+        platform_chassis = sonic_platform.platform.Platform().get_chassis()
+        if platform_chassis is None:
+            log.log_error("Platform chassis is not loaded")
+            return False
+        return True
+    except Exception as e:
+        log.log_error("Failed to instantiate Chassis: {}".format(str(e)))
+        return False
+
+
+def load_and_verify(module_name):
+    """
+    Load the platform chassis and verify the required parameters.
+
+    Args:
+        module_name (str): The name of the module to verify.
+
+    Returns:
+        bool: True if platform chassis is successfully loaded and required parameters are verified, False otherwise.
+    """
+    # Load the platform chassis if not already loaded
+    if not load_platform_chassis():
+        return False
+
+    if not is_smartswitch():
+        log.log_error("Platform is not a smartswitch")
+        return False
+
+    dpu_list = get_dpu_list()
+    if module_name.lower() not in dpu_list:
+        log.log_error("Module {} not found".format(module_name))
+        return False
+
+    return True
+
+
+def reboot_dpu(module_name, reboot_type):
+    """
+    Reboot the specified module by invoking the platform API.
+
+    Args:
+        module_name (str): The name of the module to reboot.
+        reboot_type (str): The type of reboot requested for the module.
+
+    Returns:
+        bool: True if the reboot command was successfully sent, False otherwise.
+    """
+    if not load_and_verify(module_name):
+        return False
+
+    # Attempt to reboot the module
+    if not hasattr(platform_chassis, 'reboot'):
+        log.log_error("Reboot method not found in platform chassis")
+        return False
+
+    log.log_info("Rebooting module {} with reboot_type {}...".format(module_name, reboot_type))
+    try:
+        status = platform_chassis.reboot(module_name, reboot_type)
+        if not status:
+            log.log_error("Reboot status for module {}: {}".format(module_name, status))
+            return False
+        return True
+    except Exception as e:
+        log.log_error("Unexpected error occurred while rebooting module {}: {}".format(module_name, e))
+        return False
+
+
+def pci_detach_module(module_name):
+    """
+    Detach the specified module by invoking the platform API.
+
+    Args:
+        module_name (str): The name of the module to detach.
+
+    Returns:
+        bool: True if the detach command was successfully sent, False otherwise.
+    """
+    if not load_and_verify(module_name):
+        return False
+
+    # Attempt to detach the module
+    if not hasattr(platform_chassis, 'pci_detach'):
+        log.log_error("PCI detach method not found in platform chassis")
+        return False
+
+    log.log_info("Detaching module {}...".format(module_name))
+    try:
+        status = platform_chassis.pci_detach(module_name)
+        if not status:
+            log.log_error("PCI detach status for module {}: {}".format(module_name, status))
+            return False
+        return True
+    except Exception as e:
+        log.log_error("Unexpected error occurred while detaching module {}: {}".format(module_name, e))
+        return False
+
+
+def pci_reattach_module(module_name):
+    """
+    Rescan the specified module by invoking the platform API.
+
+    Args:
+        module_name (str): The name of the module to rescan.
+
+    Returns:
+        bool: True if the rescan command was successfully sent, False otherwise.
+    """
+    if not load_and_verify(module_name):
+        return False
+
+    # Attempt to detach the module
+    if not hasattr(platform_chassis, 'pci_reattach'):
+        log.log_error("PCI reattach method not found in platform chassis")
+        return False
+
+    log.log_info("Rescaning module {}...".format(module_name))
+    try:
+        status = platform_chassis.pci_reattach(module_name)
+        if not status:
+            log.log_error("PCI detach status for module {}: {}".format(module_name, status))
+            return False
+        return True
+    except Exception as e:
+        log.log_error("Unexpected error occurred while detaching module {}: {}".format(module_name, e))
+        return False