[Python] Python] Bound the memory used for fnapi outbound data messages and receiving messages.#38407
[Python] Python] Bound the memory used for fnapi outbound data messages and receiving messages.#38407scwhittle wants to merge 6 commits into
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a capacity constraint on the data transmission queue within the gRPC data channel. By limiting the queue size, the system can better manage backpressure when the SDK produces data faster than the runner can consume it, preventing unbounded memory growth. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![gemini-code-assist[bot]](https://avatars.gh.lumenfield.work/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a fixed size for the _to_send queue in the _GrpcDataChannel class. Feedback suggests that the hardcoded maxsize=10 should be made configurable to avoid potential performance bottlenecks in high-throughput scenarios and notes a discrepancy between the PR description and the implementation of the receive queue size.
|
Assigning reviewers: R: @tvalentyn for label python. Note: If you would like to opt out of this review, comment Available commands:
The PR bot will only process comments in the main thread (not review comments). |
|
thanks, marking as draft for now |
|
waiting on author |
scwhittle
changed the title
scwhittle
force-pushed
the
queue_limit
branch
3 times, most recently
from
2a8a8d6 to
e8f5bf7
Compare
Previously an unbounded queue was used for pending data outputs to be sent over the fnapi to the runner. If outputs were being generated faster than the runner was consuming them, this would lead to memory growth and possible OOMs. This PR introduces a byte-limited queue data structure that is used instead to limit the # of bytes in the queue. This was preferred to just using a queue with max number of elements because the size of elements can vary greatly. For batch pipelines they are likely large while for stremaing pipelines there may be more small outputs.
There was a problem hiding this comment.
Code Review
This pull request introduces a ByteLimitedQueue to the Apache Beam Python SDK to limit queue capacity by both element count and total byte weight, which is then utilized in the _GrpcDataChannel to manage memory usage. The review feedback highlights several critical improvements for the new queue implementation: using time.monotonic() instead of time.time() for more robust timeout handling, ensuring compatibility with Python 3.13's queue.Queue.shutdown() to prevent deadlocks, and switching from notify() to notify_all() in the _get method to avoid potential producer starvation caused by heterogeneous item weights.
|
The failures I looked through were unrelated. Are the tests flaky? |
|
When backpressure happens, would that look to users as a processing lull? If so, that is less than ideal, since we probably should be throttling reading inputs, and we might mislead users that processing is slow. Ideally the Runner would be smart enough to not send them if it cannot consume outputs fast enough... Have you considered tracking unwritten data at the Data Channel? Something like #38422 perhaps (gemini assisted). I am a bit hesitant with us rewriting queue internals, but I can take a closer look if you think this approach would be better. |
tvalentyn
left a comment
tvalentyn
left a comment
There was a problem hiding this comment.
If we want to be conservative we could use regular queue and limit by 10k elements if that resolves this failure mode for the customer. but it might be helpful to have this utility class.
| return True | ||
| return False | ||
|
|
||
| def put(self, item, block=True, timeout=None): |
There was a problem hiding this comment.
It may be worth calling out in the docstring that we don't guarantee that the element will land as soon as enough space opens up, since https://github.com/python/cpython/blob/45c47d26c230086163ac1ef0aa9f955f794fb69c/Lib/queue.py#L214-L215 will wake up one random thread that is waiting, which might not be the one that can fit. this is fine as long as we are continuously emptying the queue
There was a problem hiding this comment.
I changed the queue to be fair as I was concerned that if we are limiting we could end up starving a bundle producing large items if there was a bundle producing small items continuously.
| """A queue.Queue that limits by both element count and total weight. | ||
|
|
||
| A single element is allowed to exceed the maxweight to avoid deadlock. | ||
| Note that shutdown is only supported after there are no more put calls. |
There was a problem hiding this comment.
is this because it is simpler or because we don't want to shutdown earlier to avoid data loss?
we could also raise NotImplemented error on shutdown if we don't want to implement this contract.
There was a problem hiding this comment.
I changed to not subclass queue.Queue. We weren't using the interface and relying overriding some of internals was messy, in particular with shutdown which was added within the range of python versions beam supports. I also wanted the queue to be fair for reasons above.
I removed shutdown and changed the test to use a poison pill like data_plane does.
Yes it does show in the logs as long bundle processing. However perhaps that is correct in that the bundle is slow to process due to this. I think that instead of trying to hide that with the runner we should just improve the lull logger to note that is the reason specifically instead of the stack trace.
A downside of throttling the input is that it isn't sufficient if a single input which is already being processed is outputting a lot of elements. I think that can be the case if we are processing a state-backed iterable as the input element would be KV<k, iterable> where the iterable may be large. Or the input could be filenames and outputs are the file contents etc. So I think blocking here is ok but we can try to improve the logging. |
|
/gemini review |
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
Since the size of the items can vary a lot, I think that it is better to try to count the bytes. |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces ByteLimitedQueue, a thread-safe queue that limits capacity based on both element count and total byte size, and integrates it into the _GrpcDataChannel to manage memory more effectively. The implementation includes a Cython interface and comprehensive unit tests. Feedback was provided regarding the use of int in the Cython definition, suggesting Py_ssize_t to prevent potential overflow issues when handling large data volumes.
|
@tvalentyn Could you take another look? I changed it to no-longer subclass since that was a little gross and tricky to support shutdown properly as it was added within the python version range beam supports. I made publishing fair to deal with the wakeup issue you noticed as well as preventing starvation if we have a large output bundle and small output bundle continually outputting. Since it's no longer a subclass, a pxd file was added to help performance. |
|
errors appear unrelated. |
| raise ValueError('Unexpected output element type %s' % type(stream)) | ||
| yield beam_fn_api_pb2.Elements(data=data_stream, timers=timer_stream) | ||
|
|
||
| def _get_element_bytes(self, element): |
There was a problem hiding this comment.
_get_element_bytes sounds as though it returns raw bytes, let's name _get_element_size_bytes ?
| if self.max_weight > 0 and self._byte_size + item_size > self.max_weight: | ||
| return True | ||
| return False | ||
| def put(self, item, item_bytes, block=True, timeout=None): |
There was a problem hiding this comment.
nit: consider marking block and timeout as keyword-only args.
def put(self, item, item_bytes, *, block=True, timeout=None)
| self.assertIn(('t4', 'success'), status) | ||
| self.assertIn(('t5', 'timeout'), status) | ||
|
|
||
|
|
There was a problem hiding this comment.
nit: you could also have a test for get with a timeout
| with self._mutex: | ||
| return len(self._queue) | ||
|
|
||
| def _is_full_locked(self, item_bytes): |
There was a problem hiding this comment.
nit: since it is used negated in the code, consider instead naming it _can_fit, reversing the logic below and removing negations.
|
|
||
| # cython: overflowcheck=True | ||
|
|
||
| cdef class ByteLimitedQueue(object): |
There was a problem hiding this comment.
I think we need to add the .py counterpart to
Line 360 in fba639a
you can also test performance diffs with regular queue to make sure no surprises there.
2 participants
Previously an unbounded queue was used for pending data outputs to
be sent over the fnapi to the runner. If outputs were being generated
faster than the runner was consuming them, this would lead to memory
growth and possible OOMs. This PR introduces a byte-limited queue
data structure that is used instead to limit the # of bytes in the
queue. This was preferred to just using a queue with max number of
elements because the size of elements can vary greatly. For batch
pipelines they are likely large while for stremaing pipelines there
may be more small outputs.
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, commentfixes #<ISSUE NUMBER>instead.CHANGES.mdwith noteworthy changes.See the Contributor Guide for more tips on how to make review process smoother.
To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md
GitHub Actions Tests Status (on master branch)
See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.