14
14
import importlib
15
15
import inspect
16
16
import itertools
17
+ import json
17
18
import linecache
18
19
import logging
19
20
import math
61
62
62
63
import torch
63
64
import torch ._functorch .config
64
- import torch ._inductor .config as inductor_config
65
65
import torch .fx .experimental .symbolic_shapes
66
66
import torch .utils ._pytree as pytree
67
67
from torch import fx
@@ -865,6 +865,7 @@ class CompilationMetrics:
865
865
post_grad_pass_time_us : Optional [int ] = None
866
866
joint_graph_pass_time_us : Optional [int ] = None
867
867
log_format_version : int = LOG_FORMAT_VERSION
868
+ inductor_config : Optional [str ] = None
868
869
869
870
870
871
DEFAULT_COMPILATION_METRICS_LIMIT = 64
@@ -912,9 +913,48 @@ def add_compilation_metrics_to_chromium(c: Dict[str, Any]) -> None:
912
913
)
913
914
914
915
916
+ def _scrubbed_inductor_config_for_logging () -> Optional [str ]:
917
+ """
918
+ Method to parse and scrub unintersting configs from inductor config
919
+ """
920
+
921
+ # TypeSafeSerializer for json.dumps()
922
+ # Skips complex types as values in config dict
923
+ class TypeSafeSerializer (json .JSONEncoder ):
924
+ def default (self , o ):
925
+ try :
926
+ return super ().default (o )
927
+ except Exception :
928
+ return "Value is not JSON serializable"
929
+
930
+ configs_to_scrub_re = r"((^TYPE_CHECKING$)|(.*_progress$)|(.*TESTING.*)|(.*(rocm|halide).*)|(^trace\..*)|(^_))"
931
+ keys_to_scrub = set ()
932
+ inductor_conf_str = None
933
+ inductor_config_copy = (
934
+ torch ._inductor .config .get_config_copy () if torch ._inductor .config else None
935
+ )
936
+ if inductor_config_copy is not None :
937
+ try :
938
+ for key , val in inductor_config_copy .items ():
939
+ if not isinstance (key , str ) or re .search (configs_to_scrub_re , key ):
940
+ keys_to_scrub .add (key )
941
+ # Convert set() to list for json.dumps()
942
+ if isinstance (val , set ):
943
+ inductor_config_copy [key ] = list (val )
944
+ # Evict unwanted keys
945
+ for key in keys_to_scrub :
946
+ del inductor_config_copy [key ]
947
+ # Stringify Inductor config
948
+ inductor_conf_str = json .dumps (
949
+ inductor_config_copy , cls = TypeSafeSerializer , skipkeys = True
950
+ )
951
+ except Exception :
952
+ # Don't crash because of runtime logging errors
953
+ inductor_conf_str = "Inductor Config is not JSON serializable"
954
+ return inductor_conf_str
955
+
956
+
915
957
def record_compilation_metrics (metrics : Dict [str , Any ]):
916
- # TODO: Temporary; populate legacy fields from their replacements.
917
- # Remove when we decide we can really deprecate them.
918
958
def us_to_s (field ):
919
959
metric = metrics .get (field , None )
920
960
return metric / 1e6 if metric is not None else None
@@ -923,7 +963,12 @@ def us_to_ms(field):
923
963
metric = metrics .get (field , None )
924
964
return metric // 1000 if metric is not None else None
925
965
926
- legacy_metrics = {
966
+ common_metrics = {
967
+ "inductor_config" : _scrubbed_inductor_config_for_logging (),
968
+ # -------- Any future common metircs go here --------
969
+ #
970
+ # Legacy metircs go here(TODO: Temporary; populate legacy fields from their replacements.)
971
+ # Remove when we decide we can really deprecate them.
927
972
"entire_frame_compile_time_s" : us_to_s ("dynamo_cumulative_compile_time_us" ),
928
973
"backend_compile_time_s" : us_to_s ("aot_autograd_cumulative_compile_time_us" ),
929
974
"inductor_compile_time_s" : us_to_s ("inductor_cumulative_compile_time_us" ),
@@ -937,7 +982,7 @@ def us_to_ms(field):
937
982
),
938
983
}
939
984
940
- compilation_metrics = CompilationMetrics (** {** metrics , ** legacy_metrics })
985
+ compilation_metrics = CompilationMetrics (** {** metrics , ** common_metrics })
941
986
_compilation_metrics .append (compilation_metrics )
942
987
if compilation_metrics .is_forward :
943
988
name = "compilation_metrics"
@@ -2036,7 +2081,7 @@ def to_tensor(t):
2036
2081
and math .isnan (res_error )
2037
2082
# Some unit test for the accuracy minifier relies on
2038
2083
# returning false in this case.
2039
- and not inductor_config .cpp .inject_relu_bug_TESTING_ONLY
2084
+ and not torch . _inductor . config .cpp .inject_relu_bug_TESTING_ONLY
2040
2085
):
2041
2086
passes_test = True
2042
2087
if not passes_test :
0 commit comments