Skip to content

core is_computed / value invariant can be violated, and root cause is difficult to debug. #427

Description

@nrfulton

mellea/mellea/core/base.py

Lines 284 to 338 in f0b1346

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:

  1. 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.
  2. 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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    area/streamingStreaming chunks, events, per-chunk validationbugSomething isn't workingp1High: important bugs (workaround exists) or high-value core features. Do soon, not on fire.

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions