vv_agent/runtime/backends/distributed/
execution.rs

1use crate::budget::{BudgetUsageSnapshot, RunBudgetLimits};
2use crate::runtime::checkpoint_resume::CheckpointController;
3use crate::runtime::CancellationToken;
4use crate::types::{AgentResult, AgentTask, CycleRecord, Message, Metadata};
5
6use super::super::{execute_cycle_loop, execute_cycle_loop_with_state, failed_backend_result};
7use super::backend::DistributedBackend;
8use super::r#loop::DistributedRunContext;
9
10impl DistributedBackend {
11    pub fn execute<F>(
12        &self,
13        task: &AgentTask,
14        initial_messages: Vec<Message>,
15        shared_state: Metadata,
16        cycle_executor: F,
17        cancellation_token: Option<&CancellationToken>,
18        max_cycles: u32,
19    ) -> AgentResult
20    where
21        F: FnMut(
22            u32,
23            &mut Vec<Message>,
24            &mut Vec<CycleRecord>,
25            &mut Metadata,
26            Option<&CancellationToken>,
27        ) -> Option<AgentResult>,
28    {
29        match (
30            &self.runtime_recipe,
31            &self.state_store,
32            &self.cycle_dispatcher,
33        ) {
34            (Some(recipe), Some(state_store), Some(cycle_dispatcher)) => {
35                return self.execute_distributed(
36                    initial_messages,
37                    shared_state,
38                    DistributedRunContext {
39                        task,
40                        recipe,
41                        state_store,
42                        cycle_dispatcher,
43                        cancellation_token,
44                        max_cycles,
45                        budget_limits: None,
46                        initial_budget_usage: None,
47                    },
48                );
49            }
50            (Some(_), _, _) => {
51                return failed_backend_result(
52                    initial_messages,
53                    Vec::new(),
54                    shared_state,
55                    "Distributed backend requires a state_store and cycle_dispatcher".to_string(),
56                );
57            }
58            (None, _, _) => {}
59        }
60        execute_cycle_loop(
61            initial_messages,
62            shared_state,
63            cycle_executor,
64            cancellation_token,
65            max_cycles,
66        )
67    }
68
69    #[allow(clippy::too_many_arguments)]
70    pub(crate) fn execute_with_state<F>(
71        &self,
72        task: &AgentTask,
73        initial_messages: Vec<Message>,
74        initial_cycles: Vec<CycleRecord>,
75        shared_state: Metadata,
76        cycle_executor: F,
77        cancellation_token: Option<&CancellationToken>,
78        cycle_index_start: u32,
79        cycle_count: u32,
80        budget_limits: Option<RunBudgetLimits>,
81        initial_budget_usage: Option<BudgetUsageSnapshot>,
82        checkpoint_controller: Option<CheckpointController>,
83    ) -> AgentResult
84    where
85        F: FnMut(
86            u32,
87            &mut Vec<Message>,
88            &mut Vec<CycleRecord>,
89            &mut Metadata,
90            Option<&CancellationToken>,
91        ) -> Option<AgentResult>,
92    {
93        if self.runtime_recipe.is_some()
94            && checkpoint_controller.is_none()
95            && (!initial_cycles.is_empty() || cycle_index_start != 1)
96        {
97            return failed_backend_result(
98                initial_messages,
99                initial_cycles,
100                shared_state,
101                "Checkpoint resume cannot recursively dispatch a distributed backend".to_string(),
102            );
103        }
104        if self.runtime_recipe.is_some() {
105            if let Some(checkpoint_controller) = checkpoint_controller {
106                return self.execute_distributed_v2(
107                    task,
108                    cycle_index_start,
109                    cycle_count,
110                    budget_limits,
111                    cancellation_token,
112                    checkpoint_controller,
113                );
114            }
115            return match (
116                &self.runtime_recipe,
117                &self.state_store,
118                &self.cycle_dispatcher,
119            ) {
120                (Some(recipe), Some(state_store), Some(cycle_dispatcher)) => self
121                    .execute_distributed(
122                        initial_messages,
123                        shared_state,
124                        DistributedRunContext {
125                            task,
126                            recipe,
127                            state_store,
128                            cycle_dispatcher,
129                            cancellation_token,
130                            max_cycles: cycle_count,
131                            budget_limits,
132                            initial_budget_usage,
133                        },
134                    ),
135                _ => failed_backend_result(
136                    initial_messages,
137                    initial_cycles,
138                    shared_state,
139                    "Distributed backend requires a state_store and cycle_dispatcher".to_string(),
140                ),
141            };
142        }
143        execute_cycle_loop_with_state(
144            initial_messages,
145            initial_cycles,
146            shared_state,
147            cycle_executor,
148            cancellation_token,
149            cycle_index_start,
150            cycle_count,
151        )
152    }
153}
vv_agent/runtime/backends/distributed/execution.rs

vv_agent/runtime/backends/distributed/
execution.rs
