jason810496 commented on code in PR #66633:
URL: https://github.com/apache/airflow/pull/66633#discussion_r3443507349


##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -2405,9 +2422,30 @@ def _configure_logging(log_path: str, client: Client) -> 
tuple[FilteringBoundLog
 
     with _remote_logging_conn(client):
         processors = logging_processors(json_output=json_logs)
+
     logger = structlog.wrap_logger(underlying_logger, processors=processors, 
logger_name="task").bind()
 
-    return logger, log_file_descriptor
+    try:
+        yield logger
+    finally:
+        # Flush and close the remote handler now — AFTER the supervisor has
+        # drained all task log messages from the subprocess pipe (i.e. after
+        # process.wait() has returned).
+        #
+        # Without this, the only thing that ever closes the handler is
+        # Python's logging.shutdown() at process exit, which fires after
+        # supervise_task() returns. Any messages still queued in the handler
+        # at that point are silently dropped. For example, the AWS CloudWatch
+        # logger will emit:
+
+        # WatchtowerWarning: "Received message after logging system shutdown"
+        remote_handler = load_remote_log_handler()

Review Comment:
   Please keep the TODO instead of dropping it silently.
   ```suggestion
           # TODO: Use logging providers to handle the chunked upload for us 
etc.
           remote_handler = load_remote_log_handler()
   ```



##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -2385,11 +2382,31 @@ def ensure_secrets_backend_loaded() -> 
list[BaseSecretsBackend]:
     return ensure_secrets_loaded(default_backends=fallback_backends)
 
 
-def _configure_logging(log_path: str, client: Client) -> 
tuple[FilteringBoundLogger, BinaryIO | TextIO]:
+def _close_remote_log_handler(handler: RemoteLogIO) -> None:
+    """
+    Close the remote log handler explicitly after all task log messages have 
been drained from the subprocess pipe.
+
+    Called after process.wait() returns, before process exit triggers
+    logging.shutdown(). This ensures the remote handler's internal batch
+    queue is flushed before the process tears down. For example, the AWS
+    CloudWatch logger will emit:
+
+        WatchtowerWarning: "Received message after logging system shutdown"
+
+    if this is not done before process exit.
+    """
+    try:
+        handler.close()

Review Comment:
   It would be better to have more defensive compatible guard against the 
`handler.close` call.
   Since not all remote logging handler support `close` method when using the 
Task-SDK version including this patch.
   
   ```suggestion
           if hasattr(handler, "close", None):
               handler.close()
   ```



##########
providers/amazon/src/airflow/providers/amazon/aws/log/cloudwatch_task_handler.py:
##########
@@ -30,6 +30,8 @@
 from typing import TYPE_CHECKING, Any
 
 import attrs
+import structlog
+import structlog.typing

Review Comment:
   ```suggestion
   ```



##########
shared/logging/src/airflow_shared/logging/remote.py:
##########
@@ -64,6 +64,10 @@ def read(self, relative_path: str, ti: RuntimeTI) -> 
LogResponse:
         """Read logs from the given remote log path."""
         ...
 
+    def close(self) -> None:
+        """Flush and close the remote log handler."""
+        ...

Review Comment:
   I don't think we need to introduce the `close` method at `RemoteLogIO` 
protocol in this PR.
   Since it's still nop for ES and OS remote IO implementation.



##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -2405,9 +2422,30 @@ def _configure_logging(log_path: str, client: Client) -> 
tuple[FilteringBoundLog
 
     with _remote_logging_conn(client):
         processors = logging_processors(json_output=json_logs)
+
     logger = structlog.wrap_logger(underlying_logger, processors=processors, 
logger_name="task").bind()
 
-    return logger, log_file_descriptor
+    try:
+        yield logger
+    finally:
+        # Flush and close the remote handler now — AFTER the supervisor has
+        # drained all task log messages from the subprocess pipe (i.e. after
+        # process.wait() has returned).
+        #
+        # Without this, the only thing that ever closes the handler is
+        # Python's logging.shutdown() at process exit, which fires after
+        # supervise_task() returns. Any messages still queued in the handler
+        # at that point are silently dropped. For example, the AWS CloudWatch
+        # logger will emit:
+
+        # WatchtowerWarning: "Received message after logging system shutdown"
+        remote_handler = load_remote_log_handler()
+        if remote_handler is not None:

Review Comment:
   nit:
   ```suggestion
           if (remote_handler := load_remote_log_handler()) is not None:
   ```



##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -1774,8 +1773,6 @@ def _handle_request(self, msg: ToSupervisor, log: 
FilteringBoundLogger, req_id:
             dump_opts = {"exclude_unset": True}
         elif isinstance(msg, GetPrevSuccessfulDagRun):
             resp, dump_opts = handle_get_prev_successful_dag_run(self.client, 
self.id)
-        elif isinstance(msg, GetXComCount):
-            resp, dump_opts = handle_get_xcom_count(self.client, msg)

Review Comment:
   Could we revert these non-related changes? (the parameters style of `def 
send_msg(` and the order of `elif isinstance(msg, GetXComCount):` ).



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to