mstar.utils.ipc_format

Classes

AbortRequest(request_id)
ConductorMessage(message_type, body)
ConductorMessageType(*values)
DrainRequest(request_id, source)
FailRequests(errors)	A worker reporting requests it can no longer serve.
InputSignals(request_id, inputs, ...)
MessageBody()
MessageSource(*values)
NewRequest(request_id, ...)
NewRequestConductor(request_id, ...)
OffloadDelta([rids, is_offload])	A rank's offloads and reloads, in the order it made them.
ReadsDone(request_id, entity_id)	An entity confirming it has no in-flight reads for a request and will start none — the conductor's gate before sending the hard RemoveRequest.
RemoveRequest(request_id, source, after_tp_seq)
ScheduleTPNode(node_name, graph_walk, ...)
SetupDone(worker_id)
Status(*values)
StopLoops(request_id, loop_names, ...)
TPNoSpeculation(node_name, graph_walk, ...)
TensorReceived(request_id, ...)
UnpersistTensors(request_id, uuid_to_ref_count)
WorkerGraphsDone(request_id, ...)
WorkerMessage(message_type, body)
WorkerMessageType(*values)

class mstar.utils.ipc_format.AbortRequest(request_id: str)

Bases: MessageBody

Parameters:

request_id (str)

request_id: str

class mstar.utils.ipc_format.ConductorMessage(message_type: mstar.utils.ipc_format.ConductorMessageType, body: mstar.utils.ipc_format.MessageBody)

Bases: object

Parameters:

• message_type (ConductorMessageType)

• body (MessageBody)

body: MessageBody

message_type: ConductorMessageType

class mstar.utils.ipc_format.ConductorMessageType(*values)

Bases: Enum

ABORT_REQUEST = 'abort_request'

FAIL_REQUESTS = 'fail_requests'

NEW_REQUEST = 'new_request'

READS_DONE = 'reads_done'

SETUP_DONE = 'setup_done'

WORKER_GRAPHS_DONE = 'worker_graphs_done'

class mstar.utils.ipc_format.DrainRequest(request_id: str, source: int = <MessageSource.CONDUCTOR: 0>)

Bases: MessageBody

Parameters:

• request_id (str)

• source (int)

request_id: str

source: int = 0

class mstar.utils.ipc_format.FailRequests(errors)

Bases: MessageBody

A worker reporting requests it can no longer serve.

errors maps request_id -> message. It’s a dict rather than a (rids, message) pair because per-rid stages (prepare_inputs, postprocess) attribute a distinct error to each request, and one step can fail several of them for different reasons.

Parameters:

errors (dict[str, str])

errors: dict[str, str]

class mstar.utils.ipc_format.InputSignals(request_id: str, inputs: list[mstar.graph.base.GraphEdge], request_info: mstar.conductor.request_info.CurrentForwardPassInfo, partition_name: str = 'default', producer_done: set[str] = <factory>)

Bases: MessageBody

Parameters:

• request_id (str)

• inputs (list[GraphEdge])

• request_info (CurrentForwardPassInfo)

• partition_name (str)

• producer_done (set[str])

inputs: list[GraphEdge]

partition_name: str = 'default'

producer_done: set[str]

request_id: str

request_info: CurrentForwardPassInfo

class mstar.utils.ipc_format.MessageBody

Bases: object

from_dict(input)

Parameters:

input (dict)

to_dict()

class mstar.utils.ipc_format.MessageSource(*values)

Bases: IntEnum

CONDUCTOR = 0

SELF = 2

TP_RANK_0 = 1

class mstar.utils.ipc_format.NewRequest(request_id: str, partition_worker_graph_ids: list[int], worker_graph_to_workers: dict[int, list[str]], initial_inputs: list[mstar.graph.base.GraphEdge], request_info: mstar.conductor.request_info.CurrentForwardPassInfo)

Bases: MessageBody

Parameters:

• request_id (str)

• partition_worker_graph_ids (list[int])

• worker_graph_to_workers (dict[int, list[str]])

• initial_inputs (list[GraphEdge])

• request_info (CurrentForwardPassInfo)

initial_inputs: list[GraphEdge]

partition_worker_graph_ids: list[int]

request_id: str

request_info: CurrentForwardPassInfo

worker_graph_to_workers: dict[int, list[str]]

class mstar.utils.ipc_format.NewRequestConductor(request_id: str, initial_signals: dict[str, list[mstar.graph.base.TensorPointerInfo]], initial_input_modalities: list[str], initial_output_modalities: list[str], input_metadata: dict[str, list[dict]], model_kwargs: dict)

Bases: MessageBody

Parameters:

• request_id (str)

• initial_signals (dict[str, list[TensorPointerInfo]])

• initial_input_modalities (list[str])

• initial_output_modalities (list[str])

• input_metadata (dict[str, list[dict]])

• model_kwargs (dict)

initial_input_modalities: list[str]

initial_output_modalities: list[str]

initial_signals: dict[str, list[TensorPointerInfo]]

input_metadata: dict[str, list[dict]]

model_kwargs: dict

request_id: str

class mstar.utils.ipc_format.OffloadDelta(rids=<factory>, is_offload=<factory>)

Bases: object

A rank’s offloads and reloads, in the order it made them.

Ordered so the follower replays the leader’s decisions in the same order: a sequence that worked for the leader works for the followers too.

Holds WIRE request ids, not worker-local handles: a handle is minted per worker, so rank 0’s would name different requests on the follower.

A dataclass of lists rather than a NamedTuple of deques so that the wire codec can automatically encode it. pop_left is O(n) on a list, which costs nothing at the handful of entries a delta ever holds.

__len__ is the queue depth, so an empty delta is falsy – several callers lean on that.

Parameters:

• rids (list[str])

• is_offload (list[bool])

add_offloaded(rid)

Parameters:

rid (str)

add_reloaded(rid)

Parameters:

rid (str)

copy()

An independent queue holding the same moves.

One per follower: each replays at its own pace, popping as it goes.

Return type:

OffloadDelta

extend(other)

Parameters:

other (OffloadDelta)

Return type:

None

is_offload: list[bool]

classmethod new()

Return type:

OffloadDelta

peek_left()

Return type:

tuple[str, bool] | None

pop_left()

Return type:

tuple[str, bool] | None

rids: list[str]

take()

Hand the queue over and leave this one empty.

Copies and clears rather than rebinding, so a holder of this delta (the engine keeps its journal in a dict) sees it emptied.

Return type:

OffloadDelta

class mstar.utils.ipc_format.ReadsDone(request_id, entity_id)

Bases: MessageBody

An entity confirming it has no in-flight reads for a request and will start none — the conductor’s gate before sending the hard RemoveRequest.

Parameters:

• request_id (str)

• entity_id (str)

entity_id: str

request_id: str

class mstar.utils.ipc_format.RemoveRequest(request_id: str, source: int = <MessageSource.CONDUCTOR: 0>, after_tp_seq: int = -1)

Bases: MessageBody

Parameters:

• request_id (str)

• source (int)

• after_tp_seq (int)

after_tp_seq: int = -1

request_id: str

source: int = 0

class mstar.utils.ipc_format.ScheduleTPNode(node_name: str, graph_walk: str, request_ids: list[str], speculative: bool = False, spec_seq: int = -1, spec_from_seq: int = -1, resident_delta: mstar.utils.ipc_format.OffloadDelta = <factory>)

Bases: MessageBody

Parameters:

• node_name (str)

• graph_walk (str)

• request_ids (list[str])

• speculative (bool)

• spec_seq (int)

• spec_from_seq (int)

• resident_delta (OffloadDelta)

graph_walk: str

node_name: str

request_ids: list[str]

resident_delta: OffloadDelta

spec_from_seq: int = -1

spec_seq: int = -1

speculative: bool = False

class mstar.utils.ipc_format.SetupDone(worker_id: str)

Bases: MessageBody

Parameters:

worker_id (str)

worker_id: str

class mstar.utils.ipc_format.Status(*values)

Bases: Enum

DONE = 'done'

IN_PROGRESS = 'in_progress'

READY = 'ready'

WAITING = 'waiting'

class mstar.utils.ipc_format.StopLoops(request_id: str, loop_names: set[str], partition_name: str, loop_stop_times: dict[str, mstar.graph.loop_indices.NestedLoopIndices] = <factory>)

Bases: MessageBody

Parameters:

• request_id (str)

• loop_names (set[str])

• partition_name (str)

• loop_stop_times (dict[str, NestedLoopIndices])

loop_names: set[str]

loop_stop_times: dict[str, NestedLoopIndices]

partition_name: str

request_id: str

class mstar.utils.ipc_format.TPNoSpeculation(node_name: str, graph_walk: str, spec_from_seq: int)

Bases: MessageBody

Parameters:

• node_name (str)

• graph_walk (str)

• spec_from_seq (int)

graph_walk: str

node_name: str

spec_from_seq: int

class mstar.utils.ipc_format.TensorReceived(request_id: str, successful_tensors: dict[int, int], failed_tensor_ids: list[int])

Bases: MessageBody

Parameters:

• request_id (str)

• successful_tensors (dict[int, int])

• failed_tensor_ids (list[int])

failed_tensor_ids: list[int]

request_id: str

successful_tensors: dict[int, int]

class mstar.utils.ipc_format.UnpersistTensors(request_id: str, uuid_to_ref_count: dict[int, int])

Bases: MessageBody

Parameters:

• request_id (str)

• uuid_to_ref_count (dict[int, int])

request_id: str

uuid_to_ref_count: dict[int, int]

class mstar.utils.ipc_format.WorkerGraphsDone(request_id: str, worker_graph_ids: list[int], is_first_tp_rank: bool, persist_signals: dict[str, list[mstar.graph.base.TensorPointerInfo]] = <factory>, new_token_counts: dict[str, int] = <factory>, output_signal_names: list[str] = <factory>, resource_publish_info: dict[str, mstar.engine.resources.base.PublishedInfo] = <factory>, partition_name: str = 'default', partition_done: bool = False, stream_tokens_consumed: dict[str, int] = <factory>, output_loop_indices: dict[str, mstar.graph.loop_indices.NestedLoopIndices] = <factory>, graph_timings: dict[tuple[str, str], mstar.profile.format.GraphTiming] = <factory>, rx_info: list[mstar.profile.format.RxInfo] = <factory>, tx_info: list[mstar.profile.format.TxInfo] = <factory>)

Bases: MessageBody

Parameters:

• request_id (str)

• worker_graph_ids (list[int])

• is_first_tp_rank (bool)

• persist_signals (dict[str, list[TensorPointerInfo]])

• new_token_counts (dict[str, int])

• output_signal_names (list[str])

• resource_publish_info (dict[str, PublishedInfo])

• partition_name (str)

• partition_done (bool)

• stream_tokens_consumed (dict[str, int])

• output_loop_indices (dict[str, NestedLoopIndices])

• graph_timings (dict[tuple[str, str], GraphTiming])

• rx_info (list[RxInfo])

• tx_info (list[TxInfo])

graph_timings: dict[tuple[str, str], GraphTiming]

is_first_tp_rank: bool

new_token_counts: dict[str, int]

output_loop_indices: dict[str, NestedLoopIndices]

output_signal_names: list[str]

partition_done: bool = False

partition_name: str = 'default'

persist_signals: dict[str, list[TensorPointerInfo]]

request_id: str

resource_publish_info: dict[str, PublishedInfo]

rx_info: list[RxInfo]

stream_tokens_consumed: dict[str, int]

tx_info: list[TxInfo]

worker_graph_ids: list[int]

class mstar.utils.ipc_format.WorkerMessage(message_type: mstar.utils.ipc_format.WorkerMessageType, body: mstar.utils.ipc_format.MessageBody)

Bases: object

Parameters:

• message_type (WorkerMessageType)

• body (MessageBody)

body: MessageBody

message_type: WorkerMessageType

class mstar.utils.ipc_format.WorkerMessageType(*values)

Bases: Enum

DRAIN_REQUEST = 'drain_request'

INPUT_SIGNALS = 'input_signals'

NEW_REQUEST = 'new_request'

REMOVE_REQUEST = 'remove_request'

SCHEDULE_TP = 'schedule_tp'

STOP_LOOPS = 'stop_loops'

TENSOR_RECEIVED = 'tensor_received'

TP_NO_SPEC = 'tp_no_spec'

UNPERSIST_TENSORS = 'unpersist'