[None][fix] repair KeyboardInterrupt drain loop in iteration logging

The handler read 'message' which is unbound when the interrupt lands
before the first recv_json(), raising UnboundLocalError. It also compared
the decoded JSON object against b'None', a sentinel the sender never
emits (LlmManager sends {'end': True}), so the drain loop could never
terminate normally. Initialize the variable, guard the socket, and drain
until the actual 'end' sentinel.

Signed-off-by: lonexreb <reach2shubhankar@gmail.com>

Author: lonexreb

tensorrt_llm/bench/benchmark/utils/processes.py

@@ -133,6 +133,7 @@ def run(address: str, log_path: Path, stop_event: MpEvent) -> None:
         """
         context = None
         socket = None
+        message = None
 
         try:
             # Create a ZeroMQ context and socket for inter-process communication
@@ -156,12 +157,13 @@ def run(address: str, log_path: Path, stop_event: MpEvent) -> None:
                     message = socket.recv_json()
                 logger.debug(f"Iteration logging: Received end message")
         except KeyboardInterrupt:
-            # Handle keyboard interrupt by continuing to receive
-            # messages until "None" is received. LlmManager will
-            # send "None" when it is finished.
+            # Handle keyboard interrupt by continuing to receive messages until
+            # the "end" sentinel is received. LlmManager sends {"end": True}
+            # when it is finished.
             logger.info("Keyboard interrupt, exiting iteration logging...")
-            while message != b"None":
-                message = socket.recv_json()
+            if socket is not None:
+                while message is None or "end" not in message:
+                    message = socket.recv_json()
         finally:
             # Finalize the logging process by closing the socket and terminating
             # the context