ts_pd/
batch.rs

1// Copyright (c) 2026 LightSeek Foundation
2//
3// Permission is hereby granted, free of charge, to any person obtaining a copy
4// of this software and associated documentation files (the "Software"), to deal
5// in the Software without restriction, including without limitation the rights
6// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
7// copies of the Software, and to permit persons to whom the Software is
8// furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
14// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
18// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
19// SOFTWARE.
20
21//! A plan's batch as a transfer reads it: the rows of a remote stream or a
22//! forward, with their block tables.
23
24use std::sync::{Arc, Mutex};
25
26use crate::layout::manifest::BlockTables;
27
28/// A batch's rows (the scheduler's `TransferBatchView`); extend rows first.
29#[derive(Debug, Clone, Default, PartialEq, Eq)]
30pub struct BatchRows {
31    /// Request ids.
32    pub request_ids: Vec<String>,
33    /// Each row's request slot.
34    pub request_pool_indices: Vec<i32>,
35    /// Tokens each row feeds this forward.
36    pub input_lengths: Vec<i32>,
37    /// Each row's prompt length.
38    pub prefill_lengths: Vec<i32>,
39    /// Extend rows: tokens already in the cache ahead of the chunk.
40    pub extend_prefix_lens: Vec<i32>,
41    /// Extend rows, a Decode admission: every state group's prompt-end
42    /// checkpoint was found locally.
43    pub state_prefix_cached: Vec<bool>,
44    /// Decode rows: on a Prefill node's remote decode, each request's first
45    /// output token.
46    pub decode_input_ids: Vec<i32>,
47    /// Decode rows: the drafter's candidates, empty for none.
48    pub spec_candidate_ids: Vec<Vec<i32>>,
49    /// Extend rows.
50    pub num_extends: usize,
51    /// Every cache group's table.
52    pub tables: BlockTables,
53}
54
55/// A plan's remote stream, staged by the forward thread when it hands the
56/// plan to the data plane, until the data plane hands it over once the
57/// stream's device prerequisites are met.
58#[derive(Debug, Clone, Default)]
59pub(crate) struct Staged(Arc<Mutex<Option<BatchRows>>>);
60
61impl Staged {
62    pub(crate) fn stage(&self, rows: BatchRows) -> Result<(), String> {
63        let mut staged = self.0.lock().expect("the staging lock is poisoned");
64        if staged.is_some() {
65            return Err(
66                "a plan's remote stream was staged before the previous was handed over".to_owned(),
67            );
68        }
69        *staged = Some(rows);
70        Ok(())
71    }
72
73    pub(crate) fn take(&self) -> Result<BatchRows, String> {
74        self.0
75            .lock()
76            .expect("the staging lock is poisoned")
77            .take()
78            .ok_or_else(|| {
79                "the data plane handed over a remote stream the plan has none of".to_owned()
80            })
81    }
82
83    pub(crate) fn is_pending(&self) -> bool {
84        self.0
85            .lock()
86            .expect("the staging lock is poisoned")
87            .is_some()
88    }
89}
ts_pd/batch.rs

ts_pd/
batch.rs
