skippy-server 0.77.0

Embedded Skippy staged runtime server
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

use std::time::Instant;

use openai_frontend::{ChatCompletionRequest, OpenAiError, OpenAiResult};
use skippy_runtime::{ModelWorkload, SamplingConfig};

use crate::frontend::generation::{
    GenerationCacheStats, OpenAiGenerationIds, StageOpenAiBackend, TokenControl,
    tool_calls_requested,
};
use crate::frontend::util::openai_backend_error;

use super::LocalSessionCleanupGuard;

impl StageOpenAiBackend {
    #[allow(clippy::too_many_arguments)]
    /// Encode once, decode bounded tokens, and clean up the exclusive session on every exit.
    pub(super) fn generate_encoder_decoder_tokens(
        &self,
        prompt_token_ids: &[i32],
        max_tokens: u32,
        sampling: &SamplingConfig,
        chat_request: Option<&ChatCompletionRequest>,
        cancellation: Option<&openai_frontend::CancellationToken>,
        ids: &OpenAiGenerationIds,
        mut emit_token: impl FnMut(i32) -> OpenAiResult<TokenControl>,
    ) -> OpenAiResult<GenerationCacheStats> {
        if !self.has_unsplit_full_model_topology() {
            return Err(OpenAiError::unsupported(
                "encoder-decoder models currently require an unsplit local runtime",
            ));
        }
        self.ensure_local_workload(ModelWorkload::EncoderDecoder)?;
        if chat_request.is_some_and(tool_calls_requested) {
            return Err(OpenAiError::unsupported(
                "tool calls are not supported by encoder-decoder models",
            ));
        }
        let session_id = ids.session_label.clone();
        let (result, mut cleanup) = LocalSessionCleanupGuard::run(
            || self.cleanup_local_generation_session(&session_id, ids),
            || {
                let mut runtime = self
                    .runtime
                    .lock()
                    .map_err(|_| OpenAiError::backend("runtime lock poisoned"))?;
                let prompt_started = Instant::now();
                let decoder_start = runtime
                    .encode_prompt(&session_id, prompt_token_ids)
                    .map_err(openai_backend_error)?;
                let prompt_ms = prompt_started.elapsed().as_secs_f64() * 1_000.0;
                let predicted_started = Instant::now();
                if max_tokens > 0 {
                    let mut predicted = runtime
                        .decode_frame_sampled(
                            &session_id,
                            decoder_start,
                            sampling.enabled.then_some(sampling),
                            None,
                            0,
                        )
                        .map_err(openai_backend_error)?
                        .0;
                    for generated in 0..max_tokens {
                        if cancellation
                            .is_some_and(openai_frontend::CancellationToken::is_cancelled)
                        {
                            return Err(OpenAiError::backend("request cancelled"));
                        }
                        if emit_token(predicted)? == TokenControl::Stop {
                            break;
                        }
                        if generated + 1 < max_tokens {
                            predicted = runtime
                                .decode_frame_sampled(
                                    &session_id,
                                    predicted,
                                    sampling.enabled.then_some(sampling),
                                    None,
                                    0,
                                )
                                .map_err(openai_backend_error)?
                                .0;
                        }
                    }
                }
                Ok(GenerationCacheStats {
                    prompt_ms,
                    predicted_ms: predicted_started.elapsed().as_secs_f64() * 1_000.0,
                    ..GenerationCacheStats::default()
                })
            },
        );
        cleanup.cleanup();
        result
    }
}