[BEAM -7741] Implement SetState for Python SDK #9090

rakeshcusat · 2019-07-17T15:04:36Z

SetState was implemented in Java sdk, to ensure the feature parity SetState is also implemented in Python SDK.

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

Choose reviewer(s) and mention them in a comment (R: @username).
Format the pull request title like [BEAM-XXX] Fixes bug in ApproximateQuantiles, where you replace BEAM-XXX with the appropriate JIRA issue, if applicable. This will automatically link the pull request to the issue.
If this contribution is large, please file an Apache Individual Contributor License Agreement.

Post-Commit Tests Status (on master branch)

Lang	Apex	Dataflow	Gearpump	Samza
Go	---	---	---	---
Java
Python	---		---	---

Pre-Commit Tests Status (on master branch)

---	Java	Python	Go	Website
Non-portable
Portable	---		---	---

See .test-infra/jenkins/README for trigger phrase, status and link of all Jenkins jobs.

merging from upstream

rakeshcusat · 2019-07-17T15:05:15Z

R: @robertwb @pabloem

rakeshcusat · 2019-07-17T15:08:58Z

Run Portable_Python PreCommit

rakeshcusat · 2019-07-17T15:18:49Z

R: @angoenka

rakeshcusat · 2019-07-17T16:10:39Z

Run Python_PVR_Flink PreCommit

robertwb

Thanks for looking at this.

robertwb · 2019-07-18T12:59:41Z

sdks/python/apache_beam/transforms/userstate.py

+  def __init__(self, state_spec, state_tag, current_value_accessor):
+    super(SetRuntimeState, self).__init__(
+        state_spec, state_tag, current_value_accessor)
+    # TODO: What is current_value_accessor? where does cached value is stored?


I think it gets the underlying value from the runner. It looks like it's only used in the direct runner.

robertwb · 2019-07-18T13:00:28Z

sdks/python/apache_beam/transforms/userstate_test.py

+      def process(self, element, buffer=DoFn.StateParam(BUFFER_STATE),
+                  timer1=DoFn.TimerParam(EXPIRY_TIMER)):
+        unused_key, value = element
+        buffer.add(str(value).encode('latin1'))


Can we just use StringCoder and avoid the explicit encoding? Or if the input is integers, just store the int itself?

sure, I will take care of this.

robertwb · 2019-07-18T13:02:56Z

sdks/python/apache_beam/transforms/userstate_test.py

+
+      @on_timer(EXPIRY_TIMER)
+      def expiry_callback(self, buffer=DoFn.StateParam(BUFFER_STATE)):
+        yield b''.join(sorted(buffer.read()))


Just yield sorted(buffer.read()).

robertwb · 2019-07-18T13:05:43Z

sdks/python/apache_beam/transforms/userstate_test.py

+      def emit_values(self, bag_state=beam.DoFn.StateParam(SET_STATE)):
+        for value in bag_state.read():
+          yield value
+        yield 'extra'


This shouldn't be needed.

robertwb · 2019-07-18T13:06:01Z

sdks/python/apache_beam/transforms/userstate_test.py

+       | test_stream
+       | beam.Map(lambda x: ('mykey', x))
+       | beam.ParDo(SimpleTestSetStatefulDoFn())
+       | beam.ParDo(self.record_dofn()))


Is there a reason assert_that doesn't work here?

I haven't tried the assert_that here. I took reference from already existing test cases. I will update the test cases.

assert_that doesn't work in case of state and timer. Here is the jira ticket filed long back: https://issues.apache.org/jira/browse/BEAM-5295 I can take care of this in a separate ticket.

robertwb · 2019-07-18T13:06:21Z

sdks/python/apache_beam/transforms/userstate_test.py

+
+      @on_timer(CLEAR_TIMER)
+      def clear_values(self, bag_state=beam.DoFn.StateParam(SET_STATE)):
+        bag_state.clear()


Maybe do an add, clear, add to make sure it really works.

robertwb · 2019-07-18T13:07:33Z

sdks/python/apache_beam/transforms/userstate_test.py

+    self.assertEqual(
+        ['extra'],
+        StatefulDoFnOnDirectRunnerTest.all_records)
+
  def test_stateful_dofn_nonkeyed_input(self):


Can you add tests without using TestStream to ensure that the FnApiRunner paths also work?

@robertwb do you have an example of the test case?

It'd be pretty much the same as the existing test, but using Create rather than TestStream. Create just emits all elements at MIN_TIMESTAMP, and then you can set timers to perform future actions.

Added some cases to make sure FnApiRunner path is also executed

pabloem · 2019-07-18T17:33:46Z

Run Python PreCommit

robertwb · 2019-07-23T09:55:09Z

sdks/python/apache_beam/transforms/userstate_test.py

+        yield 'extra'
+
+      @on_timer(CLEAR_TIMER)
+      def clear_values(self, bag_state=beam.DoFn.StateParam(SET_STATE)):


s/bag_state/set_state/g

rakeshcusat · 2019-07-24T00:44:31Z

R: @robertwb Can you take another look? I have taken care of your comments.

rakeshcusat · 2019-07-24T04:36:25Z

Run Portable_Python PreCommit

robertwb

Thanks. Just a couple more comments.

robertwb · 2019-07-24T13:08:01Z

sdks/python/apache_beam/runners/worker/bundle_processor.py

+
+  def read(self):
+    return _ConcatIterable(
+        {} if self._cleared else _StateBackedIterable(


{} is a dict, not a set. Also, should we be converting this to a frozen set here?

I still need to take care of this comment

robertwb · 2019-07-24T13:12:51Z

sdks/python/apache_beam/runners/worker/bundle_processor.py

+  def _commit(self):
+    if self._cleared:
+      self._state_handler.blocking_clear(self._state_key)
+    if self._added_elements:


One downside of this way of doing things is that if we add the same element many times from different bundles, the underlying bag state keeps growing. Perhaps we should (periodically?) "compact" by reading everything in, eeduplicating via set(), and then writing it all back out. (A TODO is fine here if you don't want to implement it now.)

I would prefer to have the compaction logic but I will do some research on my end before making the final decision.

@robertwb Check the latest implementation of read() method. I think this is what you recommended. Let me know if you want me to implement it in different way.

Looks good, thanks.

robertwb · 2019-07-24T13:13:45Z

sdks/python/apache_beam/transforms/userstate_test.py

@@ -424,18 +426,18 @@ def clear_values(self, bag_state=beam.DoFn.StateParam(BAG_STATE)):

  def test_simple_set_stateful_dofn(self):
    class SimpleTestSetStatefulDoFn(DoFn):
-      BUFFER_STATE = SetStateSpec('buffer', BytesCoder())
+      BUFFER_STATE = SetStateSpec('buffer', FastPrimitivesCoder())


Use the more specific VarIntCoder here instead.

robertwb · 2019-07-24T13:14:49Z

sdks/python/apache_beam/transforms/userstate_test.py

-        bag_state.clear()
+      def clear_values(self, set_state=beam.DoFn.StateParam(SET_STATE)):
+        set_state.clear()
+        set_state.add('value2')


Add an int, something "different" like 100, for consistency.

Wouldn't SetState give us error because we are using StrUtf8Coder and would be passing an int?

I can pass some different string value instead.

robertwb · 2019-07-24T13:15:31Z

sdks/python/apache_beam/transforms/userstate_test.py

-        StatefulDoFnOnDirectRunnerTest.all_records)
+    self.assertEqual(['value2'], StatefulDoFnOnDirectRunnerTest.all_records)
+
+  def test_stateful_set_state_fn_runner(self):


test_stateful_set_state_portably

robertwb · 2019-07-24T13:17:29Z

sdks/python/apache_beam/transforms/userstate_test.py

+        yield aggregated_value
+
+    p = TestPipeline()
+    values = p | beam.Create([('key', 1), ('key', 2), ('key', 3), ('key', 4)])


Duplicates?

Sure, I will add it.

robertwb · 2019-07-24T13:18:12Z

sdks/python/apache_beam/transforms/userstate_test.py

+        set_state.add(value)
+        for saved_value in set_state.read():
+          aggregated_value += saved_value
+        yield aggregated_value


Maybe set a timer and test clear as well?

I tried to add this test case:

def test_stateful_set_state_clean_portably(self): class SetStateClearingStatefulDoFn(beam.DoFn): SET_STATE = SetStateSpec('buffer', VarIntCoder()) EMIT_TIMER = TimerSpec('emit_timer', TimeDomain.WATERMARK) CLEAR_TIMER = TimerSpec('clear_timer', TimeDomain.WATERMARK) def process(self, element, set_state=beam.DoFn.StateParam(SET_STATE), emit_timer=beam.DoFn.TimerParam(EMIT_TIMER), clear_timer=beam.DoFn.TimerParam(CLEAR_TIMER)): _, value = element set_state.add(value) if value == 2: clear_timer.set(3) elif value == 3: emit_timer.set(4) @on_timer(EMIT_TIMER) def emit_values(self, set_state=beam.DoFn.StateParam(SET_STATE)): yield sorted(set_state.read()) @on_timer(CLEAR_TIMER) def clear_values(self, set_state=beam.DoFn.StateParam(SET_STATE)): set_state.clear() set_state.add(100) p = TestPipeline() values = p | beam.Create(range(1, 5)) actual_values = (values | beam.Map(lambda t: window.TimestampedValue(t, t)) | beam.WindowInto(window.FixedWindows(1)) | beam.ParDo(SetStateClearingStatefulDoFn())) assert_that(actual_values, equal_to([100])) result = p.run() result.wait_until_finish()

but it didn't work it complained that timer is not supported.

(beam) ➜ python git:(BEAM-7741-implement-setstate) ✗ python -m unittest apache_beam.transforms.userstate_test.StatefulDoFnOnDirectRunnerTest.test_stateful_set_state_clean_portably WARNING:root:Key coder FastPrimitivesCoder for transform <ParDo(PTransform) label=[ParDo(SetStateClearingStatefulDoFn)]> with stateful DoFn may not be deterministic. This may cause incorrect behavior for complex key types. Consider adding an input type hint for this transform. E ====================================================================== ERROR: test_stateful_set_state_clean_portably (apache_beam.transforms.userstate_test.StatefulDoFnOnDirectRunnerTest) ---------------------------------------------------------------------- Traceback (most recent call last): File "apache_beam/transforms/userstate_test.py", line 567, in test_stateful_set_state_clean_portably result = p.run() File "apache_beam/testing/test_pipeline.py", line 107, in run else test_runner_api)) File "apache_beam/pipeline.py", line 406, in run self._options).run(False) File "apache_beam/pipeline.py", line 419, in run return self.runner.run_pipeline(self, self._options) File "apache_beam/runners/direct/direct_runner.py", line 128, in run_pipeline return runner.run_pipeline(pipeline, options) File "apache_beam/runners/portability/fn_api_runner.py", line 319, in run_pipeline default_environment=self._default_environment)) File "apache_beam/runners/portability/fn_api_runner.py", line 323, in run_via_runner_api stage_context, stages = self.create_stages(pipeline_proto) File "apache_beam/runners/portability/fn_api_runner.py", line 385, in create_stages use_state_iterables=self._use_state_iterables) File "apache_beam/runners/portability/fn_api_runner_transforms.py", line 489, in create_and_optimize_stages stages = list(phase(stages, pipeline_context)) File "apache_beam/runners/portability/fn_api_runner_transforms.py", line 1168, in inject_timer_pcollections raise NotImplementedError('Timers and side inputs.') NotImplementedError: Timers and side inputs. ---------------------------------------------------------------------- Ran 1 test in 0.207s FAILED (errors=1)

This is a little bit confusing for me because I have used Timer on Flink runner and it works fine. This could be a bug on the portable runner side.

I still have to debug this, I might need some help to debug this.

Hmm... for some reason it looks like it's assuming you're using a timer and a side input in the same DoFn.

@robertwb I spent some time debugging this. It seems like FnApiRunner is running into issues when there are two timers uses.

class TwoTimersDoFn(beam.DoFn): EMIT_TIMER = TimerSpec('emit_timer', TimeDomain.WATERMARK) CLEAR_TIMER = TimerSpec('clear_timer', TimeDomain.WATERMARK) def process(self, element, emit_timer=beam.DoFn.TimerParam(EMIT_TIMER), clear_timer=beam.DoFn.TimerParam(CLEAR_TIMER) ): _, value = element if value == 2: clear_timer.set(3) pass elif value == 3: emit_timer.set(4) pass @on_timer(EMIT_TIMER) def emit_values(self): yield 'emit1' @on_timer(CLEAR_TIMER) def clear_values(self): yield 'clear1' p = TestPipeline() values = p | beam.Create([('key', 1), ('key', 2), ('key', 3), ('key', 4), ('key', 5), ('key', 6)]) actual_values = (values | beam.Map(lambda t: window.TimestampedValue(t, t[1])) | beam.WindowInto(window.FixedWindows(1)) | beam.ParDo(TwoTimersDoFn())) assert_that(actual_values, equal_to([100])) result = p.run() result.wait_until_finish()

I tested it on the master branch and it failed consistently and I get the same error message. I couldn't figure out the root cause but it is caused by two timers for sure. It works fine when there is only one timer used.
I feel I should file a separate jira ticket to address this issue, thoughts?

Thanks for the analysis. Filing a JIRA for the multiple-timer issue sounds like the right thing to do here.

I checked, there is already a Jira ticket filed by @tweise (https://issues.apache.org/jira/browse/BEAM-7074)

rakeshcusat · 2019-07-31T17:29:35Z

@robertwb can you take a another look of this PR? also check this comment: https://github.com/apache/beam/pull/9090/files#r309342482

robertwb

Mostly looks good. Just a couple of small comments.

robertwb · 2019-08-02T11:34:53Z

sdks/python/apache_beam/runners/worker/bundle_processor.py

+  def _commit(self):
+    if self._cleared:
+      self._state_handler.blocking_clear(self._state_key)
+    if self._added_elements:


Looks good, thanks.

robertwb · 2019-08-02T11:37:28Z

sdks/python/apache_beam/transforms/userstate.py

+  """Specification for a user DoFn Set State cell"""
+
+  def __init__(self, name, coder):
+    assert isinstance(name, str)


TypeError would be a better thing to raise here.

I can change this. But other StateSpec are doing the same way. Probably change their assert in a separate PR.

robertwb · 2019-08-02T11:39:31Z

sdks/python/apache_beam/transforms/userstate_test.py

+        set_state.add(value)
+        for saved_value in set_state.read():
+          aggregated_value += saved_value
+        yield aggregated_value


Thanks for the analysis. Filing a JIRA for the multiple-timer issue sounds like the right thing to do here.

robertwb · 2019-08-02T11:41:42Z

sdks/python/apache_beam/transforms/userstate_test.py

+        _, value = element
+        set_state.add(value)
+
+        if value == 5:


We can't, in general, depend on elements being received in order. Perhaps trigger based on the size of the set?

that's good idea. I will change the test case.

rakeshcusat · 2019-08-08T05:56:43Z

@robertwb I have taken care of your comments. Can you approve this now?

rakeshcusat · 2019-08-08T05:57:00Z

R: @robertwb PTAL

robertwb

Thanks!

rakeshcusat and others added 3 commits April 9, 2019 15:55

Merge pull request #1 from apache/master

665570d

merging from upstream

Merge remote-tracking branch 'upstream/master'

f70a239

Merge remote-tracking branch 'upstream/master'

6bba28a

rakeshcusat changed the title ~~[BEAM -7741] Implement Setstate for Python SDK~~ [BEAM -7741] Implement SetState for Python SDK Jul 17, 2019

Merge remote-tracking branch 'upstream/master'

a93ed31

rakeshcusat force-pushed the BEAM-7741-implement-setstate branch from c990453 to 754f4ea Compare July 17, 2019 15:51

robertwb reviewed Jul 18, 2019

View reviewed changes

robertwb reviewed Jul 23, 2019

View reviewed changes

robertwb reviewed Jul 24, 2019

View reviewed changes

rakeshcusat added 9 commits July 31, 2019 22:47

Merge remote-tracking branch 'upstream/master'

7773faf

[BEAM-7741] Implement Setstate for Py SDK

8c68fd1

remove comment

ed797fd

take care of code review comments

9e608f9

improved test cases

4130dd5

fix lint warning

7841b2c

compating the set state regularly

0325de4

fix the unit test case

ede7ed0

fix lint warning

e0dfae1

robertwb reviewed Aug 2, 2019

View reviewed changes

taken care of code review comment

ef1dc33

rakeshcusat force-pushed the BEAM-7741-implement-setstate branch from 78c8a84 to ef1dc33 Compare August 8, 2019 05:55

robertwb approved these changes Aug 8, 2019

View reviewed changes

robertwb merged commit c9e5ea8 into apache:master Aug 8, 2019

robertwb mentioned this pull request Aug 8, 2019

Move DirectRunner specific classes into direct runner package. #9297

Merged

3 tasks

[BEAM -7741] Implement SetState for Python SDK #9090

[BEAM -7741] Implement SetState for Python SDK #9090

Conversation

rakeshcusat commented Jul 17, 2019

Post-Commit Tests Status (on master branch)

Pre-Commit Tests Status (on master branch)

rakeshcusat commented Jul 17, 2019

rakeshcusat commented Jul 17, 2019

rakeshcusat commented Jul 17, 2019

rakeshcusat commented Jul 17, 2019

robertwb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pabloem commented Jul 18, 2019

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rakeshcusat commented Jul 24, 2019

rakeshcusat commented Jul 24, 2019

robertwb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rakeshcusat Jul 31, 2019 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rakeshcusat Aug 5, 2019 • edited Loading

Choose a reason for hiding this comment

rakeshcusat commented Jul 31, 2019

robertwb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rakeshcusat commented Aug 8, 2019

rakeshcusat commented Aug 8, 2019

robertwb left a comment

Choose a reason for hiding this comment

rakeshcusat Jul 31, 2019 •

edited

Loading

rakeshcusat Aug 5, 2019 •

edited

Loading