pytorch/torch/_dynamo
Sam Larsen cb15c15157 [logging] Overhaul dynamo_timed and CompilationMetrics logging. (#139849)
Here's the overview:

There's a new contextmanager singleton called MetricsContext. Entering the MetricsContext is how we demarcate the boundary on which we'll create a single CompilationMetrics object, and therefore, a single dynamo_compile log entry. While we're inside the MetricsContext, we can update/set many different metrics. Most importantly: `dynamo_timed` can also update the in-progress MetricsContext. In the proposal here, we tell `dynamo_timed` that we want it to do so by providing the name of the MetricsContext field to increment. There can be many `dynamo_timed` calls in different parts of the code updating different fields. Then when the MetricsContext exits, that's when the logging of everything gathered finally happens. One potential footgun is trying to use `dynamo_timed` when we haven't entered the MetricsContext, but we assert on that problem. Another problem is that we re-enter the context recursively, but we watch for that and do the logging only when the outermost exits.

Some specifics:
* Introduce MetricsContext - a context manager that on exit, records the CompilationMetrics (which also logs to dynamo_compile).
* Completely remove the concept of frame_phase_timing. Instead, update the MetricsContext during compilation, either directly or via dynamo_timed.
* Remove some globals we previously used to accumulate counters to later populate a CompilationMetrics. We use CompilationMetrics set/update/increment APIs instead.
* `record_compilation_metrics` is now called on exit from MetricsContext.
* Populate legacy CompilationMetrics fields right before logging, inside `record_compilation_metrics`.
* Remove the one-off `add_remote_cache_time_saved` helper; capture that timing directly into the MetricsContext.

And specifically, several changes to dynamo_timed:
* "Modernize" the parameters and update all callsites accordingly.
* Move the backwards logging of the CompilationMetrics to the backwards compile location.
* Add a parameter for which CompilationMetrics field to update

Pull Request resolved: https://github.com/pytorch/pytorch/pull/139849
Approved by: https://github.com/ezyang
ghstack dependencies: #140094
2024-11-11 14:24:23 +00:00
..
backends Class rename (#139490) 2024-11-02 00:10:17 +00:00
polyfills [Dynamo] Add itertools.compress() support (#139061) 2024-10-29 10:25:55 +00:00
repro AOTI Minifier (#139351) 2024-11-07 21:43:44 +00:00
variables Make size a property with an assertion (#139794) 2024-11-09 03:39:41 +00:00
__init__.py Profile guided optimization for automatic_dynamic (#139001) 2024-11-03 06:29:57 +00:00
_trace_wrapped_higher_order_op.py [FlexAttention] Support training bias for eager (#136910) (#137526) 2024-10-15 18:55:22 +00:00
bytecode_analysis.py [BE][Easy][15/19] enforce style for empty lines in import segments in torch/_d*/ (#129767) 2024-07-31 21:18:11 +00:00
bytecode_transformation.py [dynamo, 3.13] fix bytecode nop tests (#139323) 2024-11-02 00:39:36 +00:00
cache_size.py [dynamo][refactor][config-cleanp] Use guard_manager consistently instead of check_fn (#138896) 2024-10-26 15:14:46 +00:00
callback.py turn CompilationCallbackHandler into dataclass (#137312) 2024-10-05 19:03:28 +00:00
code_context.py add types to _dynamo/code_context.py (#136665) 2024-09-27 18:27:42 +00:00
codegen.py [dynamo] Document codegen and clean up some code paths (#139670) 2024-11-07 03:14:16 +00:00
compiled_autograd.py support nesting of suppress_guards, suppress guards when generated compiled autograd graph (#138968) 2024-10-31 00:13:39 +00:00
comptime.py Improve print stack/locals printing in comptime (#133651) 2024-08-27 01:29:50 +00:00
config.py [dynamo][guards] Consider tensors as immutable for dict tag matches (#139560) 2024-11-05 21:48:07 +00:00
convert_frame.py [logging] Overhaul dynamo_timed and CompilationMetrics logging. (#139849) 2024-11-11 14:24:23 +00:00
create_parameter_op.py type _dynamo/create_parameter_op.py (#136958) 2024-10-02 10:23:37 +00:00
current_scope_id.py add typing to _dynamo/current_scope_id.py (#136676) 2024-09-27 04:09:15 +00:00
debug_utils.py Ensure TORCH_TRACE is run for Dynamo/Distributed tests (#139786) 2024-11-07 01:58:05 +00:00
decorators.py Revert "Deprecate torch._utils.is_compiling() and torch._dynamo.external_utils.is_compiling() (#127690)" 2024-11-05 23:10:38 +00:00
device_interface.py Have Triton custom extension test use privateuseone device (#137611) 2024-10-11 21:27:29 +00:00
distributed.py [BE][Easy][15/19] enforce style for empty lines in import segments in torch/_d*/ (#129767) 2024-07-31 21:18:11 +00:00
eval_frame.py always specialize float on export path (#139486) 2024-11-03 04:47:12 +00:00
exc.py [tlparse] Add dynamo_graph_break_reason logging to trace_structured (#138778) 2024-10-25 02:00:04 +00:00
external_utils.py Revert "Deprecate torch._utils.is_compiling() and torch._dynamo.external_utils.is_compiling() (#127690)" 2024-11-05 23:10:38 +00:00
funcname_cache.py [BE][Easy][15/19] enforce style for empty lines in import segments in torch/_d*/ (#129767) 2024-07-31 21:18:11 +00:00
guards.py [dynamo][guards] Profile guard manager in C++ (#140110) 2024-11-08 18:44:08 +00:00
hooks.py [BE][Easy][15/19] enforce style for empty lines in import segments in torch/_d*/ (#129767) 2024-07-31 21:18:11 +00:00
logging.py Make automatic_dynamic state live per CodeId, rather than on code object (#138740) 2024-10-27 03:08:41 +00:00
metrics_context.py [logging] Overhaul dynamo_timed and CompilationMetrics logging. (#139849) 2024-11-11 14:24:23 +00:00
mutation_guard.py Allow Lazy Module to be modelled as UnspecializedNNModuleVariable (#138639) 2024-10-26 02:17:07 +00:00
output_graph.py [logging] Overhaul dynamo_timed and CompilationMetrics logging. (#139849) 2024-11-11 14:24:23 +00:00
pgo.py Make PGO work correctly with NJT inputs (#140046) 2024-11-08 04:27:39 +00:00
profiler.py type _dynamo/profiler.py (#137351) 2024-10-07 18:54:33 +00:00
replay_record.py type _dynamo/replay_record.py (#137183) 2024-10-04 20:29:24 +00:00
resume_execution.py [Dynamo] Trace enter/exit of TorchFunctionModes (#135422) (#137114) 2024-10-09 02:29:40 +00:00
side_effects.py [dynamo] Document codegen and clean up some code paths (#139670) 2024-11-07 03:14:16 +00:00
source.py [Dynamo] Disable torch function compilation during guard execution and in compiled bytecode (#137669) 2024-10-19 04:12:45 +00:00
symbolic_convert.py [dynamo] Remove NestedUserFunctionVariable.closure_scope (#140033) 2024-11-09 01:03:24 +00:00
tensor_version_op.py [BE][Easy][15/19] enforce style for empty lines in import segments in torch/_d*/ (#129767) 2024-07-31 21:18:11 +00:00
test_case.py Ensure TORCH_TRACE is run for Dynamo/Distributed tests (#139786) 2024-11-07 01:58:05 +00:00
test_minifier_common.py AOTI Minifier (#139351) 2024-11-07 21:43:44 +00:00
testing.py [hop free symbols] lift free symbols in example_value when create_graph_input (#138363) 2024-11-07 04:44:32 +00:00
trace_rules.py fix dynamo tracking numpy 2 ops (#138686) 2024-11-08 23:38:53 +00:00
types.py [dynamo][refactor][config-cleanp] Use guard_manager consistently instead of check_fn (#138896) 2024-10-26 15:14:46 +00:00
utils.py [logging] Overhaul dynamo_timed and CompilationMetrics logging. (#139849) 2024-11-11 14:24:23 +00:00