#

# Licensed to the Apache Software Foundation (ASF) under one or more

# contributor license agreements. See the NOTICE file distributed with

# this work for additional information regarding copyright ownership.

# The ASF licenses this file to You under the Apache License, Version 2.0

# (the "License"); you may not use this file except in compliance with

# the License. You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

#

"""A profiler context manager based on cProfile.Profile and guppy.hpy objects.

For internal use only; no backwards-compatibility guarantees.

"""

# pytype: skip-file

# mypy: check-untyped-defs

import cProfile

import io

import logging

import os

import pstats

import random

import tempfile

import time

from typing import Callable

from typing import Optional

from apache_beam.io import filesystems

_LOGGER = logging.getLogger(__name__)

class Profile(object):

"""cProfile and Heapy wrapper context for saving and logging profiler

results."""

SORTBY = 'cumulative'

profile_output: str

stats: pstats.Stats

def __init__(

self,

profile_id: str,

profile_location: Optional[str] = None,

log_results: bool = False,

file_copy_fn: Optional[Callable[[str, str], None]] = None,

time_prefix: str = '%Y-%m-%d_%H_%M_%S-',

enable_cpu_profiling: bool = False,

enable_memory_profiling: bool = False,

):

"""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.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.enable_cpu_profiling = enable_cpu_profiling

self.enable_memory_profiling = enable_memory_profiling

def __enter__(self):

_LOGGER.info('Start profiling: %s', self.profile_id)

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:

_LOGGER.info("Unable to import guppy for memory profiling")

self.hpy = None

return self

def __exit__(self, *args):

_LOGGER.info('Stop profiling: %s', self.profile_id)

if self.profile_location:

if self.enable_cpu_profiling:

self.profile.create_stats()

self.profile_output = self._upload_profile_data(

# typing: seems stats attr is missing from typeshed

self.profile_location,

'cpu_profile',

self.profile.stats)

if self.enable_memory_profiling:

if not self.hpy:

pass

else:

h = self.hpy.heap()

heap_dump_data = '%s\n%s' % (h, h.more)

self._upload_profile_data(

self.profile_location,

'memory_profile',

heap_dump_data,

write_binary=False)

if self.log_results:

if self.enable_cpu_profiling:

s = io.StringIO()

self.stats = pstats.Stats(

self.profile, stream=s).sort_stats(Profile.SORTBY)

self.stats.print_stats()

_LOGGER.info('Cpu profiler data: [%s]', s.getvalue())

if self.enable_memory_profiling and self.hpy:

_LOGGER.info('Memory profiler data: \n%s' % self.hpy.heap())

@staticmethod

def default_file_copy_fn(src, dest):

dest_handle = filesystems.FileSystems.create(dest + '.tmp')

try:

with open(src, 'rb') as src_handle:

dest_handle.write(src_handle.read())

finally:

dest_handle.close()

filesystems.FileSystems.rename([dest + '.tmp'], [dest])

@staticmethod

def factory_from_options(options) -> Optional[Callable[..., 'Profile']]:

if options.profile_cpu or options.profile_memory:

def create_profiler(profile_id, **kwargs):

if random.random() < options.profile_sample_rate:

return Profile(

profile_id,

options.profile_location,

enable_cpu_profiling=options.profile_cpu,

enable_memory_profiling=options.profile_memory,

**kwargs)

return create_profiler

return None

def _upload_profile_data(

self, profile_location, dir, data, write_binary=True) -> str:

dump_location = os.path.join(

profile_location,

dir,

time.strftime(self.time_prefix + self.profile_id))

fd, filename = tempfile.mkstemp()

try:

os.close(fd)

if write_binary:

with open(filename, 'wb') as fb:

import marshal

marshal.dump(data, fb)

else:

with open(filename, 'w') as f:

f.write(data)

_LOGGER.info('Copying profiler data to: [%s]', dump_location)

self.file_copy_fn(filename, dump_location)

finally:

os.remove(filename)

return dump_location