eval-protocol 0.0.3__py3-none-any.whl

eval_protocol/rewards/apps_testing_util.py

@@ -0,0 +1,559 @@
+# Copyright 2024 PRIME team and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Borrowed from: https://huggingface.co/spaces/codeparrot/apps_metric/blob/main/utils.py
+# Adapted for reward-kit: Removed pyext.RuntimeModule, other minor adjustments may be needed.
+
+import ast
+import faulthandler
+import importlib.util  # Added for dynamic module loading
+import json
+import platform
+import re  # Added for re.search
+
+# to run the solution files we're using a timing based approach
+import signal
+import sys
+import textwrap  # Added for dedenting model output
+import traceback
+
+# used for debugging to time steps
+from datetime import datetime
+from enum import Enum
+
+# for capturing the stdout
+from io import StringIO
+
+# used for testing the code that reads from input
+from unittest.mock import mock_open, patch
+
+import numpy as np
+
+# from pyext import RuntimeModule # Removed this problematic import
+
+
+def truncatefn(s, length=300):
+    assert isinstance(s, str)
+    if len(s) <= length:
+        return s
+
+    return s[: length // 2] + "...(truncated) ..." + s[-length // 2 :]
+
+
+class CODE_TYPE(Enum):
+    call_based = 0
+    standard_input = 1
+
+
+# used to capture stdout as a list
+# from https://stackoverflow.com/a/16571630/6416660
+# alternative use redirect_stdout() from contextlib
+class Capturing(list):
+    def __enter__(self):
+        self._stdout = sys.stdout
+        sys.stdout = self._stringio = StringIO()
+        # Make closing the StringIO a no-op
+        self._stringio.close = lambda: None  # Changed lambda x: 1 to lambda: None
+        return self
+
+    def __exit__(self, *args):
+        self.append(self._stringio.getvalue())
+        del self._stringio  # free up some memory
+        sys.stdout = self._stdout
+
+
+def only_int_check(val):
+    return isinstance(val, int)
+
+
+def string_int_check(val):
+    return isinstance(val, str) and val.isdigit()
+
+
+def combined_int_check(val):
+    return only_int_check(val) or string_int_check(val)
+
+
+def clean_traceback(error_traceback):
+    file_start = error_traceback.find('File "<string>"')
+    if file_start == -1:  # Check if "<string>" is not found, common if exec is used directly
+        file_start = error_traceback.find('File "<dynamic_module>"')  # Fallback for our dynamic module name
+
+    if file_start != -1:
+        error_traceback = "Traceback (most recent call last):\n  " + error_traceback[file_start:]
+    return error_traceback
+
+
+def _load_module_from_string(module_name, code_string):
+    """Loads a Python module from a string using importlib."""
+    spec = importlib.util.spec_from_loader(module_name, loader=None, origin="<generated_code>")
+    if spec is None:
+        raise ImportError(f"Could not create spec for dynamic module '{module_name}'")
+
+    module = importlib.util.module_from_spec(spec)
+
+    # Execute the code in the new module's namespace
+    # Ensure that the module is usable by adding it to sys.modules temporarily if needed,
+    # or by ensuring its __dict__ is correctly populated.
+    try:
+        exec(code_string, module.__dict__)
+        # sys.modules[module_name] = module # Optional: if other parts of the code expect it in sys.modules
+    except Exception as e:
+        raise
+    return module
+
+
+def run_test(in_outs, test=None, debug=False, timeout=15):
+    """
+    if test(generated_code) is not None it'll try to run the code.
+    otherwise it'll just return an input and output pair.
+    """
+    # Disable functionalities that can make destructive changes to the test.
+    reliability_guard()
+
+    if debug:
+        print(f"start = {datetime.now().time()}")
+
+    if in_outs:
+        if in_outs.get("fn_name") is None:
+            which_type = CODE_TYPE.standard_input  # Standard input
+            method_name = None
+        else:
+            which_type = CODE_TYPE.call_based  # Call-based
+            method_name = in_outs["fn_name"]
+
+    if debug:
+        print(f"loaded input_output = {datetime.now().time()}")
+
+    if test is None:
+        raise AssertionError("should not happen: test code is none")
+    elif test is not None:
+        results = []
+        # Standard library imports prepended to the solution
+        sol = "from string import *\nfrom re import *\nfrom datetime import *\nfrom collections import *\nfrom heapq import *\nfrom bisect import *\nfrom copy import *\nfrom math import *\nfrom random import *\nfrom statistics import *\nfrom itertools import *\nfrom functools import *\nfrom operator import *\nfrom io import *\nfrom sys import *\nfrom json import *\nfrom builtins import *\nfrom typing import *\nimport string\nimport re\nimport datetime\nimport collections\nimport heapq\nimport bisect\nimport copy\nimport math\nimport random\nimport statistics\nimport itertools\nimport functools\nimport operator\nimport io\nimport sys\nimport json\nsys.setrecursionlimit(6*10**5)\n"  # noqa: E501
+        if debug:
+            print(f"loading test code = {datetime.now().time()}")
+
+        if which_type == CODE_TYPE.call_based:
+            sol += test
+            if debug:
+                print(f"sol = {sol}")
+            signal.alarm(timeout)  # This is Unix-specific
+            try:
+                # Replace RuntimeModule.from_string
+                tmp_sol = _load_module_from_string("tmp_sol_call_based", sol)
+                tmp = tmp_sol if "class Solution" not in test else tmp_sol.Solution()
+                signal.alarm(0)
+            except Exception as e:
+                signal.alarm(0)
+                error_traceback = traceback.format_exc()
+                if debug:
+                    print(f"type 0 compilation error = {e}")
+                results.append(-2)
+                return results, {
+                    "error": repr(e),
+                    "traceback": clean_traceback(error_traceback),
+                }
+            signal.alarm(0)
+
+        elif which_type == CODE_TYPE.standard_input:
+            try:
+                astree = ast.parse(test)
+                last_block = astree.body[-1]
+                if isinstance(last_block, ast.If):
+                    condition = last_block.test
+                    if ast.unparse(condition).strip() == "__name__ == '__main__'":
+                        test = ast.unparse(astree.body[:-1]) + "\n" + ast.unparse(last_block.body)
+            except Exception:
+                pass
+
+            # `test` is the user's generated code string at this point.
+            # Preprocessing for `if __name__ == "__main__"` is already done.
+
+            # Dedent the entire model-generated code block first
+            try:
+                dedented_test_code = textwrap.dedent(test)
+            except Exception as e:
+                # In case dedent fails (e.g. on empty or malformed string), use original
+                if debug:
+                    print(f"Warning: textwrap.dedent failed on model code: {e}. Using original code.")
+                dedented_test_code = test
+
+            # Check if 'def main(' is in the dedented code and if 'main()' call is missing.
+            main_defined = "def main(" in dedented_test_code
+            main_called_at_toplevel = re.search(r"^\s*main\s*\(\s*\)", dedented_test_code, re.MULTILINE) is not None
+            # Also consider if it's guarded by if __name__ == "__main__": which was removed by AST.
+            # If the AST modification removed an if __name__ block that called main,
+            # the original `test` string would be different from the AST-unparsed one.
+            # This is complex to track perfectly here.
+            # For now, a simpler heuristic: if `def main` is there, and no obvious `main()` call.
+
+            user_code_lines = dedented_test_code.split("\n")
+
+            # Imports from user code should be top-level in the module `sol`
+            # Other lines form the body of `def code():`
+            code_body_lines = []
+
+            for line in user_code_lines:
+                stripped_line = line.strip()
+                if stripped_line.startswith("from ") or stripped_line.startswith("import "):
+                    sol += stripped_line + "\n"  # Add stripped import directly to sol module scope
+                else:
+                    # Add original line from (potentially dedented) user code to be tab-indented
+                    code_body_lines.append("\t" + line)
+
+            if main_defined and not main_called_at_toplevel:
+                # If system prompt asks for main(), and model provides def main() but no call, add it.
+                # This assumes main() takes no arguments if called this way.
+                # This is appended to be *inside* the `def code():` wrapper.
+                code_body_lines.append("\tmain()")
+                if debug:
+                    print("Appended main() call as it was defined but not found called at top level.")
+
+            # Construct the `def code():` wrapper string
+            code_wrapper_str = "stdin = sys.stdin\nstdout = sys.stdout\ndef code():\n"
+            code_wrapper_str += "\n".join(code_body_lines)
+
+            sol += code_wrapper_str  # Add the "def code(): ..." to sol
+
+            if debug:
+                print(f"Constructed sol for standard_input: {sol}")
+            method_name = "code"  # We will call the code() function
+            signal.alarm(timeout)  # Unix-specific
+            try:
+                # Replace RuntimeModule.from_string
+                tmp_sol = _load_module_from_string("tmp_sol_std_input", sol)
+                tmp = tmp_sol
+                signal.alarm(0)
+            except Exception as e:
+                signal.alarm(0)
+                error_traceback = traceback.format_exc()
+                if debug:
+                    print(f"type 1 compilation error = {e}")
+                results.append(-2)
+                return results, {
+                    "error": repr(e),
+                    "traceback": clean_traceback(error_traceback),
+                }
+            signal.alarm(0)
+
+        if debug:
+            print(f"get method = {datetime.now().time()}")
+
+        try:
+            method = getattr(tmp, method_name)
+        except AttributeError:  # More specific exception
+            signal.alarm(0)
+            error_traceback = traceback.format_exc()
+            # error_info = sys.exc_info() # sys.exc_info() is less clear than repr(e)
+            results.append(-2)
+            return results, {
+                "error": f"AttributeError: Method '{method_name}' not found in dynamically loaded module.",
+                "traceback": clean_traceback(error_traceback),
+            }
+        except Exception as e:  # Catch other potential errors during getattr
+            signal.alarm(0)
+            error_traceback = traceback.format_exc()
+            results.append(-2)
+            return results, {
+                "error": repr(e),
+                "traceback": clean_traceback(error_traceback),
+            }
+
+        for index, inputs_str in enumerate(in_outs["inputs"]):  # Renamed inputs to inputs_str
+            raw_inputs = inputs_str
+            raw_outputs = in_outs["outputs"][index]
+
+            current_inputs = []  # Variable to hold processed inputs for the current test case
+
+            if which_type == CODE_TYPE.call_based:
+                # Assuming inputs_str is a string where each line is a separate JSON object for an argument
+                current_inputs = [json.loads(line) for line in inputs_str.split("\n") if line.strip()]
+                # Ensure in_outs["outputs"][index] is loaded if it's a string
+                if isinstance(in_outs["outputs"][index], str):
+                    in_outs["outputs"][index] = json.loads(in_outs["outputs"][index])
+
+                truncate_line_size = 300 // (raw_inputs.count("\n") + 1) if raw_inputs.count("\n") > 0 else 300
+                raw_inputs_truncated = "\n".join(
+                    [truncatefn(line, truncate_line_size) for line in raw_inputs.strip().split("\n")]
+                )
+                raw_outputs_truncated = (
+                    truncatefn(json.dumps(in_outs["outputs"][index]), 200)
+                    if not isinstance(in_outs["outputs"][index], str)
+                    else truncatefn(in_outs["outputs"][index], 200)
+                )
+
+            else:  # standard_input
+                current_inputs = inputs_str  # For standard input, inputs might be a single string block
+                raw_inputs_truncated = truncatefn(raw_inputs)
+                raw_outputs_truncated = truncatefn(in_outs["outputs"][index], 200)
+
+            # JSON forces dictionaries to have string keys; this undoes this (assuming a singleton list)
+            # This part seems specific and might need careful handling if inputs are not always lists of dicts
+            try:
+                if which_type == CODE_TYPE.call_based and current_inputs and isinstance(current_inputs[0], dict):
+                    current_inputs = [
+                        {int(k) if isinstance(k, str) and k.isdigit() else k: v for k, v in current_inputs[0].items()}
+                    ]
+            except Exception:
+                pass  # Ignore if conversion fails, proceed with original
+
+            # Similar conversion for outputs
+            try:
+                if isinstance(in_outs["outputs"][index], dict):
+                    in_outs["outputs"][index] = {
+                        int(k) if isinstance(k, str) and k.isdigit() else k: v
+                        for k, v in in_outs["outputs"][index].items()
+                    }
+                elif (
+                    isinstance(in_outs["outputs"][index], list)
+                    and in_outs["outputs"][index]
+                    and isinstance(in_outs["outputs"][index][0], dict)
+                ):
+                    in_outs["outputs"][index][0] = {
+                        int(k) if isinstance(k, str) and k.isdigit() else k: v
+                        for k, v in in_outs["outputs"][index][0].items()
+                    }
+            except Exception:
+                pass
+
+            if debug:
+                print(
+                    f"time: {datetime.now().time()} testing index = {index}  inputs = {current_inputs}, type = {which_type}"
+                )
+
+            if which_type == CODE_TYPE.call_based:
+                signal.alarm(timeout)  # Unix-specific
+                faulthandler.enable()
+                try:
+                    output = method(*current_inputs)
+
+                    # For comparison, ensure output format matches expected (e.g. list vs tuple)
+                    # ground truth sequences are not tuples
+                    if isinstance(output, tuple):
+                        output = list(output)
+
+                    # Comparison logic
+                    tmp_result = output == in_outs["outputs"][index]
+                    # Handle cases where expected output might be a list containing the actual output
+                    if isinstance(in_outs["outputs"][index], list) and len(in_outs["outputs"][index]) == 1:
+                        tmp_result = tmp_result or (output == in_outs["outputs"][index][0])
+
+                    # Further comparison for list of tuples vs list of lists
+                    try:
+                        if (
+                            isinstance(output, list)
+                            and output
+                            and isinstance(output[0], tuple)
+                            and isinstance(in_outs["outputs"][index], list)
+                            and in_outs["outputs"][index]
+                            and isinstance(in_outs["outputs"][index][0], list)
+                        ):
+                            output_list_of_lists = [list(x) for x in output]
+                            tmp_result = tmp_result or (output_list_of_lists == in_outs["outputs"][index])
+                            if isinstance(in_outs["outputs"][index][0], list):  # If expected is list of lists
+                                tmp_result = tmp_result or (output_list_of_lists == in_outs["outputs"][index][0])
+
+                    except Exception:
+                        pass
+
+                    results.append(tmp_result)
+
+                    if tmp_result is not True:
+                        return results, {
+                            "output": truncatefn(json.dumps(output), 200),
+                            "expected": raw_outputs_truncated,
+                            "inputs": raw_inputs_truncated,
+                            "error_message": "Wrong Answer",
+                        }
+                    signal.alarm(0)
+                except Exception as e:
+                    signal.alarm(0)
+                    error_traceback = traceback.format_exc()
+                    faulthandler.disable()
+                    if debug:
+                        print(f"Call-based runtime error or time limit exceeded error = {e}")
+                    results.append(-1)  # Indicate error
+                    return results, {
+                        "error": repr(e),
+                        "traceback": clean_traceback(error_traceback),
+                    }
+                faulthandler.disable()
+                signal.alarm(0)
+
+            elif which_type == CODE_TYPE.standard_input:
+                faulthandler.enable()
+
+                # Ensure inputs_str is a single string for StringIO
+                processed_inputs_str = inputs_str
+                if isinstance(inputs_str, list):
+                    processed_inputs_str = "\n".join(inputs_str)
+
+                # Ensure ground_truth is a string for comparison
+                ground_truth_str = in_outs["outputs"][index]
+                if isinstance(ground_truth_str, list):
+                    ground_truth_str = "\n".join(ground_truth_str)
+
+                signal.alarm(timeout)  # Unix-specific
+                captured_output_str = ""
+                try:
+                    with Capturing() as output_lines:
+                        call_method(method, processed_inputs_str)
+                    captured_output_str = "".join(
+                        output_lines
+                    ).rstrip()  # rstrip to remove trailing newline often added
+                    signal.alarm(0)
+                except Exception as e:
+                    signal.alarm(0)
+                    error_traceback = traceback.format_exc()
+                    faulthandler.disable()
+                    results.append(-1)  # Indicate error
+                    return results, {
+                        "error": repr(e),
+                        "traceback": clean_traceback(error_traceback),
+                    }
+                faulthandler.disable()
+                signal.alarm(0)
+
+                # Comparison for standard input
+                # Normalize by splitting lines and stripping whitespace from each line
+                output_for_compare = [line.strip() for line in captured_output_str.splitlines()]
+                expected_for_compare = [line.strip() for line in ground_truth_str.splitlines()]
+
+                tmp_result = output_for_compare == expected_for_compare
+
+                # Additional float comparison if direct string match fails
+                if not tmp_result:
+                    try:
+                        if len(output_for_compare) == len(expected_for_compare):
+                            output_float = [float(x) for x in output_for_compare]
+                            gt_float = [float(x) for x in expected_for_compare]
+                            if np.allclose(output_float, gt_float):
+                                tmp_result = True
+                    except (ValueError, TypeError):
+                        pass  # Not all are numbers, stick to string comparison
+
+                results.append(tmp_result)
+                if tmp_result is not True:
+                    return results, {
+                        "output": truncatefn(captured_output_str, 200),
+                        "expected": raw_outputs_truncated,
+                        "inputs": raw_inputs_truncated,
+                        "error_message": "Wrong Answer",
+                    }
+        # If all test cases for this sample passed
+    return results, {}
+
+
+def custom_compare_(output, ground_truth):
+    # This function seems to be part of an older comparison logic,
+    # more direct comparisons are now in run_test.
+    # Keeping it for now in case it's referenced, but likely can be simplified/removed.
+    if isinstance(output, list):
+        output_1 = "\n".join(output)
+        if stripped_string_compare(output_1, ground_truth):
+            return True
+
+    if isinstance(output, list):
+        output_2 = [o.lstrip().rstrip() for o in output]
+        output_2 = "\n".join(output_2)
+        if stripped_string_compare(output_2, ground_truth):
+            return True
+
+    return False
+
+
+def stripped_string_compare(s1, s2):
+    s1 = s1.lstrip().rstrip()
+    s2 = s2.lstrip().rstrip()
+    return s1 == s2
+
+
+def call_method(method, inputs_str_for_mock):  # Renamed inputs to avoid conflict
+    # inputs_str_for_mock is the single string containing all inputs for stdin
+
+    inputs_line_iterator = iter(inputs_str_for_mock.split("\n"))
+
+    @patch("builtins.open", mock_open(read_data=inputs_str_for_mock))
+    @patch("sys.stdin", StringIO(inputs_str_for_mock))
+    @patch("sys.stdin.readline", lambda *args: next(inputs_line_iterator) + "\n")  # Add newline as readline expects
+    @patch(
+        "sys.stdin.readlines",
+        lambda *args: [line + "\n" for line in inputs_str_for_mock.split("\n")],
+    )
+    @patch("sys.stdin.read", lambda *args: inputs_str_for_mock)
+    def _inner_call_method(_method_to_call):  # Renamed _method to avoid conflict
+        try:
+            return _method_to_call()
+        except SystemExit:  # Allow SystemExit to pass through, e.g. if code calls exit()
+            pass
+        finally:
+            pass
+
+    return _inner_call_method(method)
+
+
+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  # Moved import here as it's Unix-specific for some parts
+
+        # Check if resource module has RLIMIT_AS, etc. (for cross-platform safety)
+        if hasattr(resource, "RLIMIT_AS"):
+            resource.setrlimit(resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes))
+        if hasattr(resource, "RLIMIT_DATA"):
+            resource.setrlimit(resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes))
+        if platform.uname().system != "Darwin" and hasattr(resource, "RLIMIT_STACK"):  # RLIMIT_STACK not on macOS
+            resource.setrlimit(resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes))
+
+    faulthandler.disable()  # This is fine
+
+    # It's generally safer to avoid modifying builtins directly if possible.
+    # For a library, this can have wide-ranging effects.
+    # Consider if this level of modification is truly necessary for eval_protocol's use case
+    # or if the multiprocessing wrapper in utils.py provides sufficient isolation.
+    # Note: The original implementation had many builtins and os/shutil functions commented out.
+    # These have been removed for clarity, as the preferred method of sandboxing
+    # would be via process isolation (e.g. multiprocessing or a dedicated sandbox env).
+    # Modifying builtins directly in a library function can have unintended side effects.
+
+    import os
+
+    os.environ["OMP_NUM_THREADS"] = "1"
+
+    # Disabling os functions: Be cautious, as this makes the execution environment very restrictive.
+    # This might be too aggressive if the generated code legitimately needs some safe os interactions.
+    # The multiprocessing wrapper in utils.py already provides process isolation.
+
+    # Example of functions that were previously considered for disabling:
+    # os.kill, os.system, os.remove, os.fork, etc.
+    # shutil.rmtree, shutil.move
+    # subprocess.Popen
+    # Modifying __builtins__ or sys.modules entries.
+
+    # For eval_protocol, rely on higher-level sandboxing if untrusted code execution is a concern.
+    # The memory limits via `resource` are a good first step for resource exhaustion.
+    import shutil  # Keep import if other shutil functions are used, or remove if not.
+    import subprocess  # Keep import if other subprocess functions are used, or remove if not.
+    import sys  # Keep import for sys.stdout, sys.stdin manipulations.