Fix stateful DoFn runtime typecheck duplicate specs error

Author: lakshitbahl

sdks/python/apache_beam/typehints/typecheck.py

@@ -100,11 +100,15 @@ class RuntimeTypeCheckWrapperDoFn(AbstractDoFnWrapper):
   def __init__(self, dofn, type_hints, full_label):
     super().__init__(dofn)
     self.full_label = full_label
-    self._process_fn = self.dofn._process_argspec_fn()
+    # Note that the bound process method must not be cached on the instance:
+    # an attribute holding a bound method is visible to the stateful DoFn
+    # reflection in userstate.get_dofn_specs, and a cached copy can diverge
+    # from self.dofn.process across (de)serialization, in which case stateful
+    # DoFn validation would see duplicate StateSpecs/TimerSpecs.
     if type_hints.input_types:
       input_args, input_kwargs = type_hints.input_types
       self._input_hints = getcallargs_forhints(
-          self._process_fn, *input_args, **input_kwargs)
+          dofn._process_argspec_fn(), *input_args, **input_kwargs)
     else:
       self._input_hints = None
     # TODO(robertwb): Multi-output.
@@ -121,7 +125,8 @@ def wrapper(self, method, args, kwargs):
   def process(self, *args, **kwargs):
     try:
       if self._input_hints:
-        actual_inputs = inspect.getcallargs(self._process_fn, *args, **kwargs)  # pylint: disable=deprecated-method
+        actual_inputs = inspect.getcallargs(
+            self.dofn._process_argspec_fn(), *args, **kwargs)  # pylint: disable=deprecated-method
         for var, hint in self._input_hints.items():
           if hint is actual_inputs[var]:
             # self parameter

sdks/python/apache_beam/typehints/typecheck_test.py

@@ -38,6 +38,8 @@
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.util import assert_that
 from apache_beam.testing.util import equal_to
+from apache_beam.transforms import combiners
+from apache_beam.transforms import userstate
 from apache_beam.typehints import decorators
 from apache_beam.typehints import typecheck
 from apache_beam.typehints import with_input_types
@@ -348,5 +350,52 @@ def test_wrapper_preserves_results(self):
     self.assertEqual(list(dofn.process(1)), [2])
 
 
+def _make_stateful_dofn():
+  # Defined inside a function so that the class is serialized by value
+  # (like the DoFn created by GroupIntoBatches), which is the case where a
+  # cached bound method diverges from the reconstructed class's process.
+  count_state = userstate.CombiningValueStateSpec(
+      'count', combiners.CountCombineFn())
+
+  class _CountingStatefulDoFn(beam.DoFn):
+    def process(self, element, count=beam.DoFn.StateParam(count_state)):
+      count.add(1)
+      yield element[1]
+
+  return _CountingStatefulDoFn()
+
+
+class RuntimeTypeCheckStatefulDoFnTest(unittest.TestCase):
+  """Regression tests for runtime_type_check with stateful DoFns.
+
+  The type-check wrapper must not expose duplicate StateSpecs/TimerSpecs to
+  stateful DoFn validation, including after the runner serializes and
+  reconstructs the wrapped DoFn.
+  """
+  def test_wrapper_does_not_cache_bound_process(self):
+    dofn = _make_stateful_dofn()
+    wrapper = typecheck.RuntimeTypeCheckWrapperDoFn(
+        dofn, dofn.get_type_hints(), 'Step')
+    # A bound method cached in the instance __dict__ is visible to
+    # userstate.get_dofn_specs and can diverge from self.dofn.process
+    # across (de)serialization.
+    for attr, value in wrapper.__dict__.items():
+      self.assertFalse(
+          hasattr(value, '__self__'),
+          'wrapper caches bound method %r, which breaks stateful DoFn '
+          'validation' % attr)
+    userstate.validate_stateful_dofn(wrapper)
+
+  def test_stateful_dofn_with_runtime_type_check(self):
+    options = PipelineOptions()
+    options.view_as(TypeOptions).runtime_type_check = True
+    with TestPipeline(options=options) as p:
+      result = (
+          p
+          | beam.Create([('k', 1), ('k', 2), ('k', 3)])
+          | beam.ParDo(_make_stateful_dofn()))
+      assert_that(result, equal_to([1, 2, 3]))
+
+
 if __name__ == '__main__':
   unittest.main()