Add logs to investigate slow stream pod termination (#602)

[ML-11919](https://iguazio.atlassian.net/browse/ML-11919)

Author: gtopper

storey/flow.py

@@ -294,6 +294,9 @@ async def _do_downstream(self, event, outlets=None, select_outlets: bool = True)
         if not outlets:
             return
         if event is _termination_obj:
+            if self.logger:
+                outlet_names = ", ".join([outlet.name for outlet in outlets])
+                self.logger.info(f"Forwarding termination signal from step '{self.name}' to steps: {outlet_names}")
             # Only propagate the termination object once we received one per inlet
             outlets[0]._termination_received += 1
             if outlets[0]._should_terminate():
@@ -937,7 +940,15 @@ async def _worker(self):
                     # be 1 higher than requested. Hence, we peek.
                     job = await self._q.peek()
                     if job is _termination_obj:
+                        if self.logger:
+                            self.logger.info(
+                                f"Terminating ConcurrentJobExecution worker belonging to step '{self.name}'"
+                            )
                         await self._q.get()
+                        if self.logger:
+                            self.logger.info(
+                                f"Terminated ConcurrentJobExecution worker belonging to step '{self.name}'"
+                            )
                         break
                     event = job[0]
                     completed = await job[1]
@@ -1016,10 +1027,22 @@ async def _do(self, event):
 
         if event is _termination_obj:
             if self._queue_size > 0:
+                if self.logger:
+                    self.logger.info(
+                        f"Sending termination signal to ConcurrentJobExecution worker belonging to step '{self.name}'"
+                    )
                 await self._q.put(_termination_obj)
+                if self.logger:
+                    self.logger.info(
+                        f"Awaiting termination of ConcurrentJobExecution worker belonging to step '{self.name}'"
+                    )
                 await self._worker_awaitable
             else:
+                if self.logger:
+                    self.logger.info(f"Terminating ConcurrentJobExecution step '{self.name}' without a worker")
                 await self._cleanup()
+                if self.logger:
+                    self.logger.info(f"Terminated ConcurrentJobExecution step '{self.name}' without a worker")
             return await self._do_downstream(_termination_obj)
         else:
             coroutine = self._process_event_with_retries(event)
@@ -1219,8 +1242,14 @@ async def _terminate(self):
 
     async def _do(self, event):
         if event is _termination_obj:
+            if self.logger:
+                self.logger.info(f"Terminating Batching step '{self.name}': emitting all remaining batches")
             await self._emit_all()
+            if self.logger:
+                self.logger.info(f"Terminating Batching step '{self.name}': running custom termination code")
             await self._terminate()
+            if self.logger:
+                self.logger.info(f"Terminated Batching step '{self.name}'")
             return await self._do_downstream(_termination_obj)
 
         key = self._extract_key(event)

storey/targets.py

@@ -1198,12 +1198,29 @@ async def _worker(self):
                         continue
                     event = await self._q.get()
                     if event is _termination_obj:  # handle outstanding batches and in flight requests on termination
+                        if self.logger:
+                            self.logger.info(
+                                f"Terminating StreamTarget worker belonging to step '{self.name}': "
+                                f"awaiting {len(in_flight_reqs)} in-flight batches"
+                            )
                         for req in in_flight_reqs:
                             await self._handle_response(req)
+                        if self.logger:
+                            self.logger.info(
+                                f"Terminating StreamTarget worker belonging to step '{self.name}': "
+                                f"sending {self._shards} final batches"
+                            )
                         for shard_id in range(self._shards):
                             self._send_batch(buffers, in_flight_reqs, buffer_events, in_flight_events, shard_id)
+                        if self.logger:
+                            self.logger.info(
+                                f"Terminating StreamTarget worker belonging to step '{self.name}': "
+                                f"awaiting {len(in_flight_reqs)} final batches"
+                            )
                         for req in in_flight_reqs:
                             await self._handle_response(req)
+                        if self.logger:
+                            self.logger.info(f"Terminated StreamTarget worker belonging to step '{self.name}'")
                         break
                     sharding_func_result = self._sharding_func(event)
                     if isinstance(sharding_func_result, int):
@@ -1281,8 +1298,12 @@ async def _do(self, event):
             raise AssertionError("StreamTarget worker has already terminated")
 
         if event is _termination_obj:
+            if self.logger:
+                self.logger.info(f"Terminating StreamTarget step '{self.name}'")
             await self._q.put(_termination_obj)
             await self._worker_awaitable
+            if self.logger:
+                self.logger.info(f"Terminated StreamTarget step '{self.name}'")
             return await self._do_downstream(_termination_obj)
         else:
             await self._q.put(event)
@@ -1365,8 +1386,14 @@ async def _do(self, event):
         await self._lazy_init()
 
         if event is _termination_obj:
+            if self.logger:
+                self.logger.info(f"Terminating KafkaTarget step '{self.name}': flushing producer")
             self._producer.flush()
+            if self.logger:
+                self.logger.info(f"Terminating KafkaTarget step '{self.name}': closing producer")
             self._producer.close()
+            if self.logger:
+                self.logger.info(f"Terminated KafkaTarget step '{self.name}': closing producer")
             return await self._do_downstream(_termination_obj)
         else:
             key = event.key
@@ -1442,7 +1469,11 @@ async def _handle_completed(self, event, response):
 
     async def _do(self, event):
         if event is _termination_obj:
+            if self.logger:
+                self.logger.info(f"Terminating NoSqlTarget step '{self.name}': terminating table")
             await self._table._terminate()
+            if self.logger:
+                self.logger.info(f"Terminated NoSqlTarget step '{self.name}'")
             return await self._do_downstream(_termination_obj)
 
         if event.key is None:

tests/test_flow.py

@@ -4678,14 +4678,16 @@ def test_verbose_logs():
     controller.terminate()
     controller.await_termination()
 
-    assert len(logger.logs) == 2
+    debug_logs = [log for log in logger.logs if log[0] == "debug"]
 
-    level, args, kwargs = logger.logs[0]
+    assert len(debug_logs) == 2
+
+    level, args, kwargs = debug_logs[0]
     assert level == "debug"
     assert args == ("SyncEmitSource -> Map1 | Event(id=myid, path=/, body={})",)
     assert kwargs == {}
 
-    level, args, kwargs = logger.logs[1]
+    level, args, kwargs = debug_logs[1]
     assert level == "debug"
     assert args == ("Map1 -> Map2 | Event(id=myid, path=/, body={})",)
     assert kwargs == {}