Fix stateful DoFn runtime typecheck duplicate specs error

Author: lakshitbahl

sdks/python/apache_beam/typehints/typecheck.py

@@ -84,25 +84,81 @@ def teardown(self):
     return self.dofn.teardown()
 
 
-class OutputCheckWrapperDoFn(AbstractDoFnWrapper):
-  """A DoFn that verifies against common errors in the output type."""
-  def __init__(self, dofn, full_label):
+class RuntimeTypeCheckWrapperDoFn(AbstractDoFnWrapper):
+  """A DoFn wrapper performing runtime type-checking of inputs and outputs.
+
+  This single wrapper performs the work that was previously split between
+  two nested wrappers (``TypeCheckWrapperDoFn`` wrapped inside
+  ``OutputCheckWrapperDoFn``): type-checking of inputs and outputs against
+  declared type hints, verification against common output errors (such as
+  returning a plain ``str``, ``bytes`` or ``dict`` from a ``ParDo``), and
+  labeling of raised ``TypeCheckError`` messages with the offending
+  transform's full label. Merging the wrappers removes one level of Python
+  call indirection per processed element when ``--runtime_type_check`` is
+  enabled.
+  """
+  def __init__(self, dofn, type_hints, full_label):
     super().__init__(dofn)
     self.full_label = full_label
+    # Note that a *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. Caching the
+    # underlying (unbound) function is safe: plain functions stored on an
+    # instance are not bound methods, so DoFn reflection ignores them.
+    process_fn = dofn._process_argspec_fn()
+    if hasattr(process_fn, '__func__'):
+      self._process_fn = process_fn.__func__
+      self._process_fn_needs_self = True
+    else:
+      self._process_fn = process_fn
+      self._process_fn_needs_self = False
+    if type_hints.input_types:
+      input_args, input_kwargs = type_hints.input_types
+      self._input_hints = getcallargs_forhints(
+          process_fn, *input_args, **input_kwargs)
+    else:
+      self._input_hints = None
+    # TODO(robertwb): Multi-output.
+    self._output_type_hint = type_hints.simple_output_type(full_label)
 
   def wrapper(self, method, args, kwargs):
     try:
-      result = method(*args, **kwargs)
+      result = self._type_check_result(method(*args, **kwargs))
     except TypeCheckError as e:
-      # TODO(BEAM-10710): Remove the 'ParDo' prefix for the label name
-      error_msg = (
-          'Runtime type violation detected within ParDo(%s): '
-          '%s' % (self.full_label, e))
-      _, _, tb = sys.exc_info()
-      raise TypeCheckError(error_msg).with_traceback(tb)
+      raise self._add_label(e)
+    else:
+      return self._check_type(result)
+
+  def process(self, *args, **kwargs):
+    try:
+      if self._input_hints:
+        if self._process_fn_needs_self:
+          actual_inputs = inspect.getcallargs(
+              self._process_fn, self.dofn, *args, **kwargs)  # pylint: disable=deprecated-method
+        else:
+          actual_inputs = inspect.getcallargs(self._process_fn, *args, **kwargs)  # pylint: disable=deprecated-method
+        for var, hint in self._input_hints.items():
+          if hint is actual_inputs[var]:
+            # self parameter
+            continue
+          _check_instance_type(hint, actual_inputs[var], var, True)
+      result = self._type_check_result(self.dofn.process(*args, **kwargs))
+    except TypeCheckError as e:
+      raise self._add_label(e)
     else:
       return self._check_type(result)
 
+  def _add_label(self, e):
+    """Returns a TypeCheckError labeled with this transform's full label."""
+    # TODO(BEAM-10710): Remove the 'ParDo' prefix for the label name
+    error_msg = (
+        'Runtime type violation detected within ParDo(%s): '
+        '%s' % (self.full_label, e))
+    _, _, tb = sys.exc_info()
+    return TypeCheckError(error_msg).with_traceback(tb)
+
   @staticmethod
   def _check_type(output):
     if output is None:
@@ -120,36 +176,6 @@ def _check_type(output):
           'iterable. %s was returned instead.' % type(output))
     return output
 
-
-class TypeCheckWrapperDoFn(AbstractDoFnWrapper):
-  """A wrapper around a DoFn which performs type-checking of input and output.
-  """
-  def __init__(self, dofn, type_hints, label=None):
-    super().__init__(dofn)
-    self._process_fn = self.dofn._process_argspec_fn()
-    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)
-    else:
-      self._input_hints = None
-    # TODO(robertwb): Multi-output.
-    self._output_type_hint = type_hints.simple_output_type(label)
-
-  def wrapper(self, method, args, kwargs):
-    result = method(*args, **kwargs)
-    return self._type_check_result(result)
-
-  def process(self, *args, **kwargs):
-    if self._input_hints:
-      actual_inputs = inspect.getcallargs(self._process_fn, *args, **kwargs)  # pylint: disable=deprecated-method
-      for var, hint in self._input_hints.items():
-        if hint is actual_inputs[var]:
-          # self parameter
-          continue
-        _check_instance_type(hint, actual_inputs[var], var, True)
-    return self._type_check_result(self.dofn.process(*args, **kwargs))
-
   def _type_check_result(self, transform_results):
     if self._output_type_hint is None or transform_results is None:
       return transform_results
@@ -289,11 +315,9 @@ def visit_transform(self, applied_transform):
         if isinstance(transform.fn, core.CombineValuesDoFn):
           transform.fn.combinefn = self._wrapped_fn
       else:
-        transform.fn = transform.dofn = OutputCheckWrapperDoFn(
-            TypeCheckWrapperDoFn(
-                transform.fn,
-                transform.get_type_hints(),
-                applied_transform.full_label),
+        transform.fn = transform.dofn = RuntimeTypeCheckWrapperDoFn(
+            transform.fn,
+            transform.get_type_hints(),
             applied_transform.full_label)
 
 

sdks/python/apache_beam/typehints/typecheck_test.py

@@ -34,10 +34,14 @@
 from apache_beam import Pipeline
 from apache_beam.options.pipeline_options import PipelineOptions
 from apache_beam.options.pipeline_options import TypeOptions
+from apache_beam.pipeline import PipelineVisitor
 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
 from apache_beam.typehints import with_output_types
 
@@ -105,7 +109,7 @@ def test_wrapper_pass_through(self):
     # not the same one that actually runs in the pipeline (it is serialized
     # here and deserialized in the worker).
     with tempfile.TemporaryDirectory() as tmp_dirname:
-      path = os.path.join(tmp_dirname + "tmp_filename")
+      path = os.path.join(tmp_dirname, "tmp_filename")
       dofn = MyDoFn(path)
       result = self.p | beam.Create([1, 2, 3]) | beam.ParDo(dofn)
       assert_that(result, equal_to([1, 2, 3]))
@@ -300,5 +304,98 @@ def process(self, element, *args, **kwargs):
       self.p.run().wait_until_finish()
 
 
+@with_input_types(int)
+@with_output_types(int)
+class _AddOneDoFn(beam.DoFn):
+  def process(self, element):
+    yield element + 1
+
+
+class RuntimeTypeCheckWrapperDoFnTest(unittest.TestCase):
+  """Tests for the merged runtime type-checking wrapper (BEAM-9489)."""
+  def test_visitor_applies_single_wrapper_layer(self):
+    # A single RuntimeTypeCheckWrapperDoFn should wrap the user's DoFn
+    # directly, with no intermediate wrapper layers.
+    p = beam.Pipeline(options=PipelineOptions(runtime_type_check=True))
+    _ = (p | beam.Create([1, 2]) | 'TypedStep' >> beam.ParDo(_AddOneDoFn()))
+    p.visit(typecheck.TypeCheckVisitor())
+
+    wrapped_dofns = []
+
+    class _Collector(PipelineVisitor):
+      def visit_transform(self, applied_transform):
+        transform = applied_transform.transform
+        if isinstance(transform, beam.ParDo) and isinstance(
+            getattr(transform, 'fn', None), typecheck.AbstractDoFnWrapper):
+          wrapped_dofns.append(transform.fn)
+
+    p.visit(_Collector())
+    self.assertTrue(wrapped_dofns)
+    for wrapper in wrapped_dofns:
+      self.assertIsInstance(wrapper, typecheck.RuntimeTypeCheckWrapperDoFn)
+      # The wrapped DoFn must be the user's DoFn, not another wrapper.
+      self.assertNotIsInstance(wrapper.dofn, typecheck.AbstractDoFnWrapper)
+
+  def test_wrapper_labels_type_check_errors(self):
+    dofn = typecheck.RuntimeTypeCheckWrapperDoFn(
+        _AddOneDoFn(), _AddOneDoFn().get_type_hints(), 'MyStep')
+    with self.assertRaisesRegex(
+        typecheck.TypeCheckError,
+        r'Runtime type violation detected within ParDo\(MyStep\)'):
+      dofn.process('not-an-int')
+
+  def test_wrapper_preserves_results(self):
+    dofn = typecheck.RuntimeTypeCheckWrapperDoFn(
+        _AddOneDoFn(), _AddOneDoFn().get_type_hints(), 'MyStep')
+    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()