[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,19 @@ 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
+
+# On a smartswitch, complete the DPU reboot and exit
+is_smartswitch=$(is_smartswitch)
+if [ "$is_smartswitch" == "True" ] && [ "$REBOOT_DPU" == "yes" ]; then
+    exit $smart_switch_result
+fi
+fi
+
 check_conflict_boot_in_fw_update
 
 setup_reboot_variables
@@ -287,6 +318,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_smartswitch_helper

@@ -0,0 +1,291 @@
+#!/bin/bash
+
+declare -r GNMI_PORT=50052
+declare -r MODULE_REBOOT_DPU="DPU"
+declare -r MODULE_REBOOT_SMARTSWITCH="SMARTSWITCH"
+
+# Function to print debug message
+function log_message() {
+    local message=$1
+    echo "$(date '+%Y-%m-%d %H:%M:%S') - $message" >&2
+}
+
+# Function to check if running on smart switch
+function is_smartswitch()
+{
+    python3 -c "from utilities_common.chassis import is_smartswitch; print(is_smartswitch())" | grep -q "True"
+}
+
+# Function to check if running on DPU
+function is_dpu()
+{
+    python3 -c "from utilities_common.chassis import is_dpu; print(is_dpu())" | grep -q "True"
+}
+
+# Function to retrieve number of DPUs
+function get_num_dpus()
+{
+    python3 -c "from utilities_common.chassis import get_num_dpus; print(get_num_dpus())"
+}
+
+# Function to retrieve DPU IP from CONFIG_DB
+function get_dpu_ip()
+{
+    local DPU_NAME=$1
+    sonic-db-cli CONFIG_DB HGET "DHCP_SERVER_IPV4_PORT|bridge-midplane|${DPU_NAME}" "ips@"
+}
+
+# Function to retrieve GNMI port from CONFIG_DB
+function get_gnmi_port()
+{
+    local DPU_NAME=$1
+    sonic-db-cli CONFIG_DB HGET "DPU_PORT|$DPU_NAME" "gnmi"
+}
+
+# Function to get reboot status from DPU
+function get_reboot_status()
+{
+    local dpu_ip=$1
+    local port=$2
+    local reboot_status
+    reboot_status=$(docker exec -i gnmi gnoi_client -target "${dpu_ip}:${port}" -logtostderr -insecure -rpc RebootStatus 2>/dev/null)
+    if [ $? -ne 0 ] || [ -z "$reboot_status" ]; then
+        log_message "Error: Failed to send reboot status command to DPU ${DPU_NAME}"
+        return ${EXIT_ERROR}
+    fi
+    local is_reboot_active
+    is_reboot_active=$(echo "$reboot_status" | grep "active" | awk '{print $2}')
+    if [ "$is_reboot_active" == "false" ]; then
+        log_message "DPU ${DPU_NAME} has finished rebooting"
+        return ${EXIT_SUCCESS}
+    fi
+    return ${EXIT_ERROR}
+}
+
+# Function to detach PCI module
+function pci_detach_module()
+{
+    local DPU_NAME=$1
+    local DPU_BUS_INFO=$2
+    python3 -c "from utilities_common.module import ModuleHelper; helper = ModuleHelper(); helper.pci_detach_module('${DPU_NAME}')"
+    if [ $? -ne 0 ]; then
+        log_message "Error: PCI detach vendor API is not available"
+        echo 1 > /sys/bus/pci/devices/${DPU_BUS_INFO}/remove
+    fi
+}
+
+# Function to rescan PCI module
+function pci_reattach_module()
+{
+    local DPU_NAME=$1
+    local DPU_BUS_INFO=$2
+    python3 -c "from utilities_common.module import ModuleHelper; helper = ModuleHelper(); helper.pci_reattach_module('${DPU_NAME}')"
+    if [ $? -ne 0 ]; then
+        log_message "Error: PCI reattach vendor API is not available"
+        echo 1 > /sys/bus/pci/rescan
+    fi
+}
+
+# Function to reboot DPU
+function reboot_dpu_platform()
+{
+    local DPU_NAME=$1
+    local REBOOT_TYPE=$2
+    python3 -c "from utilities_common.module import ModuleHelper; helper = ModuleHelper(); helper.reboot_module('${DPU_NAME}', '${REBOOT_TYPE}')"
+}
+
+# Function to wait for DPU reboot status
+function wait_for_dpu_reboot_status()
+{
+    local dpu_ip=$1
+    local port=$2
+
+    if [[ -z "$PLATFORM_JSON_PATH" ]]; then
+        log_message "Error: PLATFORM_JSON_PATH is not defined"
+        exit $EXIT_ERROR
+    fi
+
+    local dpu_halt_services_timeout=$(jq -r '.dpu_halt_services_timeout' "$PLATFORM_JSON_PATH" 2>/dev/null)
+    if [ $? -ne 0 ]; then
+        log_message "Error: Failed to retrieve dpu_halt_services_timeout from ${PLATFORM_JSON_PATH}"
+        return ${EXIT_ERROR}
+    fi
+
+    local poll_interval=5
+    local waited_time=0
+    while true; do
+        local reboot_status
+        get_reboot_status "${dpu_ip}" "${port}"
+        reboot_status=$?
+        if [ $reboot_status -eq ${EXIT_SUCCESS} ]; then
+            break
+        fi
+
+        sleep "$poll_interval"
+        waited_time=$((waited_time + poll_interval))
+        if [ $waited_time -ge $dpu_halt_services_timeout ]; then
+            log_message "Error: Timeout waiting for DPU ${DPU_NAME} to finish rebooting"
+            return ${EXIT_ERROR}
+        fi
+    done
+}
+
+# Function to send reboot command to DPU
+function gnmi_reboot_dpu()
+{
+    # Retrieve DPU IP and GNMI port
+    dpu_ip=$(get_dpu_ip "${DPU_NAME}")
+    log_message "DPU IP ${DPU_NAME}: $dpu_ip"
+    port=$(get_gnmi_port "${DPU_NAME}")
+    if [ -z "$port" ]; then
+        port=$GNMI_PORT # Default GNMI port
+    fi
+    log_message "GNMI port ${DPU_NAME}: $port"
+
+    if [ -z "$dpu_ip" ]; then
+        log_message "Error: Failed to retrieve DPU IP for ${DPU_NAME}"
+        return ${EXIT_ERROR}
+    fi
+
+    docker exec -i gnmi gnoi_client -target ${dpu_ip}:${port} -logtostderr -insecure -rpc Reboot -jsonin '{"method":3}'
+    if [ $? -ne 0 ]; then
+        log_message "Error: Failed to send reboot command to DPU ${DPU_NAME}"
+        return ${EXIT_ERROR}
+    fi
+
+    wait_for_dpu_reboot_status "${dpu_ip}" "${port}"
+}
+
+function reboot_dpu()
+{
+    local DPU_NAME=$1
+    local REBOOT_TYPE=$2
+    local DPU_INDEX=${DPU_NAME//[!0-9]/}
+
+    debug "User requested rebooting device ${DPU_NAME} ..."
+
+    # Send reboot command to DPU
+    gnmi_reboot_dpu "${DPU_NAME}"
+    if [ $? -ne 0 ]; then
+        log_message "Error: Failed to send gnoi command to reboot DPU ${DPU_NAME}"
+    fi
+
+    local DPU_BUS_INFO=$(jq -r --arg DPU_NAME "$DPU_NAME" '.DPUS[$DPU_NAME].bus_info' "$PLATFORM_JSON_PATH")
+    if [ -z "$DPU_BUS_INFO" ] || [ "$DPU_BUS_INFO" = "null" ]; then
+        log_message "Error: Failed to retrieve bus info for DPU ${DPU_NAME}"
+        return ${EXIT_ERROR}
+    fi
+
+    # Update STATE_DB and handle PCIe removal and rescan
+    sonic-db-cli STATE_DB set "PCIE_DETACH_INFO|${DPU_NAME}" '{"dpu_id": "'${DPU_INDEX}'", "dpu_state": "detaching", "bus_info": "'${DPU_BUS_INFO}'"}'
+
+    pci_detach_module ${DPU_NAME} ${DPU_BUS_INFO}
+    if [ $? -ne 0 ]; then
+        log_message "Error: Failed to detach PCI module for DPU ${DPU_NAME}"
+        return ${EXIT_ERROR}
+    fi
+
+    reboot_dpu_platform ${DPU_NAME} ${REBOOT_TYPE}
+    if [ $? -ne 0 ]; then
+        log_message "Error: Failed to send platform command to reboot DPU ${DPU_NAME}"
+        return ${EXIT_ERROR}
+    fi
+
+    pci_reattach_module ${DPU_NAME} ${DPU_BUS_INFO}
+
+    sonic-db-cli STATE_DB del "PCIE_DETACH_INFO|${DPU_NAME}"
+}
+
+# Function to reboot all DPUs in parallel
+function reboot_all_dpus() {
+    local NUM_DPU=$1
+
+    if [[ -z $NUM_DPU ]]; then
+        log_message "Error: Failed to retrieve number of DPUs or no DPUs found"
+        return
+    fi
+
+    local failures=0
+    for (( i=0; i<"$NUM_DPU"; i++ )); do
+        log_message "Rebooting DPU module dpu$i"
+        reboot_dpu "dpu$i" "$MODULE_REBOOT_SMARTSWITCH" &
+        if [ $? -ne 0 ]; then
+            ((failures++))
+        fi
+    done
+    wait
+    return $failures
+}
+
+# Function to verify DPU module name
+function verify_dpu_module_name() {
+    local DPU_MODULE_NAME=$1
+    local NUM_DPU=$2
+
+    if [[ -z "$DPU_MODULE_NAME" ]]; then
+        log_message "Error: DPU module name not provided"
+        return $EXIT_ERROR
+    fi
+
+    NUM_DPU=$((NUM_DPU - 1))
+    if [[ ! "$DPU_MODULE_NAME" =~ ^dpu[0-$NUM_DPU]$ ]]; then
+        log_message "Error: Invalid DPU module name provided"
+        return $EXIT_ERROR
+    fi
+}
+
+# Function to handle scenarios on smart switch
+function handle_smart_switch() {
+    local REBOOT_DPU=$1
+    local PRE_SHUTDOWN=$2
+    local DPU_NAME=$3
+
+    NUM_DPU=$(get_num_dpus)
+
+    if is_dpu; then
+        if [[ "$PRE_SHUTDOWN" != "yes" ]]; then
+            log_message "Error: '-p' option not specified for a DPU"
+            return $EXIT_ERROR
+        elif [[ "$REBOOT_DPU" == "yes" ]]; then
+            log_message "Error: '-d' option specified for a DPU"
+            return $EXIT_ERROR
+        fi
+        return $EXIT_SUCCESS
+    fi
+
+    if [[ "$PRE_SHUTDOWN" == "yes" ]]; then
+        log_message "Error: '-p' option specified for a non-DPU"
+        return $EXIT_ERROR
+    fi
+
+    if [[ "$REBOOT_DPU" == "yes" ]]; then
+        if is_smartswitch; then
+            if [[ -z $NUM_DPU ]]; then
+                log_message "Error: Failed to retrieve number of DPUs or no DPUs found"
+                return $EXIT_ERROR
+            fi
+
+            DPU_MODULE_NAME="${DPU_NAME,,}"
+            verify_dpu_module_name "$DPU_MODULE_NAME" "$NUM_DPU"
+            result=$?
+            if [[ $result -ne $EXIT_SUCCESS ]]; then
+                return $result
+            fi
+
+            log_message "Rebooting device ${DPU_MODULE_NAME}"
+            reboot_dpu "$DPU_MODULE_NAME" "$MODULE_REBOOT_DPU"
+            result=$?
+            return $result
+        else
+            log_message "Error: '-d' option specified for a non-smart-switch"
+            return $EXIT_ERROR
+        fi
+    fi
+
+    # If the system is a smart switch, reboot all DPUs in parallel
+    if is_smartswitch; then
+        reboot_all_dpus "$NUM_DPU" "$MODULE_REBOOT_SMARTSWITCH"
+        result=$?
+        return $result
+    fi
+}

setup.py

@@ -163,6 +163,7 @@
         'scripts/psushow',
         'scripts/queuestat',
         'scripts/reboot',
+        'scripts/reboot_smartswitch_helper',
         'scripts/route_check.py',
         'scripts/route_check_test.sh',
         'scripts/vnet_route_check.py',