mstar.streaming.topology#

Functions

client_stream_limits(topology, model_config)

Each client stream's limits: the model's declaration, with the serving config's client_streams: {<edge name>: {max_items}} applied, as validate_deployment accepted it (see mstar.validation.client_streams).

find_client_stream(conn, walks, partitions)

conn's consumer: the one node, over every walk, that takes the edge (as mstar.worker.stream_buffers.consumer_nodes resolves it), in a walk of the partition it flows into; ValueError for a stream without one.

find_client_streams(model)

Each client stream's consumer (find_client_stream).

Classes

ClientStream(node, graph_walk, partition)

Where a client stream's items go: the node that consumes the edge, a walk it runs in (for routing to its workers) and its partition.

ClientStreamLimits([max_items, item_size, unit])

Bounds on what the client feeds one client stream; None is unbounded.

Connection(from_partition, to_partition, ...)

Defines a streaming connection between two partitions, or from the client (from_partition=CLIENT) into one.

PartitionTopology(partitions[, connections])

Declares how a model's computation is split into async partitions.

ProducerWalkCtx([producer_walk, node, ...])

The producer's state for one pass that emits on a walk-driving edge.

StreamingGraphEdge(next_node, name[, ...])

A graph edge that carries streaming data between partitions.

class mstar.streaming.topology.ClientStream(node, graph_walk, partition)[source]#

Bases: NamedTuple

Where a client stream’s items go: the node that consumes the edge, a walk it runs in (for routing to its workers) and its partition.

Parameters:
  • node (str)

  • graph_walk (str)

  • partition (str)

graph_walk: str#

Alias for field number 1

node: str#

Alias for field number 0

partition: str#

Alias for field number 2

class mstar.streaming.topology.ClientStreamLimits(max_items=None, item_size=None, unit='items')[source]#

Bases: NamedTuple

Bounds on what the client feeds one client stream; None is unbounded. A bound the model declares is the most its consumer can take; a deployment may lower it (client_streams: in the serving config, see client_stream_limits()).

Parameters:
item_size: Callable[[Any], int] | None#

Alias for field number 1

max_items: int | None#

Alias for field number 0

unit: str#

Alias for field number 2

class mstar.streaming.topology.Connection(from_partition, to_partition, edge_name, chunk_policy_factory, limits=None, consumer_walk=None)[source]#

Bases: object

Defines a streaming connection between two partitions, or from the client (from_partition=CLIENT) into one. Only a client stream takes limits: a partition’s output is bounded by the model itself.

Parameters:
chunk_policy_factory: Callable[[], ChunkPolicy]#
consumer_walk: Callable[[ProducerWalkCtx], str] | None = None#
edge_name: str#
from_partition: str#
limits: ClientStreamLimits | None = None#
to_partition: str#
class mstar.streaming.topology.PartitionTopology(partitions, connections=<factory>)[source]#

Bases: object

Declares how a model’s computation is split into async partitions.

Each partition has its own set of graph walks. Connections define streaming data flow between partitions via StreamBuffers.

Parameters:
check_walk_driving_connections()[source]#

A producer-triggered partition has one authority over its walk: all its incoming connections come from one producer and share one consumer_walk, so every edge’s items agree on the walk. Its chunks must not overlap either, since a chunk runs under exactly one walk.

Return type:

None

client_connections()[source]#

The streams the client feeds.

Return type:

list[Connection]

connections: list[Connection]#
partitions: list[str]#
producer_triggered_partitions()[source]#

Partitions whose graph walk the stream drives, not the conductor.

Return type:

set[str]

walk_drivers()[source]#

Producer-triggered partition -> the connection whose consumer_walk decides its walk (any one of them; check_walk_driving_connections makes them agree).

Return type:

dict[str, Connection]

class mstar.streaming.topology.ProducerWalkCtx(producer_walk='', node='', loop_iters=<factory>, fwd_info=None)[source]#

Bases: object

The producer’s state for one pass that emits on a walk-driving edge.

The producer worker refills one instance for every call, since the hook runs per request on every emitting step: valid only during the call.

Parameters:
fwd_info: CurrentForwardPassInfo | None#
loop_iters: Mapping[str, int]#
node: str#
producer_walk: str#
class mstar.streaming.topology.StreamingGraphEdge(next_node, name, tensor_info=<factory>, persist=False, conductor_new_token=False, is_streaming=False, output_modality='', _persist_for_loop=False, _final_stream_chunk=False, _total_fanin=1, _shard_dim=None, target_partition='')[source]#

Bases: GraphEdge

A graph edge that carries streaming data between partitions.

Routed like a normal GraphEdge (producer is unaware it’s streaming). On the consumer worker, the arriving tensors are buffered in a StreamBuffer and gated by a ChunkPolicy before satisfying the consuming node’s input.

Parameters:
target_partition: str = ''#
mstar.streaming.topology.client_stream_limits(topology, model_config)[source]#

Each client stream’s limits: the model’s declaration, with the serving config’s client_streams: {<edge name>: {max_items}} applied, as validate_deployment accepted it (see mstar.validation.client_streams).

Parameters:
Return type:

dict[str, ClientStreamLimits]

mstar.streaming.topology.find_client_stream(conn, walks, partitions)[source]#

conn’s consumer: the one node, over every walk, that takes the edge (as mstar.worker.stream_buffers.consumer_nodes resolves it), in a walk of the partition it flows into; ValueError for a stream without one.

Parameters:
Return type:

ClientStream

mstar.streaming.topology.find_client_streams(model)[source]#

Each client stream’s consumer (find_client_stream). The conductor reads it after the worker graphs are built (a model may shape its walks from the deployment, as BAGEL does), and names each node in its route to the data worker, whose copy of the model is unshaped.

Parameters:

model (Model)

Return type:

dict[str, ClientStream]