From e1f86d7ca0782082e99768d77486e2b38cf1d6ed Mon Sep 17 00:00:00 2001
From: you-n-g <you-n-g@users.noreply.github.com>
Date: Thu, 6 Feb 2025 23:00:15 +0800
Subject: [PATCH] fix: add stdout context length setting and improve text
 shrinking logic (#559)

* feat: Add stdout context length setting and improve text shrinking logic

* lint
---
 rdagent/core/conf.py       | 4 ++++
 rdagent/core/experiment.py | 8 +++++++-
 rdagent/log/storage.py     | 5 ++++-
 rdagent/utils/fmt.py       | 2 +-
 4 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/rdagent/core/conf.py b/rdagent/core/conf.py
index eddf5547..dc43b0ee 100644
--- a/rdagent/core/conf.py
+++ b/rdagent/core/conf.py
@@ -82,5 +82,9 @@ class RDAgentSettings(ExtendedBaseSettings):
         # executing the function multiple times
     )
 
+    # misc
+    """The limitation of context stdout"""
+    stdout_context_len: int = 400
+
 
 RD_AGENT_SETTINGS = RDAgentSettings()
diff --git a/rdagent/core/experiment.py b/rdagent/core/experiment.py
index a2a8e1a7..28f714c5 100644
--- a/rdagent/core/experiment.py
+++ b/rdagent/core/experiment.py
@@ -14,6 +14,7 @@ from typing import Any, Generic, TypeVar
 
 from rdagent.core.conf import RD_AGENT_SETTINGS
 from rdagent.utils import filter_progress_bar
+from rdagent.utils.fmt import shrink_text
 
 if typing.TYPE_CHECKING:
     from rdagent.core.proposal import Hypothesis
@@ -236,7 +237,12 @@ class FBWorkspace(Workspace):
         self.inject_files(**self.file_dict)
         # TODO: env should be not None in new design (no code can run without environment)
         if env is not None and entry is not None:
-            return filter_progress_bar(env.run(entry, str(self.workspace_path)))
+            # We believe that removing the progress bar might resolve the issue.
+            # If needed, cutting off the information in the middle should be a backup plan.
+            return shrink_text(
+                filter_progress_bar(env.run(entry, str(self.workspace_path))),
+                context_lines=RD_AGENT_SETTINGS.stdout_context_len,
+            )
         return None
 
     def __str__(self) -> str:
diff --git a/rdagent/log/storage.py b/rdagent/log/storage.py
index acf87606..f92d2ef7 100644
--- a/rdagent/log/storage.py
+++ b/rdagent/log/storage.py
@@ -142,7 +142,10 @@ class FileStorage(Storage):
                     if "Logging object in" in msg:
                         absolute_p = msg.split("Logging object in ")[1]
                         p = Path(absolute_p)
-                        p.unlink()
+                        if p.exists():
+                            p.unlink()
+                        else:
+                            print(f"Missing pickle object: {p}.")
                     continue
 
                 new_content += content[log_start:log_end]
diff --git a/rdagent/utils/fmt.py b/rdagent/utils/fmt.py
index b49c9e71..3d382fb5 100644
--- a/rdagent/utils/fmt.py
+++ b/rdagent/utils/fmt.py
@@ -5,7 +5,7 @@ Tools that support generating better formats.
 
 def shrink_text(text: str, context_lines: int = 200) -> str:
     """
-    When the context is too long, hide the part that is not important.
+    When the context is too long, hide the part in the middle.
 
         text before
         ... (XXXXX lines are hidden) ...