Trait Server 

pub trait Server:
    Any
    + Send
    + Debug
    + ServerCommunication
    + DeviceService
    + 'static {
    // Required methods
    fn initialize_memory(
        &mut self,
        memory: ManagedMemoryHandle,
        size: u64,
        stream_id: StreamId,
    ) -> Result<(), ServerError>;
    fn logger(&self) -> Arc<ServerLogger> ;
    fn utilities(&self) -> Arc<ServerUtilities> ;
    fn read(
        &mut self,
        descriptors: Vec<CopyDescriptor>,
        stream_id: StreamId,
    ) -> DynFut<Result<Vec<Bytes>, ServerError>>;
    fn write(
        &mut self,
        descriptors: Vec<(CopyDescriptor, Bytes)>,
        stream_id: StreamId,
    );
    fn sync(
        &mut self,
        handles: Vec<BufferBinding>,
        stream_id: StreamId,
    ) -> DynFut<Result<(), ServerError>>;
    fn check(
        &mut self,
        handles: Vec<BufferBinding>,
        stream_id: StreamId,
    ) -> Result<(), ServerError>;
    unsafe fn launch(
        &mut self,
        kernel: Box<dyn CubeKernel>,
        count: CubeCount,
        bindings: KernelArguments,
        stream_id: StreamId,
        launch_mode: LaunchMode,
    );
    fn flush(&mut self, stream_id: StreamId) -> Result<(), ServerError>;
    fn memory_report(
        &mut self,
        stream_id: StreamId,
    ) -> Option<StreamMemoryReport>;
    fn memory_cleanup(&mut self, stream_id: StreamId) -> Result<(), ServerError>;
    fn start_profile(
        &mut self,
        stream_id: StreamId,
    ) -> Result<ProfilingToken, ServerError>;
    fn end_profile(
        &mut self,
        stream_id: StreamId,
        token: ProfilingToken,
    ) -> Result<ProfileDuration, ProfileError>;
    fn allocation_mode(
        &mut self,
        mode: MemoryAllocationMode,
        stream_id: StreamId,
    );

    // Provided methods
    fn staging(
        &mut self,
        _sizes: &[usize],
        _stream_id: StreamId,
    ) -> Result<Vec<Bytes>, ServerError> { ... }
    fn graph_prepare(&mut self, stream_id: StreamId) -> Result<(), ServerError> { ... }
    fn begin_capture(&mut self, stream_id: StreamId) -> Result<(), ServerError> { ... }
    fn end_capture(
        &mut self,
        stream_id: StreamId,
    ) -> Result<GraphId, ServerError> { ... }
    fn replay(
        &mut self,
        graph: GraphId,
        stream_id: StreamId,
    ) -> Result<(), ServerError> { ... }
    fn graph_destroy(&mut self, graph: GraphId, stream_id: StreamId) { ... }
    fn stream_ids(&self) -> Vec<StreamId> { ... }
    fn abandon_profile(&mut self, stream_id: StreamId, token: ProfilingToken) { ... }
}

The compute server is responsible for handling resources and computations over resources.

Everything in the server is mutable, therefore it should be solely accessed through the Client for thread safety.

Required Methods

fn initialize_memory( &mut self, memory: ManagedMemoryHandle, size: u64, stream_id: StreamId, ) -> Result<(), ServerError>

Initializes memory on the given stream with the given size.

Errors

Returns a ServerError when the memory could not be initialized, e.g. because the stream cannot be created on a poisoned device. The memory then carries the error: the next sync point reading that memory reports it.

fn logger(&self) -> Arc<ServerLogger>

Retrieve the server logger.

fn utilities(&self) -> Arc<ServerUtilities>

Retrieve the server utilities.

fn read( &mut self, descriptors: Vec<CopyDescriptor>, stream_id: StreamId, ) -> DynFut<Result<Vec<Bytes>, ServerError>>

Given bindings, returns the owned resources as bytes.

Errors

ServerError::Several when the work that was supposed to write one of these buffers failed, whichever stream it ran on — copying bytes out would hand back whatever was in memory before. Every implementation asks FailureStore::ensure_written (in cubecl-server) before it copies anything.

fn write( &mut self, descriptors: Vec<(CopyDescriptor, Bytes)>, stream_id: StreamId, )

Writes the specified bytes into the buffers given

fn sync( &mut self, handles: Vec<BufferBinding>, stream_id: StreamId, ) -> DynFut<Result<(), ServerError>>

Wait for the completion of every task in the server, then answer for handles: the barrier first, so device faults count, and then the claim check a read would have made — a read without the read.

An empty handles is the plain barrier plus the device fault, which is the only failure left that no buffer can report.

fn check( &mut self, handles: Vec<BufferBinding>, stream_id: StreamId, ) -> Result<(), ServerError>

Whether the bytes the handles name can be trusted, right now and with no barrier: the claim check a read makes, without the read. Instant — enqueue-time failures only. A device fault needs sync, which drains first.

unsafe fn launch( &mut self, kernel: Box<dyn CubeKernel>, count: CubeCount, bindings: KernelArguments, stream_id: StreamId, launch_mode: LaunchMode, )

Executes the kernel over the given memory handles.

Kernels have mutable access to every resource they are given and are responsible of determining which should be read or written.

launch_mode says whether the kernel actually runs. On LaunchMode::Skip the server must still do everything a first launch does short of dispatching — expand, compile, validate, fill its caches — and then drop the launch; skipping the compilation instead would defeat the whole point of a dry run.

Safety

When executing with mode [ExecutionMode::Unchecked], out-of-bound reads and writes can happen.

fn flush(&mut self, stream_id: StreamId) -> Result<(), ServerError>

Flush all outstanding tasks in the server.

Errors

The device fault, when the context itself is broken — a launch failure is not the flush’s to report: it lives on the buffers the launch left unwritten, and surfaces on any read, sync or check of them.

fn memory_report(&mut self, stream_id: StreamId) -> Option<StreamMemoryReport>

Structured per-pool report of the given stream’s main GPU memory: each pool’s shape, usage, and high-water marks, in allocation-routing order. The read side of a measured memory plan — see MemoryManagement::memory_report in cubecl-server.

None when the stream was never created.

fn memory_cleanup(&mut self, stream_id: StreamId) -> Result<(), ServerError>

Ask the server to release memory that it can release.

Errors

Refused while a stream records a graph: releasing memory waits on the device, and a wait on a stream that records aborts its capture.

fn start_profile( &mut self, stream_id: StreamId, ) -> Result<ProfilingToken, ServerError>

Enable collecting timestamps.

fn end_profile( &mut self, stream_id: StreamId, token: ProfilingToken, ) -> Result<ProfileDuration, ProfileError>

Disable collecting timestamps.

fn allocation_mode(&mut self, mode: MemoryAllocationMode, stream_id: StreamId)

Update the memory mode of allocation in the server.

Provided Methods

fn staging( &mut self, _sizes: &[usize], _stream_id: StreamId, ) -> Result<Vec<Bytes>, ServerError>

Reserves N Bytes of the provided sizes to be used as staging to load data.

fn graph_prepare(&mut self, stream_id: StreamId) -> Result<(), ServerError>

Prepare stream_id for an upcoming graph capture. Call this before the warmup run: the capture window allocates nothing, so it reuses what the warmup run left in the pools, and every page the recording touches is guarded for the graph’s lifetime.

Prefer having kernels already autotuned before this call, so the warmup run allocates what the recorded run asks for and nothing else.

A no-op by default (harmless on backends without graph support).

fn begin_capture(&mut self, stream_id: StreamId) -> Result<(), ServerError>

Begin recording the launches issued on stream_id into a graph instead of executing them, so the sequence can later be replayed without paying the launch path again. Call graph_prepare and warm up first.

Between this call and end_capture the stream must not synchronize — a read, a sync or a profile either aborts the capture or is refused — and should not allocate fresh device memory, which graph_prepare plus a warmup run is what avoids. Whether an operation the window cannot record fails the call or fails end_capture, and whether a mid-window allocation is fatal, is the backend’s to say; see StreamCapture in cubecl-server.

The default is unsupported. Two shapes of backend override it: a hardware graph (CUDA, HIP), where the driver records a replayable graph object, and a software graph (wgpu), where the runtime records fully-resolved dispatches and re-encodes them on replay.

fn end_capture(&mut self, stream_id: StreamId) -> Result<GraphId, ServerError>

Stop recording (see begin_capture), store the captured graph in the backend’s registry, and return its GraphId, ready to replay.

fn replay( &mut self, graph: GraphId, stream_id: StreamId, ) -> Result<(), ServerError>

Replay the graph identified by graph on stream_id, re-running the whole recorded launch sequence against its original buffers. A hardware graph replays as a single dispatch; a software graph re-encodes the recorded dispatches, which is still far cheaper than the launch path but stays O(n) in recorded launches.

The call enqueues the dispatch and returns without waiting for the device; what it reports is the enqueue — an unknown or destroyed graph, a refusal — since a caller replaying a graph is standing right there. A failure also leaves the graph’s write set carrying it, so a read of those buffers fails until a replay lands. Unsupported by default: a GraphId can only come from end_capture.

fn graph_destroy(&mut self, graph: GraphId, stream_id: StreamId)

Release the graph identified by graph, destroying whatever it recorded and unpinning the buffers it retained. Replay returns at enqueue time, so the backend must guarantee no in-flight replay can still read those buffers once they return to the pool — by syncing stream_id where nothing weaker will do (CUDA, HIP), or by relying on the queue ordering that already places a submitted replay ahead of any later write (wgpu). A no-op by default and for an unknown id.

fn stream_ids(&self) -> Vec<StreamId>

Stream ids the client should iterate to aggregate across the device.

Default is just the calling stream, which is correct for non-multi-stream backends; multi-stream backends override to return one id per initialized stream pool slot.

fn abandon_profile(&mut self, stream_id: StreamId, token: ProfilingToken)

Drop the window token opened without measuring it, for a caller that will never close it with end_profile.

An open window is not free: depending on the backend it retains command buffers, keeps timestamp writes on, or holds a device event.

Every backend overrides this, and should: the default closes the window and throws the measurement away, which is the most expensive way to be rid of it — end_profile is where the syncing and flushing live, and this is the one call that needs none of it.

Dyn Compatibility

This trait is dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors