From ec24a0f66d81cc59668382ebe493e17a2737e7b7 Mon Sep 17 00:00:00 2001
From: Maisam Arif <maisarif@amd.com>
Date: Mon, 16 Oct 2023 07:20:13 -0500
Subject: [PATCH] Enabled events subcommand to non-virtual systems

Signed-off-by: Maisam Arif <maisarif@amd.com>
Change-Id: Ied56ef015bba606b1bca1a1108a237d0c1cc7fdb
---
 amdsmi_cli/amdsmi_commands.py    | 30 ++++++++++++++++++++++++------
 amdsmi_cli/amdsmi_parser.py      |  4 ++--
 py-interface/amdsmi_interface.py | 15 ++++++---------
 src/amd_smi/amd_smi.cc           |  8 ++++----
 4 files changed, 36 insertions(+), 21 deletions(-)

diff --git a/amdsmi_cli/amdsmi_commands.py b/amdsmi_cli/amdsmi_commands.py
index 36f86fecb9..4e137f99ec 100644
--- a/amdsmi_cli/amdsmi_commands.py
+++ b/amdsmi_cli/amdsmi_commands.py
@@ -1526,13 +1526,31 @@ class AMDSMICommands():
         print('Not applicable to linux baremetal')
 
 
-    def event(self, args):
+    def event(self, args, gpu=None):
+        """ Get event information for target gpus
+
+        Args:
+            args (Namespace): argparser args to pass to subcommand
+            gpu (device_handle, optional): device_handle for target device. Defaults to None.
+
+        Return:
+            stdout event information for target gpus
+        """
+        if args.gpu:
+            gpu = args.gpu
+
+        if gpu == None:
+            args.gpu = self.device_handles
+
+        if not isinstance(args.gpu, list):
+            args.gpu = [args.gpu]
+
         print('EVENT LISTENING:\n')
-        print('Press q and hit ENTER when you want to stop (listening will stop inside 10 seconds)')
+        print('Press q and hit ENTER when you want to stop (listening will stop within 10 seconds)')
 
         threads = []
-        for i in range(len(self.device_handles)):
-            x = threading.Thread(target=self._event_thread, args=(self, i))
+        for gpu in range(len(args.gpu)):
+            x = threading.Thread(target=self._event_thread, args=(self, gpu))
             threads.append(x)
             x.start()
 
@@ -2058,8 +2076,8 @@ class AMDSMICommands():
             return
 
         device = devices[i]
-        listener = amdsmi_interface.AmdSmiEventReader(device, amdsmi_interface.AmdSmiEvtNotificationType.GPU_PRE_RESET,
-                            amdsmi_interface.AmdSmiEvtNotificationType.GPU_POST_RESET)
+        listener = amdsmi_interface.AmdSmiEventReader(device,
+                                        amdsmi_interface.AmdSmiEvtNotificationType)
         values_dict = {}
 
         while self.stop!='q':
diff --git a/amdsmi_cli/amdsmi_parser.py b/amdsmi_cli/amdsmi_parser.py
index 27c42f696f..76273578f6 100644
--- a/amdsmi_cli/amdsmi_parser.py
+++ b/amdsmi_cli/amdsmi_parser.py
@@ -564,8 +564,8 @@ class AMDSMIParser(argparse.ArgumentParser):
 
 
     def _add_event_parser(self, subparsers, func):
-        if self.helpers.is_linux() and not self.helpers.is_virtual_os():
-            # This subparser only applies to Linux Hypervisors, NOT Linux Guest
+        if self.helpers.is_virtual_os():
+            # This subparser doesn't only apply to guest systems
             return
 
         # Subparser help text
diff --git a/py-interface/amdsmi_interface.py b/py-interface/amdsmi_interface.py
index f97dd8e3ef..2b832eb7cd 100644
--- a/py-interface/amdsmi_interface.py
+++ b/py-interface/amdsmi_interface.py
@@ -348,7 +348,8 @@ class AmdSmiProcessorType(IntEnum):
 
 class AmdSmiEventReader:
     def __init__(
-        self, processor_handle: amdsmi_wrapper.amdsmi_processor_handle, *event_types
+        self, processor_handle: amdsmi_wrapper.amdsmi_processor_handle,
+        event_types: List[AmdSmiEvtNotificationType]
     ):
         if not isinstance(processor_handle, amdsmi_wrapper.amdsmi_processor_handle):
             raise AmdSmiParameterException(
@@ -375,8 +376,7 @@ class AmdSmiEventReader:
             processor_handle, ctypes.c_uint64(mask)))
 
     def read(self, timestamp, num_elem=10):
-        self.event_info = (
-            amdsmi_wrapper.amdsmi_evt_notification_data_t * num_elem)()
+        self.event_info = (amdsmi_wrapper.amdsmi_evt_notification_data_t * num_elem)()
         _check_res(
             amdsmi_wrapper.amdsmi_get_gpu_event_notification(
                 ctypes.c_int(timestamp),
@@ -387,15 +387,12 @@ class AmdSmiEventReader:
 
         ret = list()
         for i in range(0, num_elem):
-            if self.event_info[i].event in set(
-                event.value for event in AmdSmiEvtNotificationType
-            ):
+            unique_event_values = set(event.value for event in AmdSmiEvtNotificationType)
+            if self.event_info[i].event in unique_event_values:
                 ret.append(
                     {
                         "processor_handle": self.event_info[i].processor_handle,
-                        "event": AmdSmiEvtNotificationType(
-                            self.event_info[i].event
-                        ).name,
+                        "event": AmdSmiEvtNotificationType(self.event_info[i].event).name,
                         "message": self.event_info[i].message.decode("utf-8"),
                     }
                 )
diff --git a/src/amd_smi/amd_smi.cc b/src/amd_smi/amd_smi.cc
index 9253722628..145f4e639b 100644
--- a/src/amd_smi/amd_smi.cc
+++ b/src/amd_smi/amd_smi.cc
@@ -823,14 +823,14 @@ amdsmi_init_gpu_event_notification(amdsmi_processor_handle processor_handle) {
 }
 
 amdsmi_status_t
- amdsmi_set_gpu_event_notification_mask(amdsmi_processor_handle processor_handle,
-            uint64_t mask) {
+amdsmi_set_gpu_event_notification_mask(amdsmi_processor_handle processor_handle,
+          uint64_t mask) {
     return rsmi_wrapper(rsmi_event_notification_mask_set, processor_handle, mask);
 }
 
 amdsmi_status_t
- amdsmi_get_gpu_event_notification(int timeout_ms,
-                     uint32_t *num_elem, amdsmi_evt_notification_data_t *data) {
+amdsmi_get_gpu_event_notification(int timeout_ms,
+                    uint32_t *num_elem, amdsmi_evt_notification_data_t *data) {
     AMDSMI_CHECK_INIT();
 
     if (num_elem == nullptr || data == nullptr) {