Yicong-Huang commented on code in PR #59860:
URL: https://github.com/apache/spark/pull/59860#discussion_r4411188924
##########
python/pyspark/worker.py:
##########
@@ -76,6 +78,33 @@ def report_metrics(outfile: BinaryIO, metrics: dict[str,
int]) -> None:
write_with_length(payload, outfile)
+def report_init_response(
+ outfile: BinaryIO, error: Optional[BaseException] = None,
is_payload_error: bool = False
+) -> None:
+ """
+ Write the init response as a length-prefixed JSON object after
START_OF_INIT_RESPONSE.
+
+ It is sent exactly once per task, before any output: ``{}`` once the task
is initialized, or
+ ``{"error": {...}}`` if initialization failed, in which case no output
follows. An init error
+ carries the exception message, the traceback as the JVM reports it, the
exception class name,
+ and whether the worker failed to load the UDF payload (``isPayloadError``)
rather than to set
+ up the task. Must be called while handling ``error``.
+ """
+ response: dict[str, Any] = {}
+ if error is not None:
+ response["error"] = {
+ "message": str(error),
Review Comment:
Fixed in 84c2842: the message falls back to `<exception str() failed>`
(`worker.py:81`), and `test_payload_error_with_unprintable_exception` covers it.
##########
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:
Fixed in 6ef3fde: `workerJsonMapper` lifts `maxStringLength` the same way
JsonProtocol does (`PythonRunner.scala:480`), and a new BasePythonRunnerSuite
case decodes a traceback longer than the default limit.
##########
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:
Fixed in 6ef3fde: the docstring now says "a failure while processing data
after it" (`worker.py:126`).
##########
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:
Fixed in 6ef3fde: `test_payload_window_of_each_payload_kind`
(`test_worker.py:348`) adds the UDF unpickle failure, the UDTF non-struct
return type and the UDTF eval-signature mismatch cases; removing any of the
four UDF-list/UDTF window assignments now fails a test.
--
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]
