PB-402 Add more logs

xain_fl/coordinator/coordinator.py

@@ -7,6 +7,8 @@
 from numpy import ndarray
 from structlog import get_logger
 from xain_proto.fl.coordinator_pb2 import (
+    _RENDEZVOUSREPLY,
+    _STATE,
     EndTrainingRoundRequest,
     EndTrainingRoundResponse,
     HeartbeatRequest,
@@ -178,7 +180,8 @@ def on_message(
 
         logger.debug(
             "Received message from participant",
-            message_type=type(message),
+            message_type=message.DESCRIPTOR.name,
+            message_byte_size=message.ByteSize(),
             participant_id=participant_id,
         )
 
@@ -273,6 +276,10 @@ def _handle_rendezvous(
                 current_participants_count=self.participants.len(),
             )
 
+        logger.debug(
+            "Send RendezvousResponse",
+            reply=_RENDEZVOUSREPLY.values_by_number[reply].name,
+        )
         return RendezvousResponse(reply=reply)
 
     def _handle_heartbeat(
@@ -304,8 +311,9 @@ def _handle_heartbeat(
         logger.debug(
             "Heartbeat response",
             participant_id=participant_id,
-            message=state,
+            state=_STATE.values_by_number[state].name,
             round=self.current_round,
+            current_participants_count=self.participants.len(),
         )
         return HeartbeatResponse(state=state, round=self.current_round)
 
@@ -332,6 +340,11 @@ def _handle_start_training_round(
                 "StartTrainingRoundRequest outside of a round"
             )
 
+        logger.debug(
+            "Send StartTrainingRoundResponse",
+            epochs=self.epochs,
+            epoch_base=self.epoch_base,
+        )
         return StartTrainingRoundResponse(
             weights=ndarray_to_proto(self.weights),
             epochs=self.epochs,
@@ -402,4 +415,5 @@ def _handle_end_training_round(
                 # reinitialize the round
                 self.select_participant_ids_and_init_round()
 
+        logger.debug("Send EndTrainingRoundResponse", participant_id=participant_id)
         return EndTrainingRoundResponse()