diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index 4ac6949316..43252ce3d8 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -142,6 +142,10 @@ class CompilerManager: self.compilation_config = compilation_config self.compiler = make_compiler(compilation_config) self.loaded_artifacts: dict[str, Any] = {} + self.prefix: str = "" + + def _log_prefix(self) -> str: + return f"[{self.prefix}] " if self.prefix else "" def compute_hash(self, vllm_config: VllmConfig) -> str: return self.compiler.compute_hash(vllm_config) @@ -179,6 +183,7 @@ class CompilerManager: self.disable_cache = disable_cache self.cache_dir = cache_dir self.cache_file_path = os.path.join(cache_dir, "vllm_compile_cache.py") + self.prefix = prefix if not disable_cache and os.path.exists(self.cache_file_path): # load the cache from the file @@ -287,9 +292,10 @@ class CompilerManager: # after loading the last graph for this shape, record the time. # there can be multiple graphs due to piecewise compilation. elapsed = time.perf_counter() - compilation_start_time - logger.info_once( - "Directly load the compiled graph(s) for compile range %s " - "from the cache, took %.3f s", + logger.info( + "%sDirectly load the compiled graph(s) for compile " + "range %s from the cache, took %.3f s", + self._log_prefix(), str(compile_range), elapsed, ) @@ -373,8 +379,9 @@ class CompilerManager: self.is_cache_updated = True if graph_index == 0: # adds some info logging for the first graph - logger.info_once( - "Cache the graph of compile range %s for later use", + logger.info( + "%sCache the graph of compile range %s for later use", + self._log_prefix(), str(compile_range), ) logger.debug_once( @@ -388,8 +395,9 @@ class CompilerManager: # after compiling the last graph, record the end time if graph_index == num_graphs - 1: elapsed = time.perf_counter() - compilation_start_time - logger.info_once( - "Compiling a graph for compile range %s takes %.2f s", + logger.info( + "%sCompiling a graph for compile range %s takes %.2f s", + self._log_prefix(), str(compile_range), elapsed, ) @@ -865,6 +873,9 @@ class VllmBackend: # `torch.compile` is JIT compiled, so we don't need to # do anything here + def _log_prefix(self) -> str: + return f"[{self.prefix}] " if self.prefix else "" + def collect_standalone_compile_artifacts( self, ) -> tuple[Any, dict[str, list[int]] | None, dict[str, bool] | None]: @@ -1085,10 +1096,11 @@ class VllmBackend: disable_cache = disable_cache or is_ngram_gpu_enabled if disable_cache: - logger.info_once("vLLM's torch.compile cache is disabled.") + logger.info("%svLLM's torch.compile cache is disabled.", self._log_prefix()) else: - logger.info_once( - "Using cache directory: %s for vLLM's torch.compile", + logger.info( + "%sUsing cache directory: %s for vLLM's torch.compile", + self._log_prefix(), local_cache_dir, ) @@ -1148,8 +1160,9 @@ class VllmBackend: current_perf = time.perf_counter() current_epoch = time.time() dynamo_time = current_perf - torch_compile_start_time - logger.info_once( - "Dynamo bytecode transform time: %.2f s", + logger.info( + "%sDynamo bytecode transform time: %.2f s", + self._log_prefix(), dynamo_time, ) diff --git a/vllm/compilation/decorators.py b/vllm/compilation/decorators.py index 35c0e6dee5..3aad32cc24 100644 --- a/vllm/compilation/decorators.py +++ b/vllm/compilation/decorators.py @@ -292,7 +292,12 @@ def _try_load_aot_compiled_fn( Re-raises on failure when ``VLLM_FORCE_AOT_LOAD`` is set. """ try: - with monitor_torch_compile(model.vllm_config, is_encoder=model._is_encoder): + with monitor_torch_compile( + model.vllm_config, + is_encoder=model._is_encoder, + tag=getattr(model, "_compile_tag", ""), + announce_start=False, + ): with ( set_current_vllm_config(model.vllm_config), open(aot_compilation_path, "rb") as f, @@ -309,8 +314,12 @@ def _try_load_aot_compiled_fn( with maybe_use_cudagraph_partition_wrapper(model.vllm_config): loaded_fn._artifacts.compiled_fn.finalize_loading(model.vllm_config) compilation_counter.num_aot_artifacts_loaded += 1 + tag = getattr(model, "_compile_tag", "") + log_prefix = f"[{tag}] " if tag else "" logger.info( - "Directly load AOT compilation from path %s", aot_compilation_path + "%sDirectly load AOT compilation from path %s", + log_prefix, + aot_compilation_path, ) return loaded_fn except Exception as e: @@ -319,8 +328,11 @@ def _try_load_aot_compiled_fn( message = "Compile cache file corrupted." else: message = str(e) + tag = getattr(model, "_compile_tag", "") + log_prefix = f"[{tag}] " if tag else "" logger.warning( - "Compiling model again due to a load failure from %s, reason: %s", + "%sCompiling model again due to a load failure from %s, reason: %s", + log_prefix, aot_compilation_path, message, ) @@ -402,6 +414,16 @@ def _support_torch_compile( compilation_counter.num_models_seen += 1 self.compiled = False + # Capture the active model tag (e.g. "backbone", "eagle_head") so + # later compilation/warmup log lines can identify which component + # they belong to. Mirror VllmBackend's precedence rule + # (`compile_prefix or model_tag`) so log tags line up with the + # cache-dir tag: encoders use the class name, everything else + # falls through to `model_tag`. + from vllm.compilation.backends import model_tag as _current_model_tag + + self._compile_tag = (cls.__name__ if is_encoder else "") or _current_model_tag + # Handled by monkeypatching `TorchCompileWithNoGuardsWrapper` into base class TorchCompileWithNoGuardsWrapper.__init__( self, @@ -567,7 +589,7 @@ def _support_torch_compile( self.aot_compiled_fn = loaded_fn self.was_aot_compile_fn_loaded_from_disk = True with ( - monitor_profiling_run(), + monitor_profiling_run(tag=self._compile_tag), maybe_use_cudagraph_partition_wrapper(self.vllm_config), ): output = self.aot_compiled_fn(self, *args, **kwargs) @@ -667,7 +689,9 @@ def _support_torch_compile( self._aot_compilation_path = aot_compilation_path self._aot_cache_dir = cache_dir with monitor_torch_compile( - self.vllm_config, is_encoder=self._is_encoder + self.vllm_config, + is_encoder=self._is_encoder, + tag=self._compile_tag, ): self.aot_compiled_fn = self.aot_compile(*args, **kwargs) compilation_counter.num_aot_compiles += 1 @@ -675,14 +699,15 @@ def _support_torch_compile( # AOT artifact. self.save_aot_compiled_function() - with monitor_profiling_run(): + with monitor_profiling_run(tag=self._compile_tag): output = self.aot_compiled_fn(self, *args, **kwargs) else: with monitor_torch_compile( self.vllm_config, - "torch.compile and initial profiling/warmup " + "%storch.compile and initial profiling/warmup " "run together took %.2f s in total", is_encoder=self._is_encoder, + tag=self._compile_tag, ): output = TorchCompileWithNoGuardsWrapper.__call__( self, # type: ignore[arg-type] @@ -706,6 +731,7 @@ def _support_torch_compile( self.aot_compiled_fn and self._aot_compilation_path and self._aot_cache_dir ) + log_prefix = f"[{self._compile_tag}] " if self._compile_tag else "" try: os.makedirs(self._aot_cache_dir, exist_ok=True) # File saving should be atomic, so we will save to a temporary location @@ -714,13 +740,15 @@ def _support_torch_compile( self.aot_compiled_fn.save_compiled_function(tmp_file) os.replace(tmp_file, self._aot_compilation_path) compilation_counter.num_aot_artifacts_saved += 1 - logger.info_once( - "saved AOT compiled function to %s", + logger.info( + "%ssaved AOT compiled function to %s", + log_prefix, self._aot_compilation_path, ) except Exception as e: logger.warning( - "unable to save AOT compiled function to %s: %s", + "%sunable to save AOT compiled function to %s: %s", + log_prefix, self._aot_compilation_path, e, ) diff --git a/vllm/compilation/monitor.py b/vllm/compilation/monitor.py index c23a8f6722..e38192034b 100644 --- a/vllm/compilation/monitor.py +++ b/vllm/compilation/monitor.py @@ -17,18 +17,27 @@ torch_compile_start_time: float = 0.0 @contextlib.contextmanager def monitor_torch_compile( vllm_config: VllmConfig, - message: str = "torch.compile took %.2f s in total", + message: str = "%storch.compile took %.2f s in total", is_encoder: bool = False, + tag: str = "", + announce_start: bool = True, ) -> Generator[None, None, None]: """Context manager that times torch.compile and manages depyf debugging. On normal exit: logs the compile time and exits depyf. On exception: cleans up depyf without logging (compilation failed). + + ``announce_start=False`` suppresses the "Starting torch.compile" line; + used by the AOT-load probe so a failed load attempt doesn't announce a + compile it didn't perform. """ global torch_compile_start_time torch_compile_start_time = time.perf_counter() compilation_config = vllm_config.compilation_config + log_prefix = f"[{tag}] " if tag else "" + if announce_start and compilation_config.mode == CompilationMode.VLLM_COMPILE: + logger.info("%sStarting torch.compile", log_prefix) depyf_cm = None path = vllm_config.compile_debug_dump_path() if compilation_config.mode == CompilationMode.VLLM_COMPILE and path: @@ -50,7 +59,7 @@ def monitor_torch_compile( compilation_config.encoder_compilation_time += total_compile_time else: compilation_config.compilation_time += total_compile_time - logger.info_once(message, total_compile_time) + logger.info(message, log_prefix, total_compile_time) finally: if depyf_cm is not None: try: @@ -60,7 +69,7 @@ def monitor_torch_compile( @contextlib.contextmanager -def monitor_profiling_run() -> Generator[None, None, None]: +def monitor_profiling_run(tag: str = "") -> Generator[None, None, None]: """Context manager that times the initial profiling run. Asserts that no backend compilation occurs during the profiling run @@ -78,8 +87,10 @@ def monitor_profiling_run() -> Generator[None, None, None]: "backend compilation occurred during the initial profiling run; " "all compilation should be complete before the profiling run starts." ) - logger.info_once( - "Initial profiling/warmup run took %.2f s", + log_prefix = f"[{tag}] " if tag else "" + logger.info( + "%sInitial profiling/warmup run took %.2f s", + log_prefix, elapsed, )