Skip to content

Commit 94eb695

Browse files
feat: add MAXKB_SANDBOX_PYTHON_PROCESS_LIMIT_CPU_CORES env.
1 parent 447393a commit 94eb695

File tree

1 file changed

+7
-2
lines changed

1 file changed

+7
-2
lines changed

apps/common/utils/tool_code.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import pwd
1212
import resource
1313
import getpass
14+
import random
1415
import uuid_utils.compat as uuid
1516
from common.utils.logger import maxkb_logger
1617
from django.utils.translation import gettext_lazy as _
@@ -24,6 +25,7 @@ class ToolExecutor:
2425
sandbox_path = CONFIG.get("SANDBOX_HOME", '/opt/maxkb-app/sandbox') if enable_sandbox else os.path.join(PROJECT_DIR, 'data', 'sandbox')
2526
process_timeout_seconds = int(CONFIG.get("SANDBOX_PYTHON_PROCESS_TIMEOUT_SECONDS", '3600'))
2627
process_limit_mem_mb = int(CONFIG.get("SANDBOX_PYTHON_PROCESS_LIMIT_MEM_MB", '256'))
28+
process_limit_cpu_cores = min(max(int(CONFIG.get("SANDBOX_PYTHON_PROCESS_LIMIT_CPU_CORES", '1')), 1), len(os.sched_getaffinity(0))) if sys.platform.startswith("linux") else os.cpu_count() # 只支持linux,window和mac不支持
2729

2830
def __init__(self, sandbox=False):
2931
self.sandbox = sandbox
@@ -45,7 +47,7 @@ def init_sandbox_dir():
4547
except FileExistsError:
4648
# 文件已存在 → 已初始化过
4749
return
48-
maxkb_logger.debug("init dir")
50+
maxkb_logger.info("Init sandbox dir.")
4951
try:
5052
os.system("chmod -R g-rwx /dev/shm /dev/mqueue")
5153
os.system("chmod o-rwx /run/postgresql")
@@ -225,7 +227,10 @@ def _exec(self, execute_file):
225227
text=True,
226228
capture_output=True,
227229
**kwargs,
228-
preexec_fn=lambda: (None if not self.sandbox else resource.setrlimit(resource.RLIMIT_AS, (ToolExecutor.process_limit_mem_mb * 1024 * 1024,) * 2))
230+
preexec_fn=(lambda: None if (not self.sandbox or not sys.platform.startswith("linux")) else (
231+
resource.setrlimit(resource.RLIMIT_AS, (ToolExecutor.process_limit_mem_mb * 1024 * 1024,) * 2),
232+
os.sched_setaffinity(0, set(random.sample(list(os.sched_getaffinity(0)), ToolExecutor.process_limit_cpu_cores)))
233+
))
229234
)
230235
return subprocess_result
231236
except subprocess.TimeoutExpired:

0 commit comments

Comments
 (0)