mstar.worker.stream_buffers

The consumer side of streaming edges on one worker.

Each request’s ``StreamBuffer``s live on its per-request info; this object creates them, turns their ready chunks into graph edges, takes back the chunks a step did not keep, and settles which streams a step ends.

Functions

consumer_nodes(connections, walks)

Map each incoming streaming edge to the node that consumes it.

Classes

StreamBufferManager(request_state, ...[, ...])

The stream buffers of this worker's requests, for the connections whose consumer runs here.

class mstar.worker.stream_buffers.StreamBufferManager(request_state, tensor_manager, connections, walks=None, *, graph_runtime=None, walk_driven_edges=None)

Bases: object

The stream buffers of this worker’s requests, for the connections whose consumer runs here.

Parameters:

• request_state (RequestStateManager)

• tensor_manager (TensorCommunicationManager)

• connections (list[Connection])

• walks (dict[str, GraphSection] | None)

• walk_driven_edges (dict[str, str] | None)

add_request(request_id, request_info)

Create the request’s buffers. A request with several partitions here arrives once per partition: the buffer the first arrival made is kept (it may already hold items).

Parameters:

request_id (int)

Return type:

None

chunks_for(request_id, node_name, inputs)

Chunk info for the streamed inputs this step consumes, or None.

Parameters:

• request_id (int)

• node_name (str)

• inputs (dict[str, list[Tensor]])

Return type:

dict[str, StreamChunkInfo] | None

consumption(request_id)

Items consumed so far per stream, as the producer is told.

Parameters:

request_id (int)

Return type:

dict[str, int]

consumption_batch(rids)

consumption for the rids that have stream buffers at all; the rest would report an empty dict.

Parameters:

rids (list[int])

Return type:

list[tuple[int, dict[str, int]]]

static edge_block(polled)

Popped chunks as the runtime’s ingest columns, in polled order.

Parameters:

polled (Iterable[tuple[int, StreamingEdge]])

Return type:

ColumnarEdgeSpecs

static edge_specs(polled)

Popped chunks as a speculation prep’s streaming edges, in polled order.

Parameters:

polled (Iterable[tuple[int, StreamingEdge]])

Return type:

list[EdgeSpec]

input_tensors(edges, node_name, rids, skip_missing=False)

edges.to_input_tensors, with each stream chunk into node_name collated (collate): a chunk’s edge carries one uuid per item, and the items are stacked here, when the consuming batch is built.

Parameters:

• edges (ColumnarEdgeSpecs)

• node_name (str)

• skip_missing (bool)

Return type:

InputTensors

mark_ingested(request_id, streaming_edge)

Remember the chunk the graph accepted, for the batch that consumes it.

A single ingested_chunk field works because streaming edges are ingested with can_buffer=False – except when speculating a loop-back, where the current value is already in flight in the forward, so the buffered chunk is the only one a later build can still be describing.

Parameters:

• request_id (int)

• streaming_edge (StreamingEdge)

Return type:

None

poll()

Every request’s ready chunks, popped. A buffer nothing has touched since a poll found it empty is skipped (StreamBuffer.maybe_ready).

Return type:

list[tuple[int, StreamingEdge]]

poll_for(request_id, node_name)

The ready chunks of the request’s streams into node_name, for a speculated step, which was built for the current walk.

Parameters:

• request_id (int)

• node_name (str)

Return type:

list[StreamingEdge]

pop_edge(sbuf, edge_name, request_id, allow_walk_change=True)

The buffer’s next chunk as a graph edge into its consumer: one a step refused earlier first, else a newly ready one; None when there is neither. A producer-triggered consumer first enters the chunk’s walk, if allow_walk_change and nothing of it is mid-pass.

Parameters:

• sbuf (StreamBuffer)

• edge_name (str)

• request_id (int)

• allow_walk_change (bool)

Return type:

StreamingEdge | None

possible()

Whether any live request owns a stream buffer.

Return type:

bool

remove_request(request_id)

Parameters:

request_id (int)

Return type:

None

return_edge(request_id, streaming_edge)

Hand a chunk back to its StreamBuffer after it was refused, by the graph’s ingest or by a speculation rolling its ingests back. Never one the graph still holds: its items would be released twice.

Parameters:

• request_id (int)

• streaming_edge (StreamingEdge)

Return type:

None

route_tensor(request_id, edge)

Route a streaming tensor to its request’s StreamBuffer for this edge. The buffer holds the item by its uuid: the read’s usage reference passes to it, and the tensor is looked up when the consuming batch is built (input_tensors).

Parameters:

• request_id (int)

• edge (GraphEdge)

Return type:

None

settle_final(node_name, final_edges)

Record the streams whose final chunk this step consumes, per rid.

Returns two sets, (node_done, partition_done): the rids whose streams into node_name have all ended once this step runs, and the rids whose streams into the partition have. The second is a subset of the first. Streams end in any order and over several steps, so one final chunk alone means only that its own stream ended. A continue_after_producer_done stream never sends a final chunk; it holds neither back.

Parameters:

• node_name (str)

• final_edges (dict[int, set[str]])

Return type:

tuple[set[int], set[int]]

mstar.worker.stream_buffers.consumer_nodes(connections, walks)

Map each incoming streaming edge to the node that consumes it.

Recurses through get_nodes() so a consumer nested inside a Loop / Sequential / Parallel is found too. Matching the walk’s top-level input_names only finds a walk that is itself a single GraphNode; a decode Loop whose inner node consumes the streamed edge would be missed, and its chunks routed to next_node="".

Parameters:

• connections (list[Connection])

• walks (dict[str, GraphSection])

Return type:

dict[str, str]