Skip to content

Async Post Load Test Fixes #36596

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 8 commits into
base: master
Choose a base branch
from
+37 −35
Open

Async Post Load Test Fixes #36596

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
2364455
Update async_dofn.py
dustin12 Oct 21, 2025
9e3fd6d
Formatting fixes
dustin12 Oct 22, 2025
0a491a5
extra tab
dustin12 Oct 22, 2025
cb533af
import random and make randomness deterministic per key.
dustin12 Oct 22, 2025
4192812
Fix unit test failure caused by passing timeout to a function when th…
dustin12 Oct 23, 2025
badeb17
Fix unit tests now that we don't block
dustin12 Oct 24, 2025
fe24a7f
formatting
dustin12 Oct 24, 2025
d9eff31
re add sleep so message submission doesn't race with processing in th…
dustin12 Oct 24, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 20 additions & 13 deletions sdks/python/apache_beam/transforms/async_dofn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
from __future__ import absolute_import

import logging
import random
import uuid
from concurrent.futures import ThreadPoolExecutor
from math import floor
Expand Down Expand Up @@ -55,9 +56,8 @@ class AsyncWrapper(beam.DoFn):
TIMER_SET = ReadModifyWriteStateSpec('timer_set', coders.BooleanCoder())
TO_PROCESS = BagStateSpec(
'to_process',
coders.TupleCoder([coders.StrUtf8Coder(), coders.StrUtf8Coder()]),
)
_timer_frequency = 20
coders.TupleCoder(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a backwards incompatible change, since you're swapping to a different coder

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There shouldn't be any usage of this yet so I'm OK with that.

[coders.FastPrimitivesCoder(), coders.FastPrimitivesCoder()]))
# The below items are one per dofn (not instance) so are maps of UUID to
# value.
_processing_elements = {}
Expand Down Expand Up @@ -103,7 +103,7 @@ def __init__(
self._uuid = uuid.uuid4().hex
self._parallelism = parallelism
self._max_wait_time = max_wait_time
self._timer_frequency = 20
self._timer_frequency = callback_frequency
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As best I can tell, self._timer_frequency is used but self.timer_frequency_ is not. Is there any reason to have both? Same goes for all of these duped fields

if max_items_to_buffer is None:
self._max_items_to_buffer = max(parallelism * 2, 10)
else:
Expand All @@ -114,7 +114,6 @@ def __init__(
self.max_wait_time = max_wait_time
self.timer_frequency_ = callback_frequency
self.parallelism_ = parallelism
self._next_time_to_fire = Timestamp.now() + Duration(seconds=5)
self._shared_handle = Shared()

@staticmethod
Expand Down Expand Up @@ -238,9 +237,10 @@ def schedule_item(self, element, ignore_buffer=False, *args, **kwargs):
**kwargs: keyword arguments that the wrapped dofn requires.
"""
done = False
sleep_time = 1
sleep_time = 0.01
total_sleep = 0
while not done:
timeout = 1
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the timeout duration could be configurable in __init__()

while not done and total_sleep < timeout:
done = self.schedule_if_room(element, ignore_buffer, *args, **kwargs)
if not done:
sleep_time = min(self.max_wait_time, sleep_time * 2)
Expand All @@ -256,10 +256,12 @@ def schedule_item(self, element, ignore_buffer=False, *args, **kwargs):
total_sleep += sleep_time
sleep(sleep_time)

def next_time_to_fire(self):
def next_time_to_fire(self, key):
random.seed(key)
return (
floor((time() + self._timer_frequency) / self._timer_frequency) *
self._timer_frequency)
self._timer_frequency) + (
random.random() * self._timer_frequency)
Comment on lines +259 to +264
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I feel like doing all of the work to find a round increment of _timer_frequency is wasted compute once you add the extra fuzziness of random.random() * self._timer_frequency since you're no longer on a round increment afterwards

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I started just having keys setting a timer now + 10s. That doesn't work because as new work arrives the timer firing time keeps getting pushed out. ie. an element arrives at t=1, we want to check back on it at t=11 so we set the timer, but then an element arrives at t=9 and overwrites the timer to t=19.

Next setup was having this round increment firing time. so any message that arrives between t=0 and t=10 sets the timer for 0:10. That way the element at t=9 doesn't override the timer to t=19 but keeps it at t=10.

That works but means we see a spike of timers at t=10, t=20, t=30 etc. There isn't any reason the timers all need to fire at these round increments so this is attempting to add fuzzing per key (since timers are per key). Ideally this means that any 1 key has buckets 10s apart so the overwriting problem is fixed but also means that across multiple keys the buckets don't all fire at the same time. I believe this is what the random.seed(key) on line 260 is doing but correct me if I'm wrong.

Also, let me know if you know an easier way to obtain this pattern.


def accepting_items(self):
with AsyncWrapper._lock:
Expand Down Expand Up @@ -301,7 +303,7 @@ def process(
# Set a timer to fire on the next round increment of timer_frequency_. Note
# we do this so that each messages timer doesn't get overwritten by the
# next.
time_to_fire = self.next_time_to_fire()
time_to_fire = self.next_time_to_fire(element[0])
timer.set(time_to_fire)

# Don't output any elements. This will be done in commit_finished_items.
Expand Down Expand Up @@ -346,6 +348,7 @@ def commit_finished_items(
# from local state and cancel their futures.
to_remove = []
key = None
to_reschedule = []
if to_process_local:
key = str(to_process_local[0][0])
else:
Expand Down Expand Up @@ -387,9 +390,13 @@ def commit_finished_items(
'item %s found in processing state but not local state,'
' scheduling now',
x)
self.schedule_item(x, ignore_buffer=True)
to_reschedule.append(x)
items_rescheduled += 1

# Reschedule the items not under a lock
for x in to_reschedule:
self.schedule_item(x, ignore_buffer=False)

# Update processing state to remove elements we've finished
to_process.clear()
for x in to_process_local:
Expand All @@ -408,8 +415,8 @@ def commit_finished_items(
# If there are items not yet finished then set a timer to fire in the
# future.
self._next_time_to_fire = Timestamp.now() + Duration(seconds=5)
if items_not_yet_finished > 0:
time_to_fire = self.next_time_to_fire()
if items_in_processing_state > 0:
time_to_fire = self.next_time_to_fire(key)
timer.set(time_to_fire)

# Each result is a list. We want to combine them into a single
Expand Down
39 changes: 17 additions & 22 deletions sdks/python/apache_beam/transforms/async_dofn_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -343,10 +343,15 @@ def add_item(i):
self.assertEqual(async_dofn._max_items_to_buffer, 5)
self.check_items_in_buffer(async_dofn, 5)

# After 55 seconds all items should be finished (including those which were
# waiting on the buffer).
# Wait for all buffered items to finish.
self.wait_for_empty(async_dofn, 100)
# This will commit buffered items and add new items which didn't fit in the
# buffer.
result = async_dofn.commit_finished_items(fake_bag_state, fake_timer)

# Wait for the new buffered items to finish.
self.wait_for_empty(async_dofn, 100)
result.extend(async_dofn.commit_finished_items(fake_bag_state, fake_timer))
self.check_output(result, expected_output)
self.check_items_in_buffer(async_dofn, 0)

Expand Down Expand Up @@ -414,33 +419,23 @@ def add_item(i):
# Run for a while. Should be enough to start all items but not finish them
# all.
time.sleep(random.randint(30, 50))
# Commit some stuff
pre_crash_results = []
for i in range(0, 10):
pre_crash_results.append(
async_dofn.commit_finished_items(
bag_states['key' + str(i)], timers['key' + str(i)]))

# Wait for all items to at least make it into the buffer.
done = False
results = [[] for _ in range(0, 10)]
while not done:
time.sleep(10)
done = True
for future in futures:
if not future.done():
for i in range(0, 10):
results[i].extend(
async_dofn.commit_finished_items(
bag_states['key' + str(i)], timers['key' + str(i)]))
if not bag_states['key' + str(i)].items:
self.check_output(results[i], expected_outputs['key' + str(i)])
else:
done = False
break

# Wait for all items to finish.
self.wait_for_empty(async_dofn)
time.sleep(random.randint(10, 30))

for i in range(0, 10):
result = async_dofn.commit_finished_items(
bag_states['key' + str(i)], timers['key' + str(i)])
logging.info('pre_crash_results %s', pre_crash_results[i])
logging.info('result %s', result)
self.check_output(
pre_crash_results[i] + result, expected_outputs['key' + str(i)])
self.check_output(results[i], expected_outputs['key' + str(i)])
self.assertEqual(bag_states['key' + str(i)].items, [])


Expand Down
Loading