Skip to content

Commit

Permalink
encapsulate context manager exit in a try .. catch block
Browse files Browse the repository at this point in the history
  • Loading branch information
ungarj committed Jul 5, 2024
1 parent ffe823b commit 80e105a
Showing 1 changed file with 8 additions and 4 deletions.
12 changes: 8 additions & 4 deletions mapchete/processing/profilers/memory.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,7 @@ def __init__(
self._exit_stack = ExitStack()
self._temp_dir = self._exit_stack.enter_context(TemporaryDirectory())
self._temp_file = str(
MPath(self._temp_dir) / f"{os. getpid()}-{uuid.uuid4().hex}.bin"
MPath(self._temp_dir) / f"{os.getpid()}-{uuid.uuid4().hex}.bin"
)
try:
self._memray_tracker = self._exit_stack.enter_context(
Expand Down Expand Up @@ -113,16 +113,20 @@ def __exit__(self, *args):
# close memray.Tracker before attempting to read file
if self._memray_tracker:
self._memray_tracker.__exit__(*args)
reader = FileReader(self._temp_file)
allocations = list(
reader.get_high_watermark_allocation_records(merge_threads=True)
FileReader(self._temp_file).get_high_watermark_allocation_records(
merge_threads=True
)
)
self.max_allocated = max(record.size for record in allocations)
self.total_allocated = sum(record.size for record in allocations)
self.allocations = len(allocations)
if self.output_file:
copy(self._temp_file, self.output_file, overwrite=True)
finally:
self._exit_stack.__exit__(*args)
try:
self._exit_stack.__exit__(*args)
except Exception as exc: # pragma: no cover
logger.error(exc)
# we need to set this to None, so MemoryTracker can be serialized
self._memray_tracker = None

0 comments on commit 80e105a

Please sign in to comment.