Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions docs/modules/graph/builder.md
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,15 @@ parallel superstep (the active set runs in chunks of at most `n`).
`with_node_timeout(d)` fails the run with `TinyAgentsError::Timeout` if any node
handler does not resolve within `d`.

`CompiledGraph::with_run_deadline(d)` bounds the *whole run* by a wall-clock
`d`, checked at every super-step boundary: when the elapsed run time first
reaches `d` the run stops *between* super-steps with `TinyAgentsError::Timeout`,
leaving the last committed boundary checkpoint intact and resumable. Prefer this
over wrapping `run` in an external `tokio::time::timeout`, which aborts
mid-super-step and cannot leave a clean checkpoint. It bounds scheduling, not a
single in-flight node — pair it with `with_node_timeout` to also bound
individual handlers.

Compile-time options:

```rust
Expand Down
17 changes: 17 additions & 0 deletions src/graph/compiled/executor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -418,6 +418,23 @@ where
.await;
return Err(err);
}
// Whole-run wall-clock deadline: stop *between* super-steps once the
// elapsed run time reaches it, leaving the last committed boundary
// checkpoint intact (unlike an external `tokio::time::timeout`, which
// aborts mid-super-step and cannot). The already-completed super-steps
// and their checkpoints are preserved; the run fails with `Timeout`.
if let Some(deadline) = self.run_deadline {
let elapsed = started_at.elapsed().unwrap_or_default();
if elapsed >= deadline {
let err = TinyAgentsError::Timeout(format!(
"graph run exceeded its {deadline:?} deadline after {steps} super-step(s) \
({elapsed:?} elapsed)"
));
self.fail_run(&run_id, &thread_id, started_at, steps, &err, None)
.await;
return Err(err);
}
}
// Node-loop recursion: enforce `max_visits_per_node` per activation.
for activation in &active {
if let Err(err) = recursion.record_node_visit(&mut node_visits, &activation.node) {
Expand Down
21 changes: 21 additions & 0 deletions src/graph/compiled/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -277,6 +277,7 @@ impl<State, Update> CompiledGraph<State, Update> {
parallel,
max_concurrency,
node_timeout,
run_deadline: None,
durability: crate::graph::checkpoint::DurabilityMode::default(),
node_retry: None,
}
Expand Down Expand Up @@ -352,6 +353,26 @@ impl<State, Update> CompiledGraph<State, Update> {
self
}

/// Bounds the whole run by a wall-clock `deadline`, checked at every
/// super-step boundary.
///
/// Unlike wrapping [`run`](Self::run) in an external
/// [`tokio::time::timeout`] — which aborts mid-super-step and cannot leave a
/// clean checkpoint — this stops the run *between* super-steps: when the
/// elapsed run time first reaches `deadline`, the run fails with
/// [`TinyAgentsError::Timeout`] and (on a checkpointed thread) the last
/// committed boundary checkpoint stays intact, so the run can be resumed or
/// inspected rather than lost.
///
/// The deadline bounds *scheduling*, not a single in-flight node: a
/// long-running node still runs to completion within its super-step (bound
/// it independently with a per-node timeout). `None` (the default) imposes
/// no deadline.
pub fn with_run_deadline(mut self, deadline: std::time::Duration) -> Self {
self.run_deadline = Some(deadline);
self
}

/// Sets the checkpoint namespace (used by subgraph wrappers).
pub fn with_namespace(mut self, namespace: Vec<String>) -> Self {
self.namespace = namespace;
Expand Down
97 changes: 97 additions & 0 deletions src/graph/compiled/test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1861,6 +1861,103 @@ async fn node_timeout_fails_slow_handler() {
assert!(matches!(err, TinyAgentsError::Timeout(_)));
}

// ── Whole-run wall-clock deadline ────────────────────────────────────────────

/// A per-run deadline stops the run *between* super-steps once the elapsed run
/// time reaches it, surfacing [`TinyAgentsError::Timeout`].
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn run_deadline_stops_between_supersteps() {
let graph = GraphBuilder::<i32, i32>::overwrite()
.add_node("a", |s: i32, _c: NodeContext| async move {
tokio::time::sleep(Duration::from_millis(40)).await;
Ok(NodeResult::Update(s + 1))
})
.add_node("b", |s: i32, _c: NodeContext| async move {
Ok(NodeResult::Update(s + 1))
})
.set_entry("a")
.add_edge("a", "b")
.set_finish("b")
.compile()
.unwrap()
.with_run_deadline(Duration::from_millis(20));

// The first boundary (elapsed ~0) admits node `a`; the next boundary
// (elapsed ~40ms ≥ 20ms) trips the deadline before `b` ever runs.
let err = graph.run(0).await.unwrap_err();
assert!(matches!(err, TinyAgentsError::Timeout(_)), "got {err:?}");
}

/// A run that finishes within its deadline is unaffected — no false trip.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn run_deadline_allows_a_run_that_finishes_in_time() {
let graph = GraphBuilder::<i32, i32>::overwrite()
.add_node("a", |s: i32, _c: NodeContext| async move {
Ok(NodeResult::Update(s + 1))
})
.add_node("b", |s: i32, _c: NodeContext| async move {
Ok(NodeResult::Update(s + 1))
})
.set_entry("a")
.add_edge("a", "b")
.set_finish("b")
.compile()
.unwrap()
.with_run_deadline(Duration::from_secs(30));

let run = graph.run(0).await.unwrap();
assert_eq!(run.state, 2);
}

/// On a checkpointed thread, a deadline trip leaves the last committed boundary
/// checkpoint intact — so the run can be resumed to completion rather than lost
/// (the durability win over an external `tokio::time::timeout`).
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn run_deadline_leaves_last_checkpoint_resumable() {
let cp = Arc::new(InMemoryCheckpointer::<i32>::new());
let topology = || {
GraphBuilder::<i32, i32>::overwrite()
.add_node("a", |s: i32, _c: NodeContext| async move {
tokio::time::sleep(Duration::from_millis(40)).await;
Ok(NodeResult::Update(s + 1))
})
.add_node("b", |s: i32, _c: NodeContext| async move {
Ok(NodeResult::Update(s + 1))
})
.add_node("c", |s: i32, _c: NodeContext| async move {
Ok(NodeResult::Update(s + 1))
})
.set_entry("a")
.add_edge("a", "b")
.add_edge("b", "c")
.set_finish("c")
.compile()
.unwrap()
};

// Trips after `a`'s boundary (state=1, next=[b]) but before `b` runs.
let deadlined = topology()
.with_checkpointer(cp.clone())
.with_run_deadline(Duration::from_millis(20));
let err = deadlined.run_with_thread("t", 0).await.unwrap_err();
assert!(matches!(err, TinyAgentsError::Timeout(_)), "got {err:?}");

// The boundary checkpoint from the completed super-step survived intact.
let list = cp.list("t").await.unwrap();
assert!(
!list.is_empty(),
"the pre-deadline boundary checkpoint is intact"
);

// Resuming (no deadline) continues from that checkpoint to completion.
let resumed = topology().with_checkpointer(cp.clone());
let run = resumed
.resume("t", Command::resume(json!(null)))
.await
.unwrap();
assert_eq!(run.state, 3, "resume ran the remaining super-steps b and c");
}

// ── Network resilience: node retry + resumable failures ──────────────────────

/// A single-node graph whose handler fails (with a retryable model error) the
Expand Down
7 changes: 7 additions & 0 deletions src/graph/compiled/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,12 @@ pub struct CompiledGraph<State, Update> {
pub(crate) max_concurrency: Option<usize>,
/// Default per-node handler timeout (`None` = no timeout).
pub(crate) node_timeout: Option<std::time::Duration>,
/// Optional whole-run wall-clock deadline (`None` = no deadline). Checked at
/// every super-step boundary: when the elapsed run time reaches it the run
/// stops *between* super-steps with a [`TinyAgentsError::Timeout`], leaving
/// the last committed boundary checkpoint intact and resumable. Configured
/// via [`CompiledGraph::with_run_deadline`](crate::graph::CompiledGraph::with_run_deadline).
pub(crate) run_deadline: Option<std::time::Duration>,
/// When checkpoints are persisted relative to execution (default
/// [`DurabilityMode::Sync`]).
pub(crate) durability: DurabilityMode,
Expand Down Expand Up @@ -115,6 +121,7 @@ impl<State, Update> Clone for CompiledGraph<State, Update> {
parallel: self.parallel,
max_concurrency: self.max_concurrency,
node_timeout: self.node_timeout,
run_deadline: self.run_deadline,
durability: self.durability,
node_retry: self.node_retry.clone(),
}
Expand Down