ai_batch_queue/
types.rs

1use serde::{Deserialize, Serialize};
2use stack_ids::{AttemptId, TraceCtx, TrialId};
3use std::time::Duration;
4
5/// Scheduling configuration for batch queue fairness.
6///
7/// Controls how consecutive same-resource jobs are limited to prevent
8/// starvation of other resources.
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct SchedulingConfig {
11    /// Maximum consecutive jobs for the same resource key before yielding
12    /// to a different resource. Default: 3.
13    pub max_consecutive_same_key: usize,
14    /// Cooldown duration between resource switches. Default: 0 (none).
15    #[serde(with = "duration_millis")]
16    pub resource_switch_cooldown: Duration,
17    /// Whether to enable model-aware reordering. Default: true.
18    pub enable_reordering: bool,
19}
20
21impl Default for SchedulingConfig {
22    fn default() -> Self {
23        Self {
24            max_consecutive_same_key: 3,
25            resource_switch_cooldown: Duration::ZERO,
26            enable_reordering: true,
27        }
28    }
29}
30
31/// Confidence level of an ETA estimate.
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
33#[serde(rename_all = "camelCase")]
34pub enum EtaConfidence {
35    /// No historical data; estimate is a rough guess.
36    Low,
37    /// Some data points but not enough for high confidence.
38    Medium,
39    /// Sufficient data points (>= 10) for a reliable estimate.
40    High,
41}
42
43/// Structured ETA estimate with confidence indication.
44#[derive(Debug, Clone, Serialize, Deserialize)]
45#[serde(rename_all = "camelCase")]
46pub struct EtaEstimate {
47    /// Estimated remaining time in milliseconds.
48    pub remaining_ms: u64,
49    /// Number of items remaining.
50    pub items_remaining: usize,
51    /// Average per-item duration in milliseconds (based on history).
52    pub avg_item_ms: u64,
53    /// Confidence level of the estimate.
54    pub confidence: EtaConfidence,
55    /// Number of historical data points used for the estimate.
56    pub sample_count: u64,
57}
58
59/// Serde helper for Duration as milliseconds.
60mod duration_millis {
61    use serde::{Deserialize, Deserializer, Serialize, Serializer};
62    use std::time::Duration;
63
64    pub fn serialize<S: Serializer>(d: &Duration, s: S) -> Result<S::Ok, S::Error> {
65        d.as_millis().serialize(s)
66    }
67
68    pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<Duration, D::Error> {
69        let ms = u64::deserialize(d)?;
70        Ok(Duration::from_millis(ms))
71    }
72}
73
74/// Per-item status within a batch job.
75#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
76#[serde(rename_all = "camelCase")]
77pub enum BatchItemStatus {
78    Pending,
79    Running,
80    Completed,
81    Failed,
82    Skipped,
83    Cancelled,
84}
85
86/// Overall batch job status.
87#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
88#[serde(rename_all = "camelCase")]
89pub enum BatchJobStatus {
90    Queued,
91    Running,
92    Completed,
93    CompletedWithErrors,
94    Cancelled,
95}
96
97/// Overwrite policy for batch operations.
98#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
99#[serde(rename_all = "camelCase")]
100pub enum OverwritePolicy {
101    /// Skip items that already have results.
102    Skip,
103    /// Overwrite existing results.
104    Overwrite,
105}
106
107/// Size bucket for ETA estimation — groups items by processing complexity.
108#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, Serialize, Deserialize)]
109#[serde(rename_all = "camelCase")]
110pub enum SizeBucket {
111    Small,
112    Medium,
113    Large,
114    Unknown,
115}
116
117impl SizeBucket {
118    /// Classify by pixel count. Thresholds: <500K = Small, <2M = Medium, else Large.
119    pub fn from_pixel_count(pixels: u64) -> Self {
120        if pixels < 500_000 {
121            Self::Small
122        } else if pixels < 2_000_000 {
123            Self::Medium
124        } else {
125            Self::Large
126        }
127    }
128
129    /// Classify from optional width/height dimensions.
130    pub fn from_dimensions(width: Option<u32>, height: Option<u32>) -> Self {
131        match (width, height) {
132            (Some(w), Some(h)) => Self::from_pixel_count(w as u64 * h as u64),
133            _ => Self::Unknown,
134        }
135    }
136}
137
138/// A single item within a batch job.
139///
140/// The `data` field carries user-defined per-item payload (e.g. file path,
141/// image ID, document reference). It must be serializable for event emission.
142///
143/// ## Trace context
144///
145/// Optional `trace_ctx` propagates cross-crate observability context from
146/// `stack-ids`. The `attempt_id` identifies the logical retry family for this
147/// batch-item (distinct from any outer job-level retry). The `trial_id` tracks
148/// each concrete execution attempt within that family.
149#[derive(Debug, Clone, Serialize, Deserialize)]
150#[serde(rename_all = "camelCase")]
151pub struct BatchItem<D>
152where
153    D: Clone + Send + Sync + Serialize,
154{
155    /// Unique identifier for this item within the batch.
156    pub id: String,
157    /// User-defined data payload.
158    pub data: D,
159    /// Current processing status.
160    pub status: BatchItemStatus,
161    /// Error message if status is Failed.
162    pub error: Option<String>,
163    /// Processing duration in milliseconds (set after completion).
164    pub duration_ms: Option<u64>,
165    /// Size bucket for ETA estimation.
166    pub size_bucket: SizeBucket,
167    /// Cross-crate trace context for observability propagation.
168    #[serde(default, skip_serializing_if = "Option::is_none")]
169    pub trace_ctx: Option<TraceCtx>,
170    /// Logical retry family for this batch-item.
171    ///
172    /// One `AttemptId` per retry-owner boundary. Batch-item retries (via
173    /// `retry_failed`) produce new `TrialId`s under the same `AttemptId`,
174    /// keeping them distinguishable from outer job-level retries.
175    #[serde(default, skip_serializing_if = "Option::is_none")]
176    pub attempt_id: Option<AttemptId>,
177    /// Concrete execution identifier within the current retry family.
178    ///
179    /// Minted fresh on each processing attempt. Allows correlating logs,
180    /// metrics, and errors to a specific execution without ambiguity.
181    #[serde(default, skip_serializing_if = "Option::is_none")]
182    pub trial_id: Option<TrialId>,
183}
184
185/// A batch job containing multiple items processed with the same resource.
186///
187/// The "resource key" (e.g. model name) is used for intelligent reordering:
188/// jobs sharing a resource are grouped together to minimize expensive
189/// resource swaps (GPU model loads, connection setup, etc.).
190#[derive(Debug, Clone, Serialize, Deserialize)]
191#[serde(rename_all = "camelCase")]
192pub struct BatchJob<D>
193where
194    D: Clone + Send + Sync + Serialize,
195{
196    /// Unique job identifier.
197    pub id: String,
198    /// The resource key (e.g. model name) used for grouping.
199    pub resource_key: String,
200    /// Human-readable operation label (e.g. "tag", "caption", "embed").
201    pub operation: String,
202    /// Overwrite policy for items that already have results.
203    pub overwrite_policy: OverwritePolicy,
204    /// The items in this batch.
205    pub items: Vec<BatchItem<D>>,
206    /// Current job status.
207    pub status: BatchJobStatus,
208    /// ISO 8601 timestamp when the job was created.
209    pub created_at: String,
210    /// ISO 8601 timestamp when processing started.
211    pub started_at: Option<String>,
212    /// ISO 8601 timestamp when processing finished.
213    pub completed_at: Option<String>,
214    /// Whether the job was reordered for resource optimization.
215    pub reordered: bool,
216    /// Human-readable note explaining the reorder.
217    pub reorder_note: Option<String>,
218}
219
220/// Summary of a completed batch job.
221#[derive(Debug, Clone, Serialize, Deserialize)]
222#[serde(rename_all = "camelCase")]
223pub struct BatchCompletionSummary {
224    pub job_id: String,
225    pub operation: String,
226    pub resource_key: String,
227    pub total: usize,
228    pub succeeded: usize,
229    pub failed: usize,
230    pub skipped: usize,
231    pub total_duration_ms: u64,
232    pub avg_duration_ms: u64,
233}
234
235/// Result of processing a single batch item.
236#[derive(Debug, Clone)]
237pub struct ItemResult {
238    /// Whether the item was processed successfully.
239    pub success: bool,
240    /// Optional output data (e.g. generated tags, captions).
241    pub output: Option<String>,
242    /// Error message if processing failed.
243    pub error: Option<String>,
244}
245
246impl ItemResult {
247    pub fn success() -> Self {
248        Self {
249            success: true,
250            output: None,
251            error: None,
252        }
253    }
254
255    pub fn success_with_output(output: String) -> Self {
256        Self {
257            success: true,
258            output: Some(output),
259            error: None,
260        }
261    }
262
263    pub fn failure(error: String) -> Self {
264        Self {
265            success: false,
266            output: None,
267            error: Some(error),
268        }
269    }
270}
ai_batch_queue/types.rs

ai_batch_queue/
types.rs
