jason810496 commented on code in PR #75561:
URL: https://github.com/apache/airflow/pull/75561#discussion_r4385905378
##########
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)
Review Comment:
Reused in 9b6e731: `_report_startup_output` feeds `forward_to_log` instead
of splitting and decoding the lines itself.
##########
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:
Noted in `_report_startup_output`'s docstring in 4b008cb.
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -781,28 +867,37 @@ def execute_task(
read other files of the bundle, such as the JARs on a Java classpath.
Any other task runs the
artifact found by scanning the bundle that holds the artifacts.
- :raises TaskLaunchError: when a native Dag file cannot run, before the
runtime starts.
+ :raises TaskLaunchError: when the runtime cannot start: a native Dag
file cannot run, the
+ artifact of any other task cannot be found or run, the artifact
needs a supervisor schema
+ version that this worker's Task SDK does not support, or the
runtime exits or times out
+ before it connects.
"""
task_logger = logger or log
rel_path = os.fspath(dag_rel_path)
- is_native = self._is_native_dag_file(rel_path, bundle_info.name)
- if is_native:
- bundle = _resolve_task_bundle(bundle_info, task_logger)
- else:
- bundle = self._resolve_artifact_bundle(bundle_info, task_logger)
- # Hold the version lock across start()/wait() so bundle cleanup cannot
- # rmtree a version this task is still reading from, mirroring
- # task_runner.main() for the Python task path.
- with (
- BundleVersionLock(bundle_name=bundle.name,
bundle_version=bundle.version),
- self._set_scan_roots([bundle.path]),
- ):
- if is_native:
- command, subprocess_schema_version =
self._build_native_task_command(
- what=what, bundle=bundle, rel_path=rel_path
- )
- else:
- command, subprocess_schema_version =
self._build_execute_task_command(what=what)
+ # The lock and the scan roots are entered under
_raise_as_launch_error, so failing to enter
+ # them is reported as a launch failure too, but released only once
wait() returns.
Review Comment:
Removed in 3fcc709.
##########
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
Review Comment:
Documented in 4b008cb rather than changed: the catch-all stays, and
`_build_execute_task_command` / `_build_dag_file_command` now say to raise
`TaskLaunchError(..., retryable=...)` when the implementation knows the answer.
--
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]
