Skip to content

Commit

Permalink
[torch.compile] disable logging when cache is disabled (#12043)
Browse files Browse the repository at this point in the history
Signed-off-by: youkaichao <[email protected]>
  • Loading branch information
youkaichao authored Jan 17, 2025
1 parent b5b57e3 commit 7b98a65
Showing 1 changed file with 21 additions and 9 deletions.
30 changes: 21 additions & 9 deletions vllm/compilation/backends.py
Original file line number Diff line number Diff line change
Expand Up @@ -251,15 +251,27 @@ def _check_can_cache(*args, **kwargs):
def _get_shape_env() -> AlwaysHitShapeEnv:
return AlwaysHitShapeEnv()

with patch(# for hijacking the hash of the compiled graph
"torch._inductor.codecache.compiled_fx_graph_hash",
hijack_compiled_fx_graph_hash), \
patch(# for providing a dummy shape environment
"torch._inductor.codecache.FxGraphCache._get_shape_env",
_get_shape_env), \
patch(# for forcing the graph to be cached
"torch._inductor.codecache.FxGraphCache._check_can_cache",
_check_can_cache):
with ExitStack() as stack:
if not cache_data.disabled:
# compilation cache is enabled, patch several functions

# for hijacking the hash of the compiled graph
stack.enter_context(
patch("torch._inductor.codecache.compiled_fx_graph_hash",
hijack_compiled_fx_graph_hash))

# for providing a dummy shape environment
stack.enter_context(
patch(
"torch._inductor.codecache.FxGraphCache._get_shape_env",
_get_shape_env))

# for forcing the graph to be cached
stack.enter_context(
patch(
"torch._inductor.codecache.FxGraphCache._check_can_cache",
_check_can_cache))

compiled_graph = compile_fx(graph,
example_inputs,
config_patches=current_config)
Expand Down

0 comments on commit 7b98a65

Please sign in to comment.