[GitHub] [beam] angoenka commented on a change in pull request #12562: [BEAM-10200] Respect profile_memory option and add memory profiler to…

[GitBox <gi...@apache.org>]

angoenka commented on a change in pull request #12562:
URL: https://github.com/apache/beam/pull/12562#discussion_r491759030



##########
File path: sdks/python/apache_beam/utils/profiler.py
##########
@@ -44,59 +42,91 @@
 
 
 class Profile(object):
-  """cProfile wrapper context for saving and logging profiler results."""
+  """cProfile and Heapy wrapper context for saving and logging profiler
+  results."""
 
   SORTBY = 'cumulative'
 
   def __init__(
       self,
-      profile_id,
-      profile_location=None,
-      log_results=False,
-      file_copy_fn=None,
-      time_prefix='%Y-%m-%d_%H_%M_%S-'):
+      profile_id, # type: str
+      profile_location=None, # type: Optional[str]
+      log_results=False, # type: bool
+      file_copy_fn=None, # type: Optional[Callable[[str, str], None]]
+      time_prefix='%Y-%m-%d_%H_%M_%S-', # type: str
+      enable_cpu_profiling=False, # type: bool
+      enable_memory_profiling=False, # type: bool
+  ):
+    """Creates a Profile object.
+
+    Args:
+      profile_id: Unique id of the profiling session.
+      profile_location: The file location where the profiling results will be
+        stored.
+      log_results: Log the result to console if true.
+      file_copy_fn: Lambda function for copying files.
+      time_prefix: Format of the timestamp prefix in profiling result files.
+      enable_cpu_profiling: CPU profiler will be enabled during the profiling
+        session.
+      enable_memory_profiling: Memory profiler will be enabled during the
+        profiling session, the profiler only records the newly allocated objects
+        in this session.
+    """
     self.stats = None
     self.profile_id = str(profile_id)
     self.profile_location = profile_location
     self.log_results = log_results
     self.file_copy_fn = file_copy_fn or self.default_file_copy_fn
     self.time_prefix = time_prefix
     self.profile_output = None
+    self.enable_cpu_profiling = enable_cpu_profiling
+    self.enable_memory_profiling = enable_memory_profiling
 
   def __enter__(self):
     _LOGGER.info('Start profiling: %s', self.profile_id)
-    self.profile = cProfile.Profile()
-    self.profile.enable()
+    if self.enable_cpu_profiling:
+      self.profile = cProfile.Profile()
+      self.profile.enable()
+    if self.enable_memory_profiling:
+      try:
+        from guppy import hpy
+        self.hpy = hpy()
+        self.hpy.setrelheap()
+      except ImportError:

Review comment:
       Let's log the import failure




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org