mirror of
https://github.com/TauricResearch/TradingAgents.git
synced 2026-09-19 11:15:24 +03:00
fix(cli): make --checkpoint actually resume on the CLI path
- checkpoint setup lived only inside propagate(); the CLI streamed the checkpointer-less graph with no thread_id, so --checkpoint neither saved nor resumed a run - extract the lifecycle into reusable begin_checkpoint / end_checkpoint / clear_checkpoint_on_success (checkpoint_scope wraps them for propagate) and use them around the CLI stream #1249
This commit is contained in:
17
cli/main.py
17
cli/main.py
@@ -1127,6 +1127,15 @@ def run_analysis(checkpoint: bool | None = None):
|
|||||||
# (LLM tracking is handled separately via LLM constructor)
|
# (LLM tracking is handled separately via LLM constructor)
|
||||||
args = graph.propagator.get_graph_args(callbacks=[stats_handler])
|
args = graph.propagator.get_graph_args(callbacks=[stats_handler])
|
||||||
|
|
||||||
|
# Recompile with a checkpointer and inject the thread_id so --checkpoint
|
||||||
|
# actually saves and resumes on the CLI path (#1249); a no-op when
|
||||||
|
# checkpointing is disabled. Paired with end_checkpoint after the stream.
|
||||||
|
checkpoint_tid = graph.begin_checkpoint(
|
||||||
|
selections["ticker"], selections["analysis_date"], selections["asset_type"]
|
||||||
|
)
|
||||||
|
if checkpoint_tid is not None:
|
||||||
|
args.setdefault("config", {}).setdefault("configurable", {})["thread_id"] = checkpoint_tid
|
||||||
|
|
||||||
# Stream the analysis
|
# Stream the analysis
|
||||||
trace = []
|
trace = []
|
||||||
for chunk in graph.graph.stream(init_agent_state, **args):
|
for chunk in graph.graph.stream(init_agent_state, **args):
|
||||||
@@ -1231,6 +1240,14 @@ 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
|
||||||
|
# graph (#1249). A mid-stream failure skips this, leaving the checkpoint
|
||||||
|
# in place so the next run resumes.
|
||||||
|
graph.clear_checkpoint_on_success(
|
||||||
|
selections["ticker"], selections["analysis_date"], selections["asset_type"]
|
||||||
|
)
|
||||||
|
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
|
||||||
# so every report field populated across the run is present.
|
# so every report field populated across the run is present.
|
||||||
final_state = {}
|
final_state = {}
|
||||||
|
|||||||
123
tests/test_checkpoint_lifecycle.py
Normal file
123
tests/test_checkpoint_lifecycle.py
Normal file
@@ -0,0 +1,123 @@
|
|||||||
|
"""The checkpoint lifecycle is reusable so --checkpoint works on the CLI path (#1249).
|
||||||
|
|
||||||
|
Checkpoint setup previously lived only inside ``propagate``; the CLI streamed the
|
||||||
|
checkpointer-less graph, so ``--checkpoint`` neither saved nor resumed. The
|
||||||
|
lifecycle is now ``begin_checkpoint`` / ``end_checkpoint`` /
|
||||||
|
``clear_checkpoint_on_success`` on TradingAgentsGraph, used by both paths. These
|
||||||
|
tests drive that lifecycle exactly as the CLI does (begin -> stream self.graph ->
|
||||||
|
clear/end) and prove state is saved and resumed.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import tempfile
|
||||||
|
from typing import TypedDict
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from langgraph.graph import END, StateGraph
|
||||||
|
|
||||||
|
from tradingagents.graph.checkpointer import checkpoint_step
|
||||||
|
from tradingagents.graph.trading_graph import TradingAgentsGraph
|
||||||
|
|
||||||
|
_should_crash = False
|
||||||
|
|
||||||
|
|
||||||
|
class _State(TypedDict):
|
||||||
|
count: int
|
||||||
|
|
||||||
|
|
||||||
|
def _node_a(state: _State) -> dict:
|
||||||
|
return {"count": state["count"] + 1}
|
||||||
|
|
||||||
|
|
||||||
|
def _node_b(state: _State) -> dict:
|
||||||
|
if _should_crash:
|
||||||
|
raise RuntimeError("simulated mid-stream crash")
|
||||||
|
return {"count": state["count"] + 10}
|
||||||
|
|
||||||
|
|
||||||
|
def _workflow() -> StateGraph:
|
||||||
|
b = StateGraph(_State)
|
||||||
|
b.add_node("analyst", _node_a)
|
||||||
|
b.add_node("trader", _node_b)
|
||||||
|
b.set_entry_point("analyst")
|
||||||
|
b.add_edge("analyst", "trader")
|
||||||
|
b.add_edge("trader", END)
|
||||||
|
return b
|
||||||
|
|
||||||
|
|
||||||
|
def _bare_graph(tmpdir, *, enabled=True):
|
||||||
|
g = object.__new__(TradingAgentsGraph)
|
||||||
|
g.config = {
|
||||||
|
"checkpoint_enabled": enabled, "data_cache_dir": tmpdir,
|
||||||
|
"max_debate_rounds": 1, "max_risk_discuss_rounds": 1,
|
||||||
|
}
|
||||||
|
g.selected_analysts = ("market",)
|
||||||
|
g.workflow = _workflow()
|
||||||
|
g.graph = g.workflow.compile()
|
||||||
|
g._checkpointer_ctx = None
|
||||||
|
return g
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.unit
|
||||||
|
def test_disabled_is_a_noop():
|
||||||
|
with tempfile.TemporaryDirectory() as tmp:
|
||||||
|
g = _bare_graph(tmp, enabled=False)
|
||||||
|
plain = g.graph
|
||||||
|
assert g.begin_checkpoint("AAPL", "2026-05-08", "stock") is None
|
||||||
|
assert g.graph is plain # graph not recompiled
|
||||||
|
g.end_checkpoint() # safe no-op
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.unit
|
||||||
|
def test_begin_returns_thread_id_and_recompiles():
|
||||||
|
with tempfile.TemporaryDirectory() as tmp:
|
||||||
|
g = _bare_graph(tmp)
|
||||||
|
plain = g.graph
|
||||||
|
tid = g.begin_checkpoint("AAPL", "2026-05-08", "stock")
|
||||||
|
try:
|
||||||
|
assert tid # a real thread_id
|
||||||
|
assert g.graph is not plain # recompiled with a checkpointer
|
||||||
|
finally:
|
||||||
|
g.end_checkpoint()
|
||||||
|
assert g._checkpointer_ctx is None # restored
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.unit
|
||||||
|
def test_cli_style_usage_saves_then_resumes():
|
||||||
|
global _should_crash
|
||||||
|
with tempfile.TemporaryDirectory() as tmp:
|
||||||
|
cfg_args = ("AAPL", "2026-05-08", "stock")
|
||||||
|
|
||||||
|
# Run 1 (the CLI path): begin -> stream self.graph -> crash at 'trader'.
|
||||||
|
_should_crash = True
|
||||||
|
g1 = _bare_graph(tmp)
|
||||||
|
tid = g1.begin_checkpoint(*cfg_args)
|
||||||
|
args = {"config": {"configurable": {"thread_id": tid}}}
|
||||||
|
try:
|
||||||
|
with pytest.raises(RuntimeError):
|
||||||
|
for _ in g1.graph.stream({"count": 0}, **args):
|
||||||
|
pass
|
||||||
|
finally:
|
||||||
|
g1.end_checkpoint()
|
||||||
|
|
||||||
|
# A checkpoint was saved for this run signature (so --checkpoint works).
|
||||||
|
|
||||||
|
sig = g1._run_signature("stock")
|
||||||
|
assert checkpoint_step(tmp, "AAPL", "2026-05-08", sig) is not None
|
||||||
|
|
||||||
|
# Run 2 (fresh graph, as a new CLI invocation): resume and finish.
|
||||||
|
_should_crash = False
|
||||||
|
g2 = _bare_graph(tmp)
|
||||||
|
tid2 = g2.begin_checkpoint(*cfg_args)
|
||||||
|
assert tid2 == tid # stable id -> same thread resumes
|
||||||
|
try:
|
||||||
|
result = g2.graph.invoke(None, config={"configurable": {"thread_id": tid2}})
|
||||||
|
assert result["count"] == 11 # analyst(+1) resumed into trader(+10)
|
||||||
|
g2.clear_checkpoint_on_success(*cfg_args)
|
||||||
|
finally:
|
||||||
|
g2.end_checkpoint()
|
||||||
|
|
||||||
|
# Cleared on success -> a later run starts fresh.
|
||||||
|
assert checkpoint_step(tmp, "AAPL", "2026-05-08", sig) is None
|
||||||
@@ -3,6 +3,7 @@
|
|||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
|
from contextlib import contextmanager
|
||||||
from datetime import datetime, timedelta
|
from datetime import datetime, timedelta
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any
|
from typing import Any
|
||||||
@@ -400,32 +401,61 @@ class TradingAgentsGraph:
|
|||||||
# Resolve any pending memory-log entries for this ticker before the pipeline runs.
|
# Resolve any pending memory-log entries for this ticker before the pipeline runs.
|
||||||
self._resolve_pending_entries(company_name)
|
self._resolve_pending_entries(company_name)
|
||||||
|
|
||||||
# Recompile with a checkpointer if the user opted in.
|
with self.checkpoint_scope(company_name, trade_date, asset_type) as thread_id_value:
|
||||||
if self.config.get("checkpoint_enabled"):
|
return self._run_graph(
|
||||||
self._checkpointer_ctx = get_checkpointer(
|
company_name, trade_date, asset_type=asset_type,
|
||||||
self.config["data_cache_dir"], company_name
|
checkpoint_thread_id=thread_id_value,
|
||||||
)
|
)
|
||||||
saver = self._checkpointer_ctx.__enter__()
|
|
||||||
self.graph = self.workflow.compile(checkpointer=saver)
|
|
||||||
|
|
||||||
step = checkpoint_step(
|
def begin_checkpoint(self, company_name, trade_date, asset_type: str = "stock") -> str | None:
|
||||||
|
"""Recompile the graph with a per-ticker checkpointer and return the
|
||||||
|
``thread_id`` to inject into the stream/invoke ``config`` (or ``None``
|
||||||
|
when checkpointing is disabled).
|
||||||
|
|
||||||
|
Pair every call with :meth:`end_checkpoint` in a ``finally``. Both
|
||||||
|
``propagate`` (via :meth:`checkpoint_scope`) and the CLI stream path use
|
||||||
|
this so ``--checkpoint`` actually resumes (#1249); previously the setup
|
||||||
|
lived only inside ``propagate`` and the CLI streamed the checkpointer-less
|
||||||
|
graph, making the flag a no-op.
|
||||||
|
"""
|
||||||
|
if not self.config.get("checkpoint_enabled"):
|
||||||
|
return None
|
||||||
|
signature = self._run_signature(asset_type)
|
||||||
|
self._checkpointer_ctx = get_checkpointer(self.config["data_cache_dir"], company_name)
|
||||||
|
saver = self._checkpointer_ctx.__enter__()
|
||||||
|
self.graph = self.workflow.compile(checkpointer=saver)
|
||||||
|
|
||||||
|
step = checkpoint_step(
|
||||||
|
self.config["data_cache_dir"], company_name, str(trade_date), signature
|
||||||
|
)
|
||||||
|
if step is not None:
|
||||||
|
logger.info("Resuming from step %d for %s on %s", step, company_name, trade_date)
|
||||||
|
else:
|
||||||
|
logger.info("Starting fresh for %s on %s", company_name, trade_date)
|
||||||
|
return thread_id(company_name, str(trade_date), signature)
|
||||||
|
|
||||||
|
def end_checkpoint(self):
|
||||||
|
"""Restore the plain uncheckpointed graph after a checkpointed run."""
|
||||||
|
if self._checkpointer_ctx is not None:
|
||||||
|
self._checkpointer_ctx.__exit__(None, None, None)
|
||||||
|
self._checkpointer_ctx = None
|
||||||
|
self.graph = self.workflow.compile()
|
||||||
|
|
||||||
|
@contextmanager
|
||||||
|
def checkpoint_scope(self, company_name, trade_date, asset_type: str = "stock"):
|
||||||
|
"""Context-manager form of begin/end_checkpoint for the propagate path."""
|
||||||
|
try:
|
||||||
|
yield self.begin_checkpoint(company_name, trade_date, asset_type)
|
||||||
|
finally:
|
||||||
|
self.end_checkpoint()
|
||||||
|
|
||||||
|
def clear_checkpoint_on_success(self, company_name, trade_date, asset_type: str = "stock"):
|
||||||
|
"""Drop a completed run's checkpoint so a later run starts fresh (#1249)."""
|
||||||
|
if self.config.get("checkpoint_enabled"):
|
||||||
|
clear_checkpoint(
|
||||||
self.config["data_cache_dir"], company_name, str(trade_date),
|
self.config["data_cache_dir"], company_name, str(trade_date),
|
||||||
self._run_signature(asset_type),
|
self._run_signature(asset_type),
|
||||||
)
|
)
|
||||||
if step is not None:
|
|
||||||
logger.info(
|
|
||||||
"Resuming from step %d for %s on %s", step, company_name, trade_date
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
logger.info("Starting fresh for %s on %s", company_name, trade_date)
|
|
||||||
|
|
||||||
try:
|
|
||||||
return self._run_graph(company_name, trade_date, asset_type=asset_type)
|
|
||||||
finally:
|
|
||||||
if self._checkpointer_ctx is not None:
|
|
||||||
self._checkpointer_ctx.__exit__(None, None, None)
|
|
||||||
self._checkpointer_ctx = None
|
|
||||||
self.graph = self.workflow.compile()
|
|
||||||
|
|
||||||
def save_reports(self, final_state, ticker, save_path=None) -> Path:
|
def save_reports(self, final_state, ticker, save_path=None) -> Path:
|
||||||
"""Write the markdown report tree for a completed run, like the CLI does.
|
"""Write the markdown report tree for a completed run, like the CLI does.
|
||||||
@@ -442,7 +472,8 @@ class TradingAgentsGraph:
|
|||||||
)
|
)
|
||||||
return write_report_tree(final_state, ticker, save_path)
|
return write_report_tree(final_state, ticker, save_path)
|
||||||
|
|
||||||
def _run_graph(self, company_name, trade_date, asset_type: str = "stock"):
|
def _run_graph(self, company_name, trade_date, asset_type: str = "stock",
|
||||||
|
checkpoint_thread_id: str | None = None):
|
||||||
"""Execute the graph and write the resulting state to disk and memory log."""
|
"""Execute the graph and write the resulting state to disk and memory log."""
|
||||||
# Initialize state — inject memory log context for PM and the
|
# Initialize state — inject memory log context for PM and the
|
||||||
# deterministically resolved instrument identity for all agents.
|
# deterministically resolved instrument identity for all agents.
|
||||||
@@ -457,11 +488,10 @@ class TradingAgentsGraph:
|
|||||||
)
|
)
|
||||||
args = self.propagator.get_graph_args()
|
args = self.propagator.get_graph_args()
|
||||||
|
|
||||||
# Inject thread_id so same ticker+date+graph-shape resumes; a different
|
# Inject the checkpoint thread_id (from checkpoint_scope) so the same
|
||||||
# date or graph shape starts fresh (#1089).
|
# ticker+date+graph-shape resumes; a different one starts fresh (#1089).
|
||||||
if self.config.get("checkpoint_enabled"):
|
if checkpoint_thread_id is not None:
|
||||||
tid = thread_id(company_name, str(trade_date), self._run_signature(asset_type))
|
args.setdefault("config", {}).setdefault("configurable", {})["thread_id"] = checkpoint_thread_id
|
||||||
args.setdefault("config", {}).setdefault("configurable", {})["thread_id"] = tid
|
|
||||||
|
|
||||||
if self.debug:
|
if self.debug:
|
||||||
trace = []
|
trace = []
|
||||||
@@ -499,11 +529,7 @@ class TradingAgentsGraph:
|
|||||||
)
|
)
|
||||||
|
|
||||||
# Clear checkpoint on successful completion to avoid stale state.
|
# Clear checkpoint on successful completion to avoid stale state.
|
||||||
if self.config.get("checkpoint_enabled"):
|
self.clear_checkpoint_on_success(company_name, trade_date, asset_type)
|
||||||
clear_checkpoint(
|
|
||||||
self.config["data_cache_dir"], company_name, str(trade_date),
|
|
||||||
self._run_signature(asset_type),
|
|
||||||
)
|
|
||||||
|
|
||||||
return final_state, self.process_signal(final_state["final_trade_decision"])
|
return final_state, self.process_signal(final_state["final_trade_decision"])
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user