[Misc] Add more scoping for improved trace

[frank-wei]

Summary: Add more scoping in the hotspot area which can greatly help us to the find the cycle heavy area

Differential Revision: D86436375

[gemini-code-assist]

Code Review

This pull request introduces profiling scopes using record_function_or_nullcontext in key areas of the scheduler and engine core. This will help in identifying performance hotspots and analyzing cycle-heavy areas. The changes are well-targeted and correctly placed. I've found a minor typo in one of the scope names, which I've commented on. Otherwise, the changes look good.

[gemini-code-assist]

There appears to be a typo in the profiler scope name. It uses a double colon (::) which is inconsistent with the other scopes in this file that use a single colon. For consistency and to ensure easier parsing of traces, this should be corrected to LLM_ENGINE STEP: RECORD_STATS.

Suggested change

	with record_function_or_nullcontext("LLM_ENGINE STEP:: RECORD_STATS"):
	if self.logger_manager is not None and outputs.scheduler_stats is not None:
	self.logger_manager.record(
	scheduler_stats=outputs.scheduler_stats,
	iteration_stats=iteration_stats,
	mm_cache_stats=self.processor.stat_mm_cache(),
	)
	self.do_log_stats_with_interval()
	with record_function_or_nullcontext("LLM_ENGINE STEP: RECORD_STATS"):
	if self.logger_manager is not None and outputs.scheduler_stats is not None:
	self.logger_manager.record(
	scheduler_stats=outputs.scheduler_stats,
	iteration_stats=iteration_stats,
	mm_cache_stats=self.processor.stat_mm_cache(),
	)
	self.do_log_stats_with_interval()

[bangshengtang]

nit: can we try to unify the text format? we have all lower, all upper, first letter capitalized in three different layers

[frank-wei]

nit: can we try to unify the text format? we have all lower, all upper, first letter capitalized in three different layers

l prefer to use all lower case which saves the space visually. Will make corresponding changes.

[benchislett]

Typo here and above

[frank-wei]

fixed in #28584

[njhill]

I'm not a fan of this TBH. I think having these all over the code makes it really messy and harder to read and maintain. I'd hope we could instrument in a less invasive way, restructuring into functions if/where it makes sense from a code readability/modularity pov, or at minimum being more conservative/selective as to where these are placed.

[frank-wei]

I'm not a fan of this TBH. I think having these all over the code makes it really messy and harder to read. I'd hope we could instrument in a less invasive way, restructuring into functions if/where it makes sense from a code readability/modularity pov, or at minimum being more conservative/selective as to where these are placed.

thanks @njhill , do you have any code pointer for a clean way? I can follow it.

[njhill]

@frank-wei FYI I'm reverting the ones in core.py as part of #28768.