Document that autograd::Profiler::RecordFunction is available in Python

[heiner]

🚀 Feature

torch.autograd.profiler.profile().export_chrome_trace is great. Adding information to it from Python would make it even better.

Motivation

In multi-threaded PyTorch code, it can be useful to further annotate the tracing output by adding additional information to it.

Pitch

One easy / cheap way of doing this:

PYBIND11_MODULE(torch, m) {
  py::class_<torch::autograd::profiler::RecordFunction>(m, "Record")
      .def(py::init<std::string>(), py::arg("name"));
}

And then in Python

record = torch.Record("model")
outputs = model(frame)
del record  # Alternatively, turn it into a context manager.

[ezyang]

We'll accept a PR that does this. A context manager would be better than a faux-RAII object, as you can't easily control scope of objects in Python. How about calling it torch.autograd.profiler.record_function(label) for consistency with the C++ name?

[heiner]

It seems f1c8e01 has changed the RecordFunction object in a way that's not immediatly discernible to me. (Thereby also breaking my C++ extension, but I guess there's just no expectation of API stability for those parts.)

I'm still happy to supply this PR once I've understood what the second argument to the new RECORD_FUNCTION macro is about.

[ezyang]

cc @ilia-cher, could you answer the question above?

[ilia-cher]

I would definitely prefer a context manager instead of manual 'del' which is easy to forget to call.
Regarding #18717 - motivation for updates to RecordFunction code can be found in #18966. Specifically, #18717 allows us to pass inputs of the profiled function into the callbacks registered in RecordFunction.
Since you just want to pass scope name, you can use the code similar to:

torch::autograd::profiler::RecordFunction guard;
guard.before(scope_name);

when entering the context manager, and destroy guard when exiting.

[apaszke]

We already have torch.autograd.profiler.range. It's just not very well documented.

with torch.autograd.profiler.range('name'):
  ...

And yes, it definitely shouldn't depend on del. Context-manager is the way to go.

[zhangguanheng66]

Hey, I will give this a shot.