From 3a90458e1ec3aab8627bb745ac85e7d709e0ee1f Mon Sep 17 00:00:00 2001 From: Sam Larsen Date: Wed, 16 Oct 2024 20:59:52 -0700 Subject: [PATCH] Add logger logging for remote fx graph cache get + put (#2512) Summary: X-link: https://github.com/pytorch/pytorch/pull/138164 Capture the timing for the remote fx graph cache get and put operations and add them to the logger logging. Reviewed By: oulgen Differential Revision: D64484025 --- userbenchmark/dynamo/dynamobench/_dynamo/utils.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/userbenchmark/dynamo/dynamobench/_dynamo/utils.py b/userbenchmark/dynamo/dynamobench/_dynamo/utils.py index a41d92724..b82ca9bdf 100644 --- a/userbenchmark/dynamo/dynamobench/_dynamo/utils.py +++ b/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, + remote_fx_graph_cache_get_time, + remote_fx_graph_cache_put_time, ) record_compilation_metrics(metrics) @@ -810,6 +820,8 @@ class CompilationMetrics: config_inline_inbuilt_nn_modules: Optional[bool] specialize_float: Optional[bool] dynamo_config: Optional[str] + remote_fx_graph_cache_get_time_s: Optional[float] + remote_fx_graph_cache_put_time_s: Optional[float] @dataclasses.dataclass @@ -821,6 +833,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_s: Optional[float] + remote_fx_graph_cache_put_time_s: Optional[float] DEFAULT_COMPILATION_METRICS_LIMIT = 64