KEMBAR78
[ONNX] Hide draft export under a flag by justinchuby · Pull Request #162225 · pytorch/pytorch · GitHub
Skip to content

[ONNX] Hide draft export under a flag #162225

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
wants to merge 5 commits into from
Closed

[ONNX] Hide draft export under a flag #162225

wants to merge 5 commits into from

Conversation

justinchuby
Copy link
Collaborator

@justinchuby justinchuby commented Sep 4, 2025 •
edited by pytorch-bot bot
Loading

Use TORCH_ONNX_ENABLE_DRAFT_EXPORT to control whether draft_export should be used as a strategy in onnx export.

Follow up of #161454

cc @titaiwangms

@justinchuby
[ONNX] Hide draft export under a flag
ad260db 
Use `TORCH_ONNX_ENABLE_DRAFT_EXPORT` to control whether draft_export should be used as a strategy in onnx export.

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
@pytorch-bot
Copy link

pytorch-bot bot commented Sep 4, 2025 •
edited
Loading

🔗 Helpful Links

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

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

⏳ No Failures, 1 Pending

As of commit f422da9 with merge base 0d84ff3 (image):
💚 Looks good so far! There are no failures yet. 💚

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

@pytorch-bot pytorch-bot bot added the release notes: onnx torch.onnx related changes that should show up in the release notes label Sep 4, 2025
@justinchuby justinchuby requested a review from Copilot September 4, 2025 22:59
@justinchuby justinchuby added module: onnx Related to torch.onnx topic: bc breaking topic category labels Sep 4, 2025
Copilot
Copilot AI reviewed Sep 4, 2025
View reviewed changes
Copy link
Contributor

@Copilot Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull Request Overview

This PR introduces a feature flag TORCH_ONNX_ENABLE_DRAFT_EXPORT to conditionally enable the TorchExportDraftExportStrategy in ONNX export. The draft export functionality is now hidden behind this flag, preventing it from being used by default while still allowing opt-in access.

  • Adds ENABLE_DRAFT_EXPORT boolean flag controlled by environment variable
  • Conditionally includes TorchExportDraftExportStrategy in capture strategies based on flag
  • Updates export status tracking to handle the draft export strategy

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.

File Description
torch/onnx/_flags.py Replaces placeholder flag with ENABLE_DRAFT_EXPORT flag (default: False)
torch/onnx/_internal/exporter/_capture_strategies.py Conditionally adds draft export strategy and fixes import path
torch/onnx/_internal/exporter/_core.py Adds status tracking for draft export strategy

Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.

@justinchuby
wording
f1d7d13 
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
@justinchuby
api
38f7256 
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
@justinchuby justinchuby added this to the 2.9.0 milestone Sep 4, 2025
@justinchuby justinchuby added the ciflow/trunk Trigger trunk jobs on your pull request label Sep 4, 2025
@justinchuby
Copy link
Collaborator Author

@pytorchbot merge

@pytorchmergebot
Copy link
Collaborator

Merge started

Your 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

Advanced Debugging
Check the merge workflow status
here

titaiwangms
titaiwangms approved these changes Sep 5, 2025
View reviewed changes
@justinchuby
Copy link
Collaborator Author

@pytorchbot merge -f "trunk/win-vs2022 is taking too long and it's unrelated"

@pytorchmergebot
Copy link
Collaborator

The merge job was canceled or timed out. This most often happen if two merge requests were issued for the same PR, or if merge job was waiting for more than 6 hours for tests to finish. In later case, please do not hesitate to reissue the merge command
For more information see pytorch-bot wiki.

@pytorchmergebot
Copy link
Collaborator

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

@justinchuby justinchuby deleted the justinchu/draft-export-flag branch September 5, 2025 20:50
daisyden pushed a commit to daisyden/pytorch that referenced this pull request Sep 8, 2025
@justinchuby @daisyden
[ONNX] Hide draft export under a flag (pytorch#162225)
6b5a874 
Use `TORCH_ONNX_ENABLE_DRAFT_EXPORT` to control whether draft_export should be used as a strategy in onnx export.

Follow up of pytorch#161454

Pull Request resolved: pytorch#162225
Approved by: https://github.com/xadupre, https://github.com/titaiwangms
markc-614 pushed a commit to markc-614/pytorch that referenced this pull request Sep 17, 2025
@justinchuby @markc-614
[ONNX] Hide draft export under a flag (pytorch#162225)
4dcab1b 
Use `TORCH_ONNX_ENABLE_DRAFT_EXPORT` to control whether draft_export should be used as a strategy in onnx export.

Follow up of pytorch#161454

Pull Request resolved: pytorch#162225
Approved by: https://github.com/xadupre, https://github.com/titaiwangms
mansiag05 pushed a commit to mansiag05/pytorch that referenced this pull request Sep 22, 2025
@justinchuby @mansiag05
[ONNX] Hide draft export under a flag (pytorch#162225)
2c9dbe3 
Use `TORCH_ONNX_ENABLE_DRAFT_EXPORT` to control whether draft_export should be used as a strategy in onnx export.

Follow up of pytorch#161454

Pull Request resolved: pytorch#162225
Approved by: https://github.com/xadupre, https://github.com/titaiwangms
cleonard530 pushed a commit to cleonard530/pytorch that referenced this pull request Sep 22, 2025
@justinchuby @cleonard530
[ONNX] Hide draft export under a flag (pytorch#162225)
d849430 
Use `TORCH_ONNX_ENABLE_DRAFT_EXPORT` to control whether draft_export should be used as a strategy in onnx export.

Follow up of pytorch#161454

Pull Request resolved: pytorch#162225
Approved by: https://github.com/xadupre, https://github.com/titaiwangms
dsashidh pushed a commit to dsashidh/pytorch that referenced this pull request Sep 26, 2025
@justinchuby @dsashidh
[ONNX] Hide draft export under a flag (pytorch#162225)
9471e29 
Use `TORCH_ONNX_ENABLE_DRAFT_EXPORT` to control whether draft_export should be used as a strategy in onnx export.

Follow up of pytorch#161454

Pull Request resolved: pytorch#162225
Approved by: https://github.com/xadupre, https://github.com/titaiwangms
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

Copilot code review Copilot Copilot left review comments

@titaiwangms titaiwangms titaiwangms approved these changes

+1 more reviewer

@xadupre xadupre xadupre approved these changes

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

ciflow/trunk Trigger trunk jobs on your pull request Merged module: onnx Related to torch.onnx open source release notes: onnx torch.onnx related changes that should show up in the release notes topic: bc breaking topic category

Projects

None yet

Milestone

2.9.0

Development

Successfully merging this pull request may close these issues.

5 participants

@justinchuby @pytorchmergebot @titaiwangms @xadupre @pytorchbot