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,
    );
    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_usage(&mut self, stream_id: StreamId) -> MemoryUsage;
    fn memory_report(&mut self, stream_id: StreamId) -> MemoryReport;
    fn memory_cleanup(&mut self, stream_id: StreamId);
    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 install_memory_pools(
        &mut self,
        config: MemoryConfiguration,
        stream_id: StreamId,
    ) -> Result<(), InstallMemoryPoolsError> { ... }
    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, )

Initializes memory on the given stream with the given size.

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_usage(&mut self, stream_id: StreamId) -> MemoryUsage

Memory usage of the given stream.

fn memory_report(&mut self, stream_id: StreamId) -> MemoryReport

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.

fn memory_cleanup(&mut self, stream_id: StreamId)

Ask the server to release memory that it can release.

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: route allocations into a stable pool and snapshot it, so every buffer allocated between here and end_capture can be pinned for the graph’s lifetime. Call this before the warmup run so the capture window reuses the slices warmup left in the pool rather than allocating its own — which a hardware-graph backend cannot do at all (a device malloc inside the capture is illegal there), and which on any backend would grow the memory a graph pins beyond what it replays against.

Prefer having kernels already autotuned before this call: any transient benchmark buffers autotune allocates while the window is armed are forced into the persistent pool and pinned to the graph, so a graph captured over a cold autotune cache retains more device memory than it replays against. Warm the autotune cache first, then graph_prepare and warm up only to populate the pool.

A no-op by default (harmless on backends without graph support); a backend with graph support enables its persistent pool + capture recording.

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 install_memory_pools( &mut self, config: MemoryConfiguration, stream_id: StreamId, ) -> Result<(), InstallMemoryPoolsError>

Install a new dynamic-pool layout for the device’s main GPU memory.

The calling stream’s pools are rebuilt in place (see MemoryManagement::install_pools in cubecl-server — a rebuild only happens when nothing is live in them), and the layout becomes the one every stream created afterwards is built with. Pool layouts are a purely programmatic, runtime setting — there is no config-file pathway — so callers size them per workload (e.g. per model, just before loading it).

Errors

PoolsInUse when the calling stream kept its old layout because something was still live in its pools — e.g. a garbage-collection task that has not released its cross-stream pins yet, which can lag behind an explicit memory_cleanup. The layout still applies to streams created afterwards; retry to rebuild the calling stream too.

Unsupported from servers without configurable pools, which is the default implementation.

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