mstar.utils.ipc_format

Classes

AbortRequest(request_id)
ConductorMessage(message_type, body)
ConductorMessageType(*values)
FailRequests(errors)	A worker reporting requests it can no longer serve.
InputSignals(request_id, inputs, ...)
MessageBody()
MessageSource(*values)
NewRequest(request_id, ...)
NewRequestConductor(request_id, ...)
RemoveRequest(request_id, source)
ScheduleTPNode(node_name, graph_walk, ...)
SetupDone(worker_id)
Status(*values)
StopLoops(request_id, loop_names, ...)
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'

SETUP_DONE = 'setup_done'

WORKER_GRAPHS_DONE = 'worker_graphs_done'

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 = <factory>)

Bases: MessageBody

Parameters:

• request_id (str)

• inputs (list[GraphEdge])

• request_info (CurrentForwardPassInfo)

• partition_name (str)

• producer_done (set)

inputs: list[GraphEdge]

partition_name: str = 'default'

producer_done: set

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[str], worker_graph_to_workers: dict[str, 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[str])

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

• initial_inputs (list[GraphEdge])

• request_info (CurrentForwardPassInfo)

initial_inputs: list[GraphEdge]

partition_worker_graph_ids: list[str]

request_id: str

request_info: CurrentForwardPassInfo

worker_graph_to_workers: dict[str, 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.RemoveRequest(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.ScheduleTPNode(node_name: str, graph_walk: str, request_ids: list[str])

Bases: MessageBody

Parameters:

• node_name (str)

• graph_walk (str)

• request_ids (list[str])

graph_walk: str

node_name: str

request_ids: list[str]

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.TensorReceived(request_id: str, successful_tensors: dict[str, int], failed_tensor_ids: list[str])

Bases: MessageBody

Parameters:

• request_id (str)

• successful_tensors (dict[str, int])

• failed_tensor_ids (list[str])

failed_tensor_ids: list[str]

request_id: str

successful_tensors: dict[str, int]

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

Bases: MessageBody

Parameters:

• request_id (str)

• uuid_to_ref_count (dict[str, int])

request_id: str

uuid_to_ref_count: dict[str, int]

class mstar.utils.ipc_format.WorkerGraphsDone(request_id: str, worker_graph_ids: list[str], 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: int = 0, 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[str])

• is_first_tp_rank (bool)

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

• new_token_counts (dict[str, int])

• output_signal_names (int)

• 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: int = 0

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[str]

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

INPUT_SIGNALS = 'input_signals'

NEW_REQUEST = 'new_request'

REMOVE_REQUEST = 'remove_request'

SCHEDULE_TP = 'schedule_tp'

STOP_LOOPS = 'stop_loops'

TENSOR_RECEIVED = 'tensor_received'

UNPERSIST_TENSORS = 'unpersist'