Skip to content

Commit 86f21bd

Browse files
authored
Make the recursion limit workaround dependent on the aiida-core version (#72)
With aiida-core v2.8.0 we switched from nest-asyncio to greenlet and therefore do not need to dynamically increase the recursion limit anymore.
1 parent 4df7350 commit 86f21bd

File tree

1 file changed

+25
-16
lines changed

1 file changed

+25
-16
lines changed

src/aiida_pythonjob/decorator.py

Lines changed: 25 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -7,16 +7,24 @@
77
import typing as t
88
from typing import List
99

10-
from aiida.engine.processes.functions import FunctionType, get_stack_size
10+
import aiida
11+
from aiida.engine.processes.functions import FunctionType
1112
from aiida.manage import get_manager
1213
from aiida.orm import ProcessNode
1314
from node_graph.socket_spec import SocketSpec
15+
from packaging.version import parse as parse_version
1416

1517
from aiida_pythonjob.calculations.pyfunction import PyFunction
1618
from aiida_pythonjob.launch import create_inputs, prepare_pyfunction_inputs
1719

1820
LOGGER = logging.getLogger(__name__)
1921

22+
_AIIDA_VERSION = parse_version(aiida.__version__)
23+
_NEEDS_RECURSION_LIMIT_WORKAROUND = _AIIDA_VERSION < parse_version("2.8.0")
24+
25+
if _NEEDS_RECURSION_LIMIT_WORKAROUND:
26+
from aiida.engine.processes.functions import get_stack_size
27+
2028

2129
# The following code is modified from the aiida-core.engine.processes.functions module
2230
def pyfunction(
@@ -42,25 +50,26 @@ def run_get_node(*args, **kwargs) -> tuple[dict[str, t.Any] | None, "ProcessNode
4250
:param kwargs: input keyword arguments to construct the FunctionProcess
4351
:return: tuple of the outputs of the process and the process node
4452
"""
45-
frame_delta = 1000
46-
frame_count = get_stack_size()
47-
stack_limit = sys.getrecursionlimit()
48-
LOGGER.info("Executing process function, current stack status: %d frames of %d", frame_count, stack_limit)
49-
50-
# If the current frame count is more than 80% of the stack limit, or comes within 200 frames, increase the
51-
# stack limit by ``frame_delta``.
52-
if frame_count > min(0.8 * stack_limit, stack_limit - 200):
53-
LOGGER.warning(
54-
"Current stack contains %d frames which is close to the limit of %d. Increasing the limit by %d",
55-
frame_count,
56-
stack_limit,
57-
frame_delta,
53+
if _NEEDS_RECURSION_LIMIT_WORKAROUND:
54+
frame_delta = 1000
55+
frame_count = get_stack_size()
56+
stack_limit = sys.getrecursionlimit()
57+
LOGGER.info(
58+
"Executing process function, current stack status: %d frames of %d", frame_count, stack_limit
5859
)
59-
sys.setrecursionlimit(stack_limit + frame_delta)
60+
61+
if frame_count > min(0.8 * stack_limit, stack_limit - 200):
62+
LOGGER.warning(
63+
"Current stack contains %d frames which is close to the limit of %d. Increasing the limit by %d",
64+
frame_count,
65+
stack_limit,
66+
frame_delta,
67+
)
68+
sys.setrecursionlimit(stack_limit + frame_delta)
6069

6170
manager = get_manager()
6271
runner = manager.get_runner()
63-
# # Remove all the known inputs from the kwargs
72+
# Remove all the known inputs from the kwargs
6473
outputs_spec = kwargs.pop("outputs_spec", None) or outputs
6574
inputs_spec = kwargs.pop("inputs_spec", None) or inputs
6675
metadata = kwargs.pop("metadata", None)

0 commit comments

Comments
 (0)