gemini-code-assist[bot] commented on code in PR #39384:
URL: https://github.com/apache/beam/pull/39384#discussion_r3610946279
##########
sdks/python/apache_beam/runners/worker/sdk_worker.py:
##########
@@ -1372,20 +1373,34 @@ def _partially_cached_iterable(
return self.ContinuationIterable(
head,
functools.partial(
- self._lazy_iterator, state_key, coder, continuation_token))
+ self._lazy_iterator, state_key, coder, continuation_token),
+ on_failure=(
+ functools.partial(self._state_cache.invalidate, cache_key)
+ if cache_key is not None else None))
class ContinuationIterable(Generic[T], CacheAware):
- def __init__(self, head, continue_iterator_fn):
- # type: (Iterable[T], Callable[[], Iterable[T]]) -> None
+ def __init__(self, head, continue_iterator_fn, on_failure=None):
+ # type: (Iterable[T], Callable[[], Iterable[T]], Optional[Callable[[],
None]]) -> None
self.head = head
self.continue_iterator_fn = continue_iterator_fn
+ self.on_failure = on_failure
def __iter__(self):
# type: () -> Iterator[T]
for item in self.head:
yield item
- for item in self.continue_iterator_fn():
- yield item
+ try:
+ for item in self.continue_iterator_fn():
+ yield item
+ except Exception:
+ # The continuation token bound into this iterable may be permanently
+ # invalid (e.g. the runner revoked it along with the work item that
+ # created it). If this instance stays cached it will fail every
+ # subsequent read of this state, so drop it from the cache and let
+ # the next read fetch fresh state.
+ if self.on_failure is not None:
+ self.on_failure()
Review Comment:
If `self.on_failure()` raises an exception, it will mask the original
exception raised by `self.continue_iterator_fn()`. To prevent masking the
primary failure and ensure it is correctly propagated, we should wrap the
`on_failure()` call in a `try...except` block and log any invalidation failures.
```suggestion
if self.on_failure is not None:
try:
self.on_failure()
except Exception:
_LOGGER.warning('Failed to invalidate cache on continuation
failure.', exc_info=True)
```
--
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]
