77import typing as t
88from 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
1112from aiida .manage import get_manager
1213from aiida .orm import ProcessNode
1314from node_graph .socket_spec import SocketSpec
15+ from packaging .version import parse as parse_version
1416
1517from aiida_pythonjob .calculations .pyfunction import PyFunction
1618from aiida_pythonjob .launch import create_inputs , prepare_pyfunction_inputs
1719
1820LOGGER = 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
2230def 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