mstar.worker.node_manager_utils

mstar.worker.node_manager_utils#

Classes

PerPartitionInfo(current_fwd_info, ...)

PerRequestInfo([stream_buffers, ...])

What the worker still owns for a request.

RequestStateManager([per_request_info, ...])

Per-request forward-pass info and stream buffers.
class mstar.worker.node_manager_utils.PerPartitionInfo(current_fwd_info: mstar.conductor.request_info.CurrentForwardPassInfo, pending_resource_publish_info: dict[str, mstar.engine.resources.base.PublishedInfo]=<factory>)[source]#

Bases: object

Parameters:
current_fwd_info: CurrentForwardPassInfo#
pending_resource_publish_info: dict[str, PublishedInfo]#
class mstar.worker.node_manager_utils.PerRequestInfo(stream_buffers=<factory>, stream_buffers_by_consumer=<factory>, ended_streams=<factory>, per_partition_info=<factory>)[source]#

Bases: object

What the worker still owns for a request.

Everything graph-shaped – node/loop routing, the sharding config, the worker graph ids, the live queues – belongs to the graph runtime. What is left is the forward-pass info (a wire object) and the stream buffers.

Parameters:
ended_streams: set[str]#
per_partition_info: dict[str, PerPartitionInfo]#
stream_buffers: dict[str, StreamBuffer]#
stream_buffers_by_consumer: dict[str, dict[str, StreamBuffer]]#
class mstar.worker.node_manager_utils.RequestStateManager(per_request_info=<factory>, node_to_partition=<factory>, _fwd_by_partition=<factory>)[source]#

Bases: object

Per-request forward-pass info and stream buffers.

What remains here is request state the graph runtime deliberately does not take: CurrentForwardPassInfo is a wire object it treats as opaque, and a StreamBuffer holds tensors.

Parameters:
add_request(request_id, current_fwd_info)[source]#

Register a partition for a request, adding the request if new.

The conductor sends one NewRequest per partition, so this is called once per partition for a given request_id.

Parameters:
buffer_publish_info(request_id, partition_name, published)[source]#
Parameters:
Return type:

None

flush_publish_info(request_id, partition_name)[source]#
Parameters:

  • request_id (int)

  • partition_name (str)

Return type:

dict[str, PublishedInfo]

fwd_infos(request_ids, partition_name)[source]#

get_fwd_info for each rid, in order; KeyError if one has none.

Parameters:
Return type:

list[CurrentForwardPassInfo]

get_fwd_info(request_id, partition_name)[source]#
Parameters:

  • request_id (int)

  • partition_name (str)

get_partition_for_node(node_name)[source]#
Parameters:

node_name (str)

Return type:

str | None

get_pending_publish_info(request_id, partition_name)[source]#
Parameters:

  • request_id (int)

  • partition_name (str)

Return type:

dict[str, PublishedInfo]

ingested_stream_chunks(request_id, node_name)[source]#

Edge name -> the chunk each of node_name’s input streams last handed the graph: the one the request’s next step there consumes.

Parameters:

  • request_id (int)

  • node_name (str)

Return type:

dict[str, StreamChunkInfo]

node_to_partition: dict[str, str]#
per_request_info: dict[int, PerRequestInfo]#
remove_request(request_id)[source]#
Parameters:

request_id (int)

update_request_info(request_id, partition_name, current_fwd_info=None, resource_publish_info=None)[source]#
Parameters: