|
try: |
|
# Type of the chunk depends on the backend. |
|
chunks: list[Any | None] = [] |
|
while True: |
|
try: |
|
item = self._async_queue.get_nowait() |
|
chunks.append(item) |
|
except asyncio.QueueEmpty: |
|
# We've exhausted the current items in the queue. |
|
break |
|
|
|
# Make sure we always get the minimum chunk size. |
|
while len(chunks) <= self._chunk_size: |
|
if len(chunks) > 0: |
|
if chunks[-1] is None or isinstance(chunks[-1], Exception): |
|
break # Hit sentinel value or an error. |
|
# We could switch to relying on the `done` / `finish_reason` field of chunks, |
|
# but that forces us to know about the chunk type here. Prefer sentinel values |
|
# for now. |
|
|
|
item = await self._async_queue.get() |
|
chunks.append(item) |
|
|
|
# Process the sentinel value if it's there. |
|
if chunks[-1] is None: |
|
chunks.pop() # Remove the sentinel value. |
|
self._computed = True |
|
|
|
# Shouldn't be needed, but cancel the Tasks this ModelOutputThunk relied on. |
|
if self._generate is not None: |
|
self._generate.cancel() |
|
if self._generate_extra is not None: |
|
# Covers an hf edge case. The task is done generating anything useful but isn't `done` yet. |
|
await self._generate_extra |
|
self._generate_extra.cancel() |
|
|
|
# If ModelOutputThunks get too bulky, we can do additional cleanup here |
|
# and set fields to None. |
|
|
|
elif isinstance(chunks[-1], Exception): |
|
# Mark as computed so post_process runs in finally block |
|
self._computed = True |
|
# Store exception to re-raise after cleanup |
|
exception_to_raise = chunks[-1] |
|
|
|
for chunk in chunks: |
|
assert self._process is not None |
|
await self._process(self, chunk) |
|
|
|
finally: |
|
# Always call post_process if computed, even on exception |
|
# This ensures telemetry spans are properly closed |
|
if self._computed: |
|
assert self._post_process is not None |
|
await self._post_process(self) |
The temporary violation occurs when the process hook fails but the last chunk is processed successfully. We then have the computed bit set to true but haven't actually set the _underlying_value successfully yet. We then end up trying to access the value in the finally block, because computed is true, even though the value is None.
To make matters worse, the exception that caused the original violation gets swallowed up by the finally, which creates a new exception, which makes debugging the root cause quite annoying.
This is the combination of two issues with the current state of core.py:
- We set _process value callbacks instead of directly (not sure how avoidable this is, given how gross the backend code is). And we do those callbacks only after setting the computed bit to true.
- We wrap that in a finally.
A couple of fixes are in order.
First, we should never set computed to True until we actually have a final underlying value. This means that we should introduce a do_set_computed bool and use that in the if chunks[-1] is None block (and anywhere else we set _computed = True).
Then, we should:
if do_set_computed:
assert self._underlying_value is not None
self._computed = True
only after calling the process hooks.
Second, we should guard against that finally block every throwing any exceptions, probably via heavy auditing or figuring out a different solution to the logging situation.
mellea/mellea/core/base.py
Lines 284 to 338 in f0b1346
The temporary violation occurs when the process hook fails but the last chunk is processed successfully. We then have the computed bit set to true but haven't actually set the
_underlying_valuesuccessfully yet. We then end up trying to access the value in the finally block, becausecomputedis true, even though thevalueis None.To make matters worse, the exception that caused the original violation gets swallowed up by the finally, which creates a new exception, which makes debugging the root cause quite annoying.
This is the combination of two issues with the current state of core.py:
A couple of fixes are in order.
First, we should never set computed to True until we actually have a final underlying value. This means that we should introduce a do_set_computed bool and use that in the if chunks[-1] is None block (and anywhere else we set _computed = True).
Then, we should:
only after calling the process hooks.
Second, we should guard against that finally block every throwing any exceptions, probably via heavy auditing or figuring out a different solution to the logging situation.