[easy] [Precompile] Refactor guards, improve typing#160530
[easy] [Precompile] Refactor guards, improve typing#160530jamesjwu wants to merge 3 commits intogh/jamesjwu/182/basefrom
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/160530
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 4f826dd with merge base ee1b041 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
| current_device: Optional[torch.device] | ||
| global_state_guard: torch._C._dynamo.guards.GlobalStateGuard | ||
| name_of_builtins_dict_key_in_fglobals: Optional[str] = None | ||
| _guards: torch._guards.GuardsSet |
There was a problem hiding this comment.
This type is basically never None. The only time it was None was in OutputGraph.init initially, and gets properly initialized within the same function. So we refactor it so it's always non null (saving an empty value initially in that one callsite)
Lucaskabela
left a comment
There was a problem hiding this comment.
Thanks for working to clean up the typing and refactor the serialization; one minor suggestion which may or may not have already tried, otherwise LGTM!
torch/_dynamo/guards.py
Outdated
| for source in output_graph.guard_on_key_order: | ||
| prune_variable(source) | ||
|
|
||
| def normalize_create_fn(x: Any) -> Any: |
There was a problem hiding this comment.
At the very least, this always takes in an returns a callable right?
|
Add better typing to create_fn, as well as normalize_create_fn |
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Purely a refactor, improve typing and get rid of some type errors. Make certain fields as nonnull, since in general it's not empty. The goal of this stack of PRs is to move the save/load logic of guard serialization into separate, flat phases, instead of being embedded in guard creation. This way, we can put a try/catch around it and fail safely if certain guards are not serializable. Pull Request resolved: pytorch#160530 Approved by: https://github.com/Lucaskabela, https://github.com/Skylion007
Purely a refactor, improve typing and get rid of some type errors. Make certain fields as nonnull, since in general it's not empty. The goal of this stack of PRs is to move the save/load logic of guard serialization into separate, flat phases, instead of being embedded in guard creation. This way, we can put a try/catch around it and fail safely if certain guards are not serializable. Pull Request resolved: pytorch#160530 Approved by: https://github.com/Lucaskabela, https://github.com/Skylion007
Stack from ghstack (oldest at bottom):
Purely a refactor, improve typing and get rid of some type errors. Make certain fields as nonnull, since in general it's not empty.
The goal of this stack of PRs is to move the save/load logic of guard serialization into separate, flat phases, instead of being embedded in guard creation. This way, we can put a try/catch around it and fail safely if certain guards are not serializable.
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames @Lucaskabela