prodriguezdefino opened a new pull request, #40624:
URL: https://github.com/apache/beam/pull/40624
When the SDK accepts a channel split of a bundle of a stateful stage,
stageState.splitBundle removed the key of every residual element from
inprogressKeys and inprogressKeysByBundle, even when elements of the primary
still carry that key. The residual of that key was then scheduled at once in a
second bundle that ran concurrently with the primary. Prism gives each bundle a
copy of the key's state at bundle start and replaces the key's state with the
bundle's copy at commit, so the last bundle to commit wins. State updates, and
the outputs that depend on them, are lost without any error.
Serializing the processing of a key is the runner's job. The model defines
state per key and window: ReadModifyWriteState "holds a single value per key
and window. The value can be read, modified, and written back atomically"
(beam_runner_api.proto). The stateful processing post says the runner "cannot
parallelize processing of elements for a particular key+window", and the
programming guide says "State is always fully scoped only to the current
processing key". Splitting is valid SDK behaviour (beam_fn_api.proto), and the
Java and Python harnesses accept channel splits on stateful stages.
splitBundle now collects the keys of the primary elements and releases a
residual's key only if no primary element has it. A held key is released when
the primary bundle is persisted, as before. The handling of residual timer
holds is unchanged.
Reproduction with the Python SDK: all elements share one key, and a
per-element sleep makes Prism split the bundle.
class CountOnOneKey(beam.DoFn):
COUNT = ReadModifyWriteStateSpec('count', VarIntCoder())
DONE = TimerSpec('done', TimeDomain.WATERMARK)
def process(self, element, count=beam.DoFn.StateParam(COUNT),
done=beam.DoFn.TimerParam(DONE)):
time.sleep(0.25)
count.write((count.read() or 0) + 1)
done.set(GlobalWindow().max_timestamp())
@on_timer(DONE)
def report(self, count=beam.DoFn.StateParam(COUNT)):
yield count.read()
options = PipelineOptions(['--runner=PrismRunner',
'--prism_location=<prism binary>',
'--environment_type=LOOPBACK'])
with beam.Pipeline(options=options) as p:
(p | beam.Create(range(40)) | beam.Map(lambda i: ('k', i))
| beam.ParDo(CountOnOneKey()) | beam.Map(print))
A correct runner prints 40. Prism built from master printed 20 in 10 of 10
runs; with this change it printed 40 in 10 of 10 runs. With
--prism_log_level=debug, the same pipeline on master logs one split at element
20 of 40, and the residual bundle starts before the primary commits.
TestElementManager_SplitStatefulBundle covers a split inside the only key, a
split inside one of two keys, and a split between keys. The first two fail
without this change.
Not fixed here, follow-ups:
- splitBundle slices the bundle's elements by the data channel index. A
stateful bundle can hold data of one key and timers of another, and timers are
not counted in the data channel, so the slice can misattribute elements. See
#35771.
- startTriggeredBundle, for aggregate stages, does not check inprogressKeys
before it injects a bundle.
--
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]
