Add logger logging for remote fx graph cache get + put (#2512)

Summary:
Pull Request resolved: #2512

X-link: pytorch/pytorch#138164

Capture the timing for the remote fx graph cache get and put operations and add them to the logger logging.

Reviewed By: ezyang, oulgen

Differential Revision: D64484025

userbenchmark/dynamo/dynamobench/_dynamo/utils.py

@@ -351,10 +351,18 @@ def dynamo_timed(
                                 remote_cache_time_saved = frame_phase_timing[
                                     compile_id
                                 ].get("remote_cache_time_saved", None)
+                                remote_fx_graph_cache_get_time = frame_phase_timing[
+                                    compile_id
+                                ].get("remote_fx_graph_cache_get", None)
+                                remote_fx_graph_cache_put_time = frame_phase_timing[
+                                    compile_id
+                                ].get("remote_fx_graph_cache_put", None)
                             else:
                                 inductor_compile_time = None
                                 code_gen_time = None
                                 remote_cache_time_saved = None
+                                remote_fx_graph_cache_get_time = None
+                                remote_fx_graph_cache_put_time = None
                             structured_logging_overhead_s = (
                                 torch._logging.get_structured_logging_overhead()
                             )
@@ -366,6 +374,8 @@ def dynamo_timed(
                                 fail_reason,
                                 remote_cache_time_saved,
                                 structured_logging_overhead_s,
+                                to_int_ms(remote_fx_graph_cache_get_time),
+                                to_int_ms(remote_fx_graph_cache_put_time),
                             )
                             record_compilation_metrics(metrics)
 
@@ -772,6 +782,10 @@ def proxy_args_kwargs(args, kwargs):
         )
 
 
+def to_int_ms(v: Optional[float]) -> Optional[int]:
+    return None if v is None else int(v * 1000)
+
+
 @dataclasses.dataclass
 class CompilationMetrics:
     is_forward: bool = dataclasses.field(default=True, init=False)
@@ -811,6 +825,8 @@ class CompilationMetrics:
     config_inline_inbuilt_nn_modules: Optional[bool]
     specialize_float: Optional[bool]
     dynamo_config: Optional[str]
+    remote_fx_graph_cache_get_time_ms: Optional[int]
+    remote_fx_graph_cache_put_time_ms: Optional[int]
 
 
 @dataclasses.dataclass
@@ -823,6 +839,8 @@ class BwdCompilationMetrics:
     fail_reason: Optional[str]
     remote_cache_time_saved_s: Optional[float]
     structured_logging_overhead_s: Optional[float]
+    remote_fx_graph_cache_get_time_ms: Optional[int]
+    remote_fx_graph_cache_put_time_ms: Optional[int]
 
 
 DEFAULT_COMPILATION_METRICS_LIMIT = 64