[schema_upgrader] add C++ upgrader for json based upgrading by ydwu4 · Pull Request #156761 · pytorch/pytorch

ydwu4 · 2025-06-24T22:41:34Z

Stack from ghstack (oldest at bottom):

-> [schema_upgrader] add C++ upgrader for json based upgrading #156761

Differential Revision: D77459912

[ghstack-poisoned]

pytorch-bot · 2025-06-24T22:41:38Z

🔗 Helpful Links

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

📄 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.

✅ You can merge normally! (1 Unrelated Failure)

As of commit f466c82 with merge base 3ee75b7 ():

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.

[ghstack-poisoned]

ghstack-source-id: dd45748 Pull Request resolved: #156761

[ghstack-poisoned]

ghstack-source-id: 07be309 Pull Request resolved: #156761

angelayi · 2025-06-25T22:58:23Z

torch/csrc/export/upgrader.cpp

+
+// NOTE: The following version_0 and version_1 upgraders are for testing
+// purposes only. They demonstrate the upgrader system functionality and are
+// used in test/export/test_upgrader.py.


should these be put inside of a testing folder? you could follow what I did to test AOTI custom ops and just added a cpp file in the testing directory.

Good idea. I put the test upgraders in a seperate file, add a python binding for it and the python test can dynamically register/deregister it. I don't want to add a legit C++ test, the build system is complicated lol

angelayi · 2025-06-25T22:59:15Z

torch/csrc/export/upgrader.cpp

+  return keypath < other.keypath;
+}
+
+static void registerUpgrader(


how come this function isn't in the .h? it seems like this function is the most user facing so we should add proper docs

Yeah, I put it into the interface now. The initial idea is all regsiteration can be just done inside this file but it doesn't hurt to add some flexibility.

angelayi · 2025-06-25T22:59:49Z

torch/csrc/export/upgrader.cpp

+}
+
+static void registerUpgrader(
+    int version,


nit: add docstring that version is the version to upgrade from?

angelayi · 2025-06-25T23:00:47Z

torch/csrc/export/upgrader.cpp

+
+static void registerUpgrader(
+    int version,
+    const std::vector<std::string>& keypath,


what if we just passed in a string keypath and split it with .? ex. instead of passing in {"graph_module", "graph", "nodes"} we can just pass in "graph_module.graph.nodes"?

Add a separate overload for this purpose. My initial thought is that we don't want to hardcode some deliminator into the key path, which can be troublesome somtimes (e.g. the nn_module_stack deliminator).

angelayi · 2025-06-25T23:12:03Z

torch/csrc/export/upgrader.cpp

+  int current_version = current_artifact["schema_version"]["major"];
+
+  // Iteratively apply upgraders until no more are available
+  while (true) {


should we also stop if current_version == the current highest schema version? Maybe that's a little hard to find, since it's in python...

just wondering what if there aren't any more upgraders, but current_version hasn't reached the version consumable by the deserializer yet.

Yeah, we could query the schema version and pass it to the upgrader function.

torch/csrc/export/upgrader.h

[ghstack-poisoned]

angelayi · 2025-06-26T22:50:23Z

torch/csrc/export/upgrader.cpp

+  }
+
+  // Validate that we reached the target version if requested
+  if (validate_target && current_version != target_version) {


Wouldn't we always want to validate this?

sometimes, it could be easier to use not provide a target_version e.g. in testing lol.

angelayi · 2025-06-26T22:51:06Z

torch/csrc/export/upgrader.cpp

+}
+
+nlohmann::json upgrade(const nlohmann::json& artifact) {
+  return upgradeImpl(artifact, std::numeric_limits<int>::max(), false);


In what case would we want to use max? I feel like we would just always want the existing schema version?

easier for testing. Could change to always upgrade to a target version though

People might not want to always upgrade to the latest version e.g. for debug reasons maybe? It doesn't hurt to provide some flexibility i feel

hm but are you using this anywhere? all your test cases are specifying the version

angelayi · 2025-06-26T22:56:36Z

torch/csrc/export/pybind.cpp

+
+        // Query the current Python schema version as target
+        py::module_ schema_module =
+            py::module_::import("torch._export.serde.schema");


I wonder if we could add the schema version to generated_serialization_types.h so that we can query the schema version directly in the upgrade cpp function rather than here. cc @zhxchen17

can leave as a TODO for now

Right now it's inside the exported program class, we would need to deserialize ExportedProgram first in order to access the field but we cannot because we need to first do upgrade. Either we could have a top-level class that wraps up ExportedProgram and the version and make sure the schema version can be accessed all the time or we store it somewhere else.

edit: oh, you mean we expose the schema_version in generated_serialization_types.h as a constant?

edit: oh, you mean we expose the schema_version in generated_serialization_types.h as a constant?

yup

torch/csrc/export/example_upgraders.cpp

[ghstack-poisoned]

ghstack-source-id: 8168aea Pull Request resolved: #156761

angelayi · 2025-06-27T17:35:24Z

torch/csrc/export/example_upgraders.cpp

@@ -0,0 +1,89 @@
+#include <torch/csrc/export/example_upgraders.h>


I still wish this file was in test/ but I guess if the build system is too annoying to figure out it's ok

lol, it's indeed annoying

angelayi · 2025-06-27T17:36:59Z

torch/csrc/export/upgrader.cpp

+}
+
+nlohmann::json upgrade(const nlohmann::json& artifact) {
+  return upgradeImpl(artifact, std::numeric_limits<int>::max(), false);


hm but are you using this anywhere? all your test cases are specifying the version

angelayi

letsgoo

you might want to import to internal before landing, just to make sure none of the cpp changes will break

ydwu4 · 2025-06-27T19:47:57Z

@ydwu4 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

ydwu4 · 2025-06-27T23:43:00Z

@pytorchbot merge

pytorchmergebot · 2025-06-27T23:44:45Z

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

ydwu4 · 2025-06-28T03:56:44Z

@pytorchbot revert -m "break linter test, which doesn't show up in the pr" -c nosignal

pytorchmergebot · 2025-06-28T03:58:15Z

@pytorchbot successfully started a revert job. Check the current status here.
Questions? Feedback? Please reach out to the PyTorch DevX Team

…156761)" This reverts commit 61712e6. Reverted #156761 on behalf of https://github.com/ydwu4 due to break linter test, which doesn't show up in the pr ([comment](#156761 (comment)))

pytorchmergebot · 2025-06-28T03:58:28Z

@ydwu4 your PR has been successfully reverted.

Differential Revision: [D77459912](https://our.internmc.facebook.com/intern/diff/D77459912) [ghstack-poisoned]

ghstack-source-id: c020e24 Pull Request resolved: #156761

ydwu4 · 2025-06-28T18:07:49Z

@pytorchbot merge

pytorchmergebot · 2025-06-28T18:09:29Z

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

[schema_upgrader] add C++ upgrader for json based upgrading

c4de49c

[ghstack-poisoned]

Update on "[schema_upgrader] add C++ upgrader for json based upgrading"

549140c

[ghstack-poisoned]

ydwu4 added a commit that referenced this pull request Jun 24, 2025

[schema_upgrader] add C++ upgrader for json based upgrading

72bb7b8

ghstack-source-id: dd45748 Pull Request resolved: #156761

ydwu4 added topic: not user facing topic category ciflow/trunk Trigger trunk jobs on your pull request labels Jun 24, 2025

Update on "[schema_upgrader] add C++ upgrader for json based upgrading"

eb3781e

[ghstack-poisoned]

ydwu4 added a commit that referenced this pull request Jun 25, 2025

[schema_upgrader] add C++ upgrader for json based upgrading

f9f878b

ghstack-source-id: 07be309 Pull Request resolved: #156761

ydwu4 requested review from angelayi and zhxchen17 June 25, 2025 17:16

angelayi reviewed Jun 25, 2025

View reviewed changes

ydwu4 added 2 commits June 26, 2025 11:31

Update on "[schema_upgrader] add C++ upgrader for json based upgrading"

fbe8e2b

[ghstack-poisoned]

Update on "[schema_upgrader] add C++ upgrader for json based upgrading"

bda1b3d

[ghstack-poisoned]

Update on "[schema_upgrader] add C++ upgrader for json based upgrading"

20a5ff6

[ghstack-poisoned]

Update on "[schema_upgrader] add C++ upgrader for json based upgrading"

4a75fc3

[ghstack-poisoned]

Update on "[schema_upgrader] add C++ upgrader for json based upgrading"

7c29481

[ghstack-poisoned]

angelayi reviewed Jun 26, 2025

View reviewed changes

Update on "[schema_upgrader] add C++ upgrader for json based upgrading"

bd727c3

[ghstack-poisoned]

ydwu4 added a commit that referenced this pull request Jun 27, 2025

[schema_upgrader] add C++ upgrader for json based upgrading

a8c3287

ghstack-source-id: 8168aea Pull Request resolved: #156761

angelayi reviewed Jun 27, 2025

View reviewed changes

angelayi approved these changes Jun 27, 2025

View reviewed changes

pytorchmergebot added the merging label Jun 27, 2025

pytorchmergebot closed this in 61712e6 Jun 27, 2025

pytorchmergebot added Merged and removed merging labels Jun 27, 2025

pytorchmergebot added Reverted ci-no-td Do not run TD on this PR labels Jun 28, 2025

pytorchmergebot reopened this Jun 28, 2025

Update on "[schema_upgrader] add C++ upgrader for json based upgrading"

f466c82

Differential Revision: [D77459912](https://our.internmc.facebook.com/intern/diff/D77459912) [ghstack-poisoned]

ydwu4 added a commit that referenced this pull request Jun 28, 2025

[schema_upgrader] add C++ upgrader for json based upgrading

90ea1c2

ghstack-source-id: c020e24 Pull Request resolved: #156761

pytorchmergebot added the merging label Jun 28, 2025

pytorchmergebot closed this in aeffb68 Jun 28, 2025

pytorchmergebot removed the merging label Jun 28, 2025

github-actions bot deleted the gh/ydwu4/267/head branch July 29, 2025 02:20

		@@ -0,0 +1,89 @@
		#include <torch/csrc/export/example_upgraders.h>

Conversation

ydwu4 commented Jun 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jun 24, 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/156761

✅ You can merge normally! (1 Unrelated Failure)

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!

Uh oh!

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!

ydwu4 Jun 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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!

angelayi left a comment

Choose a reason for hiding this comment

Uh oh!

ydwu4 commented Jun 27, 2025

Uh oh!

ydwu4 commented Jun 27, 2025

Uh oh!

pytorchmergebot commented Jun 27, 2025

Merge started

Uh oh!

ydwu4 commented Jun 28, 2025

Uh oh!

pytorchmergebot commented Jun 28, 2025

Uh oh!

pytorchmergebot commented Jun 28, 2025

Uh oh!

ydwu4 commented Jun 28, 2025

Uh oh!

pytorchmergebot commented Jun 28, 2025

Merge started

Uh oh!

Reviewers

Assignees

ydwu4 commented Jun 24, 2025 •

edited

Loading

pytorch-bot bot commented Jun 24, 2025 •

edited

Loading

ydwu4 Jun 26, 2025 •

edited

Loading