Add a registry for GraphModuleSerializer #126550

jiashenC · 2024-05-17T16:28:41Z

This PR adds a registration function and a global registry for GraphModuleSerializer. After this PR, custom serialization methods can be done through registration instead of subclassing for ease of maintenance.

Changes

Add a test case where it injects custom op to test serialization.
Add custom op handler
Change allowed op for verifier

pytorch-bot · 2024-05-17T16:28:45Z

🔗 Helpful Links

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

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

✅ No Failures

As of commit 3acfed8 with merge base 9117779 ():
💚 Looks good so far! There are no failures yet. 💚

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

torch/_export/serde/serialize.py

zhxchen17 · 2024-05-22T16:27:22Z

torch/_export/verifier.py

@@ -187,7 +187,7 @@ def _allowed_op_types() -> Tuple[Type[Any], ...]:
            )

            if not isinstance(op, _allowed_op_types()):
-                if op not in _allowed_builtin_ops() and op not in _allowed_torch_functions:
+                if op not in _allowed_builtin_ops() and op not in _allowed_torch_functions and not isinstance(op, torch._export.serde.serialize.CustomOpHandler):


what's the error you saw? I'm confused by the change here.

Suggested change

if op not in _allowed_builtin_ops() and op not in _allowed_torch_functions and not isinstance(op, torch._export.serde.serialize.CustomOpHandler):

if op not in _allowed_builtin_ops() and op not in _allowed_torch_functions:

I think without this change, there is no easy way to test the custom op handler. If I inject the custom op in the graph, it will throw error during verification.

torch._export.verifier.SpecViolationError: Operator '<export.test_serialize.TestSerialize.test_export_with_custom_op_serialization.<locals>.FooCustomOp object at 0x7f713ef391b0>' is not an allowed operator type: (<class 'torch._ops.OpOverload'>, <class 'torch._ops.HigherOrderOperator'>)

then sounds like we need to change _allowed_op_types()?

zhxchen17 · 2024-05-22T16:27:42Z

torch/_export/serde/serialize.py

+    if namespace not in _serialization_registry:
+        _serialization_registry[namespace] = {}
+    _serialization_registry[namespace][op_name] = serialization_fn


Suggested change

if namespace not in _serialization_registry:

_serialization_registry[namespace] = {}

_serialization_registry[namespace][op_name] = serialization_fn

_deserialization_registry[namespace] = op_handler

torch/_export/serde/serialize.py

zhxchen17 · 2024-05-22T20:08:21Z

torch/_export/verifier.py

@@ -187,7 +187,7 @@ def _allowed_op_types() -> Tuple[Type[Any], ...]:
            )

            if not isinstance(op, _allowed_op_types()):
-                if op not in _allowed_builtin_ops() and op not in _allowed_torch_functions:
+                if op not in _allowed_builtin_ops() and op not in _allowed_torch_functions and not isinstance(op, torch._export.serde.serialize.CustomOpHandler):


then sounds like we need to change _allowed_op_types()?

zhxchen17

looks great!

Please address linter issues and other comments before landing this.

torch/_export/serde/serialize.py

Co-authored-by: Zhengxu Chen <zhxchen17@outlook.com>

facebook-github-bot · 2024-05-24T22:31:14Z

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

facebook-github-bot · 2024-05-28T16:28:13Z

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

facebook-github-bot · 2024-05-29T03:10:14Z

@pytorchbot merge -f 'Landed internally'

(Initiating merge automatically since Phabricator Diff has merged, using force because this PR might not pass merge_rules.json but landed internally)

pytorchmergebot · 2024-05-29T03:12:38Z

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

This PR adds a registration function and a global registry for GraphModuleSerializer. After this PR, custom serialization methods can be done through registration instead of subclassing for ease of maintenance. ## Changes - Add a test case where it injects custom op to test serialization. - Add custom op handler - Change allowed op for verifier Co-authored-by: Zhengxu Chen <zhxchen17@outlook.com> Pull Request resolved: pytorch#126550 Approved by: https://github.com/zhxchen17

pytorch-bot bot added the ciflow/inductor label May 17, 2024

jiashenC requested review from angelayi and zhxchen17 May 17, 2024 16:37

jiashenC force-pushed the hook_for_custom_serializer branch from fb2250a to 1e4b4da Compare May 21, 2024 20:46

jiashenC marked this pull request as ready for review May 21, 2024 20:47

jiashenC requested review from avikchaudhuri, gmagogsfm and tugsbayasgalan as code owners May 21, 2024 20:47

zhxchen17 requested changes May 22, 2024

View reviewed changes

jiashenC requested a review from zhxchen17 May 23, 2024 23:28

zhxchen17 approved these changes May 23, 2024

View reviewed changes

torch/_export/serde/serialize.py Outdated Show resolved Hide resolved

jiashenC and others added 15 commits May 24, 2024 15:19

draft for serialization func registry

211a991

add custom handler and test case

a48fb7b

Update torch/_export/serde/serialize.py

8142f80

Co-authored-by: Zhengxu Chen <zhxchen17@outlook.com>

Update torch/_export/serde/serialize.py

71c8129

Co-authored-by: Zhengxu Chen <zhxchen17@outlook.com>

Update torch/_export/serde/serialize.py

aca42bf

Co-authored-by: Zhengxu Chen <zhxchen17@outlook.com>

Update torch/_export/serde/serialize.py

61956da

Co-authored-by: Zhengxu Chen <zhxchen17@outlook.com>

address comments and fix lint erros

f37d866

Update torch/_export/serde/serialize.py

6005eae

Co-authored-by: Zhengxu Chen <zhxchen17@outlook.com>

Update torch/_export/serde/serialize.py

116e92c

Co-authored-by: Zhengxu Chen <zhxchen17@outlook.com>

Update torch/_export/serde/serialize.py

1768c91

Co-authored-by: Zhengxu Chen <zhxchen17@outlook.com>

improve PR based on discussions

e490af7

fix abs class

66c245a

fix lint error

0caa3ab

fix class method

7dd8d6d

fix lint errors

b7689d4

jiashenC force-pushed the hook_for_custom_serializer branch from 25302d3 to b7689d4 Compare May 24, 2024 22:30

fix import errors

3acfed8

jiashenC added the topic: not user facing topic category label May 28, 2024

pytorchmergebot added the merging label May 29, 2024

pytorchmergebot closed this in 10d2373 May 29, 2024

pytorchmergebot added Merged and removed merging labels May 29, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add a registry for GraphModuleSerializer #126550

Add a registry for GraphModuleSerializer #126550

jiashenC commented May 17, 2024 •

edited

pytorch-bot bot commented May 17, 2024 •

edited

zhxchen17 May 22, 2024

jiashenC May 22, 2024

zhxchen17 May 22, 2024

zhxchen17 May 22, 2024

zhxchen17 May 22, 2024

zhxchen17 left a comment

facebook-github-bot commented May 24, 2024

facebook-github-bot commented May 28, 2024

facebook-github-bot commented May 29, 2024

pytorchmergebot commented May 29, 2024

	if op not in _allowed_builtin_ops() and op not in _allowed_torch_functions and not isinstance(op, torch._export.serde.serialize.CustomOpHandler):
	if op not in _allowed_builtin_ops() and op not in _allowed_torch_functions:

Add a registry for GraphModuleSerializer #126550

Add a registry for GraphModuleSerializer #126550

Conversation

jiashenC commented May 17, 2024 • edited

Changes

pytorch-bot bot commented May 17, 2024 • edited

🔗 Helpful Links

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

✅ No Failures

zhxchen17 May 22, 2024

Choose a reason for hiding this comment

jiashenC May 22, 2024

Choose a reason for hiding this comment

zhxchen17 May 22, 2024

Choose a reason for hiding this comment

zhxchen17 May 22, 2024

Choose a reason for hiding this comment

zhxchen17 May 22, 2024

Choose a reason for hiding this comment

zhxchen17 left a comment

Choose a reason for hiding this comment

facebook-github-bot commented May 24, 2024

facebook-github-bot commented May 28, 2024

facebook-github-bot commented May 29, 2024

pytorchmergebot commented May 29, 2024

Merge started

jiashenC commented May 17, 2024 •

edited

pytorch-bot bot commented May 17, 2024 •

edited