Yicong-Huang commented on code in PR #59860:
URL: https://github.com/apache/spark/pull/59860#discussion_r4418524875
##########
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:
Added in b0f004a: `test_init_error_traceback_follows_worker_settings`
(`test_worker.py:318`) checks that the default rendering keeps the
payload-loading frame and that `SPARK_HIDE_TRACEBACK` yields the exception-only
line.
--
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]
