spark-test commented on code in PR #59860:
URL: https://github.com/apache/spark/pull/59860#discussion_r4418248356
##########
python/pyspark/tests/test_worker.py:
##########
@@ -126,6 +146,252 @@ 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, command: Optional[bytes] = None) -> bytes:
+ """The udf_info of a task running one single-argument UDF ``func``, or the
pickled
+ ``command`` in its place."""
+ if command is None:
+ command = CloudPickleSerializer().dumps((func, LongType()))
+ 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(command, buf)
+ write_long(0, buf) # result_id
+ return buf.getvalue()
+
+
+def _udtf_info(udtf_cls: type, return_type: str) -> bytes:
+ """The udf_info of a task running ``udtf_cls`` on one argument with
``return_type``."""
+ buf = io.BytesIO()
+ write_int(1, buf) # num_args
+ write_int(0, buf) # arg offset
+ buf.write(b"\x00") # is_kwarg
+ write_int(0, buf) # num_partition_child_indexes
+ buf.write(b"\x00") # has_pickled_analyze_result
+ write_with_length(CloudPickleSerializer().dumps(udtf_cls), buf)
+ _write_utf8(return_type, buf)
+ _write_utf8("test_udtf", buf) # name
+ return buf.getvalue()
+
+
+class _TwoArgUDTF:
+ def eval(self, a, b):
+ yield (a,)
+
+
+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):
+ """Tests for the init response the worker sends before any output, run
without a JVM."""
+
+ def setUp(self):
+ # The worker runs in this process and sets up process-wide state as a
worker would.
+ saved_files = (SparkFiles._root_directory,
SparkFiles._is_running_on_worker)
+ saved_path = list(sys.path)
+ saved_accumulators = dict(_accumulatorRegistry)
+
+ def restore():
+ TaskContext._setTaskContext(None)
+ SparkFiles._root_directory, SparkFiles._is_running_on_worker =
saved_files
+ sys.path[:] = saved_path
+ _accumulatorRegistry.clear()
+ _accumulatorRegistry.update(saved_accumulators)
+
+ self.addCleanup(restore)
+
+ def _run(self, worker_input: bytes, expect_exit: bool = False) ->
io.BytesIO:
+ outfile = io.BytesIO()
+ if expect_exit:
+ with self.assertRaises(SystemExit):
+ worker_main(io.BytesIO(worker_input), outfile)
+ else:
+ worker_main(io.BytesIO(worker_input), outfile)
+ outfile.seek(0)
+ return outfile
+
+ def _read_init_response(self, outfile: io.BytesIO) -> dict:
+ self.assertEqual(read_int(outfile),
SpecialLengths.START_OF_INIT_RESPONSE)
+ return json.loads(outfile.read(read_int(outfile)).decode("utf-8"))
+
+ def test_init_response_precedes_output(self):
+ outfile = self._run(_worker_input(_command(lambda _, it: (x * 2 for x
in it)), [1, 2]))
+ self.assertEqual(self._read_init_response(outfile), {})
+ ser = CPickleSerializer()
+ self.assertEqual([ser._read_with_length(outfile) for _ in range(2)],
[2, 4])
+ self.assertEqual(read_int(outfile), SpecialLengths.METRICS_DATA)
+
+ def test_payload_error_is_reported_in_init_response(self):
+ outfile = self._run(_worker_input(_framed(b"not a pickle"), [1]),
expect_exit=True)
+ error = self._read_init_response(outfile)["error"]
+ self.assertTrue(error["isPayloadError"])
+ self.assertEqual(error["errorClass"], "UnpicklingError")
+ self.assertIn("UnpicklingError", error["traceback"])
Review Comment:
**Nit (P3):** This check, like the other traceback assertions in
`WorkerInitResponseTests` (lines 324, 333 and 379), only looks for the
exception class name or error code, and every case runs with the default
environment. So nothing pins how the init route renders the traceback.
`report_init_response` builds it with `format_worker_exception(error)`
(`worker.py:107`). That function honors `SPARK_HIDE_TRACEBACK`,
`SPARK_SIMPLIFIED_TRACEBACK` and `SPARK_TRACEBACK_WITH_LOCALS`, but it takes
its frames from `sys.exc_info()`, so it is only correct while it runs inside
the `except` arm of `invoke_udf` (`worker.py:363`). `handleInitResponse` then
forwards the string verbatim as the `PythonException` traceback.
Two regressions on this route would pass every test:
- Rendering with `traceback.format_exc()`, or any formatter that ignores
those variables. `spark.sql.execution.pyspark.udf.hideTraceback.enabled` would
then silently stop applying to init failures such as an unpicklable UDF, a
Python version mismatch or a UDTF `eval` signature mismatch.
- Calling `report_init_response` after the handler has exited, which drops
every frame. The output then collapses to exactly the exception-only line that
hide-traceback produces, and all four substring checks still match.
`HandleWorkerExceptionTests` pins these modes only through
`handle_worker_exception`. The end-to-end traceback-mode tests
(`test_udf_traceback_with_locals`, `test_udtf_analyze_traceback_with_locals`)
cover data-phase and analyze errors, not worker init failures. Before this PR,
init failures were rendered by `handle_worker_exception` itself, the function
those tests pin.
The current rendering is correct, so this is only about keeping a regression
signal for the PR's promise that init errors carry the same traceback as
before. This route is also what the gRPC-alignment follow-ups are most likely
to rework. A couple of mode-specific cases in `WorkerInitResponseTests` would
close the gap.
**Verification:**
- **Regression:** An init failure run in-process with SPARK_HIDE_TRACEBACK
set yields an init error whose traceback is the hide-traceback rendering. The
assertion fails if the init-error route renders with a formatter that ignores
that variable.
- **Regression:** An init failure run with the default environment yields an
init error whose traceback includes the frame that raised during payload
loading. The assertion fails if the init-error route renders the exception
without the handled exception's frames, for example when it is called outside
the exception handler.
- **Behavior:** The existing init-response cases still pass unchanged:
success before output, each payload window, the unprintable exception, setup
and handler-construction errors, and a data-phase error after a successful init
response.
--
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]
