[Log] add a hook for recompile user context #157961

BoyuanFeng · 2025-07-09T20:35:45Z

Users may want compile-related but customized logging info to dynamo_compile. One example is to logging the current training iteration index when recompilation happens. In general, current training iteration index is not available to compiler, since the same compiled function may be called multiple times in the same training iteration. The user could provide the training iteration index in a user hook where torch.compile logs it when recompilation happens.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames

pytorch-bot · 2025-07-09T20:35:50Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/157961

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 Cancelled Job, 1 Unrelated Failure

As of commit 2c27345 with merge base a3ec6d6 ():

CANCELLED JOB - The following job was cancelled. Please retry:

inductor / cuda12.8-py3.10-gcc9-sm86 / test (inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh)

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

pull / cuda12.8-py3.10-gcc9-sm75 / test (pr_time_benchmarks, 1, 1, linux.g4dn.metal.nvidia.gpu, unstable) (gh) (#153987)
MISSING REGRESSION TEST

This comment was automatically generated by Dr. CI and updates every 15 minutes.

BoyuanFeng · 2025-07-09T20:39:16Z

torch/_dynamo/convert_frame.py

 output_codes = Tracker()

 initial_global_state: Optional[GlobalStateGuard] = None
+recompile_user_contexts: Optional[list[Callable[[], str]]] = None


users may want log multiple different user contexts so use a list here

williamwen42

add a test?

masnesral · 2025-07-09T21:25:31Z

torch/_dynamo/convert_frame.py

            )
        metrics_context.update_outer({"recompile_reason": recompile_reason})

+        if recompile_user_contexts:


Do you only want to do this if it's a recompile? It looks like this is called unconditionally?

It should only happen once for the first compile and once for every recompile. It should not happen for other cases. I added a test for that.

We probably don't want to distinguish the first compile and recompile here since other info (e.g., compile id) on dynamo_compile table tells us whether it's first compile or recompile.

Ok, I got fooled by the name. I assumed that you only wanted it called on the first recompile (but not the first compile).

masnesral · 2025-07-09T21:26:46Z

torch/_dynamo/convert_frame.py

        metrics_context.update_outer({"recompile_reason": recompile_reason})

+        if recompile_user_contexts:
+            user_contexts_msg = "\n".join(


I don't think "\n" is going to render very well in scuba. Maybe you want to separate by semi-colon instead (or some other separator). Alternatively, it's possible for scuba fields to be lists (called normvectors). I commented on the associated diff about that.

We probably can use set[str] following other fields (e.g., compliant_custom_ops)?

Yeah, that's another option. I think they're called "tags" or "tagsets" in scuba?? https://www.internalfb.com/wiki/Scuba/Quick_Start_Guide/Data_Types/

set becomes tagset here. Originally added by Richard in D50884828

masnesral · 2025-07-09T21:30:30Z

torch/_dynamo/convert_frame.py


+        if recompile_user_contexts:
+            user_contexts_msg = "\n".join(
+                [user_context() for user_context in recompile_user_contexts]


It's kind of dangerous to allow an unbounded string here because if a user decides to add a very long string (by purpose or by mistake), it could cause retention problems. I wonder if there's a more restricted API we could use here. At the very least, we should probably truncate whatever they return to some maximum. Do you have a sense of how big these contexts need to be? I don't know what scuba recommends, but if we don't expect something very large then we could cap to 512, 1k, or something like that?

it should be short. One use case is just "training epoch: some_integer". But I am hesitating to add an arbitrarily selected cap to user message. Since users have full control of the log message, we should allow users to log what they want?

Uh... I wouldn't give them complete freedom because the scuba table is a shared resource. So if one customer logs a ton of data, it affects retention for the whole table and we can't debug other jobs. I think it's probably fine to allow the field to be fairly large if you think it's needed. But I wouldn't allow it to be entirely unbounded.

sure we can cap each user context at 256 characters.

masnesral · 2025-07-09T21:31:04Z

torch/_dynamo/convert_frame.py

+            user_contexts_msg = "\n".join(
+                [user_context() for user_context in recompile_user_contexts]
+            )
+            metrics_context.set("recompile_user_contexts", user_contexts_msg)


What if the user context raises an exception? Do we just want to crash?

Yes. Since a user registers the context only for himself/herself, we should allow users to do whatever they want. It's like a user can raise an exception in other pytorch code.

masnesral · 2025-07-09T21:32:46Z

torch/_dynamo/convert_frame.py

+            user_contexts_msg = "\n".join(
+                [user_context() for user_context in recompile_user_contexts]
+            )
+            metrics_context.set("recompile_user_contexts", user_contexts_msg)


I think you'll also need to add a field to this thing: https://github.com/pytorch/pytorch/blob/main/torch/_dynamo/utils.py#L1234

masnesral · 2025-07-10T01:48:10Z

torch/_dynamo/convert_frame.py

+recompile_user_contexts: Optional[list[Callable[[], str]]] = None
+
+
+def register_hook_for_recompile_user_context(hook: Callable[[], str]) -> None:


nit: Isn't this kind of a weird place to put this API? Most of the logging-related stuff is currently in _dynamo/utils.py. Maybe that would be a better location?

facebook-github-bot · 2025-07-10T20:18:06Z

@BoyuanFeng has imported this pull request. If you are a Meta employee, you can view this in D78115274.

masnesral

LGTM

BoyuanFeng · 2025-07-11T03:39:48Z

@pytorchbot merge -f "skip unrelated fail"

pytorchmergebot · 2025-07-11T03:41:19Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

Summary: Users may want compile-related but customized logging info to dynamo_compile. One example is to logging the current training iteration index when recompilation happens. In general, current training iteration index is not available to compiler, since the same compiled function may be called multiple times in the same training iteration. The user could provide the training iteration index in a user hook where torch.compile logs it when recompilation happens. X-link: pytorch/pytorch#157961 Approved by: https://github.com/masnesral Reviewed By: masnesral Differential Revision: D78115274 fbshipit-source-id: 6e2023146772734515a356139aaac9263a9ece2a

BoyuanFeng added 2 commits July 9, 2025 13:32

add a hook for recompile user context

8206a8c

nit

6c68f0d

BoyuanFeng requested review from anijain2305, masnesral and williamwen42 July 9, 2025 20:35

BoyuanFeng added ciflow/trunk Trigger trunk jobs on your pull request topic: not user facing topic category module: dynamo labels Jul 9, 2025

pytorch-bot bot added the ciflow/inductor label Jul 9, 2025

BoyuanFeng commented Jul 9, 2025

View reviewed changes

williamwen42 reviewed Jul 9, 2025

View reviewed changes

masnesral reviewed Jul 9, 2025

View reviewed changes

masnesral requested changes Jul 9, 2025

View reviewed changes

BoyuanFeng added 4 commits July 9, 2025 15:59

add test that user_context is called only once per recompile

8653704

nit

e105898

nit

4c6f4fd

add cap to user context strings

45330f6

masnesral reviewed Jul 10, 2025

View reviewed changes

BoyuanFeng added 3 commits July 9, 2025 21:42

nit

9c8b74a

add more tests

0c9a1df

Merge branch 'main' into bf/recompile-context

2c27345

BoyuanFeng requested review from masnesral and williamwen42 July 10, 2025 20:17

masnesral approved these changes Jul 10, 2025

View reviewed changes

pytorchmergebot added the merging label Jul 11, 2025

pytorchmergebot closed this in 94995eb Jul 11, 2025

pytorchmergebot added Merged and removed merging labels Jul 11, 2025

github-actions bot deleted the bf/recompile-context branch August 11, 2025 02:19

		recompile_user_contexts: Optional[list[Callable[[], str]]] = None


		def register_hook_for_recompile_user_context(hook: Callable[[], str]) -> None:

[Log] add a hook for recompile user context #157961

[Log] add a hook for recompile user context #157961

Uh oh!

Conversation

BoyuanFeng commented Jul 9, 2025 • edited by pytorch-bot bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jul 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/157961

❌ 1 Cancelled Job, 1 Unrelated Failure

Uh oh!

Choose a reason for hiding this comment

Uh oh!

williamwen42 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot commented Jul 10, 2025

Uh oh!

masnesral left a comment

Choose a reason for hiding this comment

Uh oh!

BoyuanFeng commented Jul 11, 2025

Uh oh!

pytorchmergebot commented Jul 11, 2025

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

BoyuanFeng commented Jul 9, 2025 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Jul 9, 2025 •

edited

Loading