Update python_interpreter_tool.py

python_interpreter_tool.py

@@ -1,12 +1,9 @@
-import io
-import sys
-import platform
-import faulthandler
-from contextlib import contextmanager
-import signal
-
+import tempfile
+import subprocess
 from transformers import Tool
-
+import sys
+import os
+import re
 
 class PythonInterpreter(Tool):
     name = "python_interpreter_tool"
@@ -14,75 +11,41 @@ class PythonInterpreter(Tool):
 
     inputs = ["text"]
     outputs = ["text"]
-    timeout = 10.0
+    timeout = 1.0
 
     def __call__(self, task: str):
-        import os
-        import shutil
 
         if os.getenv("HF_ALLOW_CODE_EVAL", 0) != "1":
             return "Can't execute code, set code evaluation flag for PythonInterpreter"
 
-        rmtree = shutil.rmtree
-        rmdir = os.rmdir
-        chdir = os.chdir
-
-        #reliability_guard()
-        try:
-            exec_globals = {}
-            with capture_stdout() as output:
-                #with time_limit(self.timeout):
-                exec(task, exec_globals)
-            
-            captured_output = output.getvalue().strip()
-        except Exception as e:
-            captured_output = str(e)
-
-        # Needed for cleaning up.
-        shutil.rmtree = rmtree
-        os.rmdir = rmdir
-        os.chdir = chdir
-
-        return captured_output
-
-
-@contextmanager
-def capture_stdout():
-    output_buffer = io.StringIO()
-    original_stdout = sys.stdout
-    try:
-        sys.stdout = output_buffer
-        yield output_buffer
-    finally:
-        sys.stdout = original_stdout
+        with tempfile.TemporaryDirectory() as temp_dir:
 
+            code = "from safeguard import reliability_guard\nreliability_guard()\n" + task
+            code_file = os.path.join(temp_dir, 'code.py')
+            with open(code_file, 'w') as f:
+                f.write(code)
 
-@contextmanager
-def time_limit(seconds):
-    def signal_handler(signum, frame):
-        raise TimeoutException("Timed out!")
-
-    signal.setitimer(signal.ITIMER_REAL, seconds)
-    signal.signal(signal.SIGALRM, signal_handler)
-    try:
-        yield
-    finally:
-        signal.setitimer(signal.ITIMER_REAL, 0)
-
+            safeguard_file = os.path.join(temp_dir, 'safeguard.py')
+            with open(safeguard_file, 'w') as f:
+                f.write(reliability_guard_code)
 
+            cmd = f"python {code_file}"
+            
+            try:
+                output = subprocess.check_output(cmd.split(), stderr=subprocess.STDOUT, timeout=self.timeout)
+                output_text = output.decode(sys.stdout.encoding).strip()
+            except subprocess.TimeoutExpired:
+                output_text = "Code execution timed out."
+            except subprocess.CalledProcessError as e:
+                output_text = e.output.decode(sys.stdout.encoding).strip()
+                output_text = output_text.replace(temp_dir + "/", "/tmp/")
+                output_text = fix_stacktrace_linenumber(output_text)                
+        output_text = output_text.replace(temp_dir + "/", "/tmp/")
+        output_text = fix_warning_linenumber(output_text)
+        return output_text
+
+reliability_guard_code = """\
 def reliability_guard(maximum_memory_bytes=None):
-    """
-    This disables various destructive functions and prevents the generated code
-    from interfering with the test (e.g. fork bomb, killing other processes,
-    removing filesystem files, etc.)
-
-    WARNING
-    This function is NOT a security sandbox. Untrusted code, including, model-
-    generated code, should not be blindly executed outside of one. See the
-    Codex paper for more information about OpenAI's code sandbox, and proceed
-    with caution.
-    """
-
     if maximum_memory_bytes is not None:
         import resource
 
@@ -91,6 +54,8 @@ def reliability_guard(maximum_memory_bytes=None):
         if not platform.uname().system == "Darwin":
             resource.setrlimit(resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes))
 
+    import faulthandler
+
     faulthandler.disable()
 
     import builtins
@@ -99,13 +64,12 @@ def reliability_guard(maximum_memory_bytes=None):
     builtins.quit = None
 
     import os
-    try:
-        os.environ["OMP_NUM_THREADS"] = "1"
-    except:
-        pass
+
+    os.environ["OMP_NUM_THREADS"] = "1"
+
     os.kill = None
     os.system = None
-    os.putenv = None
+    # os.putenv = None # breaks e.g. numpy
     os.remove = None
     os.removedirs = None
     os.rmdir = None
@@ -141,7 +105,7 @@ def reliability_guard(maximum_memory_bytes=None):
 
     subprocess.Popen = None  # type: ignore
 
-    #__builtins__["help"] = None
+    __builtins__["help"] = None
 
     import sys
 
@@ -150,13 +114,21 @@ def reliability_guard(maximum_memory_bytes=None):
     sys.modules["resource"] = None
     sys.modules["psutil"] = None
     sys.modules["tkinter"] = None
-
-class TimeoutException(Exception):
-    pass
-
 """
-import os
-tool = PythonInterpreter()
-print(tool("import os; os.getcwd()"))
-print(os.getcwd())
-"""
+
+def fix_stacktrace_linenumber(text):
+    start = '  File "/tmp/code.py", line '
+    lines = text.split("\n")
+    fixed_lines = []
+    for line in lines:
+        if line.startswith(start):
+            number, end = line.split(start)[1].split(",")
+            new_number = int(number)-2
+            line = start + str(new_number) + end
+        fixed_lines.append(line)
+    return "\n".join(fixed_lines)
+
+def fix_warning_linenumber(string):
+    pattern = r'code\.py:(\d+):'
+    replaced_string = re.sub(pattern, lambda match: 'code.py:' + str(int(match.group(1))-2) + ':', string)
+    return replaced_string