Tests to use fast experimental communication strategy (SWE-agent#230)

Author: klieret

sweagent/environment/swe_env.py

@@ -19,6 +19,8 @@
 from simple_parsing.helpers.serialization.serializable import FrozenSerializable
 import yaml
 from sweagent.environment.utils import (
+    PROCESS_DONE_MARKER_END,
+    PROCESS_DONE_MARKER_START,
     copy_anything_to_container,
     copy_file_to_container,
     format_trajectory_markdown,
@@ -28,6 +30,7 @@
     parse_gh_issue_url,
     read_with_timeout,
     LOGGER_NAME,
+    read_with_timeout_experimental,
 )
 from swebench import (
     get_environment_yml,
@@ -445,11 +448,40 @@ def _init_scripts(self):
             error_msg="Failed to add commands directory to PATH",
         )
 
+    def _communicate_experimental(
+        self,
+        input: str,
+        timeout_duration=25,
+    ) -> str:
+        """Experimental version of `_communicate`"""
+
+        command_suffix = f"echo {PROCESS_DONE_MARKER_START}$?{PROCESS_DONE_MARKER_END}\n"
+        try:
+            self.returncode = None
+            cmd = input if input.endswith("\n") else input + "\n"
+            cmd += command_suffix
+            print(cmd)
+            os.write(self.container.stdin.fileno(), cmd.encode())
+            time.sleep(0.03)
+            self.container.stdin.flush()
+        except BrokenPipeError:
+            traceback.print_exc()
+            self.logger.error(
+                "Failed to communicate with container. Check docker logs for more information."
+            )
+            raise RuntimeError("Failed to communicate with container")
+
+        buffer, exit_code = read_with_timeout_experimental(self.container, timeout_duration) 
+        self.returncode = int(exit_code)
+        return buffer
+
     def _communicate(
         self,
         input: str,
         timeout_duration=25,
     ) -> str:
+        if "SWE_AGENT_EXPERIMENTAL_COMMUNICATE" in os.environ:
+            return self._communicate_experimental(input, timeout_duration)
         try:
             self.returncode = None
             cmd = input if input.endswith("\n") else input + "\n"

sweagent/environment/utils.py

@@ -161,6 +161,57 @@ def read_with_timeout(container, pid_func, timeout_duration):
     return buffer.decode()
 
 
+PROCESS_DONE_MARKER_START = "///PROCESS-DONE:"
+PROCESS_DONE_MARKER_END= ":PROCESS-DONE///"
+PROCESS_DONE_REGEX = re.compile(rf"{PROCESS_DONE_MARKER_START}(.+?){PROCESS_DONE_MARKER_END}")
+
+
+def read_with_timeout_experimental(container, timeout_duration):
+    """
+    Read data from a subprocess with a timeout.
+    This function uses a file descriptor to read data from the subprocess in a non-blocking way.
+
+    NOTE: This is an experimental implementation that is faster than `read_with_timeout`, but
+    has not been thoroughly tested.
+
+    Args:
+        container (subprocess.Popen): The subprocess container.
+        timeout_duration (int): The timeout duration in seconds.
+
+    Returns:
+        str: The data read from the subprocess, stripped of trailing newline characters.
+
+    Raises:
+        TimeoutError: If the timeout duration is reached while reading from the subprocess.
+    """
+    buffer = b""
+    fd = container.stdout.fileno()
+    end_time = time.time() + timeout_duration
+
+    while time.time() < end_time:
+        ready_to_read, _, _ = select.select([fd], [], [], 0.01)
+        if ready_to_read:
+            data = os.read(fd, 4096)
+            if data:
+                buffer += data
+        if PROCESS_DONE_MARKER_START in buffer.decode():
+            break
+        time.sleep(0.01)  # Prevents CPU hogging
+
+    if container.poll() is not None:
+        raise RuntimeError("Subprocess exited unexpectedly.\nCurrent buffer: {}".format(buffer.decode()))
+    if time.time() >= end_time:
+        raise TimeoutError("Timeout reached while reading from subprocess.\nCurrent buffer: {}".format(buffer.decode()))
+    decoded = buffer.decode()
+    body = "\n".join(line for line in decoded.splitlines() if not line.startswith(PROCESS_DONE_MARKER_START))
+    last_line = decoded.splitlines()[-1]
+    _results = PROCESS_DONE_REGEX.search(last_line)
+    if _results is None:
+        raise ValueError(f"Could not find process done marker in last line: {last_line=}, {body=}")
+    exit_code = _results.group(1)
+    return body, exit_code
+
+
 class timeout:
     def __init__(self, seconds=TIMEOUT_DURATION, error_message="Timeout"):
         self.seconds = seconds

tests/conftest.py

@@ -1,3 +1,4 @@
+import os
 import json
 import sys
 from pathlib import Path
@@ -10,7 +11,7 @@
 package_dir = root_dir / "sweagent"
 sys.path.insert(0, str(root_dir))
 sys.path.insert(1, str(package_dir))
-print("Adjusted path: ", sys.path)
+os.environ["SWE_AGENT_EXPERIMENTAL_COMMUNICATE"] = "1"
 
 
 @fixture

tests/test_env.py

@@ -1,4 +1,5 @@
 import dataclasses
+import os
 from pathlib import Path
 import pytest
 import yaml
@@ -80,3 +81,15 @@ def test_interrupt_close(test_env_args):
     env = SWEEnv(test_env_args)
     env.interrupt()
     env.close()
+
+
+@pytest.mark.slow
+def test_communicate_old(test_env_args):
+    del os.environ["SWE_AGENT_EXPERIMENTAL_COMMUNICATE"]
+    try:
+        env = SWEEnv(test_env_args)
+        env.reset()
+    except:
+        raise
+    finally:
+        os.environ["SWE_AGENT_EXPERIMENTAL_COMMUNICATE"] = "1"