imagegen_bridge_runtime/fanout/
mod.rs

1//! Reusable bounded multi-output emulation for providers with a lower native count.
2
3use std::{
4    collections::{BTreeSet, VecDeque},
5    sync::Arc,
6    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
7};
8
9use async_trait::async_trait;
10use futures_util::{FutureExt as _, StreamExt as _, future::BoxFuture, stream::FuturesUnordered};
11use imagegen_bridge_core::{
12    BatchCapabilities, BatchExecution, BatchMode, BridgeError, ErrorCode, ImageFailure,
13    ImageProvider, ImageRequest, ImageResponse, MultiImageFailurePolicy, ProviderCapabilities,
14    ProviderContext, ProviderDescriptor, ProviderEvent, RequestLimits, SessionMetadata, Timings,
15    U8Range, Usage, negotiate_request, validate_request,
16};
17use sha2::{Digest as _, Sha256};
18use tokio::sync::{Semaphore, mpsc};
19use tokio_util::sync::CancellationToken;
20
21const FAILED_SIBLING_GRACE: Duration = Duration::from_secs(1);
22
23/// Bounded output fan-out policy shared by provider adapters.
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub struct OutputFanoutConfig {
26    /// Effective maximum output count advertised by the decorated provider.
27    pub max_outputs: u8,
28    /// Provider-wide maximum simultaneous upstream operations.
29    pub max_parallel_outputs: u8,
30}
31
32impl OutputFanoutConfig {
33    fn validate(self) -> Result<Self, BridgeError> {
34        if self.max_outputs == 0
35            || self.max_parallel_outputs == 0
36            || self.max_parallel_outputs > self.max_outputs
37        {
38            return Err(BridgeError::new(
39                ErrorCode::Configuration,
40                "output fan-out limits are invalid",
41            ));
42        }
43        Ok(self)
44    }
45}
46
47/// Provider decorator that turns one logical multi-output request into bounded calls.
48pub struct OutputFanoutProvider {
49    inner: Arc<dyn ImageProvider>,
50    config: OutputFanoutConfig,
51    permits: Arc<Semaphore>,
52}
53
54impl std::fmt::Debug for OutputFanoutProvider {
55    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56        formatter
57            .debug_struct("OutputFanoutProvider")
58            .field("provider", &self.inner.descriptor().name)
59            .field("config", &self.config)
60            .finish_non_exhaustive()
61    }
62}
63
64impl OutputFanoutProvider {
65    /// Decorates a provider without changing its registry identity.
66    pub fn new(
67        inner: Arc<dyn ImageProvider>,
68        config: OutputFanoutConfig,
69    ) -> Result<Self, BridgeError> {
70        let config = config.validate()?;
71        Ok(Self {
72            inner,
73            config,
74            permits: Arc::new(Semaphore::new(usize::from(config.max_parallel_outputs))),
75        })
76    }
77
78    fn effective_capabilities(
79        &self,
80        mut capabilities: ProviderCapabilities,
81    ) -> Result<ProviderCapabilities, BridgeError> {
82        capabilities.validate()?;
83        if self.config.max_outputs <= capabilities.count.max {
84            return Ok(capabilities);
85        }
86        let native_count = capabilities.count;
87        capabilities.count.max = self.config.max_outputs;
88        capabilities.batching = BatchCapabilities {
89            mode: BatchMode::FanOut,
90            native_count,
91            max_parallel_outputs: self.config.max_parallel_outputs,
92        };
93        capabilities.validate()?;
94        Ok(capabilities)
95    }
96
97    async fn execute_fanout(
98        &self,
99        request: ImageRequest,
100        context: ProviderContext,
101        native_count: U8Range,
102    ) -> Result<ImageResponse, BridgeError> {
103        let started = Instant::now();
104        let requested_count = request.parameters.n;
105        let failure_policy = request.parameters.failure_policy;
106        let batch_parallel = match request.policies.batch_execution {
107            BatchExecution::Parallel => usize::from(self.config.max_parallel_outputs),
108            BatchExecution::Auto
109                if request.session.mode == imagegen_bridge_core::SessionMode::Isolated =>
110            {
111                usize::from(self.config.max_parallel_outputs)
112            }
113            BatchExecution::Sequential | BatchExecution::Auto => 1,
114        };
115        let batch_cancellation = context.cancellation.child_token();
116        let mut chunks = chunks(requested_count, native_count.max);
117        let mut active = FuturesUnordered::<BoxFuture<'static, ChunkResult>>::new();
118        let mut results = Vec::new();
119        for _ in 0..batch_parallel {
120            push_next(
121                &mut chunks,
122                &mut active,
123                &self.inner,
124                &self.permits,
125                &request,
126                &context,
127                &batch_cancellation,
128            );
129        }
130
131        while let Some(result) = active.next().await {
132            if failure_policy == MultiImageFailurePolicy::FailFast
133                && let Err(error) = &result.result
134            {
135                let siblings_may_have_run = !active.is_empty() || !results.is_empty();
136                batch_cancellation.cancel();
137                let _ = tokio::time::timeout(FAILED_SIBLING_GRACE, async {
138                    while active.next().await.is_some() {}
139                })
140                .await;
141                let mut error = error
142                    .clone()
143                    .with_detail("output_index", result.offset)
144                    .with_detail("generation_ms", result.elapsed_ms);
145                if siblings_may_have_run {
146                    error = error
147                        .retryable(false)
148                        .with_detail("outcome", "unknown")
149                        .with_detail("retry_scope", "do_not_retry_full_batch");
150                }
151                return Err(error);
152            }
153            results.push(result);
154            push_next(
155                &mut chunks,
156                &mut active,
157                &self.inner,
158                &self.permits,
159                &request,
160                &context,
161                &batch_cancellation,
162            );
163        }
164        results.sort_by_key(|result| result.offset);
165        aggregate(
166            request,
167            context.request_id,
168            requested_count,
169            failure_policy,
170            results,
171            started,
172        )
173    }
174}
175
176#[async_trait]
177impl ImageProvider for OutputFanoutProvider {
178    fn descriptor(&self) -> ProviderDescriptor {
179        self.inner.descriptor()
180    }
181
182    async fn capabilities(&self, model: Option<&str>) -> Result<ProviderCapabilities, BridgeError> {
183        self.effective_capabilities(self.inner.capabilities(model).await?)
184    }
185
186    async fn execute(
187        &self,
188        request: ImageRequest,
189        context: ProviderContext,
190    ) -> Result<ImageResponse, BridgeError> {
191        validate_request(&request, RequestLimits::default())?;
192        let native = self
193            .inner
194            .capabilities(request.routing.model.as_deref())
195            .await?;
196        let effective = self.effective_capabilities(native.clone())?;
197        let negotiated = negotiate_request(&request, &effective)?;
198        let request = negotiated.effective_request;
199        if request.parameters.n <= native.count.max {
200            let _permit = acquire(&self.permits, &context).await?;
201            let mut response = self.inner.execute(request.clone(), context).await?;
202            response.requested = negotiated.requested;
203            response.effective = request.parameters;
204            response
205                .normalizations
206                .splice(0..0, negotiated.normalizations);
207            return Ok(response);
208        }
209        let mut response = self
210            .execute_fanout(request.clone(), context, native.count)
211            .await?;
212        response.requested = negotiated.requested;
213        response.effective = request.parameters;
214        response
215            .normalizations
216            .splice(0..0, negotiated.normalizations);
217        Ok(response)
218    }
219
220    async fn check_ready(&self) -> Result<(), BridgeError> {
221        self.inner.check_ready().await
222    }
223
224    async fn get_session(&self, key: &str) -> Result<SessionMetadata, BridgeError> {
225        self.inner.get_session(key).await
226    }
227
228    async fn delete_session(&self, key: &str) -> Result<(), BridgeError> {
229        self.inner.delete_session(key).await
230    }
231
232    fn restart_count(&self) -> Option<u64> {
233        self.inner.restart_count()
234    }
235
236    async fn shutdown(&self) -> Result<(), BridgeError> {
237        self.inner.shutdown().await
238    }
239}
240
241struct ChunkSpec {
242    offset: u8,
243    count: u8,
244}
245
246struct ChunkResult {
247    offset: u8,
248    count: u8,
249    elapsed_ms: u64,
250    result: Result<ImageResponse, BridgeError>,
251}
252
253fn chunks(total: u8, native_max: u8) -> VecDeque<ChunkSpec> {
254    let mut chunks = VecDeque::new();
255    let mut offset = 0_u8;
256    while offset < total {
257        let count = native_max.min(total - offset);
258        chunks.push_back(ChunkSpec { offset, count });
259        offset += count;
260    }
261    chunks
262}
263
264#[allow(clippy::too_many_arguments)]
265fn push_next(
266    chunks: &mut VecDeque<ChunkSpec>,
267    active: &mut FuturesUnordered<BoxFuture<'static, ChunkResult>>,
268    inner: &Arc<dyn ImageProvider>,
269    permits: &Arc<Semaphore>,
270    request: &ImageRequest,
271    context: &ProviderContext,
272    cancellation: &CancellationToken,
273) {
274    let Some(chunk) = chunks.pop_front() else {
275        return;
276    };
277    let mut chunk_request = request.clone();
278    chunk_request.parameters.n = chunk.count;
279    let chunk_context = ProviderContext {
280        request_id: subrequest_id(&context.request_id, chunk.offset),
281        deadline: context.deadline,
282        cancellation: cancellation.clone(),
283        events: context.events.clone(),
284    };
285    active.push(
286        execute_chunk(
287            Arc::clone(inner),
288            Arc::clone(permits),
289            chunk_request,
290            chunk_context,
291            chunk,
292        )
293        .boxed(),
294    );
295}
296
297async fn execute_chunk(
298    inner: Arc<dyn ImageProvider>,
299    permits: Arc<Semaphore>,
300    request: ImageRequest,
301    mut context: ProviderContext,
302    chunk: ChunkSpec,
303) -> ChunkResult {
304    let started = Instant::now();
305    let result = acquire(&permits, &context).await;
306    let result = match result {
307        Ok(_permit) => {
308            let outer_events = context.events.take();
309            if let Some(outer_events) = outer_events {
310                let (sender, mut receiver) = mpsc::channel(4);
311                context.events = Some(sender);
312                let execution = inner.execute(request, context);
313                tokio::pin!(execution);
314                loop {
315                    tokio::select! {
316                        result = &mut execution => {
317                            while let Ok(event) = receiver.try_recv() {
318                                if let Some(event) = map_event(event, chunk.offset, chunk.count) {
319                                    let _ = outer_events.send(event).await;
320                                }
321                            }
322                            break result;
323                        },
324                        event = receiver.recv() => {
325                            let Some(event) = event else { continue; };
326                            if let Some(event) = map_event(event, chunk.offset, chunk.count) {
327                                let _ = outer_events.send(event).await;
328                            }
329                        }
330                    }
331                }
332            } else {
333                inner.execute(request, context).await
334            }
335        }
336        Err(error) => Err(error),
337    };
338    ChunkResult {
339        offset: chunk.offset,
340        count: chunk.count,
341        elapsed_ms: elapsed_ms(started),
342        result,
343    }
344}
345
346async fn acquire(
347    permits: &Arc<Semaphore>,
348    context: &ProviderContext,
349) -> Result<tokio::sync::OwnedSemaphorePermit, BridgeError> {
350    let remaining = context.deadline.saturating_duration_since(Instant::now());
351    tokio::select! {
352        permit = Arc::clone(permits).acquire_owned() => permit.map_err(|_| {
353            BridgeError::new(ErrorCode::Cancelled, "output fan-out is shutting down")
354        }),
355        () = context.cancellation.cancelled() => Err(BridgeError::new(
356            ErrorCode::Cancelled,
357            "output fan-out was cancelled while waiting for capacity",
358        )),
359        () = tokio::time::sleep(remaining) => Err(BridgeError::new(
360            ErrorCode::Timeout,
361            "output fan-out timed out while waiting for capacity",
362        ).retryable(true)),
363    }
364}
365
366fn map_event(event: ProviderEvent, offset: u8, count: u8) -> Option<ProviderEvent> {
367    match event {
368        ProviderEvent::Progress { stage } => Some(ProviderEvent::Progress { stage }),
369        ProviderEvent::PartialImage {
370            index,
371            partial_index,
372            b64_json,
373        } if index < count => Some(ProviderEvent::PartialImage {
374            index: offset + index,
375            partial_index,
376            b64_json,
377        }),
378        ProviderEvent::Started
379        | ProviderEvent::Completed { .. }
380        | ProviderEvent::PartialImage { .. } => None,
381    }
382}
383
384fn aggregate(
385    request: ImageRequest,
386    request_id: String,
387    requested_count: u8,
388    failure_policy: MultiImageFailurePolicy,
389    results: Vec<ChunkResult>,
390    started: Instant,
391) -> Result<ImageResponse, BridgeError> {
392    let mut successful = Vec::new();
393    let mut failures = Vec::new();
394    for result in results {
395        match result.result {
396            Ok(mut response) => {
397                validate_chunk_indices(&response, result.offset, result.count)?;
398                for image in &mut response.data {
399                    image.index += result.offset;
400                    image.generation_ms.get_or_insert(result.elapsed_ms);
401                }
402                for failure in &mut response.failures {
403                    failure.index += result.offset;
404                }
405                successful.push(response);
406            }
407            Err(error) => {
408                for index in result.offset..result.offset + result.count {
409                    failures.push(ImageFailure {
410                        index,
411                        error: error.clone().with_detail("output_index", index),
412                        generation_ms: result.elapsed_ms,
413                    });
414                }
415            }
416        }
417    }
418    if successful.is_empty() {
419        let mut error = failures.first().map_or_else(
420            || BridgeError::new(ErrorCode::Upstream, "all fan-out outputs failed"),
421            |failure| failure.error.clone(),
422        );
423        error = error
424            .with_detail("all_outputs_failed", true)
425            .with_detail("failed_outputs", requested_count);
426        return Err(error);
427    }
428
429    let provider = successful[0].provider.clone();
430    let model = successful[0].model.clone();
431    if successful
432        .iter()
433        .any(|response| response.provider != provider || response.model != model)
434    {
435        return Err(BridgeError::new(
436            ErrorCode::Protocol,
437            "fan-out provider identity changed within one batch",
438        ));
439    }
440    let mut data = Vec::new();
441    let mut warnings = BTreeSet::from(["emulated_multi_image_fanout".to_owned()]);
442    if request.policies.batch_execution == BatchExecution::Sequential {
443        warnings.insert("sequential_multi_image_fanout".to_owned());
444    }
445    let mut usage = Usage::default();
446    let mut has_usage = false;
447    let mut revised_prompt = None;
448    let session = merge_sessions(&successful);
449    for response in successful {
450        data.extend(response.data);
451        failures.extend(response.failures);
452        warnings.extend(response.warnings);
453        if revised_prompt.is_none() {
454            revised_prompt = response.revised_prompt;
455        }
456        if let Some(next) = response.usage {
457            merge_usage(&mut usage, next);
458            has_usage = true;
459        }
460    }
461    data.sort_by_key(|image| image.index);
462    failures.sort_by_key(|failure| failure.index);
463    if !failures.is_empty() {
464        if failure_policy != MultiImageFailurePolicy::BestEffort {
465            return Err(BridgeError::new(
466                ErrorCode::Protocol,
467                "fan-out provider returned partial failures under fail-fast policy",
468            ));
469        }
470        warnings.insert("partial_output_failure".to_owned());
471    }
472    Ok(ImageResponse {
473        id: request_id,
474        created: SystemTime::now()
475            .duration_since(UNIX_EPOCH)
476            .map_or(0, |duration| duration.as_secs()),
477        provider,
478        model,
479        requested: request.parameters.clone(),
480        effective: request.parameters,
481        normalizations: Vec::new(),
482        attempts: Vec::new(),
483        data,
484        failures,
485        revised_prompt,
486        usage: has_usage.then_some(usage),
487        session,
488        timings: Timings {
489            provider_ms: elapsed_ms(started),
490            total_ms: elapsed_ms(started),
491            ..Timings::default()
492        },
493        warnings: warnings.into_iter().collect(),
494    })
495}
496
497fn validate_chunk_indices(
498    response: &ImageResponse,
499    offset: u8,
500    count: u8,
501) -> Result<(), BridgeError> {
502    let mut indices = response
503        .data
504        .iter()
505        .map(|image| image.index)
506        .chain(response.failures.iter().map(|failure| failure.index))
507        .collect::<Vec<_>>();
508    indices.sort_unstable();
509    let expected = (0..count).collect::<Vec<_>>();
510    if indices != expected {
511        return Err(BridgeError::new(
512            ErrorCode::Protocol,
513            "fan-out subrequest returned an invalid output set",
514        )
515        .with_detail("output_offset", offset));
516    }
517    Ok(())
518}
519
520fn merge_usage(target: &mut Usage, value: Usage) {
521    target.input_tokens = sum_option(target.input_tokens, value.input_tokens);
522    target.output_tokens = sum_option(target.output_tokens, value.output_tokens);
523    target.total_tokens = sum_option(target.total_tokens, value.total_tokens);
524    for (key, value) in value.provider {
525        let entry = target.provider.entry(key).or_default();
526        *entry = entry.saturating_add(value);
527    }
528}
529
530fn sum_option(left: Option<u64>, right: Option<u64>) -> Option<u64> {
531    match (left, right) {
532        (None, None) => None,
533        (left, right) => Some(
534            left.unwrap_or_default()
535                .saturating_add(right.unwrap_or_default()),
536        ),
537    }
538}
539
540fn merge_sessions(responses: &[ImageResponse]) -> Option<SessionMetadata> {
541    let first = responses.first()?.session.clone()?;
542    responses
543        .iter()
544        .skip(1)
545        .all(|response| {
546            response.session.as_ref().is_some_and(|session| {
547                session.key == first.key && session.thread_id == first.thread_id
548            })
549        })
550        .then_some(first)
551}
552
553fn subrequest_id(request_id: &str, offset: u8) -> String {
554    let suffix = format!(":output:{offset}");
555    if request_id.len().saturating_add(suffix.len()) <= 128 {
556        return format!("{request_id}{suffix}");
557    }
558    let digest = base16ct::lower::encode_string(&Sha256::digest(request_id.as_bytes()));
559    format!("fanout:{}:{offset}", &digest[..32])
560}
561
562fn elapsed_ms(started: Instant) -> u64 {
563    u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX)
564}
565
566#[cfg(test)]
567mod tests;
imagegen_bridge_runtime/fanout/mod.rs

imagegen_bridge_runtime/fanout/
mod.rs
