From 243cc6ccd2e9a80cb6ea3266b7e6d7ad4ca49630 Mon Sep 17 00:00:00 2001 From: liqiang-fit2cloud Date: Wed, 3 Dec 2025 14:39:34 +0800 Subject: [PATCH] feat: catch memory error. --- apps/common/utils/tool_code.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/apps/common/utils/tool_code.py b/apps/common/utils/tool_code.py index e7fe93e27..5cafa1471 100644 --- a/apps/common/utils/tool_code.py +++ b/apps/common/utils/tool_code.py @@ -32,6 +32,7 @@ class ToolExecutor: self.user = None self.sandbox_so_path = f'{self.sandbox_path}/lib/sandbox.so' self.process_timeout_seconds = int(CONFIG.get("SANDBOX_PYTHON_PROCESS_TIMEOUT_SECONDS", '3600')) + self.process_limit_mem_mb = int(CONFIG.get("SANDBOX_PYTHON_PROCESS_LIMIT_MEM_MB", '256')) try: self._init_sandbox_dir() except Exception as e: @@ -102,6 +103,7 @@ try: exec_result=f(**keywords) builtins.print("\\n{_id}:"+base64.b64encode(json.dumps({success}, default=str).encode()).decode(), flush=True) except Exception as e: + if isinstance(e, MemoryError): e = Exception("Cannot allocate more memory: exceeded the limit of {self.process_limit_mem_mb} MB.") builtins.print("\\n{_id}:"+base64.b64encode(json.dumps({err}, default=str).encode()).decode(), flush=True) """ maxkb_logger.debug(f"Sandbox execute code: {_exec_code}") @@ -227,8 +229,8 @@ exec({dedent(code)!a}) def set_resource_limit(): if not self.sandbox: return - mem_limit = int(CONFIG.get("SANDBOX_PYTHON_PROCESS_LIMIT_MEM_MB", '128')) * 1024 * 1024 - resource.setrlimit(resource.RLIMIT_RSS, (mem_limit, mem_limit)) + mem_limit = self.process_limit_mem_mb * 1024 * 1024 + resource.setrlimit(resource.RLIMIT_AS, (mem_limit, mem_limit)) subprocess_result = subprocess.run( [python_directory, execute_file], preexec_fn=set_resource_limit,