[dynamo] Better error for invalid @contextlib.contextmanager usage (#156924)

Fixes #156716

Pull Request resolved: https://github.com/pytorch/pytorch/pull/156924
Approved by: https://github.com/williamwen42

torch/_dynamo/graph_break_registry.json

@@ -2177,5 +2177,16 @@
         "INFO"
       ]
     }
+  ],
+  "GB0221": [
+    {
+      "Gb_type": "non-generator contextlib.contextmanager",
+      "Context": "str(self.vt.get_code())",
+      "Explanation": "Cannot compile function decorated with `@contextlib.contextmanager` that is not a generator, i.e. does not use `yield`",
+      "Hints": [
+        "Use `yield` in the function body instead of `return`.",
+        "Remove the `@contextlib.contextmanager` decorator."
+      ]
+    }
   ]
 }

torch/_dynamo/variables/functions.py

@@ -942,7 +942,17 @@ class LocalGeneratorFunctionVariable(BaseUserFunctionVariable):
         args: "list[VariableTracker]",
         kwargs: "dict[str, VariableTracker]",
     ) -> "VariableTracker":
-        assert is_generator(self.vt.get_code())
+        if not is_generator(self.vt.get_code()):
+            unimplemented_v2(
+                gb_type="non-generator contextlib.contextmanager",
+                context=str(self.vt.get_code()),
+                explanation="Cannot compile function decorated with `@contextlib.contextmanager` that is not a generator"
+                ", i.e. does not use `yield`",
+                hints=[
+                    "Use `yield` in the function body instead of `return`.",
+                    "Remove the `@contextlib.contextmanager` decorator.",
+                ],
+            )
 
         inline_tracer = self._build_inline_tracer(tx, args, kwargs)
         code = self.vt.get_code()