[internal] Expose additional metadata to compilation callbacks #153596
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/153596
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 5d56d1e with merge base 82dc345 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
| def on_compile_start( | ||
| callback: Callable[[CallbackArgs], None], | ||
| ) -> Callable[[CallbackArgs], None]: | ||
| """ | ||
| Decorator to register a callback function for the start of the compilation. | ||
| """ | ||
| callback_handler.register_start_callback(callback) | ||
| return callback | ||
|
|
||
|
|
||
| def on_compile_end(callback: Callable[[], None]) -> Callable[[], None]: | ||
| def on_compile_end( | ||
| callback: Callable[[CallbackArgs], None], | ||
| ) -> Callable[[CallbackArgs], None]: |
There was a problem hiding this comment.
Would it make sense to introduce new APIs? I am guessing this change will cause type errors as the callback function being passed needs to have CallbackArgs argument.
There was a problem hiding this comment.
Is the concern about backward compatibility? This is a private API and I didn't find any usage on github, we can address the internal type changes on the import diff. The "new API" offers a superset of the old, so I didn't see much value keeping it around.
Stack from ghstack (oldest at bottom):
These hooks are used by internal stuck job detection to associate compilation events with the compile lease. Previously, we only had events for Dynamo and Inductor compilation. And recently, the callback handler was updated to ignore nested events. So the Inductor event was only really used by lazy backward.
Here, I remove the inductor event, and add an explicit lazy backward one. Additionally, 10000 I add other runtime compilation events: autotuning and cudagraphs. I also expose the CompileId as a string to avoid imports, this will let internal UIs track each graph's contribution to the timeout.
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov