From b8cb28579264f54d612baf3605c8ca698a40511e Mon Sep 17 00:00:00 2001
From: Ahmed Gheith <ahmedgheith@meta.com>
Date: Mon, 11 Nov 2024 08:49:29 -0800
Subject: [PATCH] Log REMOTE_CACHE_VERSION

Summary:
X-link: https://github.com/pytorch/pytorch/pull/140174

title

Differential Revision: D65667421
---
 .../dynamo/dynamobench/_dynamo/utils.py       | 23 +++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/userbenchmark/dynamo/dynamobench/_dynamo/utils.py b/userbenchmark/dynamo/dynamobench/_dynamo/utils.py
index 6e3549b6a1..e0cd43cd07 100644
--- a/userbenchmark/dynamo/dynamobench/_dynamo/utils.py
+++ b/userbenchmark/dynamo/dynamobench/_dynamo/utils.py
@@ -149,6 +149,27 @@
 timer_counter = itertools.count()
 
 
+class FBCodeSpecific:
+    @property
+    def remote_cache_version(self) -> Optional[int]:
+        """ Returns the remote cache version if it is available, otherwise None.
+            Needs to be defensive because:
+                (1) REMOTE_CACHE_VERSION is fbcode only
+                (2) There is a circular dependency (through dynamic_typed) between
+                    dynamo and inductor. This is why the imports are not top-level
+        """
+        try:
+            from torch._environment import is_fbcode
+            if not is_fbcode():
+                return None
+            from torch._inductor.fb.remote_cache import REMOTE_CACHE_VERSION
+            return REMOTE_CACHE_VERSION
+        except Exception:
+            return None
+
+fbCodeSpecific: FBCodeSpecific = FBCodeSpecific()
+
+
 # Abstraction on top of counters.
 class ReInplaceTrigger(enum.Enum):
     AUTO_FUNC_V1 = 1
@@ -419,6 +440,7 @@ def dynamo_timed(
                                 fail_type=fail_type,
                                 fail_reason=fail_reason,
                                 remote_cache_time_saved_s=remote_cache_time_saved,
+                                remote_cache_version=fbCodeSpecific.remote_cache_version,
                                 structured_logging_overhead_s=structured_logging_overhead_s,
                                 is_forward=False,  # is_forward
                                 num_triton_bundles=codecache_metrics.get(
@@ -898,6 +920,7 @@ class CompilationMetrics:
     # a compiled frame
     has_guarded_code: Optional[bool] = None
     remote_cache_time_saved_s: Optional[float] = None
+    remote_cache_version: Optional[int] = None
     structured_logging_overhead_s: Optional[float] = None
     config_suppress_errors: Optional[bool] = None
     config_inline_inbuilt_nn_modules: Optional[bool] = None