gavindiaz/NexQuant
1
0
Fork 0
mirror of https://github.com/NicolasBohn/NexQuant.git synced 2026-08-09 04:57:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add do_truncate control for the load function (#656)

Browse Source 
feat: add do_truncate control for the load function (#656)
This commit is contained in:
Roland Minrui
2025-03-04 15:18:36 +08:00
committed by GitHub
parent 30b52d7d57
commit 2ed6ec53d9
2 changed files with 14 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rdagent/utils/workflow.py
+7 -3
View File
@@ -161,7 +161,9 @@ class LoopBase:
pickle.dump(self, f)
@classmethod
def load(cls, path: Union[str, Path], output_path: Optional[Union[str, Path]] = None) -> "LoopBase":
def load(
cls, path: Union[str, Path], output_path: Optional[Union[str, Path]] = None, do_truncate: bool = True
) -> "LoopBase":
path = Path(path)
with path.open("rb") as f:
session = cast(LoopBase, pickle.load(f))
@@ -175,8 +177,10 @@ class LoopBase:
# set trace path
logger.set_trace_path(session.session_folder.parent)
max_loop = max(session.loop_trace.keys())
logger.storage.truncate(time=session.loop_trace[max_loop][-1].end)
# truncate future message
if do_truncate:
max_loop = max(session.loop_trace.keys())
logger.storage.truncate(time=session.loop_trace[max_loop][-1].end)
return session