add enabled API to autograd profiler (pytorch#31380)

Summary: Pull Request resolved: pytorch#31380 For being able to profile async RPCs, we attach a `RecordFunction` object to the future that is created during the RPC to persist it across the lifetime of the RPC (this is implemented in the next PR: ). Since we'd only like to do this when profiling is enabled, this PR adds an enabled API to the autograd profiler. ghstack-source-id: 96053933 Test Plan: Modified unit test. Differential Revision: D19050391 fbshipit-source-id: aa382110e69d06b4a84c83b31d2bec2d8a81ba10

Author: rohan-varma

test/test_autograd.py

@@ -2523,8 +2523,11 @@ def test_profiler(self):
  x = torch.randn(10, 10)
 
  with profile() as p:
+ self.assertTrue(torch.autograd._profiler_enabled())
  y = x * 2 + 4
 
+ self.assertFalse(torch.autograd._profiler_enabled())
+
  last_end = 0
  names = ['mul', 'add']
  self.assertEqual(len(p.function_events), len(names))

torch/csrc/autograd/init.cpp

@@ -50,6 +50,7 @@ PyObject* THPAutograd_initExtension(PyObject* _unused, PyObject *unused) {
 
  m.def("_enable_profiler", enableProfiler);
  m.def("_disable_profiler", disableProfiler);
+ m.def("_profiler_enabled", profilerEnabled);
 
  m.def("_push_range", [](std::string name) { pushRange(std::move(name)); });
  m.def("_pop_range", []() { popRange(); });

torch/csrc/autograd/profiler.cpp

@@ -54,6 +54,10 @@ void mark(std::string name, bool include_cuda /* = true */) {
  }
 }
 
+bool profilerEnabled() {
+ return state != ProfilerState::Disabled;
+}
+
 void pushRangeImpl(
  const StringView& name,
  const char* msg = "",

torch/csrc/autograd/profiler.h

@@ -230,6 +230,7 @@ using thread_event_lists = std::vector<std::vector<Event>>;
 // there no autograd functions are being executed when these function are used.
 TORCH_API void enableProfiler(ProfilerConfig);
 TORCH_API thread_event_lists disableProfiler();
+TORCH_API bool profilerEnabled();
 
 
 // Usage: