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


##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -883,11 +856,7 @@ def _serialize_response(self, msg: BaseModel | 
ErrorResponse, **dump_opts) -> di
         return msg.model_dump(**dump_opts)
 
     def send_msg(
-        self,
-        msg: BaseModel | None,
-        request_id: int,
-        error: ErrorResponse | None = None,
-        **dump_opts,
+        self, msg: BaseModel | None, request_id: int, error: ErrorResponse | 
None = None, **dump_opts
     ):

Review Comment:
   Would prefer not change this part.



##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -2437,14 +2423,36 @@ 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"
+
+        # TODO: Use logging providers to handle the chunked upload for us etc.
+        if (remote_handler := load_remote_log_handler()) is not None:
+            _close_remote_log_handler(remote_handler)
+
+        if log_file_descriptor is not None:
+            with contextlib.suppress(Exception):
+                log_file_descriptor.close()

Review Comment:
   Should we switch the order of closing `log_file_descriptor` and the `remote 
log handler`? IIUC, closing the file first then flush the renaming logs to 
remote IO makes more sense to me.



##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -2437,14 +2423,36 @@ 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"
+
+        # TODO: Use logging providers to handle the chunked upload for us etc.
+        if (remote_handler := load_remote_log_handler()) is not None:
+            _close_remote_log_handler(remote_handler)

Review Comment:
   The scope and the purpose of the function is small and clear enough, we 
don't need the additional function for this. I'd like to collapse it back.
   ```suggestion
               try:
                   if hasattr(handler, "close"):
                       handler.close()
               except Exception:
                   log.warning("Failed to close remote log handler", 
exc_info=True)
   ```



##########
task-sdk/src/airflow/sdk/log.py:
##########
@@ -133,6 +139,15 @@ def configure_logging(
         callsite_parameters=callsite_params,
     )
 
+    # dictConfig has now run — safe to build the watchtower handler.
+    # We append remote processors into the already-configured structlog chain,
+    # inserting before the final renderer (last processor in the chain).
+    if (remote := load_remote_log_handler()) and (remote_processors := 
getattr(remote, "processors")):
+        current_processors = list(structlog.get_config()["processors"])
+        # Insert before the final renderer
+        updated_processors = current_processors[:-1] + list(remote_processors) 
+ [current_processors[-1]]
+        structlog.configure(processors=updated_processors)

Review Comment:
   Unlike the old `extra_processors` path, stdlib-routed records 
(`ProcessorFormatter`'s `foreign_pre_chain`) no longer pass through the remote 
processors. Neutral today — `ProcessorFormatter` invokes processors with 
`logger=None`, so the CloudWatch proc skipped them anyway — but worth 
documenting so the asymmetry isn't "fixed" later:
   
   ```suggestion
           # NOTE: unlike the old extra_processors path, stdlib-routed records 
(ProcessorFormatter's
           # foreign_pre_chain) intentionally do NOT pass through the remote 
processors.
           structlog.configure(processors=updated_processors)
   ```
   



##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -2417,11 +2382,32 @@ 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:
+        if hasattr(handler, "close"):
+            handler.close()
+    except Exception:
+        log.warning("Failed to close remote log handler", exc_info=True)

Review Comment:
   ```suggestion
   ```



##########
task-sdk/src/airflow/sdk/log.py:
##########
@@ -133,6 +139,15 @@ def configure_logging(
         callsite_parameters=callsite_params,
     )
 
+    # dictConfig has now run — safe to build the watchtower handler.
+    # We append remote processors into the already-configured structlog chain,
+    # inserting before the final renderer (last processor in the chain).
+    if (remote := load_remote_log_handler()) and (remote_processors := 
getattr(remote, "processors")):
+        current_processors = list(structlog.get_config()["processors"])
+        # Insert before the final renderer
+        updated_processors = current_processors[:-1] + list(remote_processors) 
+ [current_processors[-1]]
+        structlog.configure(processors=updated_processors)

Review Comment:
   Two hardening points, both cheap while this is being rewritten:
   
   1. This also runs in the task subprocess 
(`_configure_logs_over_json_channel` → `configure_logging(..., 
sending_to_supervisor=True)`), whose chain deliberately **omits `mask_logs`** — 
masking happens supervisor-side. Any future `RemoteLogIO` whose processors 
don't require a file-backed logger would ship **unmasked** events from the 
subprocess *and* duplicate them (the supervisor re-logs the same events through 
its own remote-processor chain). It's inert for CloudWatch today, but this is 
exactly the trap the old code had, so let's gate it while we're here.
   2. `getattr(remote, "processors")` without a default raises `AttributeError` 
for a third-party `REMOTE_TASK_LOG` object missing the property, and that takes 
down logging setup for the whole process (same pre-existing pattern in 
`logging_processors`).
   
   ```suggestion
       if (
           not sending_to_supervisor
           and (remote := load_remote_log_handler())
           and (remote_processors := getattr(remote, "processors", None))
       ):
           current_processors = list(structlog.get_config()["processors"])
           # Insert before the final renderer
           updated_processors = current_processors[:-1] + 
list(remote_processors) + [current_processors[-1]]
           structlog.configure(processors=updated_processors)
   ```
   



##########
task-sdk/src/airflow/sdk/log.py:
##########
@@ -133,6 +139,15 @@ def configure_logging(
         callsite_parameters=callsite_params,
     )
 
+    # dictConfig has now run — safe to build the watchtower handler.
+    # We append remote processors into the already-configured structlog chain,
+    # inserting before the final renderer (last processor in the chain).

Review Comment:
   The injection itself is inert for CloudWatch: every global-chain logger is 
stdout/pipe-backed, so `relative_path_from_logger()` returns `None` and the 
processor early-returns. Actual streaming uses the file-backed task logger from 
`logging_processors()`, which loads `remote.processors` lazily — also after 
`dictConfig`.
   
   ```suggestion
       # dictConfig has now run, so it is safe to build the watchtower handler. 
Re-inject the remote
       # processors into the global structlog chain (before the final renderer) 
for parity with the old
       # extra_processors layout. Task-log streaming itself does not rely on 
this: it uses the
       # file-backed logger built from logging_processors(), which loads 
remote.processors lazily.
   ```
   



-- 
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