[Profiler] Memory profiler part 2: Config validation

Taylor Robie · Taylor Robie · commit 2c110b0b29f4 · 2022-10-31T14:46:42.000-07:00
Pull Request resolved: #86853 Memory profiling requires `record_shapes`, `profile_memory`, and `with_stack`. This PR just adds a skeleton endpoint with a good error message if certain flags are missing. ghstack-source-id: 172169616 Differential Revision: [D39920801](https://our.internmc.facebook.com/intern/diff/D39920801/)
diff --git a/test/profiler/test_memory_profiler.py b/test/profiler/test_memory_profiler.py
@@ -12,6 +12,28 @@
     torch.profiler.profile, record_shapes=True, profile_memory=True, with_stack=True
 )
 
+@skipIfTorchDynamo("TorchDynamo removes profiler altogether.")
+class TestMemoryProfiler(TestCase):
+    def test_config_check(self) -> None:
+        with torch.profiler.profile() as prof:
+            pass
+
+        pattern = r"record_shapes=True, profile_memory=True, with_stack=True"
+        with self.assertRaisesRegex(ValueError, pattern):
+            prof._memory_profile()
+
+        with torch.profiler.profile(record_shapes=True, with_stack=True) as prof:
+            pass
+
+        pattern = r"^profile_memory=True required for memory profiling\.$"
+        with self.assertRaisesRegex(ValueError, pattern):
+            prof._memory_profile()
+
+        with profile() as prof:
+            pass
+
+        self.assertIsInstance(prof._memory_profile(), _memory_profiler.MemoryProfile)
+
 
 class ScaleLayer(torch.nn.Module):
     def __init__(self) -> None:
diff --git a/torch/profiler/_memory_profiler.py b/torch/profiler/_memory_profiler.py
@@ -2,7 +2,13 @@
 from typing import Any, Iterator, Optional, Tuple
 
 import torch
-from torch._C._profiler import _EventType, _ProfilerEvent, _TensorMetadata, RecordScope
+from torch._C._autograd import _ProfilerResult
+from torch._C._profiler import (
+    _EventType,
+    _ProfilerEvent,
+    _TensorMetadata,
+    RecordScope,
+)
 
 
 @dataclasses.dataclass
@@ -112,3 +118,8 @@ def extract_gradients(
                 p_grad_key = TensorKey.from_tensor(p_grad)
                 if p_grad_key is not None:
                     yield TensorKey.from_tensor(p), p_grad_key
+
+
+class MemoryProfile:
+    def __init__(self, result: _ProfilerResult) -> None:
+        pass
diff --git a/torch/profiler/profiler.py b/torch/profiler/profiler.py
@@ -16,10 +16,19 @@
     _ExperimentalConfig,
     _remove_execution_graph_observer,
 )
-from torch.autograd import ProfilerActivity, kineto_available
+from torch.autograd import kineto_available, ProfilerActivity
+from torch.profiler import _memory_profiler
+
+
+__all__ = [
+    "supported_activities",
+    "ProfilerAction",
+    "schedule",
+    "tensorboard_trace_handler",
+    "profile",
+    "ExecutionGraphObserver",
+]
 
-__all__ = ['supported_activities', 'ProfilerAction', 'schedule', 'tensorboard_trace_handler', 'profile',
-           'ExecutionGraphObserver']
 
 def supported_activities():
     """
@@ -208,6 +217,15 @@ def _get_distributed_info(self):
             "world_size": dist.get_world_size()
         }
 
+    def _memory_profile(self) -> _memory_profiler.MemoryProfile:
+        required = ("record_shapes", "profile_memory", "with_stack")
+        missing = [f"{i}=True" for i in required if not getattr(self, i)]
+        if missing:
+            raise ValueError(f"{', '.join(missing)} required for memory profiling.")
+
+        assert self.profiler is not None and self.profiler.kineto_results is not None
+        return _memory_profiler.MemoryProfile(self.profiler.kineto_results)
+
 
 class ProfilerAction(Enum):
     """
