Move put_code_state to only trigger on successful compile (#140042)

ezyang · pytorchmergebot · commit af682f3cd712 · 2024-11-08T04:19:50.000Z
Signed-off-by: Edward Z. Yang <ezyang@meta.com> Differential Revision: [D65623081](https://our.internmc.facebook.com/intern/diff/D65623081) Pull Request resolved: #140042 Approved by: https://github.com/markkm
diff --git a/torch/_dynamo/convert_frame.py b/torch/_dynamo/convert_frame.py
@@ -977,6 +977,18 @@ def format_guard_failures() -> str:
         codecache_metrics.clear()
         try:
             guarded_code = compile_inner(code, one_graph, hooks, transform)
+
+            # NB: We only put_code_state in success case.  Success case here
+            # does include graph breaks; specifically, if a graph break still
+            # resulted in a partially compiled graph, we WILL return here.  An
+            # Unsupported exception will only bubble to the top level if we
+            # are unable to compile the frame at all.  In this case, there's
+            # no point in uploading the code state, because we will always
+            # fail exactly the same way even without the update.  (It's useful
+            # to upload for graph break though, because this can prevent
+            # extra graph break compilations.)
+            put_code_state()
+
             return guarded_code
         except Exception as e:
             fail_type = type(e).__qualname__
@@ -1017,7 +1029,13 @@ def format_guard_failures() -> str:
                     f"{type(e).__qualname__}: {str(e)}"
                 ).with_traceback(e.__traceback__) from None
         finally:
-            put_code_state()
+            # === WARNING WARNING WARNING ===
+            # If you commit a bug here, it will suppress writing to
+            # dynamo_compile table, and we will not have telemetry.
+            # Be extra careful when making changes here!
+            #
+            # TODO to masnesral: feel free to delete these comments
+            # to resolve any merge conflict you have
 
             if tracer:
                 tracer.output.local_scope = {}
@@ -1182,6 +1200,7 @@ def clean_for_json(d: Dict[str, Any]) -> Dict[str, Any]:
             chromium_event_log.log_event_end(
                 "dynamo", time.time_ns(), {}, chromium_start_time, True
             )
+            # === END WARNING WARNING WARNING ===
 
 
 class ConvertFrame:
