NVTX Push and Pop Context Manager?

[wkaisertexas]

🚀 The feature, motivation and pitch

Could the two functions:

• torch.cuda.nvtx.range_push("msg")
• torch.cuda.nvtx.range_pop()

Be combined into a easy-to-use context manager?

"""
from chatgpt
"""
import torch.cuda.nvtx
from contextlib import contextmanager

@contextmanager
def nvtx_context(message):
    try:
        torch.cuda.nvtx.range_push(message)
        yield
    finally:
        torch.cuda.nvtx.range_pop()

# Example usage:
with nvtx_context("my_message"):
    # Code block where NVTX range is active
    pass

Doing this would be dead simple and make using range push and pop less error-prone.

Alternatives

No response

Additional context

No response

cc @ptrblck

[Aidyn-A]

Hi @wkaisertexas,
Thanks for submitting the feature request. Fortunately, there is already a context manager that does exactly what you described:

pytorch/torch/cuda/nvtx.py

Lines 77 to 91 in 5b5d423

	@contextmanager
	def range(msg, *args, **kwargs):
	"""
	Context manager / decorator that pushes an NVTX range at the beginning
	of its scope, and pops it at the end. If extra arguments are given,
	they are passed as arguments to msg.format().
	Args:
	msg (str): message to associate with the range
	"""
	range_push(msg.format(*args, **kwargs))
	try:
	yield
	finally:
	range_pop()

Although, I am afraid it is not documented.

[Aidyn-A]

@wkaisertexas, as the torch.cuda.nvtx.range context manager is now being documented https://pytorch.org/docs/main/generated/torch.cuda.nvtx.range.html I am closing this issue. Please fell free to comment and re-open this issue if necessary.

[wkaisertexas]

Yeah, I saw that here:

#121699