Skip to content

Commit

Permalink
Cap large exception messages (#1952)
Browse files Browse the repository at this point in the history
  • Loading branch information
kramstrom authored Jul 1, 2024
1 parent a37a408 commit ce3b6b5
Show file tree
Hide file tree
Showing 3 changed files with 36 additions and 1 deletion.
12 changes: 11 additions & 1 deletion modal/_container_io_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -514,13 +514,23 @@ async def handle_input_exception(self, input_id, started_at: float) -> AsyncGene
# serializing the exception, which may have some issues (there
# was an earlier note about it that it might not be possible
# to unpickle it in some cases). Let's watch out for issues.

repr_exc = repr(exc)
if len(repr_exc) >= MAX_OBJECT_SIZE_BYTES:
# We prevent large exception messages to avoid
# unhandled exceptions causing inf loops
# and just send backa trimmed version
trimmed_bytes = len(repr_exc) - MAX_OBJECT_SIZE_BYTES - 1000
repr_exc = repr_exc[: MAX_OBJECT_SIZE_BYTES - 1000]
repr_exc = f"{repr_exc}...\nTrimmed {trimmed_bytes} bytes from original exception"

await self._push_output(
input_id,
started_at=started_at,
data_format=api_pb2.DATA_FORMAT_PICKLE,
status=api_pb2.GenericResult.GENERIC_STATUS_FAILURE,
data=self.serialize_exception(exc),
exception=repr(exc),
exception=repr_exc,
traceback=traceback.format_exc(),
serialized_tb=serialized_tb,
tb_line_cache=tb_line_cache,
Expand Down
19 changes: 19 additions & 0 deletions test/container_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
serialize_data_format,
)
from modal._utils import async_utils
from modal._utils.blob_utils import MAX_OBJECT_SIZE_BYTES
from modal.app import _App
from modal.exception import InvalidError
from modal.partial_function import enter, method
Expand Down Expand Up @@ -1423,6 +1424,24 @@ async def custom_heartbeater(self):
assert "Traceback" not in caplog.text # should not print a full traceback - don't scare users!


@skip_github_non_linux
@pytest.mark.usefixtures("server_url_env")
def test_container_doesnt_send_large_exceptions(servicer):
# Tests that large exception messages (>2mb are trimmed)
ret = _run_container(
servicer,
"test.supports.functions",
"raise_large_unicode_exception",
inputs=_get_inputs(((), {})),
)

assert len(ret.items) == 1
assert len(ret.items[0].SerializeToString()) < MAX_OBJECT_SIZE_BYTES * 1.5
assert ret.items[0].result.status == api_pb2.GenericResult.GENERIC_STATUS_FAILURE
assert "UnicodeDecodeError" in ret.items[0].result.exception
assert servicer.task_result is None # should not cause a failure result


@skip_github_non_linux
@pytest.mark.usefixtures("server_url_env")
def test_sigint_termination_input_concurrent(servicer):
Expand Down
6 changes: 6 additions & 0 deletions test/supports/functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -433,3 +433,9 @@ def sandbox_f(x):
@app.function()
def is_local_f(x):
return is_local()


@app.function()
def raise_large_unicode_exception():
byte_str = (b"k" * 120_000_000) + b"\x99"
byte_str.decode("utf-8")

0 comments on commit ce3b6b5

Please sign in to comment.