mirror of
https://github.com/TauricResearch/TradingAgents.git
synced 2026-09-27 15:02:39 +03:00
fix(cli): resume a checkpoint without duplicating messages or leaking the saver
- on resume, the CLI and propagate re-passed the initial state to a thread with an existing checkpoint; nodes do not re-run, but the message reducer appended the initial messages again, duplicating them in the resumed state - feed None on resume (checkpoint_input) so LangGraph continues the interrupted run, and wrap the CLI stream in try/finally so the checkpointer tears down even if the stream raises - correct the _fetch_returns docstring to the 4-tuple return #1249
This commit is contained in:
+10
-6
@@ -1129,16 +1129,19 @@ def run_analysis(checkpoint: bool | None = None):
|
|||||||
|
|
||||||
# Recompile with a checkpointer and inject the thread_id so --checkpoint
|
# Recompile with a checkpointer and inject the thread_id so --checkpoint
|
||||||
# actually saves and resumes on the CLI path (#1249); a no-op when
|
# actually saves and resumes on the CLI path (#1249); a no-op when
|
||||||
# checkpointing is disabled. Paired with end_checkpoint after the stream.
|
# checkpointing is disabled. Torn down in the finally below.
|
||||||
checkpoint_tid = graph.begin_checkpoint(
|
checkpoint_tid = graph.begin_checkpoint(
|
||||||
selections["ticker"], selections["analysis_date"], selections["asset_type"]
|
selections["ticker"], selections["analysis_date"], selections["asset_type"]
|
||||||
)
|
)
|
||||||
if checkpoint_tid is not None:
|
if checkpoint_tid is not None:
|
||||||
args.setdefault("config", {}).setdefault("configurable", {})["thread_id"] = checkpoint_tid
|
args.setdefault("config", {}).setdefault("configurable", {})["thread_id"] = checkpoint_tid
|
||||||
|
|
||||||
# Stream the analysis
|
# Stream the analysis. On resume, feed None so LangGraph continues the
|
||||||
|
# interrupted run instead of re-appending the initial state (#1249); the
|
||||||
|
# try/finally tears the checkpointer down even if the stream raises.
|
||||||
trace = []
|
trace = []
|
||||||
for chunk in graph.graph.stream(init_agent_state, **args):
|
try:
|
||||||
|
for chunk in graph.graph.stream(graph.checkpoint_input(init_agent_state), **args):
|
||||||
# Process all messages in chunk, deduplicating by message ID
|
# Process all messages in chunk, deduplicating by message ID
|
||||||
for message in chunk.get("messages", []):
|
for message in chunk.get("messages", []):
|
||||||
msg_id = getattr(message, "id", None)
|
msg_id = getattr(message, "id", None)
|
||||||
@@ -1240,12 +1243,13 @@ def run_analysis(checkpoint: bool | None = None):
|
|||||||
|
|
||||||
trace.append(chunk)
|
trace.append(chunk)
|
||||||
|
|
||||||
# The stream completed: drop this run's checkpoint and restore the plain
|
# Clean run: drop this run's checkpoint so a later run starts fresh.
|
||||||
# graph (#1249). A mid-stream failure skips this, leaving the checkpoint
|
# A mid-stream failure skips this, keeping the checkpoint for resume.
|
||||||
# in place so the next run resumes.
|
|
||||||
graph.clear_checkpoint_on_success(
|
graph.clear_checkpoint_on_success(
|
||||||
selections["ticker"], selections["analysis_date"], selections["asset_type"]
|
selections["ticker"], selections["analysis_date"], selections["asset_type"]
|
||||||
)
|
)
|
||||||
|
finally:
|
||||||
|
# Always restore the plain uncheckpointed graph, even on failure.
|
||||||
graph.end_checkpoint()
|
graph.end_checkpoint()
|
||||||
|
|
||||||
# Streamed chunks are per-node deltas, not full state. Merge them
|
# Streamed chunks are per-node deltas, not full state. Merge them
|
||||||
|
|||||||
@@ -84,6 +84,38 @@ def test_begin_returns_thread_id_and_recompiles():
|
|||||||
assert g._checkpointer_ctx is None # restored
|
assert g._checkpointer_ctx is None # restored
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.unit
|
||||||
|
def test_checkpoint_input_is_none_only_when_resuming():
|
||||||
|
global _should_crash
|
||||||
|
with tempfile.TemporaryDirectory() as tmp:
|
||||||
|
init = {"count": 0}
|
||||||
|
args = ("AAPL", "2026-05-08", "stock")
|
||||||
|
# Fresh run: no checkpoint yet -> stream the initial state, then crash.
|
||||||
|
g1 = _bare_graph(tmp)
|
||||||
|
tid = g1.begin_checkpoint(*args)
|
||||||
|
try:
|
||||||
|
assert g1._resuming is False
|
||||||
|
assert g1.checkpoint_input(init) is init # not resuming -> initial state
|
||||||
|
_should_crash = True
|
||||||
|
with pytest.raises(RuntimeError):
|
||||||
|
for _ in g1.graph.stream(init, config={"configurable": {"thread_id": tid}}):
|
||||||
|
pass
|
||||||
|
finally:
|
||||||
|
g1.end_checkpoint()
|
||||||
|
assert g1.checkpoint_input(init) is init # reset after teardown
|
||||||
|
|
||||||
|
# A later run finds the checkpoint -> resume by feeding None, not the
|
||||||
|
# initial state (re-passing it would duplicate messages, #1249).
|
||||||
|
_should_crash = False
|
||||||
|
g2 = _bare_graph(tmp)
|
||||||
|
g2.begin_checkpoint(*args)
|
||||||
|
try:
|
||||||
|
assert g2._resuming is True
|
||||||
|
assert g2.checkpoint_input(init) is None
|
||||||
|
finally:
|
||||||
|
g2.end_checkpoint()
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.unit
|
@pytest.mark.unit
|
||||||
def test_cli_style_usage_saves_then_resumes():
|
def test_cli_style_usage_saves_then_resumes():
|
||||||
global _should_crash
|
global _should_crash
|
||||||
|
|||||||
@@ -163,6 +163,7 @@ class TradingAgentsGraph:
|
|||||||
self.workflow = self.graph_setup.setup_graph(selected_analysts)
|
self.workflow = self.graph_setup.setup_graph(selected_analysts)
|
||||||
self.graph = self.workflow.compile()
|
self.graph = self.workflow.compile()
|
||||||
self._checkpointer_ctx = None
|
self._checkpointer_ctx = None
|
||||||
|
self._resuming = False
|
||||||
|
|
||||||
def _get_provider_kwargs(self) -> dict[str, Any]:
|
def _get_provider_kwargs(self) -> dict[str, Any]:
|
||||||
"""Get provider-specific kwargs for LLM client creation."""
|
"""Get provider-specific kwargs for LLM client creation."""
|
||||||
@@ -277,8 +278,10 @@ class TradingAgentsGraph:
|
|||||||
|
|
||||||
``benchmark`` is the index used as the alpha baseline (resolved by the
|
``benchmark`` is the index used as the alpha baseline (resolved by the
|
||||||
caller via ``_resolve_benchmark``). Returns ``(raw_return, alpha_return,
|
caller via ``_resolve_benchmark``). Returns ``(raw_return, alpha_return,
|
||||||
actual_holding_days)`` or ``(None, None, None)`` if price data is
|
actual_holding_days, resolution_date)`` — where ``resolution_date`` is
|
||||||
unavailable (too recent, delisted, or network error).
|
the date of the last price bar used, i.e. when the outcome became known
|
||||||
|
(#1251) — or ``(None, None, None, None)`` if price data is unavailable
|
||||||
|
(too recent, delisted, or network error).
|
||||||
"""
|
"""
|
||||||
from tradingagents.dataflows.symbol_utils import normalize_symbol
|
from tradingagents.dataflows.symbol_utils import normalize_symbol
|
||||||
|
|
||||||
@@ -433,6 +436,7 @@ class TradingAgentsGraph:
|
|||||||
lived only inside ``propagate`` and the CLI streamed the checkpointer-less
|
lived only inside ``propagate`` and the CLI streamed the checkpointer-less
|
||||||
graph, making the flag a no-op.
|
graph, making the flag a no-op.
|
||||||
"""
|
"""
|
||||||
|
self._resuming = False
|
||||||
if not self.config.get("checkpoint_enabled"):
|
if not self.config.get("checkpoint_enabled"):
|
||||||
return None
|
return None
|
||||||
signature = self._run_signature(asset_type)
|
signature = self._run_signature(asset_type)
|
||||||
@@ -443,18 +447,30 @@ class TradingAgentsGraph:
|
|||||||
step = checkpoint_step(
|
step = checkpoint_step(
|
||||||
self.config["data_cache_dir"], company_name, str(trade_date), signature
|
self.config["data_cache_dir"], company_name, str(trade_date), signature
|
||||||
)
|
)
|
||||||
|
self._resuming = step is not None
|
||||||
if step is not None:
|
if step is not None:
|
||||||
logger.info("Resuming from step %d for %s on %s", step, company_name, trade_date)
|
logger.info("Resuming from step %d for %s on %s", step, company_name, trade_date)
|
||||||
else:
|
else:
|
||||||
logger.info("Starting fresh for %s on %s", company_name, trade_date)
|
logger.info("Starting fresh for %s on %s", company_name, trade_date)
|
||||||
return thread_id(company_name, str(trade_date), signature)
|
return thread_id(company_name, str(trade_date), signature)
|
||||||
|
|
||||||
|
def checkpoint_input(self, init_state):
|
||||||
|
"""The value to stream/invoke: ``None`` to resume an existing checkpoint,
|
||||||
|
else the initial state for a fresh run.
|
||||||
|
|
||||||
|
LangGraph resumes an interrupted thread when invoked with ``None``;
|
||||||
|
re-passing the initial state instead appends it through the message
|
||||||
|
reducer, duplicating messages in the resumed state (#1249).
|
||||||
|
"""
|
||||||
|
return None if self._resuming else init_state
|
||||||
|
|
||||||
def end_checkpoint(self):
|
def end_checkpoint(self):
|
||||||
"""Restore the plain uncheckpointed graph after a checkpointed run."""
|
"""Restore the plain uncheckpointed graph after a checkpointed run."""
|
||||||
if self._checkpointer_ctx is not None:
|
if self._checkpointer_ctx is not None:
|
||||||
self._checkpointer_ctx.__exit__(None, None, None)
|
self._checkpointer_ctx.__exit__(None, None, None)
|
||||||
self._checkpointer_ctx = None
|
self._checkpointer_ctx = None
|
||||||
self.graph = self.workflow.compile()
|
self.graph = self.workflow.compile()
|
||||||
|
self._resuming = False
|
||||||
|
|
||||||
@contextmanager
|
@contextmanager
|
||||||
def checkpoint_scope(self, company_name, trade_date, asset_type: str = "stock"):
|
def checkpoint_scope(self, company_name, trade_date, asset_type: str = "stock"):
|
||||||
@@ -512,10 +528,12 @@ class TradingAgentsGraph:
|
|||||||
if checkpoint_thread_id is not None:
|
if checkpoint_thread_id is not None:
|
||||||
args.setdefault("config", {}).setdefault("configurable", {})["thread_id"] = checkpoint_thread_id
|
args.setdefault("config", {}).setdefault("configurable", {})["thread_id"] = checkpoint_thread_id
|
||||||
|
|
||||||
|
# None resumes an existing checkpoint; init_agent_state starts fresh (#1249).
|
||||||
|
graph_input = self.checkpoint_input(init_agent_state)
|
||||||
if self.debug:
|
if self.debug:
|
||||||
trace = []
|
trace = []
|
||||||
last_printed = None
|
last_printed = None
|
||||||
for chunk in self.graph.stream(init_agent_state, **args):
|
for chunk in self.graph.stream(graph_input, **args):
|
||||||
if chunk["messages"]:
|
if chunk["messages"]:
|
||||||
msg = chunk["messages"][-1]
|
msg = chunk["messages"][-1]
|
||||||
# Nodes after the trader don't append to messages, so the
|
# Nodes after the trader don't append to messages, so the
|
||||||
@@ -532,7 +550,7 @@ class TradingAgentsGraph:
|
|||||||
for chunk in trace:
|
for chunk in trace:
|
||||||
final_state.update(chunk)
|
final_state.update(chunk)
|
||||||
else:
|
else:
|
||||||
final_state = self.graph.invoke(init_agent_state, **args)
|
final_state = self.graph.invoke(graph_input, **args)
|
||||||
|
|
||||||
# Store current state for reflection.
|
# Store current state for reflection.
|
||||||
self.curr_state = final_state
|
self.curr_state = final_state
|
||||||
|
|||||||
Reference in New Issue
Block a user