Yicong-Huang commented on code in PR #59976:
URL: https://github.com/apache/spark/pull/59976#discussion_r4447072533
##########
python/pyspark/worker.py:
##########
@@ -111,6 +111,8 @@ def report_init_response(
payload = json.dumps(response, separators=(",", ":")).encode("utf-8")
write_int(SpecialLengths.START_OF_INIT_RESPONSE, outfile)
write_with_length(payload, outfile)
+ # The JVM timestamps this response to measure initialization.
+ outfile.flush()
Review Comment:
**Non-blocking (P2):** On the default TCP worker socket this flush turns a
short task's response into two small writes: this 10-byte init frame, and later
the end-of-task tail (output, `METRICS_DATA`, `END_OF_DATA_SECTION`,
accumulators, `END_OF_STREAM`). No worker socket sets `TCP_NODELAY`, and if the
executor has already sent all of its input it has no data to piggyback an ACK
on. Nagle can then hold the tail until the executor's delayed-ACK timer fires,
about 40 ms on Linux. I reproduced this over loopback TCP with the same write
pattern: a 10-byte flushed frame, about 1 ms of work, then a small tail,
repeated on one connection. About 11% of tasks stalled by roughly 40 ms, and
the run took about 5.6 times as long. The base branch's single end-of-task
flush showed no stalls, and neither did this flush with `TCP_NODELAY` on the
sending socket. The reproduction models only the socket pattern, not Spark's
executor reader, so frequency in real jobs may differ. Queries with many tiny
Python UDF tasks over TCP would see the added wall time; Unix domain sockets,
pipelined mode, and large-output tasks are essentially unaffected.
**Recommended change:** Disable Nagle's algorithm (TCP_NODELAY) on the
Python worker's TCP data socket to the executor. Set it both where the daemon
accepts a worker connection and where a non-daemon worker connects back to the
executor. Leave Unix domain sockets unchanged. Add a test that fails when a
small tail written after the flushed init frame is held over a TCP connection.
**Why this works:** With Nagle disabled, each flush sends immediately even
while the init segment is unacknowledged, so the end-of-task tail no longer
waits for the executor's delayed ACK. The worker's 64 KiB buffered writer still
coalesces writes, so segments still go out only at flush points: the init
response, buffer-full writes, per-batch flushes in pipelined mode, and task end.
**Scope:** Set TCP_NODELAY on the Python worker's TCP data socket in both
worker creation routes and add a regression test for the end-of-task tail.
**Compatibility:** The bytes written, the init response's early delivery,
and Unix domain socket behavior stay as in the PR.
**Risks:** TCP_NODELAY also stops the kernel from coalescing any other
back-to-back small flushes on the worker socket, such as per-batch flushes in
pipelined mode. That can add small segments, but the bytes do not change. The
change is one socket option per route and is easy to roll back. The stall was
reproduced with a loopback model of the PR's write pattern, not measured inside
a Spark executor, so its frequency in real jobs may differ.
**Constraints:** Keep the init response's early delivery to the executor. Do
not change Unix domain socket behavior.
**Exclusions:**
- Where pythonInitTime ends is a separate fix and does not depend on this
change.
**Success:** A short task over a TCP worker socket whose output fits in one
buffer delivers its end-of-task output without waiting for the executor's
delayed ACK of the init response. The init response still reaches the executor
before the worker processes input.
**Verification:**
- **Regression:** Over a TCP connection set up like the worker's, a short
task that flushes the init frame and then a small tail delivers the tail
without a delayed-ACK stall across many repeated tasks.
- **Behavior:** The init response is still readable by the executor before
the worker processes the task's input.
- **Compatibility:** Unix domain socket workers keep working unchanged, and
existing worker protocol tests pass.
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -1206,6 +1207,7 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
s"Expected the Python worker init response, but got $other")
}
} catch handleException
+ initResponseTimeMs = System.currentTimeMillis()
Review Comment:
**Non-blocking (P2):** This stamp is taken only after `readInt()` returns,
which can be well after the init response arrives. In the default
(non-pipelined) mode the read goes through `ReaderInputStream.read`. In the
same loop iteration where it reads the response bytes, that method also runs
`writeAdditionalInputToPythonWorker()` whenever the key is writable, and
`writeBufferedInputToPythonWorker` keeps pulling and writing input while
`channel.write` accepts data. For example, a slow upstream scan feeding a cheap
UDF with small output keeps that loop running while the now-initialized worker
drains input. The time to produce and write that input, plus the worker's
processing of it, is then counted as initialization. In pipelined mode the
writer thread drives the worker on its own, and nothing reads the socket until
the consumer's first `hasNext`, so any delay before that first pull is counted
too. In both modes `pythonInitTime` can exceed `pythonTotalTime`, which could
not happen
on the base branch because `initTimestampMs` was recorded by the worker. The
new grid test reads an in-memory stream, and the `PythonUDFSuite` case writes
its one-row input before the response arrives, so neither exercises this
ordering.
**Recommended change:** End pythonInitTime at the worker-reported
initTimestampMs again, keeping the start at max(startTime, bootTimestampMs),
and remove the JVM-side receipt timestamp that handleInitResponse now records.
Restore the timing tests to match: the core reader-timing test should expect
init from the reported initTimestampMs, and the reused-worker SQL timing test
should again bound init by the task's total Python time.
**Why this works:** The worker records initTimestampMs immediately before it
writes and flushes the init response, so on the same host the timestamp
precedes the response's arrival by only the local transport time and does not
depend on the JVM reader. Deriving init from it removes both delays that now
leak in: the input-write loop that ReaderInputStream.read runs after reading
the response, and a pipelined consumer's delayed first pull. Because
initTimestampMs is never later than finishTimestampMs, init cannot exceed total
again.
**Scope:** Change the init end used in the JVM runner's metric calculation
back to the worker's timestamp, drop the now-unused receipt stamp, and realign
the core and SQL Python timing tests with that definition.
**Compatibility:** A reused worker's pre-task idle time stays excluded from
initialization, and every timing other than pythonInitTime keeps its
base-branch definition.
**Risks:** This replaces the PR's stated user-facing change (initialization
ending when the executor receives the response) with the worker's own
timestamp, which the response's flush makes nearly identical to arrival on the
same host. The metric then no longer includes the time to deliver the response
to the JVM. The change is easy to roll back because it touches only the metric
calculation.
**Constraints:** Keep the start of initialization at max(startTime,
bootTimestampMs) so a reused worker's idle time stays excluded. Keep
pythonBootTime, pythonTotalTime, pythonProcessingTime, spill metrics, and the
logged finish duration unchanged.
**Exclusions:**
- The extra per-task latency that the worker-side flush of the init response
causes over TCP sockets is fixed separately; this repair neither depends on the
flush nor removes it.
**Success:** pythonInitTime equals initTimestampMs - max(startTime,
bootTimestampMs) for every Python runner, independent of how long the JVM
reader keeps writing input after the response arrives or when the consumer
first pulls the output iterator. pythonInitTime never exceeds pythonTotalTime
for a task that completes normally. Boot, total, processing, spill, and
finish-log values are unchanged from the base branch.
**Verification:**
- **Regression:** A default (non-pipelined) task whose input is still being
produced and written after the worker sends its init response, for example a
slow upstream iterator feeding a cheap UDF with small output, reports a
pythonInitTime that excludes that post-arrival writing and stays no larger than
pythonTotalTime.
- **Regression:** With pipelined Python UDF execution enabled, delaying the
consumer's first pull of the output iterator after the worker has initialized
does not increase pythonInitTime.
- **Behavior:** For a reused worker that sat idle before the task,
pythonInitTime excludes the idle period and does not exceed the task's total
Python time.
- **Compatibility:** Boot, total, processing, spill, and logged finish
values computed from a fixed worker metrics report match the base-branch
calculation.
##########
python/pyspark/tests/test_worker.py:
##########
@@ -283,6 +285,46 @@ def test_init_response_precedes_output(self):
ser = CPickleSerializer()
self.assertEqual([ser._read_with_length(outfile) for _ in range(2)],
[2, 4])
self.assertEqual(read_int(outfile), SpecialLengths.METRICS_DATA)
+ metrics = json.loads(outfile.read(read_int(outfile)))
+ self.assertIn("bootTimestampMs", metrics)
+ self.assertIn("finishTimestampMs", metrics)
+ self.assertIn("initTimestampMs", metrics)
+ self.assertNotIn("pythonInitializationDurationMs", metrics)
Review Comment:
**Nit (P3):** `pythonInitializationDurationMs` appears nowhere else in
`python/`, `core/`, or `sql/`. No worker writes it, no JVM code reads it, and
`handleMetricsData` would drop it anyway because it keeps only registered
metric names. This assertion therefore cannot fail under any reachable
regression of the boot/init/finish timestamp protocol, while it reads as though
the protocol deliberately forbids that key. The `assertIn` checks above already
cover what the report actually carries.
**Verification:**
- **Inspection:** No test asserts the presence or absence of a
metrics-report key that neither the Python worker nor the JVM runner defines.
- **Behavior:** The worker init-response test still checks that the metrics
report carries the boot, init, and finish timestamps after the task's output.
--
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]
