uranusjr commented on code in PR #75561:
URL: https://github.com/apache/airflow/pull/75561#discussion_r4385406506
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -206,31 +216,96 @@ def _build_runtime_env() -> dict[str, str]:
}
+def _read_available(soc: socket.socket, drain_size: int) -> bytes:
+ """Return the data *soc* holds now, without waiting for more from a
process that may still run."""
+ soc.setblocking(False)
+ chunks = []
+ budget = _STARTUP_DRAIN_BUDGET
+ with contextlib.suppress(BlockingIOError, OSError):
+ while budget > 0 and (chunk := soc.recv(min(drain_size, budget))):
+ budget -= len(chunk)
+ chunks.append(chunk)
+ return b"".join(chunks)
+
+
+def _report_startup_output(
+ logger: FilteringBoundLogger,
+ drains: dict[str, socket.socket],
+ drained: dict[socket.socket, bytes],
+ drain_size: int,
+) -> None:
+ """
+ Write what a runtime printed before it failed to connect to the task log,
one event per line.
+
+ The output goes to the log the way a connected runtime's output does,
rather than into the error
+ message, so every line keeps its own level and timestamp.
+ """
+ for key, soc in drains.items():
+ output = drained[soc] + _read_available(soc, drain_size)
+ # A connected runtime's stderr is logged as an error and its stdout as
info.
+ write = logger.error if key == "stderr" else logger.info
+ name = f"task.{key}"
+ if len(output) > _STARTUP_OUTPUT_LIMIT:
+ output = output[-_STARTUP_OUTPUT_LIMIT:]
+ write(_STARTUP_OUTPUT_TRUNCATED, logger=name)
+ for line in output.decode(errors="replace").splitlines():
+ write(line, logger=name)
+
+
[email protected]
+def _raise_as_launch_error(prefix: str) -> Generator[None, None, None]:
+ """
+ Raise a failure in the block as a TaskLaunchError, so the supervisor fails
the task with its reason.
+
+ Use it only before the task instance is started: the supervisor starts the
task instance itself
+ when it fails one with a TaskLaunchError. A TaskLaunchError raised in the
block is kept as it is.
+ An OSError or ValueError, such as a missing runtime or artifact, or an
artifact that cannot run,
+ is retryable. Any other failure, such as a coordinator bug, is not:
retrying cannot fix it.
+ """
+ try:
+ yield
+ except TaskLaunchError:
+ raise
+ except (OSError, ValueError) as e:
+ raise TaskLaunchError(f"{prefix}: {e}") from e
+ except Exception as e:
+ raise TaskLaunchError(f"{prefix}: {e}", retryable=False) from e
+
+
def _accept_connections(
servers: dict[str, socket.socket],
drains: dict[str, socket.socket],
proc: subprocess.Popen,
*,
+ logger: FilteringBoundLogger,
max_wait: float = 10.0,
drain_size: int = 4096,
) -> tuple[dict[socket.socket, socket.socket], dict[socket.socket, bytes]]:
- """Block until the subprocess connects to servers, draining stdout/stderr
along the way."""
+ """
+ Block until the subprocess connects to servers, draining stdout/stderr
along the way.
+
+ :raises TaskLaunchError: when the subprocess exits or *max_wait* elapses
before it connects. What
+ the subprocess wrote to the *drains* until then is written to *logger*
first.
+ """
accepted: dict[socket.socket, socket.socket] = {}
drained: dict[socket.socket, bytes] = {s: b"" for s in drains.values()}
+
+ def _fail(reason: str) -> NoReturn:
+ for s in accepted.values():
+ s.close()
+ _report_startup_output(logger, drains, drained, drain_size)
+ raise TaskLaunchError(f"The runtime {reason}")
Review Comment:
Nit: the startup output is written to the task log before
`_fail_task_before_runtime` calls `task_instances.start()`. If start is
rejected as cleared, the log still contains output for a try that never ran.
Probably acceptable, but worth mentioning in the docstring.
--
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]
