[Dynamo] Guard serialization for TUPLE_ITERATOR_LEN (pytorch#152865)

Tests serialization for TUPLE_ITERATOR_LEN; includes no non-test changes.

Passing a tuple iterator as input results in the iterator being exhausted during testing. I threw together a super janky workaround via accepting a func for kwarg generation and replacing `frame_state.f_locals` with newly-generated kwargs to get fresh iterators, but insights into a better approach are welcome!

Pull Request resolved: pytorch#152865
Approved by: https://github.com/jansel
ghstack dependencies: pytorch#152725, pytorch#152727, pytorch#152728, pytorch#152730

Author: jbschlosser

test/dynamo/test_guard_serialization.py

@@ -5,6 +5,7 @@
 import pickle
 import sys
 import types
+from collections.abc import Iterator
 from unittest.mock import patch
 
 import torch
@@ -239,6 +240,11 @@ def _tracefunc(self, frame, event, arg):
         )
 
     def _test_serialization(self, guard_type, fn, *args, **kwargs):
+        # kwargs might contain a callable that generates kwargs
+        kwarg_gen_fn = kwargs.get("_gen_fn", None)
+        if kwarg_gen_fn is not None:
+            kwargs = kwarg_gen_fn()
+
         self._frame_state = None
         sys.settrace(self._tracefunc)
         if isinstance(fn, torch.nn.Module):
@@ -250,6 +256,14 @@ def _test_serialization(self, guard_type, fn, *args, **kwargs):
 
         assert self._frame_state is not None
 
+        # Set f_locals from regenerated kwargs to handle exhausted input iterators
+        # NB: This is super janky and might cause unforeseen problems
+        if kwarg_gen_fn is not None:
+            kwargs = kwarg_gen_fn()
+            for key in self._frame_state.f_locals.keys():
+                if key in kwargs and isinstance(kwargs[key], Iterator):
+                    self._frame_state.f_locals[key] = kwargs[key]
+
         def guard_filter_fn(guards):
             ret = [
                 g.guard_type == guard_type or guard_type in g.derived_guard_types
@@ -688,6 +702,41 @@ def fn(t, x):
             False,
         )
 
+    def test_tuple_iterator_len(self):
+        def fn(t, x):
+            if len(list(t)) > 2:
+                return x * 2
+            return x + 1
+
+        tup = (1, 2, 3)
+        x = torch.randn(3)
+
+        # func to generate kwargs; useful for avoiding iterator exhaustion issues
+        def _gen_kwargs(tup=tup, x=x):
+            return {"t": iter(tup), "x": x}
+
+        ref, loaded = self._test_serialization(
+            "TUPLE_ITERATOR_LEN", fn, _gen_fn=_gen_kwargs
+        )
+
+        # same tuple
+        self._test_check_fn(ref, loaded, {"t": iter(tup), "x": x}, True)
+        self._test_check_fn(ref, loaded, {"t": iter(tup), "x": torch.randn(4)}, True)
+        # same length tuple, different contents
+        self._test_check_fn(ref, loaded, {"t": iter((3, 2, 1)), "x": x}, True)
+        self._test_check_fn(
+            ref, loaded, {"t": iter((3, 2, 1)), "x": torch.randn(4)}, True
+        )
+        # different tuple lengths
+        self._test_check_fn(ref, loaded, {"t": iter((1, 2)), "x": x}, False)
+        self._test_check_fn(
+            ref, loaded, {"t": iter((1, 2)), "x": torch.randn(4)}, False
+        )
+        self._test_check_fn(ref, loaded, {"t": iter((1, 2, 3, 4)), "x": x}, False)
+        self._test_check_fn(
+            ref, loaded, {"t": iter((1, 2, 3, 4)), "x": torch.randn(4)}, False
+        )
+
     def test_dict_version(self):
         def fn(x):
             return pytree.tree_leaves(x)[0] + 1