From 4901fdcf99b5ca7f1f2dd9a74e69380a84b90647 Mon Sep 17 00:00:00 2001
From: Dawei Huang <daweihuang@microsoft.com>
Date: Fri, 7 Aug 2026 17:23:34 -0500
Subject: [PATCH 01/16] reboot: use TLS for DPU gNOI calls

Signed-off-by: Dawei Huang <daweihuang@microsoft.com>
---
 scripts/reboot_smartswitch_helper       |  4 +--
 tests/reboot_smartswitch_helper_test.py | 44 +++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)
 create mode 100644 tests/reboot_smartswitch_helper_test.py

diff --git a/scripts/reboot_smartswitch_helper b/scripts/reboot_smartswitch_helper
index 5dfe8ae691f..ceb221c7c82 100644
--- a/scripts/reboot_smartswitch_helper
+++ b/scripts/reboot_smartswitch_helper
@@ -53,7 +53,7 @@ function get_reboot_status()
     local port=$2
     local reboot_output_file="reboot_status_${dpu_ip}.txt"
 
-    $(docker exec gnmi gnoi_client -target "${dpu_ip}:${port}" -logtostderr -notls -module System -rpc RebootStatus | tee "$reboot_output_file" &>/dev/null)
+    $(docker exec gnmi gnoi_client -target "${dpu_ip}:${port}" -logtostderr -insecure -module System -rpc RebootStatus | tee "$reboot_output_file" &>/dev/null)
     if [ $? -ne 0 ]; then
         rm -f "$reboot_output_file"
         return ${EXIT_ERROR}
@@ -188,7 +188,7 @@ function gnmi_reboot_dpu()
         return ${EXIT_ERROR}
     fi
 
-    $(docker exec gnmi gnoi_client -target ${dpu_ip}:${port} -logtostderr -notls -module System -rpc Reboot -jsonin '{"method":3, "message":"User initiated reboot"}' &>/dev/null)
+    $(docker exec gnmi gnoi_client -target ${dpu_ip}:${port} -logtostderr -insecure -module System -rpc Reboot -jsonin '{"method":3, "message":"User initiated reboot"}' &>/dev/null)
     if [ $? -ne 0 ]; then
         log_message "ERROR: Failed to send gnoi command to halt services on ${DPU_NAME}"
         log_message "ERROR: proceeding without halting the services"
diff --git a/tests/reboot_smartswitch_helper_test.py b/tests/reboot_smartswitch_helper_test.py
new file mode 100644
index 00000000000..22f2e0a93e2
--- /dev/null
+++ b/tests/reboot_smartswitch_helper_test.py
@@ -0,0 +1,44 @@
+import os
+import subprocess
+from pathlib import Path
+
+
+SCRIPT = Path(__file__).parents[1] / "scripts" / "reboot_smartswitch_helper"
+
+
+def run_helper_function(tmp_path, function_call):
+    command_log = tmp_path / "command.log"
+    env = os.environ.copy()
+    env["COMMAND_LOG"] = str(command_log)
+    script = f'''
+EXIT_SUCCESS=0
+EXIT_ERROR=1
+source "{SCRIPT}"
+docker() {{
+    printf '%s\n' "$*" > "$COMMAND_LOG"
+    printf '{{"active":false}}\n'
+}}
+jq() {{ printf 'false\n'; }}
+get_dpu_ip() {{ printf '169.254.200.1\n'; }}
+get_gnmi_port() {{ printf '8080\n'; }}
+wait_for_dpu_reboot_status() {{ return 0; }}
+{function_call}
+'''
+    subprocess.run(["bash", "-c", script], env=env, check=True)
+    return command_log.read_text()
+
+
+def test_get_reboot_status_uses_tls(tmp_path):
+    command = run_helper_function(
+        tmp_path, "get_reboot_status 169.254.200.1 8080"
+    )
+    assert "-insecure" in command
+    assert "-notls" not in command
+    assert "-rpc RebootStatus" in command
+
+
+def test_gnmi_reboot_dpu_uses_tls(tmp_path):
+    command = run_helper_function(tmp_path, "gnmi_reboot_dpu dpu0")
+    assert "-insecure" in command
+    assert "-notls" not in command
+    assert "-rpc Reboot" in command

From a3cc3fbe86db5963186f02a9fd23bfa36c317629 Mon Sep 17 00:00:00 2001
From: Dawei Huang <daweihuang@microsoft.com>
Date: Fri, 7 Aug 2026 20:40:32 -0500
Subject: [PATCH 02/16] reboot: preserve DPU gNOI failures

Signed-off-by: Dawei Huang <daweihuang@microsoft.com>
---
 scripts/reboot_smartswitch_helper       |  7 +++---
 tests/reboot_smartswitch_helper_test.py | 31 +++++++++++++++++++++----
 2 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/scripts/reboot_smartswitch_helper b/scripts/reboot_smartswitch_helper
index ceb221c7c82..59a194a35db 100644
--- a/scripts/reboot_smartswitch_helper
+++ b/scripts/reboot_smartswitch_helper
@@ -53,8 +53,9 @@ function get_reboot_status()
     local port=$2
     local reboot_output_file="reboot_status_${dpu_ip}.txt"
 
-    $(docker exec gnmi gnoi_client -target "${dpu_ip}:${port}" -logtostderr -insecure -module System -rpc RebootStatus | tee "$reboot_output_file" &>/dev/null)
-    if [ $? -ne 0 ]; then
+    docker exec gnmi gnoi_client -target "${dpu_ip}:${port}" -logtostderr -insecure -module System -rpc RebootStatus | tee "$reboot_output_file" &>/dev/null
+    local gnoi_rc=${PIPESTATUS[0]}
+    if [ $gnoi_rc -ne 0 ]; then
         rm -f "$reboot_output_file"
         return ${EXIT_ERROR}
     fi
@@ -188,7 +189,7 @@ function gnmi_reboot_dpu()
         return ${EXIT_ERROR}
     fi
 
-    $(docker exec gnmi gnoi_client -target ${dpu_ip}:${port} -logtostderr -insecure -module System -rpc Reboot -jsonin '{"method":3, "message":"User initiated reboot"}' &>/dev/null)
+    docker exec gnmi gnoi_client -target "${dpu_ip}:${port}" -logtostderr -insecure -module System -rpc Reboot -jsonin '{"method":3, "message":"User initiated reboot"}' &>/dev/null
     if [ $? -ne 0 ]; then
         log_message "ERROR: Failed to send gnoi command to halt services on ${DPU_NAME}"
         log_message "ERROR: proceeding without halting the services"
diff --git a/tests/reboot_smartswitch_helper_test.py b/tests/reboot_smartswitch_helper_test.py
index 22f2e0a93e2..2a55a6a1974 100644
--- a/tests/reboot_smartswitch_helper_test.py
+++ b/tests/reboot_smartswitch_helper_test.py
@@ -6,7 +6,7 @@
 SCRIPT = Path(__file__).parents[1] / "scripts" / "reboot_smartswitch_helper"
 
 
-def run_helper_function(tmp_path, function_call):
+def run_helper_function(tmp_path, function_call, docker_rc=0):
     command_log = tmp_path / "command.log"
     env = os.environ.copy()
     env["COMMAND_LOG"] = str(command_log)
@@ -17,6 +17,7 @@ def run_helper_function(tmp_path, function_call):
 docker() {{
     printf '%s\n' "$*" > "$COMMAND_LOG"
     printf '{{"active":false}}\n'
+    return "$DOCKER_RC"
 }}
 jq() {{ printf 'false\n'; }}
 get_dpu_ip() {{ printf '169.254.200.1\n'; }}
@@ -24,21 +25,41 @@ def run_helper_function(tmp_path, function_call):
 wait_for_dpu_reboot_status() {{ return 0; }}
 {function_call}
 '''
-    subprocess.run(["bash", "-c", script], env=env, check=True)
-    return command_log.read_text()
+    env["DOCKER_RC"] = str(docker_rc)
+    result = subprocess.run(
+        ["bash", "-c", script], env=env, capture_output=True, text=True
+    )
+    return result, command_log.read_text()
 
 
 def test_get_reboot_status_uses_tls(tmp_path):
-    command = run_helper_function(
+    result, command = run_helper_function(
         tmp_path, "get_reboot_status 169.254.200.1 8080"
     )
+    assert result.returncode == 0
     assert "-insecure" in command
     assert "-notls" not in command
     assert "-rpc RebootStatus" in command
 
 
 def test_gnmi_reboot_dpu_uses_tls(tmp_path):
-    command = run_helper_function(tmp_path, "gnmi_reboot_dpu dpu0")
+    result, command = run_helper_function(tmp_path, "gnmi_reboot_dpu dpu0")
+    assert result.returncode == 0
     assert "-insecure" in command
     assert "-notls" not in command
     assert "-rpc Reboot" in command
+
+
+def test_get_reboot_status_preserves_gnoi_failure(tmp_path):
+    result, _ = run_helper_function(
+        tmp_path, "get_reboot_status 169.254.200.1 8080", docker_rc=1
+    )
+    assert result.returncode != 0
+
+
+def test_gnmi_reboot_dpu_reports_gnoi_failure(tmp_path):
+    result, _ = run_helper_function(
+        tmp_path, "gnmi_reboot_dpu dpu0", docker_rc=1
+    )
+    assert result.returncode == 0
+    assert "Failed to send gnoi command to halt services" in result.stderr

From 1099d93335c5b849ed4cb8f5b201a5345c586ccc Mon Sep 17 00:00:00 2001
From: Dawei Huang <daweihuang@microsoft.com>
Date: Sat, 8 Aug 2026 02:33:05 -0500
Subject: [PATCH 03/16] reboot: fall back to native DPU gNMI port

Signed-off-by: Dawei Huang <daweihuang@microsoft.com>
---
 scripts/reboot_smartswitch_helper       | 44 ++++++++++++++++++-------
 tests/reboot_smartswitch_helper_test.py | 27 ++++++++++++---
 2 files changed, 56 insertions(+), 15 deletions(-)

diff --git a/scripts/reboot_smartswitch_helper b/scripts/reboot_smartswitch_helper
index 59a194a35db..938f9d23c6e 100644
--- a/scripts/reboot_smartswitch_helper
+++ b/scripts/reboot_smartswitch_helper
@@ -1,6 +1,7 @@
 #!/bin/bash
 
-declare -r GNMI_PORT=8080 # Default GNMI port
+declare -ra COMMON_GNMI_PORTS=(8080 50052)
+declare -r GNMI_PROBE_TIMEOUT=10
 declare -r MODULE_REBOOT_DPU="DPU"
 declare -r MODULE_REBOOT_SMARTSWITCH="SMARTSWITCH"
 
@@ -37,13 +38,32 @@ function get_dpu_ip()
     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()
+# Function to retrieve configured and common GNMI ports in preference order
+function get_gnmi_ports()
 {
     local DPU_NAME=${1:-dpu0}
+    local configured_port
     for k in $(sonic-db-cli CONFIG_DB keys "DPU|*$DPU_NAME"); do
-        sonic-db-cli CONFIG_DB hget "$k" 'gnmi_port'
+        configured_port=$(sonic-db-cli CONFIG_DB hget "$k" 'gnmi_port')
+        if [ -n "$configured_port" ]; then
+            break
+        fi
+    done
+    printf '%s\n' "$configured_port" "${COMMON_GNMI_PORTS[@]}" | awk 'NF && !seen[$0]++'
+}
+
+function find_working_gnmi_port()
+{
+    local dpu_ip=$1
+    shift
+    local port
+    for port in "$@"; do
+        if timeout "$GNMI_PROBE_TIMEOUT" docker exec gnmi gnoi_client -target "${dpu_ip}:${port}" -logtostderr -insecure -module System -rpc Time &>/dev/null; then
+            printf '%s\n' "$port"
+            return ${EXIT_SUCCESS}
+        fi
     done
+    return ${EXIT_ERROR}
 }
 
 # Function to get reboot status from DPU
@@ -178,19 +198,21 @@ function gnmi_reboot_dpu()
 
     # Retrieve DPU IP and GNMI port
     dpu_ip=$(get_dpu_ip "${DPU_NAME}")
-    port=$(get_gnmi_port "${DPU_NAME}")
-    if [ -z "$port" ]; then
-        port=$GNMI_PORT # Default GNMI port
-    fi
-    log_message "INFO: Rebooting ${DPU_NAME}, ip:$dpu_ip gnmi_port:$port"
+    local ports
+    mapfile -t ports < <(get_gnmi_ports "${DPU_NAME}")
+    log_message "INFO: Rebooting ${DPU_NAME}, ip:$dpu_ip gnmi_ports:${ports[*]}"
 
     if [ -z "$dpu_ip" ]; then
         log_message "ERROR: Failed to retrieve DPU IP for ${DPU_NAME}"
         return ${EXIT_ERROR}
     fi
 
-    docker exec gnmi gnoi_client -target "${dpu_ip}:${port}" -logtostderr -insecure -module System -rpc Reboot -jsonin '{"method":3, "message":"User initiated reboot"}' &>/dev/null
-    if [ $? -ne 0 ]; then
+    local port
+    port=$(find_working_gnmi_port "$dpu_ip" "${ports[@]}")
+    if [ -z "$port" ]; then
+        log_message "ERROR: Failed to find a reachable gNMI port on ${DPU_NAME}"
+        log_message "ERROR: proceeding without halting the services"
+    elif ! docker exec gnmi gnoi_client -target "${dpu_ip}:${port}" -logtostderr -insecure -module System -rpc Reboot -jsonin '{"method":3, "message":"User initiated reboot"}' &>/dev/null; then
         log_message "ERROR: Failed to send gnoi command to halt services on ${DPU_NAME}"
         log_message "ERROR: proceeding without halting the services"
     else
diff --git a/tests/reboot_smartswitch_helper_test.py b/tests/reboot_smartswitch_helper_test.py
index 2a55a6a1974..027e3c495ff 100644
--- a/tests/reboot_smartswitch_helper_test.py
+++ b/tests/reboot_smartswitch_helper_test.py
@@ -6,7 +6,7 @@
 SCRIPT = Path(__file__).parents[1] / "scripts" / "reboot_smartswitch_helper"
 
 
-def run_helper_function(tmp_path, function_call, docker_rc=0):
+def run_helper_function(tmp_path, function_call, docker_rc=0, fail_port=""):
     command_log = tmp_path / "command.log"
     env = os.environ.copy()
     env["COMMAND_LOG"] = str(command_log)
@@ -15,17 +15,22 @@ def run_helper_function(tmp_path, function_call, docker_rc=0):
 EXIT_ERROR=1
 source "{SCRIPT}"
 docker() {{
-    printf '%s\n' "$*" > "$COMMAND_LOG"
+    printf '%s\n' "$*" >> "$COMMAND_LOG"
     printf '{{"active":false}}\n'
+    if [ -n "$FAIL_PORT" ] && [[ "$*" == *":$FAIL_PORT"* ]]; then
+        return 1
+    fi
     return "$DOCKER_RC"
 }}
+timeout() {{ shift; docker "$@"; }}
 jq() {{ printf 'false\n'; }}
 get_dpu_ip() {{ printf '169.254.200.1\n'; }}
-get_gnmi_port() {{ printf '8080\n'; }}
+get_gnmi_ports() {{ printf '8080\n50052\n'; }}
 wait_for_dpu_reboot_status() {{ return 0; }}
 {function_call}
 '''
     env["DOCKER_RC"] = str(docker_rc)
+    env["FAIL_PORT"] = fail_port
     result = subprocess.run(
         ["bash", "-c", script], env=env, capture_output=True, text=True
     )
@@ -62,4 +67,18 @@ def test_gnmi_reboot_dpu_reports_gnoi_failure(tmp_path):
         tmp_path, "gnmi_reboot_dpu dpu0", docker_rc=1
     )
     assert result.returncode == 0
-    assert "Failed to send gnoi command to halt services" in result.stderr
+    assert "Failed to find a reachable gNMI port" in result.stderr
+
+
+def test_gnmi_reboot_dpu_falls_back_to_native_port(tmp_path):
+    result, commands = run_helper_function(
+        tmp_path, "gnmi_reboot_dpu dpu0", fail_port="8080"
+    )
+    assert result.returncode == 0
+    assert "-target 169.254.200.1:8080" in commands
+    assert "-target 169.254.200.1:50052" in commands
+    command_lines = commands.splitlines()
+    assert "-rpc Time" in command_lines[0]
+    assert "-rpc Time" in command_lines[1]
+    assert "-rpc Reboot" in command_lines[2]
+    assert sum("-rpc Reboot" in line for line in command_lines) == 1

From e1fd871251f32dc445de5baedeb69e51016b79db Mon Sep 17 00:00:00 2001
From: Dawei Huang <daweihuang@microsoft.com>
Date: Sat, 8 Aug 2026 02:35:02 -0500
Subject: [PATCH 04/16] test: cover DPU gNMI port ordering

Signed-off-by: Dawei Huang <daweihuang@microsoft.com>
---
 tests/reboot_smartswitch_helper_test.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/tests/reboot_smartswitch_helper_test.py b/tests/reboot_smartswitch_helper_test.py
index 027e3c495ff..c3522e0df53 100644
--- a/tests/reboot_smartswitch_helper_test.py
+++ b/tests/reboot_smartswitch_helper_test.py
@@ -82,3 +82,25 @@ def test_gnmi_reboot_dpu_falls_back_to_native_port(tmp_path):
     assert "-rpc Time" in command_lines[1]
     assert "-rpc Reboot" in command_lines[2]
     assert sum("-rpc Reboot" in line for line in command_lines) == 1
+
+
+def test_get_gnmi_ports_orders_and_deduplicates(tmp_path):
+    script = f'''
+EXIT_SUCCESS=0
+EXIT_ERROR=1
+source "{SCRIPT}"
+sonic-db-cli() {{
+    if [ "$2" = "keys" ]; then
+        printf 'DPU|dpu0\n'
+    else
+        printf '%s\n' "$CONFIGURED_PORT"
+    fi
+}}
+get_gnmi_ports dpu0
+'''
+    env = os.environ.copy()
+    env["CONFIGURED_PORT"] = "50052"
+    result = subprocess.run(
+        ["bash", "-c", script], env=env, capture_output=True, text=True, check=True
+    )
+    assert result.stdout.splitlines() == ["50052", "8080"]

From 9436884762c8a3f4ffe0986935ba056470ace731 Mon Sep 17 00:00:00 2001
From: Dawei Huang <daweihuang@microsoft.com>
Date: Sat, 8 Aug 2026 03:32:42 -0500
Subject: [PATCH 05/16] reboot: bound DPU gNOI operations

Signed-off-by: Dawei Huang <daweihuang@microsoft.com>
---
 scripts/reboot_smartswitch_helper       | 19 ++++++++++++-----
 tests/reboot_smartswitch_helper_test.py | 27 +++++++++++++++++++++----
 2 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/scripts/reboot_smartswitch_helper b/scripts/reboot_smartswitch_helper
index 938f9d23c6e..576065f44d4 100644
--- a/scripts/reboot_smartswitch_helper
+++ b/scripts/reboot_smartswitch_helper
@@ -2,6 +2,8 @@
 
 declare -ra COMMON_GNMI_PORTS=(8080 50052)
 declare -r GNMI_PROBE_TIMEOUT=10
+declare -r GNMI_REBOOT_TIMEOUT=60
+declare -r GNMI_STATUS_TIMEOUT=10
 declare -r MODULE_REBOOT_DPU="DPU"
 declare -r MODULE_REBOOT_SMARTSWITCH="SMARTSWITCH"
 
@@ -73,15 +75,19 @@ function get_reboot_status()
     local port=$2
     local reboot_output_file="reboot_status_${dpu_ip}.txt"
 
-    docker exec gnmi gnoi_client -target "${dpu_ip}:${port}" -logtostderr -insecure -module System -rpc RebootStatus | tee "$reboot_output_file" &>/dev/null
+    timeout "$GNMI_STATUS_TIMEOUT" docker exec gnmi gnoi_client -target "${dpu_ip}:${port}" -logtostderr -insecure -module System -rpc RebootStatus | tee "$reboot_output_file" &>/dev/null
     local gnoi_rc=${PIPESTATUS[0]}
     if [ $gnoi_rc -ne 0 ]; then
         rm -f "$reboot_output_file"
         return ${EXIT_ERROR}
     fi
-    local is_reboot_active
-    is_reboot_active=$(cat "$reboot_output_file" | awk '/^{.*}$/' | jq -r '.active')
+    local status_json
+    status_json=$(awk '/^{.*}$/{json=$0} END{print json}' "$reboot_output_file")
     rm -f "$reboot_output_file"
+    local is_reboot_active
+    if [ -z "$status_json" ] || ! is_reboot_active=$(jq -r 'if has("active") then .active elif has("status") then false else error("invalid RebootStatus response") end' <<< "$status_json"); then
+        return ${EXIT_ERROR}
+    fi
     if [ "$is_reboot_active" != "true" ]; then
         return ${EXIT_SUCCESS}
     fi
@@ -185,7 +191,7 @@ function wait_for_dpu_reboot_status()
         waited_time=$((waited_time + poll_interval))
         if [ $waited_time -ge $dpu_halt_services_timeout ]; then
             log_message "ERROR: Timeout waiting for ${DPU_NAME} to finish halting the services"
-            return
+            return ${EXIT_ERROR}
         fi
     done
     return
@@ -212,12 +218,15 @@ function gnmi_reboot_dpu()
     if [ -z "$port" ]; then
         log_message "ERROR: Failed to find a reachable gNMI port on ${DPU_NAME}"
         log_message "ERROR: proceeding without halting the services"
-    elif ! docker exec gnmi gnoi_client -target "${dpu_ip}:${port}" -logtostderr -insecure -module System -rpc Reboot -jsonin '{"method":3, "message":"User initiated reboot"}' &>/dev/null; then
+        return ${EXIT_ERROR}
+    elif ! timeout "$GNMI_REBOOT_TIMEOUT" docker exec gnmi gnoi_client -target "${dpu_ip}:${port}" -logtostderr -insecure -module System -rpc Reboot -jsonin '{"method":3, "message":"User initiated reboot"}' &>/dev/null; then
         log_message "ERROR: Failed to send gnoi command to halt services on ${DPU_NAME}"
         log_message "ERROR: proceeding without halting the services"
+        return ${EXIT_ERROR}
     else
         # Wait for DPU to halt services, if reboot command is successful
         wait_for_dpu_reboot_status "${dpu_ip}" "${port}" "${DPU_NAME}"
+        return $?
     fi
 }
 
diff --git a/tests/reboot_smartswitch_helper_test.py b/tests/reboot_smartswitch_helper_test.py
index c3522e0df53..eca8a1c1c2e 100644
--- a/tests/reboot_smartswitch_helper_test.py
+++ b/tests/reboot_smartswitch_helper_test.py
@@ -8,15 +8,18 @@
 
 def run_helper_function(tmp_path, function_call, docker_rc=0, fail_port=""):
     command_log = tmp_path / "command.log"
+    platform_json = tmp_path / "platform.json"
+    platform_json.write_text('{"dpu_halt_services_timeout": 1}')
     env = os.environ.copy()
     env["COMMAND_LOG"] = str(command_log)
+    env["PLATFORM_JSON_PATH"] = str(platform_json)
     script = f'''
 EXIT_SUCCESS=0
 EXIT_ERROR=1
 source "{SCRIPT}"
 docker() {{
     printf '%s\n' "$*" >> "$COMMAND_LOG"
-    printf '{{"active":false}}\n'
+    printf '%s\n' "$DOCKER_OUTPUT"
     if [ -n "$FAIL_PORT" ] && [[ "$*" == *":$FAIL_PORT"* ]]; then
         return 1
     fi
@@ -26,11 +29,11 @@ def run_helper_function(tmp_path, function_call, docker_rc=0, fail_port=""):
 jq() {{ printf 'false\n'; }}
 get_dpu_ip() {{ printf '169.254.200.1\n'; }}
 get_gnmi_ports() {{ printf '8080\n50052\n'; }}
-wait_for_dpu_reboot_status() {{ return 0; }}
 {function_call}
 '''
     env["DOCKER_RC"] = str(docker_rc)
     env["FAIL_PORT"] = fail_port
+    env["DOCKER_OUTPUT"] = '{"active":false}'
     result = subprocess.run(
         ["bash", "-c", script], env=env, capture_output=True, text=True
     )
@@ -66,7 +69,7 @@ def test_gnmi_reboot_dpu_reports_gnoi_failure(tmp_path):
     result, _ = run_helper_function(
         tmp_path, "gnmi_reboot_dpu dpu0", docker_rc=1
     )
-    assert result.returncode == 0
+    assert result.returncode != 0
     assert "Failed to find a reachable gNMI port" in result.stderr
 
 
@@ -81,7 +84,23 @@ def test_gnmi_reboot_dpu_falls_back_to_native_port(tmp_path):
     assert "-rpc Time" in command_lines[0]
     assert "-rpc Time" in command_lines[1]
     assert "-rpc Reboot" in command_lines[2]
-    assert sum("-rpc Reboot" in line for line in command_lines) == 1
+    assert "-target 169.254.200.1:50052" in command_lines[3]
+    assert "-rpc RebootStatus" in command_lines[3]
+    assert sum("-rpc Reboot " in f"{line} " for line in command_lines) == 1
+
+
+def test_get_reboot_status_rejects_malformed_output(tmp_path):
+    command_log = tmp_path / "command.log"
+    script = f'''
+EXIT_SUCCESS=0
+EXIT_ERROR=1
+source "{SCRIPT}"
+timeout() {{ shift; docker "$@"; }}
+docker() {{ printf 'not-json\n'; }}
+get_reboot_status 169.254.200.1 50052
+'''
+    result = subprocess.run(["bash", "-c", script], capture_output=True, text=True)
+    assert result.returncode != 0
 
 
 def test_get_gnmi_ports_orders_and_deduplicates(tmp_path):

From e7500b781025d4d1140976d961566ee0e5c44c34 Mon Sep 17 00:00:00 2001
From: Dawei Huang <daweihuang@microsoft.com>
Date: Sat, 8 Aug 2026 03:51:48 -0500
Subject: [PATCH 06/16] reboot: validate DPU halt status

Signed-off-by: Dawei Huang <daweihuang@microsoft.com>
---
 scripts/reboot_smartswitch_helper       | 20 +++++++++-----------
 tests/reboot_smartswitch_helper_test.py | 16 ++++++++++++++--
 2 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/scripts/reboot_smartswitch_helper b/scripts/reboot_smartswitch_helper
index 576065f44d4..1f6626218e6 100644
--- a/scripts/reboot_smartswitch_helper
+++ b/scripts/reboot_smartswitch_helper
@@ -85,7 +85,7 @@ function get_reboot_status()
     status_json=$(awk '/^{.*}$/{json=$0} END{print json}' "$reboot_output_file")
     rm -f "$reboot_output_file"
     local is_reboot_active
-    if [ -z "$status_json" ] || ! is_reboot_active=$(jq -r 'if has("active") then .active elif has("status") then false else error("invalid RebootStatus response") end' <<< "$status_json"); then
+    if [ -z "$status_json" ] || ! is_reboot_active=$(jq -r 'if .active == true then true elif (.status.status // 0) == 1 then false else error("reboot not complete") end' <<< "$status_json"); then
         return ${EXIT_ERROR}
     fi
     if [ "$is_reboot_active" != "true" ]; then
@@ -176,8 +176,8 @@ function wait_for_dpu_reboot_status()
     fi
 
     local poll_interval=5
-    local waited_time=0
-    while true; do
+    local deadline=$((SECONDS + dpu_halt_services_timeout))
+    while [ "$SECONDS" -lt "$deadline" ]; do
         sleep "$poll_interval"
 
         local reboot_status
@@ -185,16 +185,12 @@ function wait_for_dpu_reboot_status()
         reboot_status=$?
         if [ $reboot_status -eq ${EXIT_SUCCESS} ]; then
             log_message "INFO: ${DPU_NAME} halted the services successfully"
-            break
+            return ${EXIT_SUCCESS}
         fi
 
-        waited_time=$((waited_time + poll_interval))
-        if [ $waited_time -ge $dpu_halt_services_timeout ]; then
-            log_message "ERROR: Timeout waiting for ${DPU_NAME} to finish halting the services"
-            return ${EXIT_ERROR}
-        fi
     done
-    return
+    log_message "ERROR: Timeout waiting for ${DPU_NAME} to finish halting the services"
+    return ${EXIT_ERROR}
 }
 
 # Function to send reboot command to DPU
@@ -264,7 +260,8 @@ function reboot_dpu()
 
     # Send reboot command to DPU
     gnmi_reboot_dpu "${DPU_NAME}"
-    if [ $? -ne 0 ]; then
+    local gnoi_reboot_rc=$?
+    if [ "$gnoi_reboot_rc" -ne 0 ]; then
         log_message "ERROR: Failed to send gnoi command to reboot ${DPU_NAME}"
     fi
 
@@ -301,6 +298,7 @@ function reboot_dpu()
         fi
         clear_module_state_transition_flag "${DPU_NAME}"
     fi
+    return "$gnoi_reboot_rc"
 }
 
 # Function to reboot all DPUs in parallel
diff --git a/tests/reboot_smartswitch_helper_test.py b/tests/reboot_smartswitch_helper_test.py
index eca8a1c1c2e..9deb2cec254 100644
--- a/tests/reboot_smartswitch_helper_test.py
+++ b/tests/reboot_smartswitch_helper_test.py
@@ -26,14 +26,13 @@ def run_helper_function(tmp_path, function_call, docker_rc=0, fail_port=""):
     return "$DOCKER_RC"
 }}
 timeout() {{ shift; docker "$@"; }}
-jq() {{ printf 'false\n'; }}
 get_dpu_ip() {{ printf '169.254.200.1\n'; }}
 get_gnmi_ports() {{ printf '8080\n50052\n'; }}
 {function_call}
 '''
     env["DOCKER_RC"] = str(docker_rc)
     env["FAIL_PORT"] = fail_port
-    env["DOCKER_OUTPUT"] = '{"active":false}'
+    env["DOCKER_OUTPUT"] = '{"active":false,"status":{"status":1}}'
     result = subprocess.run(
         ["bash", "-c", script], env=env, capture_output=True, text=True
     )
@@ -103,6 +102,19 @@ def test_get_reboot_status_rejects_malformed_output(tmp_path):
     assert result.returncode != 0
 
 
+def test_get_reboot_status_rejects_failure_status(tmp_path):
+    script = f'''
+EXIT_SUCCESS=0
+EXIT_ERROR=1
+source "{SCRIPT}"
+timeout() {{ shift; docker "$@"; }}
+docker() {{ printf '{{"active":false,"status":{{"status":3}}}}\n'; }}
+get_reboot_status 169.254.200.1 50052
+'''
+    result = subprocess.run(["bash", "-c", script], capture_output=True, text=True)
+    assert result.returncode != 0
+
+
 def test_get_gnmi_ports_orders_and_deduplicates(tmp_path):
     script = f'''
 EXIT_SUCCESS=0

From 02418fb13e4c5f9e7b6045e4e62231948cc088c7 Mon Sep 17 00:00:00 2001
From: Dawei Huang <daweihuang@microsoft.com>
Date: Sat, 8 Aug 2026 04:00:13 -0500
Subject: [PATCH 07/16] reboot: propagate parallel DPU failures

Signed-off-by: Dawei Huang <daweihuang@microsoft.com>
---
 scripts/reboot_smartswitch_helper       | 34 ++++++++++++++++++----
 tests/reboot_smartswitch_helper_test.py | 38 ++++++++++++++++++++++++-
 2 files changed, 65 insertions(+), 7 deletions(-)

diff --git a/scripts/reboot_smartswitch_helper b/scripts/reboot_smartswitch_helper
index 1f6626218e6..cae16e5cf6c 100644
--- a/scripts/reboot_smartswitch_helper
+++ b/scripts/reboot_smartswitch_helper
@@ -73,9 +73,10 @@ function get_reboot_status()
 {
     local dpu_ip=$1
     local port=$2
+    local timeout_sec=${3:-$GNMI_STATUS_TIMEOUT}
     local reboot_output_file="reboot_status_${dpu_ip}.txt"
 
-    timeout "$GNMI_STATUS_TIMEOUT" docker exec gnmi gnoi_client -target "${dpu_ip}:${port}" -logtostderr -insecure -module System -rpc RebootStatus | tee "$reboot_output_file" &>/dev/null
+    timeout "$timeout_sec" docker exec gnmi gnoi_client -target "${dpu_ip}:${port}" -logtostderr -insecure -module System -rpc RebootStatus | tee "$reboot_output_file" &>/dev/null
     local gnoi_rc=${PIPESTATUS[0]}
     if [ $gnoi_rc -ne 0 ]; then
         rm -f "$reboot_output_file"
@@ -177,11 +178,28 @@ function wait_for_dpu_reboot_status()
 
     local poll_interval=5
     local deadline=$((SECONDS + dpu_halt_services_timeout))
-    while [ "$SECONDS" -lt "$deadline" ]; do
-        sleep "$poll_interval"
+    while true; do
+        local remaining=$((deadline - SECONDS))
+        if [ "$remaining" -le 0 ]; then
+            break
+        fi
+        local sleep_time=$poll_interval
+        if [ "$remaining" -lt "$sleep_time" ]; then
+            sleep_time=$remaining
+        fi
+        sleep "$sleep_time"
+
+        remaining=$((deadline - SECONDS))
+        if [ "$remaining" -le 0 ]; then
+            break
+        fi
+        local status_timeout=$GNMI_STATUS_TIMEOUT
+        if [ "$remaining" -lt "$status_timeout" ]; then
+            status_timeout=$remaining
+        fi
 
         local reboot_status
-        get_reboot_status "${dpu_ip}" "${port}"
+        get_reboot_status "${dpu_ip}" "${port}" "$status_timeout"
         reboot_status=$?
         if [ $reboot_status -eq ${EXIT_SUCCESS} ]; then
             log_message "INFO: ${DPU_NAME} halted the services successfully"
@@ -311,13 +329,17 @@ function reboot_all_dpus() {
     fi
 
     local failures=0
+    local pids=()
     for (( i=0; i<"$NUM_DPU"; i++ )); do
         reboot_dpu "dpu$i" "$MODULE_REBOOT_SMARTSWITCH" &
-        if [ $? -ne 0 ]; then
+        pids+=("$!")
+    done
+    local pid
+    for pid in "${pids[@]}"; do
+        if ! wait "$pid"; then
             ((failures++))
         fi
     done
-    wait
     return $failures
 }
 
diff --git a/tests/reboot_smartswitch_helper_test.py b/tests/reboot_smartswitch_helper_test.py
index 9deb2cec254..99e783ddb6b 100644
--- a/tests/reboot_smartswitch_helper_test.py
+++ b/tests/reboot_smartswitch_helper_test.py
@@ -9,7 +9,7 @@
 def run_helper_function(tmp_path, function_call, docker_rc=0, fail_port=""):
     command_log = tmp_path / "command.log"
     platform_json = tmp_path / "platform.json"
-    platform_json.write_text('{"dpu_halt_services_timeout": 1}')
+    platform_json.write_text('{"dpu_halt_services_timeout": 6}')
     env = os.environ.copy()
     env["COMMAND_LOG"] = str(command_log)
     env["PLATFORM_JSON_PATH"] = str(platform_json)
@@ -135,3 +135,39 @@ def test_get_gnmi_ports_orders_and_deduplicates(tmp_path):
         ["bash", "-c", script], env=env, capture_output=True, text=True, check=True
     )
     assert result.stdout.splitlines() == ["50052", "8080"]
+
+
+def test_reboot_dpu_continues_hardware_reboot_and_returns_gnoi_failure(tmp_path):
+    platform_json = tmp_path / "platform.json"
+    platform_json.write_text('{"DPUS":{"dpu0":{"bus_info":"0000:00:00.0"}}}')
+    marker = tmp_path / "platform-rebooted"
+    script = f'''
+EXIT_SUCCESS=0
+EXIT_ERROR=1
+PLATFORM_JSON_PATH="{platform_json}"
+source "{SCRIPT}"
+show() {{ printf '  DPU0 test Online up\n'; }}
+get_module_state_transition_flag() {{ return 1; }}
+set_module_state_transition_flag() {{ return 0; }}
+clear_module_state_transition_flag() {{ return 0; }}
+gnmi_reboot_dpu() {{ return 1; }}
+module_pre_shutdown() {{ return 0; }}
+module_post_startup() {{ return 0; }}
+reboot_dpu_platform() {{ touch "{marker}"; return 0; }}
+reboot_dpu dpu0 DPU
+'''
+    result = subprocess.run(["bash", "-c", script], capture_output=True, text=True)
+    assert result.returncode != 0
+    assert marker.exists()
+
+
+def test_reboot_all_dpus_collects_background_failures():
+    script = f'''
+EXIT_SUCCESS=0
+EXIT_ERROR=1
+source "{SCRIPT}"
+reboot_dpu() {{ [ "$1" != "dpu1" ]; }}
+reboot_all_dpus 3
+'''
+    result = subprocess.run(["bash", "-c", script], capture_output=True, text=True)
+    assert result.returncode == 1

From 15b2197b3064d65af8e695594d9f5b908f871061 Mon Sep 17 00:00:00 2001
From: Dawei Huang <daweihuang@microsoft.com>
Date: Sat, 8 Aug 2026 04:21:06 -0500
Subject: [PATCH 08/16] reboot: continue whole-switch recovery

Signed-off-by: Dawei Huang <daweihuang@microsoft.com>
---
 scripts/reboot_smartswitch_helper       |  6 ++++--
 tests/reboot_smartswitch_helper_test.py | 17 ++++++++++++++++-
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/scripts/reboot_smartswitch_helper b/scripts/reboot_smartswitch_helper
index cae16e5cf6c..ad05619a230 100644
--- a/scripts/reboot_smartswitch_helper
+++ b/scripts/reboot_smartswitch_helper
@@ -410,7 +410,9 @@ function handle_smart_switch() {
     # 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
+        if [ $? -ne 0 ]; then
+            log_message "ERROR: One or more DPU reboot preparations failed; continuing SmartSwitch reboot"
+        fi
+        return ${EXIT_SUCCESS}
     fi
 }
diff --git a/tests/reboot_smartswitch_helper_test.py b/tests/reboot_smartswitch_helper_test.py
index 99e783ddb6b..3e5ad0e8f5b 100644
--- a/tests/reboot_smartswitch_helper_test.py
+++ b/tests/reboot_smartswitch_helper_test.py
@@ -89,7 +89,6 @@ def test_gnmi_reboot_dpu_falls_back_to_native_port(tmp_path):
 
 
 def test_get_reboot_status_rejects_malformed_output(tmp_path):
-    command_log = tmp_path / "command.log"
     script = f'''
 EXIT_SUCCESS=0
 EXIT_ERROR=1
@@ -171,3 +170,19 @@ def test_reboot_all_dpus_collects_background_failures():
 '''
     result = subprocess.run(["bash", "-c", script], capture_output=True, text=True)
     assert result.returncode == 1
+
+
+def test_whole_smartswitch_reboot_continues_after_dpu_failure():
+    script = f'''
+EXIT_SUCCESS=0
+EXIT_ERROR=1
+source "{SCRIPT}"
+get_num_dpus() {{ printf '2\n'; }}
+is_dpu() {{ return 1; }}
+is_smartswitch() {{ return 0; }}
+reboot_all_dpus() {{ return 1; }}
+handle_smart_switch no no ""
+'''
+    result = subprocess.run(["bash", "-c", script], capture_output=True, text=True)
+    assert result.returncode == 0
+    assert "continuing SmartSwitch reboot" in result.stderr

From 217f882eeb6d7786b7d6e04650c7499c4b401df1 Mon Sep 17 00:00:00 2001
From: Dawei Huang <daweihuang@microsoft.com>
Date: Sat, 8 Aug 2026 04:44:28 -0500
Subject: [PATCH 09/16] reboot: propagate DPU platform failures

Signed-off-by: Dawei Huang <daweihuang@microsoft.com>
---
 scripts/reboot_smartswitch_helper       | 46 ++++++++++++++-------
 tests/reboot_smartswitch_helper_test.py | 54 +++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 15 deletions(-)

diff --git a/scripts/reboot_smartswitch_helper b/scripts/reboot_smartswitch_helper
index ad05619a230..a1551d4a40a 100644
--- a/scripts/reboot_smartswitch_helper
+++ b/scripts/reboot_smartswitch_helper
@@ -100,9 +100,11 @@ function module_pre_shutdown()
 {
     local DPU_NAME=$1
     python3 -c "from utilities_common.module import ModuleHelper; helper = ModuleHelper(); helper.module_pre_shutdown('${DPU_NAME}')"
-    if [ $? -ne 0 ]; then
+    local rc=$?
+    if [ "$rc" -ne 0 ]; then
         log_message "ERROR: Module pre-shutdown vendor API failed"
     fi
+    return "$rc"
 }
 
 # Function to rescan PCI module
@@ -110,9 +112,11 @@ function module_post_startup()
 {
     local DPU_NAME=$1
     python3 -c "from utilities_common.module import ModuleHelper; helper = ModuleHelper(); helper.module_post_startup('${DPU_NAME}')"
-    if [ $? -ne 0 ]; then
+    local rc=$?
+    if [ "$rc" -ne 0 ]; then
         log_message "ERROR: Module post-startup vendor API failed"
     fi
+    return "$rc"
 }
 
 # Function to set state_transition_in_progress flag
@@ -121,9 +125,11 @@ function set_module_state_transition_flag()
     local DPU_NAME=$1
     local FLAG_VALUE=$2
     python3 -c "from utilities_common.module import ModuleHelper; helper = ModuleHelper(); helper.set_module_state_transition('${DPU_NAME}', ${FLAG_VALUE})"
-    if [ $? -ne 0 ]; then
+    local rc=$?
+    if [ "$rc" -ne 0 ]; then
         log_message "ERROR: Setting module state transition flag failed"
     fi
+    return "$rc"
 }
 
 # Function to clear state_transition_in_progress flag
@@ -131,9 +137,11 @@ function clear_module_state_transition_flag()
 {
     local DPU_NAME=$1
     python3 -c "from utilities_common.module import ModuleHelper; helper = ModuleHelper(); helper.clear_module_state_transition('${DPU_NAME}')"
-    if [ $? -ne 0 ]; then
+    local rc=$?
+    if [ "$rc" -ne 0 ]; then
         log_message "ERROR: Clearing module state transition flag failed"
     fi
+    return "$rc"
 }
 
 # Function to get state_transition_in_progress flag
@@ -155,7 +163,9 @@ 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}')"
+    local rc=$?
     log_message "INFO: Rebooting ${DPU_NAME} with reboot_type:${REBOOT_TYPE}..."
+    return "$rc"
 }
 
 # Function to wait for DPU reboot status
@@ -253,23 +263,26 @@ function reboot_dpu()
     log_message "User requested rebooting device ${DPU_NAME} ..."
 
     # Check if the DPU operation status is online before rebooting
-    local oper_status
-    oper_status=$(show chassis modules status "${DPU_NAME^^}" | sed -n '/^ *DPU/ s/.*\s\+\(Online\|Offline\)\s\+.*/\1/p')
-    if [ $? -ne 0 ]; then
+    local module_status
+    if ! module_status=$(show chassis modules status "${DPU_NAME^^}"); then
         log_message "ERROR: Failed to retrieve DPU status."
-    else
-        log_message "INFO: DPU ${DPU_NAME} is in '$oper_status' state before reboot."
-        oper_status=$(echo "$oper_status" | tr '[:upper:]' '[:lower:]')
-        if [ "$oper_status" != "online" ]; then
-            log_message "INFO: ${DPU_NAME} is not online. Current status: $oper_status"
-            return ${EXIT_DPU_DOWN}
-        fi
+        return ${EXIT_ERROR}
+    fi
+    local oper_status
+    oper_status=$(sed -n '/^ *DPU/ s/.*\s\+\(Online\|Offline\)\s\+.*/\1/p' <<< "$module_status")
+    log_message "INFO: DPU ${DPU_NAME} is in '$oper_status' state before reboot."
+    oper_status=$(echo "$oper_status" | tr '[:upper:]' '[:lower:]')
+    if [ "$oper_status" != "online" ]; then
+        log_message "INFO: ${DPU_NAME} is not online. Current status: $oper_status"
+        return ${EXIT_DPU_DOWN}
     fi
 
     if [[ "$REBOOT_TYPE" != $MODULE_REBOOT_SMARTSWITCH ]]; then
         # get and set the state_transition_in_progress flag before reboot
         if ! get_module_state_transition_flag "${DPU_NAME}"; then
-            set_module_state_transition_flag "${DPU_NAME}" True
+            if ! set_module_state_transition_flag "${DPU_NAME}" True; then
+                return ${EXIT_ERROR}
+            fi
         else
             log_message "ERROR: state_transition_in_progress flag is already set for ${DPU_NAME}"
             return ${EXIT_ERROR}
@@ -286,6 +299,9 @@ function reboot_dpu()
     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_NAME}"
+        if [[ "$REBOOT_TYPE" != $MODULE_REBOOT_SMARTSWITCH ]]; then
+            clear_module_state_transition_flag "${DPU_NAME}"
+        fi
         return ${EXIT_ERROR}
     fi
 
diff --git a/tests/reboot_smartswitch_helper_test.py b/tests/reboot_smartswitch_helper_test.py
index 3e5ad0e8f5b..5bae83b6781 100644
--- a/tests/reboot_smartswitch_helper_test.py
+++ b/tests/reboot_smartswitch_helper_test.py
@@ -186,3 +186,57 @@ def test_whole_smartswitch_reboot_continues_after_dpu_failure():
     result = subprocess.run(["bash", "-c", script], capture_output=True, text=True)
     assert result.returncode == 0
     assert "continuing SmartSwitch reboot" in result.stderr
+
+
+def test_reboot_dpu_propagates_platform_reboot_failure(tmp_path):
+    platform_json = tmp_path / "platform.json"
+    platform_json.write_text('{"DPUS":{"dpu0":{"bus_info":"0000:00:00.0"}}}')
+    script = f'''
+EXIT_SUCCESS=0
+EXIT_ERROR=1
+PLATFORM_JSON_PATH="{platform_json}"
+source "{SCRIPT}"
+show() {{ printf '  DPU0 test Online up\n'; }}
+get_module_state_transition_flag() {{ return 1; }}
+set_module_state_transition_flag() {{ return 0; }}
+clear_module_state_transition_flag() {{ return 0; }}
+gnmi_reboot_dpu() {{ return 0; }}
+module_pre_shutdown() {{ return 0; }}
+reboot_dpu_platform() {{ return 1; }}
+reboot_dpu dpu0 DPU
+'''
+    result = subprocess.run(["bash", "-c", script], capture_output=True, text=True)
+    assert result.returncode != 0
+
+
+def test_reboot_dpu_clears_transition_flag_when_bus_info_missing(tmp_path):
+    platform_json = tmp_path / "platform.json"
+    platform_json.write_text('{"DPUS":{"dpu0":{}}}')
+    marker = tmp_path / "cleared"
+    script = f'''
+EXIT_SUCCESS=0
+EXIT_ERROR=1
+PLATFORM_JSON_PATH="{platform_json}"
+source "{SCRIPT}"
+show() {{ printf '  DPU0 test Online up\n'; }}
+get_module_state_transition_flag() {{ return 1; }}
+set_module_state_transition_flag() {{ return 0; }}
+clear_module_state_transition_flag() {{ touch "{marker}"; return 0; }}
+gnmi_reboot_dpu() {{ return 0; }}
+reboot_dpu dpu0 DPU
+'''
+    result = subprocess.run(["bash", "-c", script], capture_output=True, text=True)
+    assert result.returncode != 0
+    assert marker.exists()
+
+
+def test_reboot_dpu_fails_when_chassis_status_fails(tmp_path):
+    script = f'''
+EXIT_SUCCESS=0
+EXIT_ERROR=1
+source "{SCRIPT}"
+show() {{ return 1; }}
+reboot_dpu dpu0 DPU
+'''
+    result = subprocess.run(["bash", "-c", script], capture_output=True, text=True)
+    assert result.returncode != 0

From 7186dc9f999c95a8c7541f52c1e6b4f4fc09ae77 Mon Sep 17 00:00:00 2001
From: Dawei Huang <daweihuang@microsoft.com>
Date: Sat, 8 Aug 2026 05:07:08 -0500
Subject: [PATCH 10/16] reboot: honor DPU platform API results

Signed-off-by: Dawei Huang <daweihuang@microsoft.com>
---
 scripts/reboot_smartswitch_helper       | 14 ++++++++------
 tests/reboot_smartswitch_helper_test.py | 22 ++++++++++++++++++++++
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/scripts/reboot_smartswitch_helper b/scripts/reboot_smartswitch_helper
index a1551d4a40a..369095b1d4e 100644
--- a/scripts/reboot_smartswitch_helper
+++ b/scripts/reboot_smartswitch_helper
@@ -99,7 +99,7 @@ function get_reboot_status()
 function module_pre_shutdown()
 {
     local DPU_NAME=$1
-    python3 -c "from utilities_common.module import ModuleHelper; helper = ModuleHelper(); helper.module_pre_shutdown('${DPU_NAME}')"
+    python3 -c "from utilities_common.module import ModuleHelper; raise SystemExit(0 if ModuleHelper().module_pre_shutdown('${DPU_NAME}') else 1)"
     local rc=$?
     if [ "$rc" -ne 0 ]; then
         log_message "ERROR: Module pre-shutdown vendor API failed"
@@ -111,7 +111,7 @@ function module_pre_shutdown()
 function module_post_startup()
 {
     local DPU_NAME=$1
-    python3 -c "from utilities_common.module import ModuleHelper; helper = ModuleHelper(); helper.module_post_startup('${DPU_NAME}')"
+    python3 -c "from utilities_common.module import ModuleHelper; raise SystemExit(0 if ModuleHelper().module_post_startup('${DPU_NAME}') else 1)"
     local rc=$?
     if [ "$rc" -ne 0 ]; then
         log_message "ERROR: Module post-startup vendor API failed"
@@ -124,7 +124,7 @@ function set_module_state_transition_flag()
 {
     local DPU_NAME=$1
     local FLAG_VALUE=$2
-    python3 -c "from utilities_common.module import ModuleHelper; helper = ModuleHelper(); helper.set_module_state_transition('${DPU_NAME}', ${FLAG_VALUE})"
+    python3 -c "from utilities_common.module import ModuleHelper; raise SystemExit(0 if ModuleHelper().set_module_state_transition('${DPU_NAME}', ${FLAG_VALUE}) else 1)"
     local rc=$?
     if [ "$rc" -ne 0 ]; then
         log_message "ERROR: Setting module state transition flag failed"
@@ -136,7 +136,7 @@ function set_module_state_transition_flag()
 function clear_module_state_transition_flag()
 {
     local DPU_NAME=$1
-    python3 -c "from utilities_common.module import ModuleHelper; helper = ModuleHelper(); helper.clear_module_state_transition('${DPU_NAME}')"
+    python3 -c "from utilities_common.module import ModuleHelper; raise SystemExit(0 if ModuleHelper().clear_module_state_transition('${DPU_NAME}') else 1)"
     local rc=$?
     if [ "$rc" -ne 0 ]; then
         log_message "ERROR: Clearing module state transition flag failed"
@@ -162,7 +162,7 @@ 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}')"
+    python3 -c "from utilities_common.module import ModuleHelper; raise SystemExit(0 if ModuleHelper().reboot_module('${DPU_NAME}', '${REBOOT_TYPE}') else 1)"
     local rc=$?
     log_message "INFO: Rebooting ${DPU_NAME} with reboot_type:${REBOOT_TYPE}..."
     return "$rc"
@@ -330,7 +330,9 @@ function reboot_dpu()
             clear_module_state_transition_flag "${DPU_NAME}"
             return ${EXIT_ERROR}
         fi
-        clear_module_state_transition_flag "${DPU_NAME}"
+        if ! clear_module_state_transition_flag "${DPU_NAME}"; then
+            return ${EXIT_ERROR}
+        fi
     fi
     return "$gnoi_reboot_rc"
 }
diff --git a/tests/reboot_smartswitch_helper_test.py b/tests/reboot_smartswitch_helper_test.py
index 5bae83b6781..3f16a0ac0cc 100644
--- a/tests/reboot_smartswitch_helper_test.py
+++ b/tests/reboot_smartswitch_helper_test.py
@@ -230,6 +230,28 @@ def test_reboot_dpu_clears_transition_flag_when_bus_info_missing(tmp_path):
     assert marker.exists()
 
 
+def test_reboot_dpu_reports_final_transition_cleanup_failure(tmp_path):
+    platform_json = tmp_path / "platform.json"
+    platform_json.write_text('{"DPUS":{"dpu0":{"bus_info":"0000:00:00.0"}}}')
+    script = f'''
+EXIT_SUCCESS=0
+EXIT_ERROR=1
+PLATFORM_JSON_PATH="{platform_json}"
+source "{SCRIPT}"
+show() {{ printf '  DPU0 test Online up\n'; }}
+get_module_state_transition_flag() {{ return 1; }}
+set_module_state_transition_flag() {{ return 0; }}
+clear_module_state_transition_flag() {{ return 1; }}
+gnmi_reboot_dpu() {{ return 0; }}
+module_pre_shutdown() {{ return 0; }}
+module_post_startup() {{ return 0; }}
+reboot_dpu_platform() {{ return 0; }}
+reboot_dpu dpu0 DPU
+'''
+    result = subprocess.run(["bash", "-c", script], capture_output=True, text=True)
+    assert result.returncode != 0
+
+
 def test_reboot_dpu_fails_when_chassis_status_fails(tmp_path):
     script = f'''
 EXIT_SUCCESS=0

From 368c90e2ad318e60c642cd0eb9b5511062d13162 Mon Sep 17 00:00:00 2001
From: Dawei Huang <daweihuang@microsoft.com>
Date: Sat, 8 Aug 2026 05:26:40 -0500
Subject: [PATCH 11/16] reboot: restore DPU after reset failure

Signed-off-by: Dawei Huang <daweihuang@microsoft.com>
---
 scripts/reboot_smartswitch_helper       | 1 +
 tests/reboot_smartswitch_helper_test.py | 3 ++-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/scripts/reboot_smartswitch_helper b/scripts/reboot_smartswitch_helper
index 369095b1d4e..66f3c2913fb 100644
--- a/scripts/reboot_smartswitch_helper
+++ b/scripts/reboot_smartswitch_helper
@@ -318,6 +318,7 @@ function reboot_dpu()
     if [ $? -ne 0 ]; then
         log_message "ERROR: Failed to reboot the module for ${DPU_NAME}"
         if [[ "$REBOOT_TYPE" != $MODULE_REBOOT_SMARTSWITCH ]]; then
+            module_post_startup "${DPU_NAME}"
             clear_module_state_transition_flag "${DPU_NAME}"
         fi
         return ${EXIT_ERROR}
diff --git a/tests/reboot_smartswitch_helper_test.py b/tests/reboot_smartswitch_helper_test.py
index 3f16a0ac0cc..59ded6f33a2 100644
--- a/tests/reboot_smartswitch_helper_test.py
+++ b/tests/reboot_smartswitch_helper_test.py
@@ -244,12 +244,13 @@ def test_reboot_dpu_reports_final_transition_cleanup_failure(tmp_path):
 clear_module_state_transition_flag() {{ return 1; }}
 gnmi_reboot_dpu() {{ return 0; }}
 module_pre_shutdown() {{ return 0; }}
-module_post_startup() {{ return 0; }}
+    module_post_startup() {{ touch "{tmp_path / 'post-startup'}"; return 0; }}
 reboot_dpu_platform() {{ return 0; }}
 reboot_dpu dpu0 DPU
 '''
     result = subprocess.run(["bash", "-c", script], capture_output=True, text=True)
     assert result.returncode != 0
+    assert (tmp_path / "post-startup").exists()
 
 
 def test_reboot_dpu_fails_when_chassis_status_fails(tmp_path):

From 56d9bc760e294021a769303619d9005aebfec6a9 Mon Sep 17 00:00:00 2001
From: Dawei Huang <daweihuang@microsoft.com>
Date: Sat, 8 Aug 2026 05:38:55 -0500
Subject: [PATCH 12/16] reboot: restore DPU after detach failure

Signed-off-by: Dawei Huang <daweihuang@microsoft.com>
---
 scripts/reboot_smartswitch_helper       |  1 +
 tests/reboot_smartswitch_helper_test.py | 23 +++++++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/scripts/reboot_smartswitch_helper b/scripts/reboot_smartswitch_helper
index 66f3c2913fb..3bbfde2f987 100644
--- a/scripts/reboot_smartswitch_helper
+++ b/scripts/reboot_smartswitch_helper
@@ -309,6 +309,7 @@ function reboot_dpu()
     if [ $? -ne 0 ]; then
         log_message "ERROR: Failed to pre-shutdown the module for ${DPU_NAME}"
         if [[ "$REBOOT_TYPE" != $MODULE_REBOOT_SMARTSWITCH ]]; then
+            module_post_startup "${DPU_NAME}"
             clear_module_state_transition_flag "${DPU_NAME}"
         fi
         return ${EXIT_ERROR}
diff --git a/tests/reboot_smartswitch_helper_test.py b/tests/reboot_smartswitch_helper_test.py
index 59ded6f33a2..59facb4d893 100644
--- a/tests/reboot_smartswitch_helper_test.py
+++ b/tests/reboot_smartswitch_helper_test.py
@@ -253,6 +253,29 @@ def test_reboot_dpu_reports_final_transition_cleanup_failure(tmp_path):
     assert (tmp_path / "post-startup").exists()
 
 
+def test_reboot_dpu_restores_after_partial_pre_shutdown_failure(tmp_path):
+    platform_json = tmp_path / "platform.json"
+    platform_json.write_text('{"DPUS":{"dpu0":{"bus_info":"0000:00:00.0"}}}')
+    marker = tmp_path / "post-startup"
+    script = f'''
+EXIT_SUCCESS=0
+EXIT_ERROR=1
+PLATFORM_JSON_PATH="{platform_json}"
+source "{SCRIPT}"
+show() {{ printf '  DPU0 test Online up\n'; }}
+get_module_state_transition_flag() {{ return 1; }}
+set_module_state_transition_flag() {{ return 0; }}
+clear_module_state_transition_flag() {{ return 0; }}
+gnmi_reboot_dpu() {{ return 0; }}
+module_pre_shutdown() {{ return 1; }}
+module_post_startup() {{ touch "{marker}"; return 0; }}
+reboot_dpu dpu0 DPU
+'''
+    result = subprocess.run(["bash", "-c", script], capture_output=True, text=True)
+    assert result.returncode != 0
+    assert marker.exists()
+
+
 def test_reboot_dpu_fails_when_chassis_status_fails(tmp_path):
     script = f'''
 EXIT_SUCCESS=0

From 1ee145c3c91aaff296d8d7961068e708bb433f46 Mon Sep 17 00:00:00 2001
From: Dawei Huang <daweihuang@microsoft.com>
Date: Sat, 8 Aug 2026 22:31:09 -0500
Subject: [PATCH 13/16] dump: support modern protobuf descriptors

Signed-off-by: Dawei Huang <daweihuang@microsoft.com>
---
 dump/dash_util.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/dump/dash_util.py b/dump/dash_util.py
index 4e1c7bae61c..40be9690cde 100644
--- a/dump/dash_util.py
+++ b/dump/dash_util.py
@@ -79,7 +79,11 @@ def find_index(proto_obj, proto_dict=pb2_dict):
             field_type = field_descriptor.type
             if field_type == field_descriptor.TYPE_MESSAGE:
                 obj = getattr(proto_obj, field_name)
-                if field_descriptor.label == field_descriptor.LABEL_REPEATED:
+                if hasattr(field_descriptor, 'is_repeated'):
+                    is_repeated = field_descriptor.is_repeated
+                else:
+                    is_repeated = field_descriptor.label == field_descriptor.LABEL_REPEATED
+                if is_repeated:
                     process_rep_field(obj, proto_dict, field_name)
                 else:
                     process_msg_field(obj, proto_dict, field_name)

From 90d0fef0b2f1069ebfbd80f980e1bd1662af4bb7 Mon Sep 17 00:00:00 2001
From: Dawei Huang <daweihuang@microsoft.com>
Date: Sat, 8 Aug 2026 22:31:09 -0500
Subject: [PATCH 14/16] test: request all ports in multi-ASIC pfcstat

Signed-off-by: Dawei Huang <daweihuang@microsoft.com>
---
 tests/pfcstat_test.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tests/pfcstat_test.py b/tests/pfcstat_test.py
index 1e8d49732e4..ef41365e0c0 100644
--- a/tests/pfcstat_test.py
+++ b/tests/pfcstat_test.py
@@ -277,6 +277,7 @@ def test_pfc_counters_all(self):
         runner = CliRunner()
         result = runner.invoke(
             show.cli.commands["pfc"].commands["counters"],
+            ["--display", "all"]
         )
         print(result.output)
         assert result.exit_code == 0
@@ -288,7 +289,7 @@ def test_pfc_counters_all_with_clear(self):
         assert result.exit_code == 0
         result = runner.invoke(
             show.cli.commands["pfc"].commands["counters"],
-            []
+            ["--display", "all"]
         )
         print(result.output)
         show.run_command(['pfcstat', '-d'])
@@ -326,7 +327,7 @@ def test_pfc_counters_history_all(self):
         runner = CliRunner()
         result = runner.invoke(
             show.cli.commands["pfc"].commands["counters"],
-            ["--history"]
+            ["--display", "all", "--history"]
         )
         print(result.output)
         assert result.exit_code == 0
@@ -339,7 +340,7 @@ def test_pfc_counters_history_all_with_clear(self):
         assert result.exit_code == 0
         result = runner.invoke(
             show.cli.commands["pfc"].commands["counters"],
-            ["--history"]
+            ["--display", "all", "--history"]
         )
         print(result.output)
         show.run_command(['pfcstat', '-d'])

From a3e84d75ff74ced62c4f547508d026f76035ae43 Mon Sep 17 00:00:00 2001
From: Dawei Huang <daweihuang@microsoft.com>
Date: Sat, 8 Aug 2026 23:00:14 -0500
Subject: [PATCH 15/16] test: cover legacy protobuf descriptors

Signed-off-by: Dawei Huang <daweihuang@microsoft.com>
---
 tests/dump_tests/dash_util_test.py | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)
 create mode 100644 tests/dump_tests/dash_util_test.py

diff --git a/tests/dump_tests/dash_util_test.py b/tests/dump_tests/dash_util_test.py
new file mode 100644
index 00000000000..2b6a58122c1
--- /dev/null
+++ b/tests/dump_tests/dash_util_test.py
@@ -0,0 +1,26 @@
+import pytest
+
+pytest.importorskip("dash_api.types_pb2")
+
+from dump.dash_util import find_known_types_sec
+
+
+class LegacyRepeatedField:
+    name = "children"
+    type = 11
+    TYPE_MESSAGE = 11
+    label = 3
+    LABEL_REPEATED = 3
+
+
+class LegacyProto:
+    children = []
+
+    def ListFields(self):
+        return [(LegacyRepeatedField(), self.children)]
+
+
+def test_find_known_types_supports_legacy_repeated_descriptor():
+    proto_dict = {"children": []}
+
+    assert find_known_types_sec(LegacyProto(), proto_dict) == proto_dict

From dc9817955a65ed6a7a3b0b366003710b8a4e70fb Mon Sep 17 00:00:00 2001
From: Dawei Huang <daweihuang@microsoft.com>
Date: Sat, 8 Aug 2026 23:22:49 -0500
Subject: [PATCH 16/16] test: satisfy protobuf test lint

Signed-off-by: Dawei Huang <daweihuang@microsoft.com>
---
 tests/dump_tests/dash_util_test.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tests/dump_tests/dash_util_test.py b/tests/dump_tests/dash_util_test.py
index 2b6a58122c1..a7217475ac1 100644
--- a/tests/dump_tests/dash_util_test.py
+++ b/tests/dump_tests/dash_util_test.py
@@ -1,8 +1,10 @@
+import importlib
+
 import pytest
 
 pytest.importorskip("dash_api.types_pb2")
 
-from dump.dash_util import find_known_types_sec
+dash_util = importlib.import_module("dump.dash_util")
 
 
 class LegacyRepeatedField:
@@ -23,4 +25,6 @@ def ListFields(self):
 def test_find_known_types_supports_legacy_repeated_descriptor():
     proto_dict = {"children": []}
 
-    assert find_known_types_sec(LegacyProto(), proto_dict) == proto_dict
+    result = dash_util.find_known_types_sec(LegacyProto(), proto_dict)
+
+    assert result == proto_dict