Yicong-Huang opened a new pull request, #58835: URL: https://github.com/apache/spark/pull/58835
### What changes were proposed in this pull request? When Python worker logging is enabled (`spark.sql.pyspark.worker.logging.enabled=true`), a task's worker log lines are captured from the worker's stdout by a background redirect thread and saved as blocks (keyed by session id and worker PID) via `PythonWorkerLogCapture`. The `python_worker_logs()` table-valued function then reads those blocks back in a subsequent query. A block only becomes visible once the worker's end-of-logs sentinel (an empty-payload marker line emitted at the end of each task) is processed and the rolling log writer is closed. This capture happens on a channel that is independent of the task result: the worker flushes the sentinel to stdout, but the query result flows over the data socket. There is no happens-before relationship between the two, so a `python_worker_logs()` call issued right after the triggering query can race the async capture and see missing (or zero) log blocks. This PR establishes that ordering. On the worker side, the sentinel is already flushed to stdout before the task writes `END_OF_DATA_SECTION`/accumulators/`END_OF_STREAM` to the data socket, so by the time the executor reaches `handleEndOfDataSection()` the sentinel bytes are already in the pipe. The change makes the executor wait, at the end of the task, for that sentinel to be processed: - `PythonWorkerLogCapture` now tracks a per-worker count of processed sentinels and exposes `sentinelCount` / `awaitLogsFlushed`. The count is incremented only after the writer is closed (block saved), so a woken waiter is guaranteed to see the saved block. - `PythonWorkerFactory` and `SparkEnv` expose thin delegators (the wait is performed outside the `SparkEnv` lock so a slow flush never blocks other worker create/release calls). - `BasePythonRunner` snapshots the sentinel count for the worker PID before the task runs, and `ReaderIterator.handleEndOfDataSection()` waits for the count to advance past that baseline. Correlation is exact because a worker is released back to the pool only after this barrier, so tasks on a given PID are serial and each emits exactly one sentinel. The wait is bounded by a new internal config `spark.python.worker.logging.flushTimeout` (default 10s, `0` disables). On timeout the task still completes and a warning is logged. The barrier only runs when worker logging is enabled for the task, so the default (logging-off) path is unchanged. ### Why are the changes needed? `python_worker_logs()` intermittently returns incomplete results for the query that just ran, because the log capture is asynchronous with respect to task completion. This surfaces as flaky tests such as `pyspark.sql.tests.connect.arrow.test_parity_arrow_udtf` (`test_arrow_udtf_with_logging`), which is far more likely to manifest in Spark Connect remote mode where worker and capture timing are looser. The same latent race affects every test/usage that reads `python_worker_logs()` immediately after a logging-enabled query. ### Does this PR introduce _any_ user-facing change? No behavior change to query results. With worker logging enabled, `python_worker_logs()` now deterministically observes the logs emitted by the preceding query in the same session instead of racing the asynchronous capture. Adds one internal config, `spark.python.worker.logging.flushTimeout`. ### How was this patch tested? - New unit suite `PythonWorkerLogCaptureSuite` covering the sentinel-count/await semantics, per-worker independence, and that `awaitLogsFlushed` blocks until a concurrently produced sentinel arrives (4 tests, all passing). - Ran `test_arrow_udtf_with_logging` under Spark Connect (repeated) and classic mode against a local `-Phive` build; both pass with the barrier active and no regression on the logging-off path. ### Was this patch authored or co-authored using generative AI tooling? No This pull request and its description were written by Isaac. -- 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] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
