cubecl_server/stream/

capture.rs

//! A backend stream's graph capture: where it is in the capture lifecycle, and the memory its
//! recorded launches were given. Shared by every backend with graph support (see
//! [`Server::graph_prepare`](crate::server::Server::graph_prepare)); the lifecycle itself is
//! [`StreamCaptureState`], which also publishes it to the device's captures.

use crate::memory_management::{ManagedMemoryBinding, MemoryLocation, PageUpdate};
use crate::metadata_cache::CacheMode;
use crate::server::{BufferBinding, ServerError, WeakBufferBinding};
use alloc::vec::Vec;
use cubecl_common::bytes::Bytes;
use cubecl_environment::stream::StreamId;
use cubecl_runtime::server::{CaptureStatus, DeviceCaptures, StreamCaptureState};

pub use cubecl_runtime::server::CaptureEnd;

/// The graph capture of one pooled backend stream: where it sits in the
/// lifecycle, and the memory its recorded launches were given.
///
/// The two travel together because neither is meaningful alone. A launch is
/// remembered only while the window is recording, and what it remembers is
/// only ever read once the window closes — so pairing them
/// makes "buffers accumulate inside a window and nowhere else" a property of
/// the type rather than a rule each backend has to keep.
///
/// Why the memory is worth keeping: a replay runs every recorded launch or
/// none, so a replay that fails to enqueue leaves all of them exactly as they
/// were, and so does a capture that never seals. Either way a later read of one
/// has to fail rather than copy out bytes nothing wrote — which needs the list
/// the launches themselves no longer hold, as bindings, so the failure can be
/// tainted onto the allocations they resolve to.
#[derive(Debug)]
pub struct StreamCapture {
    /// Where the stream is in the capture lifecycle, published to the device's captures.
    state: StreamCaptureState,
    /// What the recorded launches write, kept without holding their memory:
    /// a held binding keeps its slice from being reused, and the recording
    /// has to reuse memory exactly as the warmup run did.
    recorded: Vec<WeakBufferBinding>,
    /// The memory the recorded launches were given, as the pages the graph
    /// will [guard](crate::memory_management::PageGuard) once it seals.
    touched: Vec<TouchedPage>,
    /// The pinned staging the warmup run reserved, held until the window
    /// opens. A recorded copy keeps its staging for the graph's life, so the
    /// recorded run needs one slice per copy where the warmup run could reuse
    /// one; holding them makes the pool that big before the window opens,
    /// where nothing may be allocated.
    primed: Vec<ManagedMemoryBinding>,
    /// The host memory the recorded copies read from, held while the window
    /// is open and handed to the graph it seals into. A recorded memcpy node
    /// keeps the raw host pointer, so the bytes must live exactly as long as
    /// the graph — whatever kind of allocation they are, pinned-pool slice or
    /// plain heap. A window that never seals drops them here: its copies
    /// never ran and now never will.
    retained_host: Vec<Bytes>,
    /// The failure that doomed the window, if one did — work inside it failed
    /// or was skipped, so the recording is missing an operation. Sealing it
    /// would hand back a graph silently missing that work, and the replay
    /// contract has the caller write fresh inputs before each replay,
    /// clearing the very taint that would explain the hole — so the window is
    /// doomed instead, and `end_capture` refuses to seal it.
    failed: Option<ServerError>,
}

impl StreamCapture {
    /// The capture of one more stream of the device whose captures `device` records.
    pub fn new(device: &DeviceCaptures) -> Self {
        Self {
            state: device.stream(),
            recorded: Vec::new(),
            touched: Vec::new(),
            primed: Vec::new(),
            retained_host: Vec::new(),
            failed: None,
        }
    }

    /// The captures of the whole device this stream belongs to, this one included.
    pub fn device(&self) -> &CaptureStatus {
        self.state.device()
    }

    /// Remember the memory a launch was given, when the stream is recording.
    ///
    /// A no-op outside a window, where a launch that fails taints its own
    /// buffers on the spot and there is no graph to answer for them later.
    pub fn record(&mut self, buffers: impl IntoIterator<Item = BufferBinding>) {
        if self.state.is_recording() {
            self.recorded
                .extend(buffers.into_iter().map(|binding| binding.downgrade()));
        }
    }

    /// Remember that a launch was given the memory at `location`, owned by
    /// `stream`, when the stream is recording: the graph replays against that
    /// page, so it has to stay where it is for the graph's life.
    pub fn touch(&mut self, stream: StreamId, location: MemoryLocation) {
        if self.state.is_recording() {
            self.touched.push(TouchedPage { stream, location });
        }
    }

    /// The pages the recording touched, each named once, taken as the window
    /// closes. A launch touches its pages again on every recorded launch, so
    /// the list is deduplicated here, once, rather than on every touch.
    pub fn take_touched(&mut self) -> Vec<TouchedPage> {
        let mut touched = core::mem::take(&mut self.touched);
        touched.sort_unstable();
        touched.dedup();
        touched
    }

    /// What this stream's reservations may do to the pages held.
    ///
    /// A recording stream allocates nothing: an allocation would become part
    /// of the recording. While any stream of the device records, no page is
    /// released or renumbered, so the pages a recording touched are still the
    /// ones it guards when it seals.
    pub fn page_update(&self) -> PageUpdate {
        match (self.is_recording(), self.device().any_recording()) {
            (true, _) => PageUpdate::Forbidden,
            (false, true) => PageUpdate::AddOnly,
            (false, false) => PageUpdate::Allow,
        }
    }

    /// Hold `staging` until the window opens, when the stream is preparing a
    /// capture: a recorded copy keeps its staging for the graph's life, so the
    /// recorded run needs a slice per copy where the warmup run could reuse
    /// one, and holding them grows the pool before nothing may be allocated.
    pub fn prime(&mut self, staging: &ManagedMemoryBinding) {
        if self.state.is_preparing() {
            self.primed.push(staging.clone());
        }
    }

    /// The memory of the capture that just closed, each claim named once —
    /// the same buffer comes back once per recorded launch that was given it.
    ///
    /// Deduplicated by [`claim_key`](BufferBinding::claim_key), because the
    /// taint bookkeeping is range-exact and this list is what gets claimed
    /// and released: two tensors carved from one batched allocation are two
    /// claims, and collapsing them to their shared memory id would leave
    /// every sibling but one unclaimed on a refusal and unreleased on a
    /// replay.
    ///
    /// A buffer whose memory is gone by then is left out: nothing will read
    /// it again.
    pub fn take_recorded(&mut self) -> Vec<BufferBinding> {
        let mut recorded: Vec<BufferBinding> = core::mem::take(&mut self.recorded)
            .iter()
            .filter_map(WeakBufferBinding::upgrade)
            .collect();
        recorded.sort_unstable_by_key(|binding| binding.claim_key());
        recorded.dedup_by_key(|binding| binding.claim_key());
        recorded
    }

    /// Doom the recording: work inside the window failed or was skipped —
    /// see [`Self::take_failure`]. The first failure wins, and a stream that
    /// is not recording has no window to doom.
    pub fn fail(&mut self, error: ServerError) {
        if self.state.is_recording() && self.failed.is_none() {
            self.failed = Some(error);
        }
    }

    /// Keep `bytes` alive for the graph this recording seals into: a
    /// recorded copy holds their raw pointer and re-reads them on every
    /// replay, so they must not return to any pool or allocator while the
    /// graph lives. Handed to the graph by [`take_retained_host`](Self::take_retained_host).
    ///
    /// Only meaningful while recording — outside a window the bytes belong
    /// in the drop queue, whose fence knows when the device is done with
    /// them.
    pub fn retain_host(&mut self, bytes: Bytes) {
        debug_assert!(
            self.state.is_recording(),
            "host bytes are the window's to retain only while it records"
        );
        self.retained_host.push(bytes);
    }

    /// The host memory the window's recorded copies read from, taken as it
    /// closes — onto the graph when it seals, or to be dropped when it does
    /// not, since a recording that never becomes a graph never runs them.
    pub fn take_retained_host(&mut self) -> Vec<Bytes> {
        core::mem::take(&mut self.retained_host)
    }

    /// The failure that doomed this window, taken as it closes. `Some` means
    /// the recording is missing at least one operation and must not seal.
    pub fn take_failure(&mut self) -> Option<ServerError> {
        self.failed.take()
    }

    /// Whether launches on the stream are being recorded into a graph right
    /// now — the window during which a host sync would abort the capture, and
    /// during which a neighbour sharing the pooled stream is refused.
    pub fn is_recording(&self) -> bool {
        self.state.is_recording()
    }

    /// Whether a capture is prepared or recording — the span over which the
    /// pooled stream is committed to one logical stream's window.
    pub fn is_active(&self) -> bool {
        self.state.is_active()
    }

    /// The logical stream that opened the window, while one is open. `None`
    /// outside a capture, which is what distinguishes a neighbour's operation
    /// from the owner's.
    pub fn owner(&self) -> Option<StreamId> {
        self.state.owner()
    }

    /// The [`CacheMode`] the metadata info cache should run in at this lifecycle
    /// position: a window pins what it builds, so a replay finds the same entries it recorded
    /// against. Both while a graph is being *prepared* (warmup, which primes
    /// the cache) and while it is being *recorded* the cache runs in
    /// [`CacheMode::Capture`] — caching every buffer and invalidating none — so
    /// the capture window finds every info buffer warm and drops none out from
    /// under a recorded launch. Normal operation uses [`CacheMode::Normal`].
    pub fn cache_mode(&self) -> CacheMode {
        match self.state.is_active() {
            true => CacheMode::Capture,
            false => CacheMode::Normal,
        }
    }

    /// Start the warmup run; [`begin`](Self::begin) may open the window
    /// afterwards. A capture starts from an empty recording, so a window that
    /// was abandoned mid-flight cannot leak its buffers into the next one.
    ///
    /// # Errors
    ///
    /// Fails when a capture is already prepared or recording, leaving both the
    /// state and the recording untouched.
    pub fn prepare(&mut self, owner: StreamId) -> Result<(), ServerError> {
        self.state.prepare(owner)?;
        self.clear();
        Ok(())
    }

    /// Open the window: launches from here until [`end`](Self::end) are
    /// recorded rather than executed.
    ///
    /// # Errors
    ///
    /// Fails when no capture is prepared, or one is already recording.
    ///
    /// The staging the warmup run primed goes back to the pool, for the
    /// recorded run to reuse.
    pub fn begin(&mut self) -> Result<(), ServerError> {
        self.state.begin()?;
        self.primed.clear();
        Ok(())
    }

    /// Close the window, saying whether `caller` owned it — see
    /// [`CaptureEnd`]. The recording survives the transition for
    /// [`take_recorded`](Self::take_recorded) to collect.
    ///
    /// # Errors
    ///
    /// Fails when no capture is recording, leaving the state untouched.
    pub fn end(&mut self, caller: StreamId) -> Result<CaptureEnd, ServerError> {
        self.state.end(caller)
    }

    /// Give up a prepared capture that never opened, restoring the stream to
    /// no-capture. Whatever it had recorded or retained goes with it.
    pub fn abort(&mut self) {
        self.state.abort();
        self.clear();
    }

    fn clear(&mut self) {
        self.recorded.clear();
        self.touched.clear();
        self.primed.clear();
        self.retained_host.clear();
        self.failed = None;
    }
}

/// A page a recorded launch was given memory on.
///
/// Ordered by stream, then the whole location: a pool whose slices are their
/// own allocations tells its pages apart by slice.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct TouchedPage {
    /// The stream whose memory holds the page.
    pub stream: StreamId,
    /// Where the memory sat when the launch was recorded. Nothing moves or
    /// renumbers pages while a stream records, so it still names the page
    /// when the window closes.
    pub location: MemoryLocation,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::memory_management::ManagedMemoryId;
    use crate::server::Handle;

    /// A service for handles that never reach a device.
    fn service() -> cubecl_common::device::ServiceId {
        cubecl_common::device::ServiceId::of::<()>(cubecl_common::device::DeviceId::new(0, 0))
    }

    const OWNER: StreamId = StreamId { value: 7 };

    /// A distinct buffer per call, on the owner's stream.
    fn buffer() -> BufferBinding {
        let binding = Handle::new(service(), OWNER, 8).binding();
        // The reference the pool's slice keeps on a real allocation, which
        // no test here has: a recording only answers for allocations
        // something other than their pool still holds.
        core::mem::forget(binding.memory.clone());
        binding
    }

    fn ids(bindings: &[BufferBinding]) -> Vec<ManagedMemoryId> {
        bindings.iter().map(|binding| binding.memory.id()).collect()
    }

    /// A graph answers for every buffer its launches were given, and names each
    /// one once however many launches shared it.
    ///
    /// A failed replay runs none of the recorded launches, so the list is what
    /// a later read of any of those buffers fails on. Repeats would make that
    /// list grow with the length of the capture rather than with the working
    /// set — a chat step records hundreds of launches over the same handful of
    /// weights.
    #[test]
    fn a_capture_names_each_buffer_its_launches_were_given_once() {
        let (a, b, c) = (buffer(), buffer(), buffer());

        let mut capture = StreamCapture::new(&DeviceCaptures::default());
        capture.prepare(OWNER).unwrap();
        capture.begin().unwrap();
        capture.record([a.clone(), b.clone()]);
        capture.record([b.clone(), c.clone()]);
        capture.record([a.clone()]);

        assert_eq!(
            capture.end(OWNER).unwrap(),
            CaptureEnd::Owned { owner: OWNER }
        );
        assert_eq!(ids(&capture.take_recorded()), ids(&[a, b, c]));
        assert!(
            capture.take_recorded().is_empty(),
            "the recording moves onto the graph, it is not left on the stream"
        );
    }

    /// Outside a window a launch answers for its own buffers as it fails, so
    /// nothing is remembered for a graph that will never exist.
    ///
    /// A stream that accumulated ids while not recording would grow one entry
    /// per launch for the life of the process, and hand the next capture a list
    /// of buffers it never touched.
    #[test]
    fn a_launch_outside_a_window_is_not_recorded() {
        let (before, warmup, recorded) = (buffer(), buffer(), buffer());
        let mut capture = StreamCapture::new(&DeviceCaptures::default());

        capture.record([before]);
        capture.prepare(OWNER).unwrap();
        // Prepared is the warmup run: it executes rather than records.
        capture.record([warmup]);

        capture.begin().unwrap();
        capture.record([recorded.clone()]);
        capture.end(OWNER).unwrap();

        assert_eq!(ids(&capture.take_recorded()), ids(&[recorded]));
    }

    /// A batched allocation is carved into sibling tensors sharing one memory
    /// id and nothing else, and the taint bookkeeping is range-exact — so the
    /// write set keeps every range. Collapsed to the id, a refusal would
    /// claim (and a replay release) only whichever sibling survived the
    /// dedup, leaving the others readable with stale bytes or unreadable with
    /// clean ones.
    #[test]
    fn a_capture_names_every_range_of_a_batched_allocation() {
        let handle = Handle::new(service(), OWNER, 8);
        let mut front = handle.clone().binding();
        front.offset_end = Some(4);
        let mut back = handle.clone().binding();
        back.offset_start = Some(4);
        assert_eq!(
            front.memory.id(),
            back.memory.id(),
            "one allocation carved in two is the case under test"
        );

        let mut capture = StreamCapture::new(&DeviceCaptures::default());
        capture.prepare(OWNER).unwrap();
        capture.begin().unwrap();
        capture.record([front.clone(), back.clone()]);
        // The same range again, from a second launch given the same tensor:
        // still one claim.
        capture.record([front.clone()]);
        capture.end(OWNER).unwrap();

        let recorded = capture.take_recorded();
        let keys: Vec<_> = recorded.iter().map(|binding| binding.claim_key()).collect();
        assert_eq!(
            keys,
            alloc::vec![front.claim_key(), back.claim_key()],
            "both siblings survive, each named once"
        );
    }

    /// The host bytes a recorded copy reads from live with the window: taken
    /// once as it closes — onto the graph, when one seals — and dropped with
    /// an aborted window, whose copies never ran and now never will.
    #[test]
    fn a_window_owns_the_host_bytes_its_copies_read() {
        let mut capture = StreamCapture::new(&DeviceCaptures::default());
        capture.prepare(OWNER).unwrap();
        capture.begin().unwrap();
        capture.retain_host(Bytes::from_bytes_vec(alloc::vec![7u8; 4]));
        capture.end(OWNER).unwrap();

        assert_eq!(capture.take_retained_host().len(), 1);
        assert!(
            capture.take_retained_host().is_empty(),
            "taken means moved onto the graph, not copied"
        );

        capture.prepare(OWNER).unwrap();
        capture.begin().unwrap();
        capture.retain_host(Bytes::from_bytes_vec(alloc::vec![7u8; 4]));
        capture.abort();
        capture.prepare(OWNER).unwrap();
        assert!(
            capture.take_retained_host().is_empty(),
            "an aborted window keeps nothing alive"
        );
    }

    /// A window that never sealed leaves nothing behind for the next one.
    ///
    /// The stream is re-capturable after an abort or an abandoned end, and a
    /// second capture that inherited the first one's buffers would pin memory
    /// its own launches never saw — failing reads that had nothing to do with
    /// it, on a graph that outlives the mistake.
    #[test]
    fn a_new_capture_starts_from_an_empty_recording() {
        let neighbour = StreamId { value: 8 };

        let (aborted, abandoned) = (buffer(), buffer());
        let mut capture = StreamCapture::new(&DeviceCaptures::default());
        capture.prepare(OWNER).unwrap();
        capture.begin().unwrap();
        capture.record([aborted]);
        capture.abort();

        capture.prepare(OWNER).unwrap();
        capture.begin().unwrap();
        capture.record([abandoned.clone()]);
        assert!(capture.end(neighbour).unwrap().is_abandoned());
        assert_eq!(ids(&capture.take_recorded()), ids(&[abandoned]));

        capture.prepare(neighbour).unwrap();
        capture.begin().unwrap();
        assert!(
            capture.take_recorded().is_empty(),
            "the abandoned window's buffers are not this capture's to answer for"
        );
    }

    /// The cache runs in capture mode for the *whole* prepare → record window,
    /// not just while recording: warmup is what makes the recorded launches hit
    /// warm info buffers, and an entry evicted between the two would be one a
    /// recorded launch dropped out from under itself.
    #[test]
    fn the_cache_captures_across_the_whole_window() {
        let mut capture = StreamCapture::new(&DeviceCaptures::default());
        assert_eq!(capture.cache_mode(), CacheMode::Normal);
        capture.prepare(OWNER).unwrap();
        assert_eq!(capture.cache_mode(), CacheMode::Capture);
        capture.begin().unwrap();
        assert_eq!(capture.cache_mode(), CacheMode::Capture);
        capture.end(OWNER).unwrap();
        assert_eq!(capture.cache_mode(), CacheMode::Normal);
    }

    /// A capture is visible to the stream that prepared it, from its warmup run
    /// until it ends however it ends, and to no other stream: code that has to
    /// decide the same way in the warmup and the recording reads it, and would
    /// keep paying for it if a finished capture lingered.
    #[test]
    fn a_capture_is_seen_by_its_stream_until_it_ends() {
        const OTHER: StreamId = StreamId { value: 8 };
        let device = DeviceCaptures::default();
        let status = device.status();
        let mut capture = StreamCapture::new(&device);

        assert!(!status.is_capturing(OWNER), "no capture before prepare");
        capture.prepare(OWNER).unwrap();
        assert!(status.is_capturing(OWNER), "the warmup run is part of it");
        assert!(
            !status.is_capturing(OTHER),
            "another stream isn't capturing"
        );
        capture.begin().unwrap();
        assert!(status.is_capturing(OWNER), "the recorded run is part of it");
        capture.end(OWNER).unwrap();
        assert!(!status.is_capturing(OWNER), "a sealed capture is over");

        capture.prepare(OWNER).unwrap();
        capture.begin().unwrap();
        capture.end(OTHER).unwrap();
        assert!(
            !status.is_capturing(OWNER),
            "an abandoned capture is over for its owner too"
        );

        capture.prepare(OWNER).unwrap();
        capture.abort();
        assert!(!status.is_capturing(OWNER), "an aborted capture is over");

        let mut other = StreamCapture::new(&device);
        other.prepare(OTHER).unwrap();
        capture.prepare(OWNER).unwrap();
        drop(capture);
        assert!(
            !status.is_capturing(OWNER),
            "a dropped stream's capture is over"
        );
        assert!(
            status.is_capturing(OTHER),
            "another stream's capture is untouched"
        );
    }

    /// Every stream of a device sees a window one of them opens, and the
    /// window leaves the count however it closes: sealed, aborted, or with
    /// its stream dropped.
    #[test]
    fn a_recording_is_seen_device_wide() {
        let device = DeviceCaptures::default();
        let mut recording = StreamCapture::new(&device);
        let neighbour = StreamCapture::new(&device);
        let status = device.status();
        let recording_now =
            |capture: &StreamCapture| (capture.page_update(), status.any_recording());

        recording.prepare(OWNER).unwrap();
        assert!(!status.any_recording(), "a warmup records nothing");
        recording.begin().unwrap();
        assert_eq!(recording_now(&recording), (PageUpdate::Forbidden, true));
        assert_eq!(recording_now(&neighbour), (PageUpdate::AddOnly, true));
        recording.end(OWNER).unwrap();
        assert_eq!(recording_now(&neighbour), (PageUpdate::Allow, false));

        recording.prepare(OWNER).unwrap();
        recording.begin().unwrap();
        recording.abort();
        assert!(
            !status.any_recording(),
            "an aborted window leaves the count"
        );

        recording.prepare(OWNER).unwrap();
        recording.abort();
        assert!(
            !status.any_recording(),
            "aborting a warmup leaves the count alone"
        );

        recording.prepare(OWNER).unwrap();
        recording.begin().unwrap();
        drop(recording);
        assert!(
            !status.any_recording(),
            "a stream dropped while recording leaves the count"
        );
    }
}