feat: catch memory error.

This commit is contained in:
liqiang-fit2cloud 2025-12-03 14:39:34 +08:00
parent 46f5a6c9ce
commit 243cc6ccd2

View File

@ -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,