Skip to content

memory leak through plumpy's PROCESS_STACK context variable #4698

@ltalirz

Description

@ltalirz

Describe the bug

plumpy uses the contextvars module to enable access to the process itself (via Process.current()) at any point in time during a task run by the process.

https://github.com/aiidateam/plumpy/blob/33f4e9f029e0573f168237b8db030de537453878/plumpy/processes.py#L516-L534

When used in combination with asyncio's call_soon, call_later or call_at methods, each individual function execution actually gets their own copy of this context.
This means that as long as a handle to these scheduled executions remains in memory, the copy of the 'process stack' context var (and thus the process itself) remain in memory [1].

One example of this happening is the Transport future opened by aiida-core for file transfers - it is passed around to several classes inside aiida-core and even outside (e.g. paramiko).

def do_open():
""" Actually open the transport """
if transport_request and transport_request.count > 0:
# The user still wants the transport so open it
_LOGGER.debug('Transport request opening transport for %s', authinfo)
try:
transport.open()
except Exception as exception: # pylint: disable=broad-except
_LOGGER.error('exception occurred while trying to open transport:\n %s', exception)
transport_request.future.set_exception(exception)
# Cleanup of the stale TransportRequest with the excepted transport future
self._transport_requests.pop(authinfo.id, None)
else:
transport_request.future.set_result(transport)
# Save the handle so that we can cancel the callback if the user no longer wants it
open_callback_handle = self._loop.call_later(safe_open_interval, do_open)

The problem is: if any of these places forgets to shed the reference to this handle, not only does the transport itself remain in memory but also the Process that was responsible for creating it (and everything it references).

One can work around this issue by explicitly passing an empty context to call_soon, call_later or call_at where the context is not needed, but I would like to point out that this behaviour may be unexpected to users of plumpy and can be very difficult to debug.

Your environment

  • aiida-core Current develop

@muhrin

[1] Not sure whether memory consumption during execution should be a concern as well - if it's a shallow copy, it should be ok.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions