Merge pull request #77 from shwinshaker/add-unique-cache

Specify a unique cache directory before each code execution

Author: terryyz

bigcodebench/gen/util/__init__.py

@@ -3,6 +3,7 @@
 import sys
 import types
 import unittest
+import tempfile
 import multiprocessing
 from multiprocessing import Array, Value, Manager
 from bigcodebench.eval.utils import (
@@ -17,55 +18,69 @@
 
 def trusted_exec(code, test_code, task_id, max_as_limit, max_data_limit, max_stack_limit, times):
     """Execute trusted code in place."""
+    # Specify a unique cache dir by modifying XDG_CONFIG_HOME
+    old_xdg = os.environ.get("XDG_CONFIG_HOME")
+    temp_xdg = tempfile.mkdtemp(prefix="xdg_config_")
+    os.environ["XDG_CONFIG_HOME"] = temp_xdg
 
-    with create_tempdir():
-        import os
-        import shutil
-        import builtins
-
-        rmtree = shutil.rmtree
-        rmdir = os.rmdir
-        chdir = os.chdir
-        module_name = "__test__"
-        new_module = types.ModuleType(module_name)
-        reliability_guard(max_as_limit, max_data_limit, max_stack_limit)
-        # Set necessary attributes for the module
-        new_module.__dict__.update({
-            '__builtins__': builtins,
-            '__file__': f"{module_name}.py",
-            '__package__': None,
-            '__doc__': None,
-            'sys': sys,
-            'os': os,
-            'environ': os.environ,
-        })
-
-        # Combine the user code and the test code
-        full_code = code + "\n" + test_code
-
-        # Compile and execute the combined code within the new module
-        exec(compile(full_code, f"{module_name}.py", 'exec'),
-             new_module.__dict__)
-        sys.modules[module_name] = new_module
-        TestCases = getattr(new_module, 'TestCases')
-        loader = unittest.TestLoader()
-        suite = loader.loadTestsFromTestCase(TestCases)
-        test_result = unittest.TestResult()
-        start = time.time()
-        with safe_environment(), swallow_io(), time_limit(seconds=TIMEOUT_LIMIT):
-            suite.run(test_result)
-        
-        errors = test_result.failures + test_result.errors
-        if len(errors) > 0:
-            print(errors)
-            times.value = -1
+    try:
+        with create_tempdir():
+            import shutil
+            import builtins
+
+            rmtree = shutil.rmtree
+            rmdir = os.rmdir
+            chdir = os.chdir
+            module_name = "__test__"
+            new_module = types.ModuleType(module_name)
+
+            reliability_guard(max_as_limit, max_data_limit, max_stack_limit)
+
+            # Set necessary attributes for the module
+            new_module.__dict__.update({
+                '__builtins__': builtins,
+                '__file__': f"{module_name}.py",
+                '__package__': None,
+                '__doc__': None,
+                'sys': sys,
+                'os': os,
+                'environ': os.environ,
+            })
+
+            # Combine the user code and the test code
+            full_code = code + "\n" + test_code
+
+            # Compile and execute the combined code within the new module
+            exec(compile(full_code, f"{module_name}.py", 'exec'),
+                 new_module.__dict__)
+            sys.modules[module_name] = new_module
+            TestCases = getattr(new_module, 'TestCases')
+            loader = unittest.TestLoader()
+            suite = loader.loadTestsFromTestCase(TestCases)
+            test_result = unittest.TestResult()
+            start = time.time()
+            with safe_environment(), swallow_io(), time_limit(seconds=TIMEOUT_LIMIT):
+                suite.run(test_result)
+
+            errors = test_result.failures + test_result.errors
+            if len(errors) > 0:
+                print(errors)
+                times.value = -1
+            else:
+                times.value = time.time() - start
+
+            # Needed for cleaning up.
+            shutil.rmtree = rmtree
+            os.rmdir = rmdir
+            os.chdir = chdir
+
+    finally:
+        # Restore the original environment variable
+        if old_xdg is None:
+            os.environ.pop("XDG_CONFIG_HOME", None)
         else:
-            times.value = time.time() - start
-        
-        # Needed for cleaning up.
-        shutil.rmtree = rmtree
-        os.rmdir = rmdir
-        os.chdir = chdir
+            os.environ["XDG_CONFIG_HOME"] = old_xdg
+        shutil.rmtree(temp_xdg, ignore_errors=True)
 
 
 def trusted_check_exec(code, inputs):