1111import pwd
1212import resource
1313import getpass
14+ import random
1415import uuid_utils .compat as uuid
1516from common .utils .logger import maxkb_logger
1617from 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