Update shared_vis_python_exe.py

shared_vis_python_exe.py

@@ -1,3 +1,5 @@
+# shared_vis_python_exe.py
+
 import os
 import io
 import regex
@@ -6,19 +8,21 @@ import traceback
 import copy
 import datetime
 import dateutil.relativedelta
-import multiprocess
-from multiprocess import Pool
+import multiprocessing
+from multiprocessing import Queue, Process
 from typing import Any, Dict, Optional, Tuple, List, Union
-from pebble import ProcessPool
 from tqdm import tqdm
 from concurrent.futures import TimeoutError
-from functools import partial
-from timeout_decorator import timeout
 from contextlib import redirect_stdout
 import base64
 from io import BytesIO
 from PIL import Image
-import pdb
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import numpy as np
+import time
+import queue
 
 def encode_image(image_path):
     with open(image_path, "rb") as image_file:
@@ -30,42 +34,212 @@ def base64_to_image(
     convert_mode: Optional[str] = "RGB"
 ) -> Union[Image.Image, None]:
     """
-    将Base64编码的图片字符串转换为PIL Image对象
-    
+    Convert a Base64-encoded image string to a PIL Image object.
+
     Args:
-        base64_str: Base64编码的图片字符串（可带data:前缀）
-        remove_prefix: 是否自动去除"data:image/..."前缀（默认True）
-        convert_mode: 转换为指定模式（如"RGB"/"RGBA"，None表示不转换）
-    
+        base64_str: Base64-encoded image string (can include data: prefix)
+        remove_prefix: Whether to automatically remove the "data:image/..." prefix (default True)
+        convert_mode: Convert to the specified mode (such as "RGB"/"RGBA", None means no conversion)
+
     Returns:
-        PIL.Image.Image 对象，解码失败时返回None
+        PIL.Image.Image object, or None if decoding fails
         
     Examples:
         >>> img = base64_to_image("data:image/png;base64,iVBORw0KGg...")
         >>> img = base64_to_image("iVBORw0KGg...", remove_prefix=False)
     """
     try:
-        # 1. 处理Base64前缀
+        # 1. Handle Base64 prefix
         if remove_prefix and "," in base64_str:
             base64_str = base64_str.split(",")[1]
 
-        # 2. 解码Base64
+        # 2. Decode Base64
         image_data = base64.b64decode(base64_str)
         
-        # 3. 转换为PIL Image
+        # 3. Convert to PIL Image
         image = Image.open(BytesIO(image_data))
         
-        # 4. 可选模式转换
+        # 4. Optional mode conversion
         if convert_mode:
             image = image.convert(convert_mode)
             
         return image
     
     except (base64.binascii.Error, OSError, Exception) as e:
-        print(f"Base64解码失败: {str(e)}")
+        print(f"Base64 decode failed: {str(e)}")
         return None
 
 
+class PersistentWorker:
+    """Persistent worker process."""
+    
+    def __init__(self):
+        self.input_queue = multiprocessing.Queue()
+        self.output_queue = multiprocessing.Queue()
+        self.process = None
+        self.start()
+    
+    def start(self):
+        """Start the worker process."""
+        self.process = Process(target=self._worker_loop)
+        self.process.daemon = True
+        self.process.start()
+    
+    def _worker_loop(self):
+        """Main loop for the worker process."""
+        runtime = None
+        runtime_class = None
+        
+        while True:
+            try:
+                # Get task
+                task = self.input_queue.get()
+                
+                if task is None:  # Termination signal
+                    break
+                
+                task_type = task.get('type')
+                
+                if task_type == 'init':
+                    # Initialize runtime
+                    messages = task.get('messages', [])
+                    runtime_class = task.get('runtime_class', ImageRuntime)
+                    runtime = runtime_class(messages)
+                    self.output_queue.put({
+                        'status': 'success',
+                        'result': 'Initialized'
+                    })
+                
+                elif task_type == 'execute':
+                    # Execute code
+                    if runtime is None:
+                        messages = task.get('messages', [])
+                        runtime_class = task.get('runtime_class', ImageRuntime)
+                        runtime = runtime_class(messages)
+                    
+                    code = task.get('code')
+                    get_answer_from_stdout = task.get('get_answer_from_stdout', True)
+                    answer_symbol = task.get('answer_symbol')
+                    answer_expr = task.get('answer_expr')
+                    
+                    try:
+                        # Record the number of images before execution
+                        pre_figures_count = len(runtime._global_vars.get("_captured_figures", []))
+                        
+                        if get_answer_from_stdout:
+                            program_io = io.StringIO()
+                            with redirect_stdout(program_io):
+                                runtime.exec_code("\n".join(code))
+                            program_io.seek(0)
+                            result = program_io.read()
+                        elif answer_symbol:
+                            runtime.exec_code("\n".join(code))
+                            result = runtime._global_vars.get(answer_symbol, "")
+                        elif answer_expr:
+                            runtime.exec_code("\n".join(code))
+                            result = runtime.eval_code(answer_expr)
+                        else:
+                            if len(code) > 1:
+                                runtime.exec_code("\n".join(code[:-1]))
+                                result = runtime.eval_code(code[-1])
+                            else:
+                                runtime.exec_code("\n".join(code))
+                                result = ""
+                        
+                        # Get newly generated images
+                        all_figures = runtime._global_vars.get("_captured_figures", [])
+                        new_figures = all_figures[pre_figures_count:]
+                        
+                        # Build result
+                        if new_figures:
+                            result = {
+                                'text': result,
+                                'images': new_figures
+                            } if result else {'images': new_figures}
+                        else:
+                            result = {'text': result} if result else {}
+                        
+                        self.output_queue.put({
+                            'status': 'success',
+                            'result': result,
+                            'report': 'Done'
+                        })
+                    
+                    except Exception as e:
+                        self.output_queue.put({
+                            'status': 'error',
+                            'error': str(e),
+                            'traceback': traceback.format_exc(),
+                            'report': f'Error: {str(e)}'
+                        })
+                
+                elif task_type == 'reset':
+                    # Reset runtime
+                    messages = task.get('messages', [])
+                    runtime_class = task.get('runtime_class', ImageRuntime)
+                    runtime = runtime_class(messages)
+                    self.output_queue.put({
+                        'status': 'success',
+                        'result': 'Reset'
+                    })
+                    
+            except Exception as e:
+                self.output_queue.put({
+                    'status': 'error',
+                    'error': f'Worker error: {str(e)}',
+                    'traceback': traceback.format_exc()
+                })
+    
+    def execute(self, code: List[str], messages: list = None, runtime_class=None, 
+                get_answer_from_stdout=True, answer_symbol=None, answer_expr=None, timeout: int = 30):
+        """Execute code."""
+        self.input_queue.put({
+            'type': 'execute',
+            'code': code,
+            'messages': messages,
+            'runtime_class': runtime_class,
+            'get_answer_from_stdout': get_answer_from_stdout,
+            'answer_symbol': answer_symbol,
+            'answer_expr': answer_expr
+        })
+        
+        try:
+            result = self.output_queue.get(timeout=timeout)
+            return result
+        except queue.Empty:
+            return {
+                'status': 'error',
+                'error': 'Execution timeout',
+                'report': 'Timeout Error'
+            }
+    
+    def init_runtime(self, messages: list, runtime_class=None):
+        """Initialize runtime."""
+        self.input_queue.put({
+            'type': 'init',
+            'messages': messages,
+            'runtime_class': runtime_class
+        })
+        return self.output_queue.get()
+    
+    def reset_runtime(self, messages: list = None, runtime_class=None):
+        """Reset runtime."""
+        self.input_queue.put({
+            'type': 'reset',
+            'messages': messages,
+            'runtime_class': runtime_class
+        })
+        return self.output_queue.get()
+    
+    def terminate(self):
+        """Terminate the worker process."""
+        if self.process and self.process.is_alive():
+            self.input_queue.put(None)
+            self.process.join(timeout=5)
+            if self.process.is_alive():
+                self.process.terminate()
+
+
 class GenericRuntime:
     GLOBAL_DICT = {}
     LOCAL_DICT = None
@@ -80,17 +254,14 @@ class GenericRuntime:
             self.exec_code(c)
 
     def exec_code(self, code_piece: str) -> None:
-        if regex.search(r"(\s|^)?input\(", code_piece) or regex.search(
-            r"(\s|^)?os.system\(", code_piece
-        ):
+        # Security check
+        if regex.search(r"(\s|^)?(input|os\.system|subprocess)\(", code_piece):
             raise RuntimeError("Forbidden function calls detected")
-
         
-        
-        # 检测并修改plt.show()调用
+        # Detect and modify plt.show() calls
         if "plt.show()" in code_piece:
             modified_code = code_piece.replace("plt.show()", """
-# 捕获当前图像
+# Capture current figure
 buf = io.BytesIO()
 plt.savefig(buf, format='png')
 buf.seek(0)
@@ -98,7 +269,7 @@ _captured_image = base64.b64encode(buf.read()).decode('utf-8')
 _captured_figures.append(_captured_image)
 plt.close()
 """)
-            # 确保_captured_figures变量存在
+            # Ensure _captured_figures variable exists
             if "_captured_figures" not in self._global_vars:
                 self._global_vars["_captured_figures"] = []
             
@@ -123,46 +294,41 @@ plt.close()
 
 
 class ImageRuntime(GenericRuntime):
-    # """支持图像处理的运行时环境"""
-    # GLOBAL_DICT = {}  # 不预加载模块，避免序列化问题
-    # LOCAL_DICT = None
-    
     HEADERS = [
         "import matplotlib",
-        "matplotlib.use('Agg')",  # 使用非交互式后端
+        "matplotlib.use('Agg')",  # Use non-interactive backend
         "import matplotlib.pyplot as plt",
         "from PIL import Image",
         "import io",
         "import base64",
         "import numpy as np",
-        "_captured_figures = []",  # 初始化图像捕获列表
+        "_captured_figures = []",  # Initialize image capture list
     ]
 
     def __init__(self, messages):
         super().__init__()
-        # pdb.set_trace()
-
-        self._global_vars = copy.copy(self.GLOBAL_DICT)
-        self._local_vars = copy.copy(self.LOCAL_DICT) if self.LOCAL_DICT else None
-        self._captured_figures = []
-
-        for c in self.HEADERS:
-            self.exec_code(c)
 
         image_var_dict = {}
         image_var_idx = 0
+        init_captured_figures = []
+        
         for message_item in messages:
-            content = message_item['content']  # {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
+            content = message_item['content']
             for item in content:
-                item_type = item['type']
-                if item_type == "image_url":
-                    item_image_url = item['image_url']['url']
-                    image = base64_to_image(item_image_url)
-                    image_var_dict[f"image_clue_{image_var_idx}"] = image
-                    image_var_idx += 1
-
+                if isinstance(item, dict):
+                    item_type = item.get('type')
+                    if item_type == "image_url":
+                        item_image_url = item['image_url']['url']
+                        image = base64_to_image(item_image_url)
+                        if image:
+                            image_var_dict[f"image_clue_{image_var_idx}"] = image
+                            init_captured_figures.append(base64.b64encode(
+                                BytesIO(image.tobytes()).getvalue()).decode('utf-8'))
+                            image_var_idx += 1
+
+        image_var_dict["_captured_figures"] = init_captured_figures
         self.inject(image_var_dict)
-                    
+
 
 class DateRuntime(GenericRuntime):
     GLOBAL_DICT = {}
@@ -190,97 +356,24 @@ class PythonExecutor:
         get_answer_expr: Optional[str] = None,
         get_answer_from_stdout: bool = True,
         timeout_length: int = 20,
+        use_process_isolation: bool = True,
     ) -> None:
         self.runtime_class = runtime_class if runtime_class else ImageRuntime
-        print(self.runtime_class)
         self.answer_symbol = get_answer_symbol
         self.answer_expr = get_answer_expr
         self.get_answer_from_stdout = get_answer_from_stdout
         self.timeout_length = timeout_length
+        self.use_process_isolation = use_process_isolation
+        self.persistent_worker = None
 
-        # Create a persistent runtime instance if messages are provided
-        self.persistent_runtime = None
+    def _ensure_worker(self):
+        """Ensure the worker process exists."""
+        if self.persistent_worker is None:
+            self.persistent_worker = PersistentWorker()
 
     def process_generation_to_code(self, gens: str):
         return [g.split("\n") for g in gens]
 
-    # def execute(
-    #     self,
-    #     code,
-    #     messages,
-    #     get_answer_from_stdout=True,
-    #     runtime_class=None,
-    #     # run_time_instance=None,
-    #     answer_symbol=None,
-    #     answer_expr=None,
-    #     # 移除 timeout_length 参数
-    # ) -> Tuple[Union[str, Dict[str, Any]], str]:
-    #     # print("dome")
-    #     # try:
-    #     # 在每个进程中创建新的运行时实例
-    #     print(runtime_class)
-    #     # runtime = runtime_class(messages)
-    #     runtime = self.persistent_runtime
-        
-    #     if get_answer_from_stdout:
-    #         program_io = io.StringIO()
-    #         with redirect_stdout(program_io):
-    #             # 移除 timeout 调用
-    #             runtime.exec_code("\n".join(code))
-    #         program_io.seek(0)
-    #         result = program_io.read()
-    #     elif answer_symbol:
-    #         # 移除 timeout 调用
-    #         runtime.exec_code("\n".join(code))
-    #         result = runtime._global_vars.get(answer_symbol, "")
-    #     elif answer_expr:
-    #         # 移除 timeout 调用
-    #         runtime.exec_code("\n".join(code))
-    #         # 移除 timeout 调用
-    #         result = runtime.eval_code(answer_expr)
-    #     else:
-    #         if len(code) > 1:
-    #             # 移除 timeout 调用
-    #             runtime.exec_code("\n".join(code[:-1]))
-    #             # 移除 timeout 调用
-    #             result = runtime.eval_code(code[-1])
-    #         else:
-    #             # 移除 timeout 调用
-    #             runtime.exec_code("\n".join(code))
-    #             result = ""
-        
-    #     # 检查是否有捕获的图像
-    #     captured_figures = runtime._global_vars.get("_captured_figures", [])
-    #     if captured_figures:
-    #         # 如果有文本输出和图像，将它们组合
-    #         if result:
-    #             result = {
-    #                 'text': result,
-    #                 'images': captured_figures
-    #             }
-    #         else:
-    #             result = {'images': captured_figures}
-    #     else:
-    #         if result:
-    #             result = {
-    #                 'text': result,
-    #             }
-            
-        
-    #     report = "Done"
-    #     # except Exception as e:
-    #     #     result = ""
-    #     #     report = f"Error: {str(e)}\n{traceback.format_exc()}"
-        
-    #     # 确保结果可序列化
-    #     try:
-    #         pickle.dumps(result)
-    #     except Exception as e:
-    #         result = f"Result serialization error: {str(e)}"
-    #         report = f"Serialization Error: {str(e)}"
-            
-    #     return result, report
-
     def execute(
         self,
         code,
@@ -290,67 +383,83 @@ class PythonExecutor:
         answer_symbol=None,
         answer_expr=None,
     ) -> Tuple[Union[str, Dict[str, Any]], str]:
-        print(runtime_class)
-        runtime = self.persistent_runtime
-    
-        try:
-            if get_answer_from_stdout:
-                program_io = io.StringIO()
-                with redirect_stdout(program_io):
-                    runtime.exec_code("\n".join(code))
-                program_io.seek(0)
-                result = program_io.read()
-            elif answer_symbol:
-                runtime.exec_code("\n".join(code))
-                result = runtime._global_vars.get(answer_symbol, "")
-            elif answer_expr:
-                runtime.exec_code("\n".join(code))
-                result = runtime.eval_code(answer_expr)
+        
+        if self.use_process_isolation:
+            # Ensure worker process exists
+            self._ensure_worker()
+            
+            # Execute code
+            result = self.persistent_worker.execute(
+                code,
+                messages,
+                runtime_class or self.runtime_class,
+                get_answer_from_stdout,
+                answer_symbol,
+                answer_expr,
+                timeout=self.timeout_length
+            )
+            
+            if result['status'] == 'success':
+                return result['result'], result.get('report', 'Done')
             else:
-                if len(code) > 1:
-                    runtime.exec_code("\n".join(code[:-1]))
-                    result = runtime.eval_code(code[-1])
-                else:
+                error_result = {
+                    'error': result.get('error', 'Unknown error'),
+                    'traceback': result.get('traceback', '')
+                }
+                return error_result, result.get('report', f"Error: {result.get('error', 'Unknown error')}")
+        else:
+            # Non-isolation mode (for backward compatibility)
+            runtime = runtime_class(messages) if runtime_class else self.runtime_class(messages)
+            
+            try:
+                if get_answer_from_stdout:
+                    program_io = io.StringIO()
+                    with redirect_stdout(program_io):
+                        runtime.exec_code("\n".join(code))
+                    program_io.seek(0)
+                    result = program_io.read()
+                elif answer_symbol:
                     runtime.exec_code("\n".join(code))
-                    result = ""
-    
-            # Check for captured figures
-            captured_figures = runtime._global_vars.get("_captured_figures", [])
-            if captured_figures:
+                    result = runtime._global_vars.get(answer_symbol, "")
+                elif answer_expr:
+                    runtime.exec_code("\n".join(code))
+                    result = runtime.eval_code(answer_expr)
+                else:
+                    if len(code) > 1:
+                        runtime.exec_code("\n".join(code[:-1]))
+                        result = runtime.eval_code(code[-1])
+                    else:
+                        runtime.exec_code("\n".join(code))
+                        result = ""
+
+                # Check for captured figures
+                captured_figures = runtime.captured_figures
+                if captured_figures:
+                    result = {
+                        'text': result,
+                        'images': captured_figures
+                    } if result else {'images': captured_figures}
+                else:
+                    result = {'text': result} if result else {}
+
+                report = "Done"
+
+            except Exception as e:
                 result = {
-                    'text': result,
-                    'images': captured_figures
-                } if result else {'images': captured_figures}
-            else:
-                result = {'text': result} if result else {}
-    
-            report = "Done"
-    
-        except Exception as e:
-            result = {
-                'error': str(e),
-                'traceback': traceback.format_exc()
-            }
-            report = f"Error: {str(e)}"
-    
-        # Ensure result is serializable
-        try:
-            pickle.dumps(result)
-        except Exception as e:
-            result = f"Result serialization error: {str(e)}"
-            report = f"Serialization Error: {str(e)}"
-    
-        return result, report
+                    'error': str(e),
+                    'traceback': traceback.format_exc()
+                }
+                report = f"Error: {str(e)}"
 
+            return result, report
 
-    
     def apply(self, code, messages):
         return self.batch_apply([code], messages)[0]
 
     @staticmethod
     def truncate(s, max_length=400):
         if isinstance(s, dict):
-            # 如果是字典（包含图像），只截断文本部分
+            # If it is a dict (with images), truncate only the text part
             if 'text' in s:
                 half = max_length // 2
                 if len(s['text']) > max_length:
@@ -362,21 +471,12 @@ class PythonExecutor:
                 s = s[:half] + "..." + s[-half:]
             return s
 
-    def update_persistent_runtime_with_messages():
-        pass
-
-    def get_persistent_runtime(self):
-        return self.persistent_runtime
-
     def batch_apply(self, batch_code, messages):
-        if not self.persistent_runtime and messages:
-            self.persistent_runtime = self.runtime_class(messages)
         all_code_snippets = self.process_generation_to_code(batch_code)
 
         timeout_cnt = 0
         all_exec_results = []
 
-        # 去掉 ProcessPool，改为单进程顺序执行
         if len(all_code_snippets) > 100:
             progress_bar = tqdm(total=len(all_code_snippets), desc="Execute")
         else:
@@ -384,7 +484,6 @@ class PythonExecutor:
 
         for code in all_code_snippets:
             try:
-                # 直接调用 self.execute，而不是用 ProcessPool
                 result = self.execute(
                     code,
                     messages=messages,
@@ -392,7 +491,6 @@ class PythonExecutor:
                     runtime_class=self.runtime_class,
                     answer_symbol=self.answer_symbol,
                     answer_expr=self.answer_expr,
-                    # timeout_length=self.timeout_length,
                 )
                 all_exec_results.append(result)
             except TimeoutError as error:
@@ -411,16 +509,16 @@ class PythonExecutor:
 
         batch_results = []
         for code, (res, report) in zip(all_code_snippets, all_exec_results):
-            # 处理结果
+            # Handle results
             if isinstance(res, dict):
-                # 如果结果包含图像，特殊处理
+                # If result contains images, special handling
                 if 'text' in res:
                     res['text'] = str(res['text']).strip()
                     res['text'] = self.truncate(res['text'])
                 report = str(report).strip()
                 report = self.truncate(report)
             else:
-                # 普通文本结果
+                # Normal text result
                 res = str(res).strip()
                 res = self.truncate(res)
                 report = str(report).strip()
@@ -428,101 +526,12 @@ class PythonExecutor:
             batch_results.append((res, report))
         return batch_results
 
+    def reset(self, messages=None):
+        """Reset executor state."""
+        if self.use_process_isolation and self.persistent_worker:
+            self.persistent_worker.reset_runtime(messages, self.runtime_class)
 
-def _test():
-    image_path = "/mnt/petrelfs/zhaoshitian/vis_tool_inference_engine/test_data/0.JPG"
-    image_base64 = encode_image(image_path)
-    messages = [
-        {
-            "role": "user",
-            "content": [{"type": "text", "text": "From the information on that advertising board, what is the type of this shop?"}]
-        },
-        {
-            "role": "user",
-            "content": [{"type": "text", "text": "image_clue_0"}] + [{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_base64}"}}]
-        }
-    ]
-    # 测试普通计算
-    math_code ="""
-a = 1
-b = 2
-c = a + b
-print(c)
-"""
-
-    batch_code = [math_code]
-
-    executor = PythonExecutor()
-    predictions = executor.apply(batch_code[0], messages)
-    print("数学计算结果:", predictions)
-    
-    # 测试图像显示
-    image_code = """
-import matplotlib.pyplot as plt
-import numpy as np
-from PIL import Image
-import io
-
-# 创建一个简单的图像
-x = np.linspace(0, 10, 100)
-y = np.sin(x)
-
-plt.figure(figsize=(8, 6))
-plt.plot(x, y, 'r-', linewidth=2)
-plt.title('Sine Wave')
-plt.grid(True)
-plt.show()
-
-# 也可以显示一个简单的图像
-# 创建一个彩色渐变图像
-arr = np.zeros((100, 100, 3), dtype=np.uint8)
-for i in range(100):
-    for j in range(100):
-        arr[i, j, 0] = i  # 红色通道
-        arr[i, j, 1] = j  # 绿色通道
-        arr[i, j, 2] = 100  # 蓝色通道
-
-img = Image.fromarray(arr)
-plt.figure()
-plt.imshow(img)
-plt.title('Gradient Image')
-plt.show()
-
-print("图像生成完成")
-    """
-
-    image_code = """
-import matplotlib.pyplot as plt
-import numpy as np
-from PIL import Image
-import io
-
-plt.imshow(image_clue_0)
-plt.title("Original Image - Locate Advertising Board")
-plt.show()
-    """
-    
-    image_result = executor.apply(image_code, messages)
-    print("\n图像结果类型:", type(image_result[0]))
-    if isinstance(image_result[0], dict) and 'images' in image_result[0]:
-        print(f"捕获到 {len(image_result[0]['images'])} 个图像")
-        print("第一个图像的base64编码前20个字符:", image_result[0]['images'][0][:20])
-        
-        # 可选：保存图像到文件
-        for i, img_data in enumerate(image_result[0]['images']):
-            img_bytes = base64.b64decode(img_data)
-            with open(f"captured_image_{i}.png", "wb") as f:
-                f.write(img_bytes)
-            print(f"图像已保存为 captured_image_{i}.png")
-            
-        if 'text' in image_result[0]:
-            print("文本输出:", image_result[0]['text'])
-    else:
-        print("未捕获到图像")
-        print("结果:", image_result[0])
-    
-    print("\n执行状态:", image_result[1])
-
-
-if __name__ == "__main__":
-    _test()
+    def __del__(self):
+        """Clean up resources."""
+        if self.persistent_worker:
+            self.persistent_worker.terminate()