[Feat] Enhance upgrade image script to support chassis devices

ansible/devutil/devices/ansible_hosts.py

@@ -741,6 +741,7 @@ def _run_ansible_module(self, *args, **kwargs):
                 }
         """
         caller_info = kwargs.pop("caller_info", None)
+        target_hosts = kwargs.pop("target_hosts", None)
         if not caller_info:
             previous_frame = inspect.currentframe().f_back
             caller_info = inspect.getframeinfo(previous_frame)
@@ -763,7 +764,10 @@ def _run_ansible_module(self, *args, **kwargs):
         self._log_modules(caller_info, module_info, verbosity)
 
         task = self.build_task(**module_info)
-        results = self.run_tasks(self.host_pattern, self.loader, self.im, self.vm, self.options, tasks=[task])
+        host_pattern = self.host_pattern
+        if target_hosts:
+            host_pattern = target_hosts
+        results = self.run_tasks(host_pattern, self.loader, self.im, self.vm, self.options, tasks=[task])
 
         self._log_results(caller_info, module_info, results, verbosity)
         self._check_results(caller_info, module_info, results, module_ignore_errors, verbosity)

ansible/devutil/devices/chassis_utils.py

@@ -0,0 +1,26 @@
+import enum
+
+
+class ChassisCardType(str, enum.Enum):
+    # Sample: lab-1111-sup-1
+    SUPERVISOR_CARD = "-sup-"
+    # Sample: lab-1111-lc1-1
+    LINE_CARD = "-lc"
+
+
+def is_chassis(sonichosts):
+    supervisor_card_exists, line_card_exists = False, False
+    for hostname in sonichosts.hostnames:
+        if ChassisCardType.SUPERVISOR_CARD.value in hostname:
+            supervisor_card_exists = True
+        if ChassisCardType.LINE_CARD.value in hostname:
+            line_card_exists = True
+    return supervisor_card_exists and line_card_exists
+
+
+def get_chassis_hostnames(sonichosts, chassis_card_type: ChassisCardType):
+    res = []
+    for hostname in sonichosts.hostnames:
+        if chassis_card_type.value in hostname:
+            res.append(hostname)
+    return res

ansible/devutil/devices/sonic.py

@@ -1,8 +1,11 @@
 import logging
+import time
+
 import yaml
 
 from .ansible_hosts import AnsibleHosts
 from .ansible_hosts import RunAnsibleModuleFailed
+from .chassis_utils import is_chassis, get_chassis_hostnames, ChassisCardType
 
 logger = logging.getLogger(__name__)
 
@@ -26,14 +29,40 @@ def sonic_version(self):
             return {}
 
 
-def upgrade_by_sonic(sonichosts, image_url, disk_used_percent):
+def upgrade_by_sonic(sonichosts, localhost, image_url, disk_used_percent):
     try:
         sonichosts.reduce_and_add_sonic_images(
             disk_used_pcent=disk_used_percent,
             new_image_url=image_url,
             module_attrs={"become": True}
         )
-        sonichosts.shell("reboot", module_attrs={"become": True, "async": 300, "poll": 0})
+        if is_chassis(sonichosts):
+            logger.info("Upgrading image on chassis device...")
+            # Chassis DUT need to firstly upgrade and reboot supervisor cards.
+            # Until supervisor cards back online, then upgrade and reboot line cards.
+            rp_hostnames = get_chassis_hostnames(sonichosts, ChassisCardType.SUPERVISOR_CARD)
+            lc_hostnames = get_chassis_hostnames(sonichosts, ChassisCardType.LINE_CARD)
+            sonichosts.shell("reboot", target_hosts=rp_hostnames,
+                             module_attrs={"become": True, "async": 300, "poll": 0})
+            logger.info("Sleep 900s to wait for supervisor card to be ready...")
+            time.sleep(900)
+            for i in range(len(sonichosts.ips)):
+                localhost.wait_for(
+                    host=sonichosts.ips[i],
+                    port=22,
+                    state="started",
+                    search_regex="OpenSSH",
+                    delay=0,
+                    timeout=600,
+                    module_attrs={"changed_when": False}
+                )
+            sonichosts.shell("reboot", target_hosts=lc_hostnames,
+                             module_attrs={"become": True, "async": 300, "poll": 0})
+            logger.info("Sleep 300s to wait for line cards to be ready...")
+            time.sleep(300)
+        else:
+            sonichosts.shell("reboot", module_attrs={"become": True, "async": 300, "poll": 0})
+
         return True
     except RunAnsibleModuleFailed as e:
         logger.error(
@@ -205,7 +234,7 @@ def upgrade_image(sonichosts, localhost, image_url, upgrade_type="sonic", disk_u
         return False
 
     if upgrade_type == "sonic":
-        upgrade_result = upgrade_by_sonic(sonichosts, image_url, disk_used_percent)
+        upgrade_result = upgrade_by_sonic(sonichosts, localhost, image_url, disk_used_percent)
     elif upgrade_type == "onie":
         upgrade_result = upgrade_by_onie(sonichosts, localhost, image_url, onie_pause_time)
     if not upgrade_result: