Add nvtx.range() context manager

[chrish42]

Small quality-of-life improvement to NVTX Python bindings, that we're using internally and that would be useful to other folks using NVTX annotations via PyTorch. (And my first potential PyTorch contribution.)

Instead of needing to be careful with try/finally to make sure all your range_push'es are range_pop'ed:

nvtx.range_push("Some event")
try:
    # Code here...
finally:
    nvtx.range_pop()

you can simply do:

with nvtx.range("Some event"):
    # Code here...

or even use it as a decorator:

class MyModel(nn.Module):

    # Other methods here...

    @nvtx.range("MyModel.forward()")
    def forward(self, *input):
        # Forward pass code here...

A couple small open questions:

1. I also added the ability to call msg.format() inside range(), with the intention that, if there is nothing listening to NVTX events, we should skip the string formatting, to lower the overhead in that case. If you like that idea, I could add the actual "skip string formatting if nobody is listening to events" parts. We can also just leave it as is. Or I can remove that if you folks don't like it. (In the first two cases, should we add that to range_push() and mark() too?) Just let me know which one it is, and I'll update the pull request.

2. I don't think there are many places for bugs to hide in that function, but I can certainly add a quick test, if you folks want.

[dr-ci]

💊 CI failures summary and remediations

As of commit 352a19e (more details on the Dr. CI page):

---

• 1/2 failures possibly* introduced in this PR
  • 1/1 non-CircleCI failure(s)
• 1/2 broken upstream at merge base bda0007 on Aug 12 from 8:35am to 10:28am PDT (9 commits; 4665f3f - 3d3752d)

---

🚧 1 fixed upstream failure:

These were probably caused by upstream breakages that were already fixed.

Please rebase on the viable/strict branch (expand for instructions)

Since your merge base is older than viable/strict, run these commands:

git fetch https://github.com/pytorch/pytorch viable/strict
git rebase FETCH_HEAD

Check out the recency history of this "viable master" tracking branch.

• pytorch_macos_10_13_py3_test on Aug 12 from 8:35am to 10:28am PDT (9 commits; 4665f3f - 3d3752d)
  • 🔁 rerun

---

ci.pytorch.org: 1 failed

• Failed: pr/caffe2-pytorch-linux-xenial-rocm3.5.1-py3.6-test

---

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 4 times.

[chrish42]

Hi. I see that the code review was approved. (Thanks!) Kinda new to PyTorch contributing. Is there anything else that's needed from my part to have this merged? I can't merge it myself ("base branch restricts merging to authorized users"), at least. Just making sure this change doesn't get lost...

[chrish42]

@peterjc123 @zhangguanheng66 This is my first PyTorch contribution, so I'm not sure what is or isn't a normal delay between pull request approval and merging. Is it one of you folks that merges this? Just making sure this hasn't slipped through the cracks... Thanks!

[facebook-github-bot]

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

[facebook-github-bot]

@ezyang merged this pull request in 511f89e.