TorchScript with record_function #44345

rohan-varma · 2020-09-09T00:07:17Z

Stack from ghstack:

Support record_shapes in RPC profiling #44419 Support record_shapes in RPC profiling
TorchScript with record_function #44345 TorchScript with record_function

As part of enhancing profiler support for RPC, when executing TorchScript functions over RPC, we would like to be able to support user-defined profiling scopes created by with record_function(...).

Since after #34705, we support with statements in TorchScript, this PR adds support for with torch.autograd.profiler.record_function to be used within TorchScript.

This can be accomplished via the following without this PR:

torch.opts.profiler._record_function_enter(...)
# Script code, such as forward pass
torch.opts.profiler._record_function_exit(....)

This is a bit hacky and it would be much cleaner to use the context manager now that we support with statements. Also, _record_function_ type operators are internal operators that are subject to change, this change will help avoid BC issues in the future.

Tested with python test/test_jit.py TestWith. test_with_record_function -v

Differential Revision: D23332074

NOTE FOR REVIEWERS: This PR has internal Facebook specific changes or comments, please review them on Phabricator!

As part of enhancing profiler support for RPC, when executing TorchScript functions over RPC, we would like to be able to support user-defined profiling scopes created by `with record_function(...)`. Since after #34705, we support `with` statements in TorchScript, this PR adds support for `with torch.autograd.profiler.record_function` to be used within TorchScript. This can be accomplished via the following without this PR: ``` torch.opts.profiler._record_function_enter(...) # Script code, such as forward pass torch.opts.profiler._record_function_exit(....) ``` This is a bit hacky and it would be much cleaner to use the context manager now that we support `with` statements. Also, `_record_function_` type operators are internal operators that are subject to change, this change will help avoid BC issues in the future. Tested with `python test/test_jit.py TestWith.test_with_foo -v` Differential Revision: [D23332074](https://our.internmc.facebook.com/intern/diff/D23332074/) **NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D23332074/)! [ghstack-poisoned]

SplitInfinity

JIT side changes look okay. I'm actually surprised it was this simple; I expected there to be some work on the JIT side in getting torch.autograd.profiler.record_function to resolve properly.

I cannot comment on the changes to autograd/profiler.py or the distributed changes.

SplitInfinity · 2020-09-09T00:25:44Z

test/jit/test_with.py

+        """
+        def with_rf(x, y):
+            # type: (Tensor, Tensor) -> Tensor
+            with torch.autograd.profiler.record_function("foo"):


Consider testing nested record_function calls to make sure the context managers are nested properly. There are tests for this already with a much simpler context manager though so it's not a must-have.

Added this test.

dr-ci · 2020-09-09T00:38:15Z

💊 CI failures summary and remediations

As of commit 2c9e63c (more details on the Dr. CI page):

2/2 failures possibly* introduced in this PR
- 2/2 non-CircleCI failure(s)

Extra GitHub checks: 1 failed

Failed: Codecov - codecov/patch

codecov.io: 1 failed

Failed: codecov/patch

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions on the GitHub issue tracker or post in the (internal) Dr. CI Users group.

See how this bot performed.

This comment has been revised 58 times.

As part of enhancing profiler support for RPC, when executing TorchScript functions over RPC, we would like to be able to support user-defined profiling scopes created by `with record_function(...)`. Since after #34705, we support `with` statements in TorchScript, this PR adds support for `with torch.autograd.profiler.record_function` to be used within TorchScript. This can be accomplished via the following without this PR: ``` torch.opts.profiler._record_function_enter(...) # Script code, such as forward pass torch.opts.profiler._record_function_exit(....) ``` This is a bit hacky and it would be much cleaner to use the context manager now that we support `with` statements. Also, `_record_function_` type operators are internal operators that are subject to change, this change will help avoid BC issues in the future. Tested with `python test/test_jit.py TestWith. test_with_record_function -v` Differential Revision: [D23332074](https://our.internmc.facebook.com/intern/diff/D23332074/) **NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D23332074/)! [ghstack-poisoned]

Pull Request resolved: #44345 As part of enhancing profiler support for RPC, when executing TorchScript functions over RPC, we would like to be able to support user-defined profiling scopes created by `with record_function(...)`. Since after #34705, we support `with` statements in TorchScript, this PR adds support for `with torch.autograd.profiler.record_function` to be used within TorchScript. This can be accomplished via the following without this PR: ``` torch.opts.profiler._record_function_enter(...) # Script code, such as forward pass torch.opts.profiler._record_function_exit(....) ``` This is a bit hacky and it would be much cleaner to use the context manager now that we support `with` statements. Also, `_record_function_` type operators are internal operators that are subject to change, this change will help avoid BC issues in the future. Tested with `python test/test_jit.py TestWith.test_with_record_function -v` ghstack-source-id: 111638601 Differential Revision: [D23332074](https://our.internmc.facebook.com/intern/diff/D23332074/) **NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D23332074/)!

As part of enhancing profiler support for RPC, when executing TorchScript functions over RPC, we would like to be able to support user-defined profiling scopes created by `with record_function(...)`. Since after #34705, we support `with` statements in TorchScript, this PR adds support for `with torch.autograd.profiler.record_function` to be used within TorchScript. This can be accomplished via the following without this PR: ``` torch.opts.profiler._record_function_enter(...) # Script code, such as forward pass torch.opts.profiler._record_function_exit(....) ``` This is a bit hacky and it would be much cleaner to use the context manager now that we support `with` statements. Also, `_record_function_` type operators are internal operators that are subject to change, this change will help avoid BC issues in the future. Tested with `python test/test_jit.py TestWith. test_with_record_function -v` Differential Revision: [D23332074](https://our.internmc.facebook.com/intern/diff/D23332074/) **NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D23332074/)! [ghstack-poisoned]

Pull Request resolved: #44345 As part of enhancing profiler support for RPC, when executing TorchScript functions over RPC, we would like to be able to support user-defined profiling scopes created by `with record_function(...)`. Since after #34705, we support `with` statements in TorchScript, this PR adds support for `with torch.autograd.profiler.record_function` to be used within TorchScript. This can be accomplished via the following without this PR: ``` torch.opts.profiler._record_function_enter(...) # Script code, such as forward pass torch.opts.profiler._record_function_exit(....) ``` This is a bit hacky and it would be much cleaner to use the context manager now that we support `with` statements. Also, `_record_function_` type operators are internal operators that are subject to change, this change will help avoid BC issues in the future. Tested with `python test/test_jit.py TestWith.test_with_record_function -v` ghstack-source-id: 111697778 Differential Revision: [D23332074](https://our.internmc.facebook.com/intern/diff/D23332074/) **NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D23332074/)!

codecov · 2020-09-09T23:04:13Z

Codecov Report

Merging #44345 into gh/rohan-varma/165/base will decrease coverage by 0.04%.
The diff coverage is 34.42%.

@@                     Coverage Diff                     @@
##           gh/rohan-varma/165/base   #44345      +/-   ##
===========================================================
- Coverage                    67.89%   67.85%   -0.05%     
===========================================================
  Files                          384      384              
  Lines                        49890    49943      +53     
===========================================================
+ Hits                         33874    33888      +14     
- Misses                       16016    16055      +39

Impacted Files	Coverage Δ
.../testing/_internal/distributed/rpc/jit/rpc_test.py	`29.13% <23.07%> (-0.52%)`	⬇️
torch/autograd/profiler.py	`78.30% <100.00%> (+0.28%)`	⬆️
torch/distributed/rpc/internal.py	`40.86% <100.00%> (+0.64%)`	⬆️
torch/testing/_internal/expecttest.py	`77.55% <0.00%> (-1.03%)`	⬇️

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update bee97d5...2c9e63c. Read the comment docs.

pritamdamania87 · 2020-09-10T01:06:17Z

torch/testing/_internal/distributed/rpc/jit/rpc_test.py

+    @dist_init
+    def test_rpc_torchscript_record_function(self):


Can we add a test where with record_function is on the caller side inside a torchscript function and we call rpc_async within the record_function context manager?

@pritamdamania87 Just to make sure we're on the same page, were you thinking of something like the following?

@torch.jit.script def foo(): with record_function(name) as rf: rpc.rpc_async(....) # Ensures profiling callbacks run on future completion and not exit of the scope. rf._call_end_callbacks_on_future()

and ensure we have the correct time for rpc_async?

@pritamdamania87 Actually it doesn't seem like the above test will work unless the return future of rpc_async in JIT is typed as Future[Any], due to current TorchScript limitations. Although, I don't anticipate a use case such as the above where a user manually has to call rf._call_end_callbacks_on_future, that should be a Python only thing which is called internally. For rpc_async in JIT, the profiling is done in the C++ operator. Or did you have a different idea of what this test should look like?

@rohan-varma I meant something similar to your example, but more concretely:

@torch.jit.script def foo(): with record_function(name) as rf: fut1 = rpc.rpc_async(....) fut2 = rpc.rpc_async(....) fut1.wait() fut2.wait() res = fut1.value() + fut2.value()

I'm guessing this should work since we're not actually using rf._call_end_callbacks_on_future?

Added a test for this. Once #44570 lands we should probably also validate with a test that rf._call_end_callbacks is torchscriptable (I patched the diff locally to verify it is)

As part of enhancing profiler support for RPC, when executing TorchScript functions over RPC, we would like to be able to support user-defined profiling scopes created by `with record_function(...)`. Since after #34705, we support `with` statements in TorchScript, this PR adds support for `with torch.autograd.profiler.record_function` to be used within TorchScript. This can be accomplished via the following without this PR: ``` torch.opts.profiler._record_function_enter(...) # Script code, such as forward pass torch.opts.profiler._record_function_exit(....) ``` This is a bit hacky and it would be much cleaner to use the context manager now that we support `with` statements. Also, `_record_function_` type operators are internal operators that are subject to change, this change will help avoid BC issues in the future. Tested with `python test/test_jit.py TestWith. test_with_record_function -v` Differential Revision: [D23332074](https://our.internmc.facebook.com/intern/diff/D23332074/) **NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D23332074/)! [ghstack-poisoned]

rohan-varma · 2020-09-12T04:36:10Z

Test will fail until #44570 is landed, we can either skip the test and enable it later or wait until it is landed.

As part of enhancing profiler support for RPC, when executing TorchScript functions over RPC, we would like to be able to support user-defined profiling scopes created by `with record_function(...)`. Since after #34705, we support `with` statements in TorchScript, this PR adds support for `with torch.autograd.profiler.record_function` to be used within TorchScript. This can be accomplished via the following without this PR: ``` torch.opts.profiler._record_function_enter(...) # Script code, such as forward pass torch.opts.profiler._record_function_exit(....) ``` This is a bit hacky and it would be much cleaner to use the context manager now that we support `with` statements. Also, `_record_function_` type operators are internal operators that are subject to change, this change will help avoid BC issues in the future. Tested with `python test/test_jit.py TestWith. test_with_record_function -v` Differential Revision: [D23332074](https://our.internmc.facebook.com/intern/diff/D23332074/) **NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D23332074/)! [ghstack-poisoned]

Pull Request resolved: #44345 As part of enhancing profiler support for RPC, when executing TorchScript functions over RPC, we would like to be able to support user-defined profiling scopes created by `with record_function(...)`. Since after #34705, we support `with` statements in TorchScript, this PR adds support for `with torch.autograd.profiler.record_function` to be used within TorchScript. This can be accomplished via the following without this PR: ``` torch.opts.profiler._record_function_enter(...) # Script code, such as forward pass torch.opts.profiler._record_function_exit(....) ``` This is a bit hacky and it would be much cleaner to use the context manager now that we support `with` statements. Also, `_record_function_` type operators are internal operators that are subject to change, this change will help avoid BC issues in the future. Tested with `python test/test_jit.py TestWith.test_with_record_function -v` ghstack-source-id: 111935620 Differential Revision: [D23332074](https://our.internmc.facebook.com/intern/diff/D23332074/) **NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D23332074/)!

As part of enhancing profiler support for RPC, when executing TorchScript functions over RPC, we would like to be able to support user-defined profiling scopes created by `with record_function(...)`. Since after #34705, we support `with` statements in TorchScript, this PR adds support for `with torch.autograd.profiler.record_function` to be used within TorchScript. This can be accomplished via the following without this PR: ``` torch.opts.profiler._record_function_enter(...) # Script code, such as forward pass torch.opts.profiler._record_function_exit(....) ``` This is a bit hacky and it would be much cleaner to use the context manager now that we support `with` statements. Also, `_record_function_` type operators are internal operators that are subject to change, this change will help avoid BC issues in the future. Tested with `python test/test_jit.py TestWith. test_with_record_function -v` Differential Revision: [D23332074](https://our.internmc.facebook.com/intern/diff/D23332074/) **NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D23332074/)! [ghstack-poisoned]

Pull Request resolved: #44345 As part of enhancing profiler support for RPC, when executing TorchScript functions over RPC, we would like to be able to support user-defined profiling scopes created by `with record_function(...)`. Since after #34705, we support `with` statements in TorchScript, this PR adds support for `with torch.autograd.profiler.record_function` to be used within TorchScript. This can be accomplished via the following without this PR: ``` torch.opts.profiler._record_function_enter(...) # Script code, such as forward pass torch.opts.profiler._record_function_exit(....) ``` This is a bit hacky and it would be much cleaner to use the context manager now that we support `with` statements. Also, `_record_function_` type operators are internal operators that are subject to change, this change will help avoid BC issues in the future. Tested with `python test/test_jit.py TestWith.test_with_record_function -v` ghstack-source-id: 111935758 Differential Revision: [D23332074](https://our.internmc.facebook.com/intern/diff/D23332074/) **NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D23332074/)!

As part of enhancing profiler support for RPC, when executing TorchScript functions over RPC, we would like to be able to support user-defined profiling scopes created by `with record_function(...)`. Since after #34705, we support `with` statements in TorchScript, this PR adds support for `with torch.autograd.profiler.record_function` to be used within TorchScript. This can be accomplished via the following without this PR: ``` torch.opts.profiler._record_function_enter(...) # Script code, such as forward pass torch.opts.profiler._record_function_exit(....) ``` This is a bit hacky and it would be much cleaner to use the context manager now that we support `with` statements. Also, `_record_function_` type operators are internal operators that are subject to change, this change will help avoid BC issues in the future. Tested with `python test/test_jit.py TestWith. test_with_record_function -v` Differential Revision: [D23332074](https://our.internmc.facebook.com/intern/diff/D23332074/) **NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D23332074/)! [ghstack-poisoned]

Pull Request resolved: #44345 As part of enhancing profiler support for RPC, when executing TorchScript functions over RPC, we would like to be able to support user-defined profiling scopes created by `with record_function(...)`. Since after #34705, we support `with` statements in TorchScript, this PR adds support for `with torch.autograd.profiler.record_function` to be used within TorchScript. This can be accomplished via the following without this PR: ``` torch.opts.profiler._record_function_enter(...) # Script code, such as forward pass torch.opts.profiler._record_function_exit(....) ``` This is a bit hacky and it would be much cleaner to use the context manager now that we support `with` statements. Also, `_record_function_` type operators are internal operators that are subject to change, this change will help avoid BC issues in the future. Tested with `python test/test_jit.py TestWith.test_with_record_function -v` ghstack-source-id: 112254463 Differential Revision: [D23332074](https://our.internmc.facebook.com/intern/diff/D23332074/) **NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D23332074/)!

As part of enhancing profiler support for RPC, when executing TorchScript functions over RPC, we would like to be able to support user-defined profiling scopes created by `with record_function(...)`. Since after #34705, we support `with` statements in TorchScript, this PR adds support for `with torch.autograd.profiler.record_function` to be used within TorchScript. This can be accomplished via the following without this PR: ``` torch.opts.profiler._record_function_enter(...) # Script code, such as forward pass torch.opts.profiler._record_function_exit(....) ``` This is a bit hacky and it would be much cleaner to use the context manager now that we support `with` statements. Also, `_record_function_` type operators are internal operators that are subject to change, this change will help avoid BC issues in the future. Tested with `python test/test_jit.py TestWith. test_with_record_function -v` Differential Revision: [D23332074](https://our.internmc.facebook.com/intern/diff/D23332074/) **NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D23332074/)! [ghstack-poisoned]

pritamdamania87

Test will fail until #44570 is landed, we can either skip the test and enable it later or wait until it is landed.

Approving since it looks like #44570 has landed.

facebook-github-bot · 2020-09-18T02:14:07Z

This pull request has been merged in 5dbcbea.

Summary: Pull Request resolved: #44345 As part of enhancing profiler support for RPC, when executing TorchScript functions over RPC, we would like to be able to support user-defined profiling scopes created by `with record_function(...)`. Since after #34705, we support `with` statements in TorchScript, this PR adds support for `with torch.autograd.profiler.record_function` to be used within TorchScript. This can be accomplished via the following without this PR: ``` torch.opts.profiler._record_function_enter(...) # Script code, such as forward pass torch.opts.profiler._record_function_exit(....) ``` This is a bit hacky and it would be much cleaner to use the context manager now that we support `with` statements. Also, `_record_function_` type operators are internal operators that are subject to change, this change will help avoid BC issues in the future. Tested with `python test/test_jit.py TestWith.test_with_record_function -v` ghstack-source-id: 112320645 Test Plan: Repro instructions: 1) Change `def script_add_ones_return_any(x) -> Any` to `def script_add_ones_return_any(x) -> Tensor` in `jit/rpc_test.py` 2) `buck test mode/dev-nosan //caffe2/test/distributed/rpc:process_group_agent -- test_record_function_on_caller_rpc_async --print-passing-details` 3) The function which ideally should accept `Future[Any]` is `def _call_end_callbacks_on_future` in `autograd/profiler.py`. python test/test_jit.py TestWith.test_with_foo -v Reviewed By: pritamdamania87 Differential Revision: D23332074 fbshipit-source-id: 61b0078578e8b23bfad5eeec3b0b146b6b35a870

rohan-varma requested review from albanD, apaszke, mrshenli, pritamdamania87 and zhaojuanmao as code owners September 9, 2020 00:07

facebook-github-bot added the oncall: jit Add this issue/PR to JIT oncall triage queue label Sep 9, 2020

rohan-varma requested review from SplitInfinity and wanchaol September 9, 2020 00:10

SplitInfinity reviewed Sep 9, 2020

View reviewed changes

SplitInfinity linked an issue Sep 9, 2020 that may be closed by this pull request

[JIT] Support with torch.autograd.profiler.record_function #40260

Closed

rohan-varma mentioned this pull request Sep 9, 2020

Support record_shapes in RPC profiling #44419

Closed

pritamdamania87 reviewed Sep 10, 2020

View reviewed changes

rohan-varma requested a review from pritamdamania87 September 12, 2020 04:36

rohan-varma added 2 commits September 11, 2020 21:38

rohan-varma added 3 commits September 14, 2020 11:07

rohan-varma requested a review from pietern as a code owner September 17, 2020 02:44

rohan-varma requested a review from SplitInfinity September 17, 2020 02:52

pritamdamania87 approved these changes Sep 18, 2020

View reviewed changes

facebook-github-bot closed this in 5dbcbea Sep 18, 2020

facebook-github-bot added the merged label Sep 18, 2020

facebook-github-bot deleted the gh/rohan-varma/165/head branch September 21, 2020 14:14

mruberry added the Merged label Oct 28, 2020

TorchScript with record_function #44345

TorchScript with record_function #44345

Uh oh!

Conversation

rohan-varma commented Sep 9, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

SplitInfinity left a comment

Choose a reason for hiding this comment

Uh oh!

SplitInfinity Sep 9, 2020

Choose a reason for hiding this comment

Uh oh!

rohan-varma Sep 9, 2020

Choose a reason for hiding this comment

Uh oh!

dr-ci bot commented Sep 9, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

💊 CI failures summary and remediations

Extra GitHub checks: 1 failed

codecov.io: 1 failed

Uh oh!

codecov bot commented Sep 9, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

pritamdamania87 Sep 10, 2020

Choose a reason for hiding this comment

Uh oh!

rohan-varma Sep 10, 2020 • 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!

rohan-varma Sep 11, 2020

Choose a reason for hiding this comment

Uh oh!

pritamdamania87 Sep 17, 2020

Choose a reason for hiding this comment

Uh oh!

rohan-varma Sep 17, 2020

Choose a reason for hiding this comment

Uh oh!

rohan-varma commented Sep 12, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pritamdamania87 left a comment

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot commented Sep 18, 2020

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

rohan-varma commented Sep 9, 2020 •

edited

Loading

dr-ci bot commented Sep 9, 2020 •

edited

Loading

codecov bot commented Sep 9, 2020 •

edited

Loading

rohan-varma Sep 10, 2020 •

edited

Loading

rohan-varma commented Sep 12, 2020 •

edited

Loading