spark-test commented on code in PR #59860:
URL: https://github.com/apache/spark/pull/59860#discussion_r4407662812
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -507,6 +507,45 @@ private[spark] object BasePythonRunner extends Logging {
}.toMap
}
+ /**
+ * Decode the init response that follows START_OF_INIT_RESPONSE: a
length-prefixed JSON object
+ * that is `{}` once the worker is initialized, or carries an `error` if
initialization failed.
+ * Unknown fields are ignored.
+ */
+ private[python] def readInitResponse(stream: DataInputStream):
Option[PythonWorkerInitError] = {
+ val length = stream.readInt()
+ if (length <= 0) {
+ throw new SparkException(s"Invalid Python worker init response length:
$length")
+ }
+ val json = PythonWorkerUtils.readUTF(length, stream)
+ val response = try {
+ workerJsonMapper.readTree(json)
Review Comment:
**Nit (P3):** `workerJsonMapper` is a plain `ObjectMapper`, so this
`readTree` enforces Jackson's default `StreamReadConstraints`, which reject any
JSON string value longer than `StreamReadConstraints.DEFAULT_MAX_STRING_LEN`.
The init response now carries the worker's full formatted traceback and the
exception message as JSON strings. Before this PR, init failures used the
`PYTHON_EXCEPTION_THROWN` frame read by `PythonWorkerUtils.readUTF`, which has
no such limit.
When an init-phase traceback or message is longer than that, `readTree`
throws and the task fails with `SparkException("Malformed Python worker init
response JSON")` in place of the `PythonException` that carries the worker
traceback. The worker sent the traceback only in this response, so the user
never sees it. The likeliest trigger is
`spark.sql.execution.pyspark.udf.tracebackWithLocals.enabled`, the setting
users turn on to debug exactly such failures. It formats with
`TracebackException(capture_locals=True)`, which embeds the full `repr()` of
every local with no truncation, so one large local in a user frame on the init
path (an unpickling reconstructor or a UDTF `__init__`) is enough. A very large
exception message works the same way. Failures after the init response still go
through `readUTF`, so only init failures are affected, and the trigger is rare.
Lifting the limit on this mapper, as `JsonProtocol` and `KVUtils` already do
for SPARK-49872, restores parity, for example:
```scala
private lazy val workerJsonMapper = {
val mapper = new ObjectMapper()
.enable(JsonParser.Feature.STRICT_DUPLICATE_DETECTION)
.enable(DeserializationFeature.FAIL_ON_TRAILING_TOKENS)
mapper.getFactory.setStreamReadConstraints(
StreamReadConstraints.builder().maxStringLength(Int.MaxValue).build())
mapper
}
```
The metrics report shares this mapper but carries only short keys and
integers, so it loses nothing from the lifted limit.
**Verification:**
- **Regression:** Decoding an init response whose error traceback exceeds
Jackson's default maximum string length returns the init error with the
complete traceback, instead of failing with the malformed-JSON SparkException.
- **Compatibility:** The existing init-response and worker-metrics decoding
cases behave as before, including success with unknown fields, error decoding,
and rejection of malformed JSON, invalid error fields and non-positive lengths.
##########
python/pyspark/worker.py:
##########
@@ -84,7 +121,12 @@ def invoke_udf(message_receiver: SparkMessageReceiver,
outfile: BinaryIO):
This method goes through three phases:
Initialization -> Processing -> Finish/Cleanup
+
+ Initialization ends with the init response (see ``report_init_response``);
a failure before
+ it is reported there, and a failure after it as PYTHON_EXCEPTION_THROWN.
Review Comment:
**Nit (P3):** This promises more than the code does. Only exceptions raised
inside the `try` block (processing and the `TaskContext` reset) reach
`handle_worker_exception`. The finish section after the `try`
(`report_metrics`, `END_OF_DATA_SECTION`, `send_accumulator_updates` and the
finish-signal check) never writes `PYTHON_EXCEPTION_THROWN`. An accumulator
value that does not pickle raises out of `main` with no frame, so the JVM sees
the worker exit while reading accumulator updates. A bad finish signal makes
the worker write `END_OF_DATA_SECTION` and exit. That behavior predates this
PR, so only the new sentence is off. It matters because this docstring now
describes the protocol that the gRPC alignment builds on. Something like "a
failure while processing data after it as PYTHON_EXCEPTION_THROWN" would match
the code.
**Verification:**
- **Inspection:** The invoke_udf docstring's failure-reporting statement
matches the code. Failures inside the processing try block go through
handle_worker_exception, and the docstring does not say that the metrics,
accumulator or finish-signal steps produce PYTHON_EXCEPTION_THROWN.
##########
python/pyspark/tests/test_worker.py:
##########
@@ -126,6 +146,195 @@ def test_accumulate_and_propagate_exception(self):
self.assertEqual(metrics.to_dict(), {"duration": 3})
+def _write_utf8(s: str, buf: io.BytesIO) -> None:
+ write_with_length(s.encode("utf-8"), buf)
+
+
+def _worker_input(
+ udf_info: bytes,
+ items: list,
+ python_version: Optional[str] = None,
+ eval_type: int = PythonEvalType.NON_UDF,
+) -> bytes:
+ """Build the socket input of a task of ``eval_type`` whose init message
ends with
+ ``udf_info``, followed by ``items`` as pickled input."""
+ init = io.BytesIO()
+ write_int(0, init) # split_index
+ _write_utf8(python_version or "%d.%d" % sys.version_info[:2], init)
+ task_context = {
+ "isBarrier": False,
+ "connInfo": None,
+ "secret": None,
+ "stageId": 0,
+ "partitionId": 0,
+ "attemptNumber": 0,
+ "taskAttemptId": 0,
+ "cpus": 1,
+ "cpuAmount": "1",
+ "resources": {},
+ "localProperties": {},
+ }
+ _write_utf8(json.dumps(task_context), init)
+ _write_utf8(tempfile.gettempdir(), init) # spark_files_dir
+ write_int(0, init) # num_python_includes
+ init.write(b"\x00") # needs_broadcast_decryption_server
+ write_int(0, init) # num_broadcast_variables
+ write_int(eval_type, init)
+ write_int(0, init) # runner conf
+ write_int(0, init) # eval conf
+ init.write(udf_info)
+
+ buf = io.BytesIO()
+ write_int(SpecialLengths.START_OF_INIT_MESSAGE, buf)
+ write_with_length(init.getvalue(), buf)
+ CPickleSerializer().dump_stream(iter(items), buf)
+ write_int(SpecialLengths.END_OF_DATA_SECTION, buf)
+ write_int(SpecialLengths.END_OF_STREAM, buf)
+ return buf.getvalue()
+
+
+def _framed(payload: bytes) -> bytes:
+ buf = io.BytesIO()
+ write_with_length(payload, buf)
+ return buf.getvalue()
+
+
+def _command(func: Any) -> bytes:
+ """The udf_info of a classic RDD (NON_UDF) task running ``func``."""
+ ser = CPickleSerializer()
+ return _framed(CloudPickleSerializer().dumps((func, None, ser, ser)))
+
+
+def _udf_info(func: Any) -> bytes:
+ """The udf_info of a task running one single-argument UDF ``func``."""
+ buf = io.BytesIO()
+ write_int(1, buf) # num_udfs
+ write_int(1, buf) # num_args
+ write_int(0, buf) # arg offset
+ buf.write(b"\x00") # is_kwarg
+ write_int(1, buf) # num_chained
+ write_with_length(CloudPickleSerializer().dumps((func, LongType())), buf)
+ write_long(0, buf) # result_id
+ return buf.getvalue()
+
+
+class _UnprintableError(Exception):
+ def __str__(self):
+ raise RuntimeError("str() failed")
+
+
+def _raise_unprintable():
+ raise _UnprintableError()
+
+
+class _RaisesOnUnpickle:
+ def __reduce__(self):
+ return _raise_unprintable, ()
+
+
+class _FailingHandler:
+ def __init__(self, **kwargs):
+ raise ImportError("pyarrow is not installed")
+
+
+class WorkerInitResponseTests(unittest.TestCase):
Review Comment:
**Nit (P3):** `isPayloadError` comes from three separate `loading_payload`
windows in `invoke_udf` (NON_UDF, UDTF and the SQL UDF list), but the
assertions here cover only the NON_UDF window and the closing edge of the
UDF-list window. They check true for the NON_UDF command (two cases), and false
for the Python-version setup error and for the handler-construction error after
the UDF-list window closes. No case fails a payload inside the UDF-list window
(`worker.py:206-221`), and no case builds a `UDTFInfo` payload, so neither edge
of the UDTF window (`worker.py:179-200`) is exercised, including the
docstring's own non-struct return type example. Deleting `loading_payload =
True` at `worker.py:206` or `:179`, or `loading_payload = False` at `:200`,
would flip the flag for real failures with every test still passing. The JVM
only debug-logs the flag, so no end-to-end suite would notice either. Three
more cases would close the gap: a `SQL_BATCHED_UDF` payload whose function does
not
unpickle (expect true), a UDTF payload with a non-struct return type (expect
true), and a UDTF whose handler construction fails, for example on an `eval`
signature mismatch (expect false).
**Verification:**
- **Behavior:** Each in-process worker run yields exactly one init error
response with the expected isPayloadError value and no further output. The runs
are a SQL UDF payload that does not unpickle, a UDTF payload with a non-struct
return type, and a UDTF whose handler construction fails.
- **Regression:** Deleting any one of the four loading_payload assignments
that open and close the SQL UDF-list and UDTF windows makes at least one
init-response test fail. The new cases cover the UDF list's opening assignment
and both UDTF assignments, and the existing handler-construction case covers
the UDF list's closing assignment.
--
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]
