chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,862 @@
|
||||
import copy
|
||||
import heapq
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from enum import Enum
|
||||
from functools import total_ordering
|
||||
from typing import Dict, List, Optional, Set, Tuple
|
||||
|
||||
import ray
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class _DAGNodeOperationType(Enum):
|
||||
"""
|
||||
There are three types of operations that a DAG node can perform:
|
||||
1. READ: Read from an input channel.
|
||||
2. COMPUTE: Execute the method corresponding to the node.
|
||||
3. WRITE: Write to an output channel.
|
||||
"""
|
||||
|
||||
READ = "READ"
|
||||
COMPUTE = "COMPUTE"
|
||||
WRITE = "WRITE"
|
||||
|
||||
def viz_str(self):
|
||||
"""
|
||||
A string representation of the operation type to be used in visualization.
|
||||
|
||||
The result string is a single character because conciseness is preferred.
|
||||
"""
|
||||
if self == _DAGNodeOperationType.READ:
|
||||
return "R"
|
||||
elif self == _DAGNodeOperationType.COMPUTE:
|
||||
return "C"
|
||||
elif self == _DAGNodeOperationType.WRITE:
|
||||
return "W"
|
||||
assert False, f"Unknown operation type: {self}"
|
||||
|
||||
|
||||
class _DAGNodeOperation:
|
||||
def __init__(
|
||||
self,
|
||||
exec_task_idx: int,
|
||||
operation_type: _DAGNodeOperationType,
|
||||
method_name: Optional[str] = None,
|
||||
):
|
||||
"""Initialize a _DAGNodeOperation.
|
||||
|
||||
Args:
|
||||
exec_task_idx: The index of the task that this operation belongs to
|
||||
in the actor's ExecutableTask list. The index is not the same
|
||||
as bind_index because there may be more tasks bound to an actor
|
||||
than tasks that appear in the current compiled DAG.
|
||||
operation_type: The type of operation to perform.
|
||||
method_name: The name of the method that this operation originates
|
||||
from. This is only for visualization and debugging purposes.
|
||||
"""
|
||||
self.exec_task_idx = exec_task_idx
|
||||
self.type = operation_type
|
||||
self.method_name = method_name
|
||||
|
||||
def __repr__(self):
|
||||
return (
|
||||
f"_DAGNodeOperation("
|
||||
f"exec_task_idx: {self.exec_task_idx}, "
|
||||
f"type: {self.type}, "
|
||||
f"method_name: {self.method_name})"
|
||||
)
|
||||
|
||||
def viz_str(self):
|
||||
"""
|
||||
A string representation of the node to be used in visualization.
|
||||
"""
|
||||
return f"[{self.exec_task_idx}] {self.method_name} {self.type.viz_str()}"
|
||||
|
||||
def __hash__(self):
|
||||
return hash((self.exec_task_idx, self.type))
|
||||
|
||||
def __eq__(self, other):
|
||||
# An operation is uniquely identified by its `exec_task_idx` and type.
|
||||
# `method_name` is only for debugging purposes.
|
||||
return self.exec_task_idx == other.exec_task_idx and self.type == other.type
|
||||
|
||||
|
||||
@total_ordering
|
||||
class _DAGOperationGraphNode:
|
||||
def __init__(
|
||||
self,
|
||||
operation: _DAGNodeOperation,
|
||||
task_idx: int,
|
||||
actor_handle: "ray.actor.ActorHandle",
|
||||
requires_accelerator: bool,
|
||||
):
|
||||
"""
|
||||
_DAGOperationGraphNode represents a node in the DAG operation graph.
|
||||
It contains information about the node's in-degree, out-degree, edges,
|
||||
and the operation it performs.
|
||||
|
||||
Args:
|
||||
operation: The operation that this node performs. The operation
|
||||
can be a READ, COMPUTE, or WRITE operation.
|
||||
task_idx: A unique index which can be used to index into
|
||||
`CompiledDAG.idx_to_task` to get the corresponding task.
|
||||
actor_handle: The actor handle to which this operation belongs.
|
||||
requires_accelerator: Whether this operation requires accelerator.
|
||||
"""
|
||||
self.operation = operation
|
||||
self.task_idx = task_idx
|
||||
self.actor_handle = actor_handle
|
||||
self.requires_accelerator = requires_accelerator
|
||||
# The in_edges and out_edges are dicts of tuples to strings.
|
||||
# Each tuple (the key) contains an integer `task_idx`, which can be
|
||||
# used to index into `idx_to_task` to get the corresponding task,
|
||||
# and a `_DAGNodeOperationType`, which can be READ, COMPUTE, or WRITE.
|
||||
# The string (the value) is the visualization information of the edge,
|
||||
# it is a tuple of a label of the edge and a boolean indicating whether
|
||||
# the edge is a control dependency.
|
||||
self.in_edges: Dict[Tuple[int, _DAGNodeOperationType], Tuple[str, bool]] = {}
|
||||
self.out_edges: Dict[Tuple[int, _DAGNodeOperationType], Tuple[str, bool]] = {}
|
||||
# The synchronous nodes are all the nodes that belong to the same accelerator
|
||||
# operation. Each node is represented by a tuple of its task idx and type.
|
||||
self.sync_idxs: Set[Tuple[int, _DAGNodeOperationType]] = set()
|
||||
# The pending synchronous nodes are the nodes that are pending to be executed,
|
||||
# i.e., their in-degrees are zero. When a synchronous node is pending, it
|
||||
# will be added to the pending synchronous nodes of all the nodes in the
|
||||
# accelerator operation.
|
||||
self.pending_sync_idxs: Set[Tuple[int, _DAGNodeOperationType]] = set()
|
||||
|
||||
def __repr__(self):
|
||||
return (
|
||||
f"_DAGOperationGraphNode("
|
||||
f"operation: {self.operation}, "
|
||||
f"task_idx: {self.task_idx}, "
|
||||
f"actor_id: {self.actor_handle._ray_actor_id}, "
|
||||
f"requires_accelerator: {self.requires_accelerator})"
|
||||
)
|
||||
|
||||
def __lt__(self, other: "_DAGOperationGraphNode"):
|
||||
"""
|
||||
This function defines the order of the nodes in the priority queue used in
|
||||
`_select_next_nodes`. The priority queue is a min-heap, so the node with
|
||||
higher priority is considered "less than" the other node.
|
||||
"""
|
||||
if self.is_accelerator_op != other.is_accelerator_op:
|
||||
# When one node is an accelerator operation and the other is not,
|
||||
# prioritize the accelerator operation.
|
||||
return self.is_accelerator_op
|
||||
else:
|
||||
# When either both nodes are accelerator operations or both nodes
|
||||
# are not accelerator operations, prioritize the earlier task within
|
||||
# the same actor and load balance tasks across actors. The tie is
|
||||
# broken by the `task_idx`.
|
||||
return (self.operation.exec_task_idx, self.task_idx) < (
|
||||
other.operation.exec_task_idx,
|
||||
other.task_idx,
|
||||
)
|
||||
|
||||
def __eq__(self, other: "_DAGOperationGraphNode"):
|
||||
"""
|
||||
Two operations are equal only when they have the same `exec_task_idx` and `type`
|
||||
and belong to the same actor.
|
||||
"""
|
||||
return (
|
||||
self.actor_handle == other.actor_handle
|
||||
and self.operation.exec_task_idx == other.operation.exec_task_idx
|
||||
and self.operation.type == other.operation.type
|
||||
)
|
||||
|
||||
def __hash__(self):
|
||||
"""
|
||||
An operation is uniquely identified by its `task_idx` and type.
|
||||
"""
|
||||
return hash((self.operation, self.task_idx))
|
||||
|
||||
@property
|
||||
def in_degree(self) -> int:
|
||||
return len(self.in_edges)
|
||||
|
||||
@property
|
||||
def is_ready(self) -> bool:
|
||||
"""
|
||||
If a node is not an accelerator operation, it is ready when it has a zero
|
||||
in-degree.
|
||||
If it is an accelerator operation, it is ready when all the nodes in the
|
||||
operation have zero in-degrees.
|
||||
"""
|
||||
return self.in_degree == 0 and (
|
||||
len(self.pending_sync_idxs) == len(self.sync_idxs)
|
||||
)
|
||||
|
||||
@property
|
||||
def is_read(self) -> bool:
|
||||
return self.operation.type == _DAGNodeOperationType.READ
|
||||
|
||||
@property
|
||||
def is_accelerator_read(self) -> bool:
|
||||
"""
|
||||
A node is an accelerator read if it is a read node and requires accelerator.
|
||||
"""
|
||||
return (
|
||||
self.operation.type == _DAGNodeOperationType.READ
|
||||
and self.requires_accelerator
|
||||
)
|
||||
|
||||
@property
|
||||
def is_accelerator_compute(self) -> bool:
|
||||
"""
|
||||
A node is an accelerator compute if it is a compute node and requires accelerator.
|
||||
"""
|
||||
return (
|
||||
self.operation.type == _DAGNodeOperationType.COMPUTE
|
||||
and self.requires_accelerator
|
||||
)
|
||||
|
||||
@property
|
||||
def is_accelerator_write(self) -> bool:
|
||||
"""
|
||||
A node is an accelerator write if it is a write node and requires accelerator.
|
||||
"""
|
||||
return (
|
||||
self.operation.type == _DAGNodeOperationType.WRITE
|
||||
and self.requires_accelerator
|
||||
)
|
||||
|
||||
@property
|
||||
def is_accelerator_op(self) -> bool:
|
||||
return (
|
||||
self.is_accelerator_read
|
||||
or self.is_accelerator_compute
|
||||
or self.is_accelerator_write
|
||||
)
|
||||
|
||||
def viz_str(self):
|
||||
"""
|
||||
A string representation of the node to be used in visualization.
|
||||
"""
|
||||
return self.operation.viz_str()
|
||||
|
||||
@property
|
||||
def _actor_id(self):
|
||||
return self.actor_handle._ray_actor_id.hex()
|
||||
|
||||
|
||||
def _add_edge(
|
||||
from_node: _DAGOperationGraphNode,
|
||||
to_node: _DAGOperationGraphNode,
|
||||
label: str = "",
|
||||
control_dependency: bool = False,
|
||||
):
|
||||
"""
|
||||
Add an edge from `from_node` to `to_node`.
|
||||
|
||||
Args:
|
||||
from_node: The node from which the edge originates.
|
||||
to_node: The node to which the edge points.
|
||||
label: The label of the edge. This will be used to annotate the edge
|
||||
in the visualization of the execution schedule.
|
||||
control_dependency: If True, the edge represents a control dependency
|
||||
(used for visualization) rather than a data dependency.
|
||||
"""
|
||||
from_node.out_edges[(to_node.task_idx, to_node.operation.type)] = (
|
||||
label,
|
||||
control_dependency,
|
||||
)
|
||||
to_node.in_edges[(from_node.task_idx, from_node.operation.type)] = (
|
||||
label,
|
||||
control_dependency,
|
||||
)
|
||||
|
||||
|
||||
def _update_pending_sync_idxs(
|
||||
graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]],
|
||||
node: _DAGOperationGraphNode,
|
||||
) -> None:
|
||||
"""
|
||||
Update the node as pending for its synchronous nodes.
|
||||
"""
|
||||
idx = (node.task_idx, node.operation.type)
|
||||
for task_idx, op_type in node.sync_idxs:
|
||||
sync_node = graph[task_idx][op_type]
|
||||
sync_node.pending_sync_idxs.add(idx)
|
||||
|
||||
|
||||
def _push_candidate_node_if_ready(
|
||||
actor_to_candidates: Dict["ray._raylet.ActorID", List[_DAGOperationGraphNode]],
|
||||
graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]],
|
||||
node: _DAGOperationGraphNode,
|
||||
) -> None:
|
||||
"""
|
||||
Push the node with a zero in-degree to the candidates if its operation is ready.
|
||||
If it has synchronous nodes, its accelerator operation is not ready until all
|
||||
the nodes are pending, then all the nodes will be pushed to the candidates.
|
||||
"""
|
||||
assert node.in_degree == 0, "Expected to have a zero in-degree"
|
||||
# For the accelerator write node, update the in-degrees of the downstream
|
||||
# accelerator read nodes and update them as pending. This is necessary because
|
||||
# the data dependency edges between accelerator write and read nodes are only
|
||||
# updated here. The accelerator P2P operation becomes ready after both the write
|
||||
# and read nodes are marked as pending.
|
||||
if node.is_accelerator_write:
|
||||
for task_idx, op_type in node.out_edges:
|
||||
read_node = graph[task_idx][op_type]
|
||||
read_node.in_edges.pop((node.task_idx, node.operation.type))
|
||||
assert read_node.is_accelerator_read and len(read_node.in_edges) == 0
|
||||
_update_pending_sync_idxs(graph, read_node)
|
||||
# For the accelerator operation node, update it as pending.
|
||||
if len(node.sync_idxs) != 0:
|
||||
_update_pending_sync_idxs(graph, node)
|
||||
# The accelerator operation is ready when all the nodes have zero in-degrees.
|
||||
# When the last node in the operation is updated as pending, push all the nodes
|
||||
# to the candidates.
|
||||
if node.is_ready:
|
||||
if len(node.sync_idxs) == 0:
|
||||
heapq.heappush(
|
||||
actor_to_candidates[node.actor_handle._actor_id],
|
||||
node,
|
||||
)
|
||||
else:
|
||||
for task_idx, op_type in node.sync_idxs:
|
||||
sync_node = graph[task_idx][op_type]
|
||||
heapq.heappush(
|
||||
actor_to_candidates[sync_node.actor_handle._actor_id],
|
||||
sync_node,
|
||||
)
|
||||
|
||||
|
||||
def _select_next_nodes(
|
||||
actor_to_candidates: Dict["ray._raylet.ActorID", List[_DAGOperationGraphNode]],
|
||||
graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]],
|
||||
) -> Optional[List[_DAGOperationGraphNode]]:
|
||||
"""
|
||||
This function selects the next nodes for the topological sort to generate
|
||||
execution schedule. If there are multiple candidate _DAGOperationGraphNodes,
|
||||
select the node with the top priority. The priority is defined in
|
||||
`_DAGOperationGraphNode.__lt__`.
|
||||
|
||||
For the implementation details, we maintain a priority queue for each actor,
|
||||
where the head of the priority queue is the node with the smallest `exec_task_idx`.
|
||||
When a node has a zero in-degree, it is added to the corresponding actor's
|
||||
priority queue. For a node other than an accelerator collective node, it is ready to be
|
||||
executed if it has a zero in-degree. For an accelerator collective node, it is ready to
|
||||
be executed when all the nodes in its collective operation have zero in-degrees.
|
||||
|
||||
If a node is an accelerator collective node, it updates the `ready_collective_nodes` of
|
||||
all the nodes in its collective operation. Unless all the nodes in its collective
|
||||
group have zero in-degrees, this node is removed from the candidate list.
|
||||
Eventually, exactly one accelerator collective node from its collective operation is
|
||||
selected from the candidate list.
|
||||
|
||||
If the selected node is an accelerator write node, select all the downstream accelerator
|
||||
read nodes. If the selected node is an accelerator collective node, select all the accelerator
|
||||
compute nodes in its collective operation.
|
||||
|
||||
Args:
|
||||
actor_to_candidates: A dictionary mapping an actor id to a list of
|
||||
candidate nodes. The list is maintained as a priority queue, so
|
||||
the head of the queue, i.e., `candidates[0]`, is the node with
|
||||
the smallest `bind_index`.
|
||||
graph: A dictionary mapping the index of a task to a dictionary of its
|
||||
_DAGOperationGraphNodes for different operations.
|
||||
|
||||
Returns:
|
||||
A list of _DAGOperationGraphNodes to be placed into the corresponding
|
||||
execution schedules.
|
||||
"""
|
||||
top_priority_node = None
|
||||
for candidates in actor_to_candidates.values():
|
||||
if len(candidates) == 0:
|
||||
continue
|
||||
if top_priority_node is None or candidates[0] < top_priority_node:
|
||||
top_priority_node = candidates[0]
|
||||
|
||||
if top_priority_node is None:
|
||||
return None
|
||||
next_nodes = [top_priority_node]
|
||||
|
||||
# Select all the synchronous nodes in the accelerator operation.
|
||||
if len(top_priority_node.sync_idxs) != 0:
|
||||
for task_idx, op_type in top_priority_node.sync_idxs:
|
||||
node = graph[task_idx][op_type]
|
||||
if node != top_priority_node:
|
||||
next_nodes.append(node)
|
||||
|
||||
# Remove the selected nodes from the candidates.
|
||||
for node in next_nodes:
|
||||
candidates = actor_to_candidates[node.actor_handle._actor_id]
|
||||
candidates.remove(node)
|
||||
heapq.heapify(candidates)
|
||||
|
||||
# Remove the selected nodes from the candidates.
|
||||
for node in next_nodes:
|
||||
candidates = actor_to_candidates[node.actor_handle._actor_id]
|
||||
# The accelerator read nodes are not added to the candidates.
|
||||
if node in candidates:
|
||||
candidates.remove(node)
|
||||
heapq.heapify(candidates)
|
||||
|
||||
return next_nodes
|
||||
|
||||
|
||||
def _build_dag_node_operation_graph(
|
||||
idx_to_task: Dict[int, "ray.dag.compiled_dag_node.CompiledTask"],
|
||||
actor_to_operation_nodes: Dict[
|
||||
"ray.actor.ActorHandle", List[List[_DAGOperationGraphNode]]
|
||||
],
|
||||
) -> Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]]:
|
||||
"""
|
||||
Generate a DAG node operation graph by adding edges based on the
|
||||
following rules:
|
||||
|
||||
#1 Add edges from READ to COMPUTE, and from COMPUTE to WRITE, which
|
||||
belong to the same task.
|
||||
#2 Add an edge from COMPUTE with bind_index i to COMPUTE with bind_index
|
||||
i+1 if they belong to the same actor.
|
||||
#3 Add an edge from WRITE of the writer task to READ of the reader task.
|
||||
|
||||
This is the step one of building an execution schedule for each actor.
|
||||
|
||||
Args:
|
||||
idx_to_task: A dictionary that maps the `task_idx` to the `CompiledTask`.
|
||||
`CompiledTask` contains information about a DAGNode and its downstream
|
||||
nodes.
|
||||
|
||||
actor_to_operation_nodes: A dictionary that maps an actor handle to
|
||||
a list of lists of _DAGOperationGraphNode. For the same actor, the
|
||||
index of the outer list corresponds to the index of the ExecutableTask
|
||||
in the list of `executable_tasks` in `actor_to_executable_tasks`. In
|
||||
the inner list, the order of operations is READ, COMPUTE, and WRITE.
|
||||
|
||||
Returns:
|
||||
A graph where each node is a _DAGOperationGraphNode. The key is `task_idx`,
|
||||
the index to retrieve its task from `idx_to_task`, and the value is a
|
||||
dictionary that maps the _DAGNodeOperationType (READ, COMPUTE, or WRITE)
|
||||
to the corresponding _DAGOperationGraphNode
|
||||
"""
|
||||
assert idx_to_task
|
||||
graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]] = {}
|
||||
|
||||
for _, operation_nodes_list in actor_to_operation_nodes.items():
|
||||
prev_compute_node = None
|
||||
for operation_nodes in operation_nodes_list:
|
||||
task_idx = operation_nodes[0].task_idx
|
||||
read_node, compute_node, write_node = (
|
||||
operation_nodes[0],
|
||||
operation_nodes[1],
|
||||
operation_nodes[2],
|
||||
)
|
||||
# Add edges from READ to COMPUTE, and from COMPUTE to WRITE, which
|
||||
# belong to the same task.
|
||||
_add_edge(read_node, compute_node)
|
||||
_add_edge(compute_node, write_node)
|
||||
# Add an edge from COMPUTE with `bind_index` i to COMPUTE with
|
||||
# `bind_index` i+1 if they belong to the same actor.
|
||||
if prev_compute_node is not None:
|
||||
_add_edge(prev_compute_node, compute_node, "", True)
|
||||
prev_compute_node = compute_node
|
||||
assert task_idx not in graph
|
||||
graph[task_idx] = {
|
||||
_DAGNodeOperationType.READ: read_node,
|
||||
_DAGNodeOperationType.COMPUTE: compute_node,
|
||||
_DAGNodeOperationType.WRITE: write_node,
|
||||
}
|
||||
|
||||
# Import `ray.dag` here to avoid circular import.
|
||||
from ray.dag import ClassMethodNode, CollectiveOutputNode, MultiOutputNode
|
||||
from ray.dag.collective_node import _CollectiveOperation
|
||||
|
||||
# Add an edge from WRITE of the writer task to READ of the reader task.
|
||||
# Set synchronous nodes for accelerator P2P operations.
|
||||
for task_idx, task in idx_to_task.items():
|
||||
if not (
|
||||
isinstance(task.dag_node, ClassMethodNode)
|
||||
or isinstance(task.dag_node, CollectiveOutputNode)
|
||||
):
|
||||
# The graph is used to generate an execution schedule for each actor.
|
||||
# The edge from the InputNode has no impact on the final execution
|
||||
# schedule.
|
||||
continue
|
||||
if (
|
||||
isinstance(task.dag_node, ClassMethodNode)
|
||||
and task.dag_node.is_class_method_output
|
||||
):
|
||||
# Class method output node dependencies are handled at its upstream:
|
||||
# i.e., class method node
|
||||
continue
|
||||
for downstream_task_idx in task.downstream_task_idxs:
|
||||
downstream_dag_node = idx_to_task[downstream_task_idx].dag_node
|
||||
if isinstance(downstream_dag_node, MultiOutputNode):
|
||||
continue
|
||||
write_node = graph[task_idx][_DAGNodeOperationType.WRITE]
|
||||
if (
|
||||
isinstance(downstream_dag_node, ClassMethodNode)
|
||||
and downstream_dag_node.is_class_method_output
|
||||
):
|
||||
consumer_idxs = idx_to_task[downstream_task_idx].downstream_task_idxs
|
||||
for consumer_idx in consumer_idxs:
|
||||
if consumer_idx in graph:
|
||||
read_node = graph[consumer_idx][_DAGNodeOperationType.READ]
|
||||
_add_edge(
|
||||
write_node,
|
||||
read_node,
|
||||
"accelerator" if write_node.requires_accelerator else "shm",
|
||||
)
|
||||
if write_node.requires_accelerator:
|
||||
idxs = {
|
||||
(task_idx, _DAGNodeOperationType.WRITE),
|
||||
(consumer_idx, _DAGNodeOperationType.READ),
|
||||
}
|
||||
for node in [write_node, read_node]:
|
||||
node.sync_idxs.update(idxs)
|
||||
continue
|
||||
read_node = graph[downstream_task_idx][_DAGNodeOperationType.READ]
|
||||
_add_edge(
|
||||
write_node,
|
||||
read_node,
|
||||
"accelerator" if write_node.requires_accelerator else "shm",
|
||||
)
|
||||
if write_node.requires_accelerator:
|
||||
idxs = {
|
||||
(task_idx, _DAGNodeOperationType.WRITE),
|
||||
(downstream_task_idx, _DAGNodeOperationType.READ),
|
||||
}
|
||||
for node in [write_node, read_node]:
|
||||
node.sync_idxs.update(idxs)
|
||||
|
||||
# Set synchronous nodes for accelerator collective operations.
|
||||
collective_op_to_idxs: Dict[
|
||||
_CollectiveOperation, Set[Tuple[int, _DAGNodeOperationType]]
|
||||
] = defaultdict(set)
|
||||
for task_idx, task in idx_to_task.items():
|
||||
if (
|
||||
isinstance(task.dag_node, CollectiveOutputNode)
|
||||
and not task.dag_node.is_class_method_output
|
||||
):
|
||||
collective_op_to_idxs[task.dag_node.collective_op].add(
|
||||
(task_idx, _DAGNodeOperationType.COMPUTE)
|
||||
)
|
||||
for idxs in collective_op_to_idxs.values():
|
||||
for task_idx, op_type in idxs:
|
||||
graph[task_idx][op_type].sync_idxs = idxs
|
||||
|
||||
return graph
|
||||
|
||||
|
||||
def _actor_viz_label(actor: "ray.actor.ActorHandle") -> str:
|
||||
"""
|
||||
Returns the label of an actor in the visualization of the execution schedule.
|
||||
|
||||
Args:
|
||||
actor: The actor to be represented.
|
||||
|
||||
Returns:
|
||||
A human-readable label combining the actor's class name and ID.
|
||||
"""
|
||||
class_name = actor._ray_actor_creation_function_descriptor.class_name
|
||||
actor_id = actor._ray_actor_id.hex()
|
||||
return f"Actor class name: {class_name}\nActor ID: {actor_id}"
|
||||
|
||||
|
||||
def _node_viz_id_and_label(
|
||||
node: _DAGOperationGraphNode, idx: int, optimized_index: int
|
||||
) -> Tuple[str, str]:
|
||||
"""
|
||||
Returns the visualization id and label of a node. The visualization id is unique
|
||||
across all nodes.
|
||||
|
||||
Args:
|
||||
node: The node to be represented.
|
||||
idx: The index of the node in the execution schedule.
|
||||
optimized_index: The index of the node in the optimized execution schedule.
|
||||
|
||||
Returns:
|
||||
A ``(node_viz_id, node_viz_label)`` tuple suitable for visualization.
|
||||
"""
|
||||
node_viz_label = node.viz_str() + f" {idx},{optimized_index}"
|
||||
node_viz_id = f"{node._actor_id}_{node_viz_label}"
|
||||
return node_viz_id, node_viz_label
|
||||
|
||||
|
||||
def _visualize_execution_schedule(
|
||||
actor_to_execution_schedule: Dict[
|
||||
"ray.actor.ActorHandle", List[_DAGOperationGraphNode]
|
||||
],
|
||||
actor_to_overlapped_schedule: Optional[
|
||||
Dict["ray.actor.ActorHandle", List[_DAGOperationGraphNode]]
|
||||
],
|
||||
graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]],
|
||||
):
|
||||
"""
|
||||
Visualize the execution schedule for each actor.
|
||||
|
||||
The visualization will be saved as a PNG file named `compiled_graph_schedule.png`.
|
||||
Details of the visualization: # noqa
|
||||
|
||||
Node description format:
|
||||
[<task_index>] <method_name> <operation> <orig_index>, <overlap_index>
|
||||
|
||||
Node description fields:
|
||||
operation: is R(READ), C(COMPUTE), or W(WRITE)
|
||||
orig_index: the index in the original execution schedule
|
||||
overlap_index: the index in the overlap-communication optimized execution schedule
|
||||
If this is different from orig_index, the node is highlighted in red color
|
||||
|
||||
Node grouping:
|
||||
The nodes belonging to the same actor are grouped in the same rectangle
|
||||
The actor class name and the actor id are shown in the rectangle
|
||||
|
||||
Edges:
|
||||
black color (without label): data dependency
|
||||
black color (annotated with "shm"): shared memory channel
|
||||
blue color (annotated with "accelerator): accelerator channel
|
||||
dashed edge: control dependency between compute operations
|
||||
|
||||
Args:
|
||||
actor_to_execution_schedule: A dictionary that maps an actor handle to
|
||||
the execution schedule which is a list of operation nodes.
|
||||
actor_to_overlapped_schedule: A dictionary that maps an actor handle to the
|
||||
optimized execution schedule which is a list of operation nodes.
|
||||
graph: A graph where each node is a _DAGOperationGraphNode. The key is
|
||||
`task_idx`, the index to retrieve its task from `idx_to_task`, and
|
||||
the value is a dictionary that maps the _DAGNodeOperationType (READ,
|
||||
COMPUTE, or WRITE) to the corresponding _DAGOperationGraphNode. It is
|
||||
generated by `_build_dag_node_operation_graph`.
|
||||
"""
|
||||
try:
|
||||
import graphviz
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Please install graphviz to visualize the execution schedule. "
|
||||
"You can install it by running `pip install graphviz`."
|
||||
)
|
||||
|
||||
dot = graphviz.Digraph(comment="DAG")
|
||||
# A dictionary that maps a node to its visualization id
|
||||
node_to_viz_id: Dict[_DAGOperationGraphNode, str] = {}
|
||||
|
||||
if actor_to_overlapped_schedule is None:
|
||||
# TODO(rui): make the visualization more concise by only displaying
|
||||
# the original schedule
|
||||
actor_to_overlapped_schedule = actor_to_execution_schedule
|
||||
for actor, execution_nodes in actor_to_execution_schedule.items():
|
||||
overlapped_schedule = actor_to_overlapped_schedule[actor]
|
||||
node_to_optimized_index = {
|
||||
node: i for i, node in enumerate(overlapped_schedule)
|
||||
}
|
||||
|
||||
actor_id = actor._ray_actor_id.hex()
|
||||
with dot.subgraph(name=f"cluster_{actor_id}") as subgraph:
|
||||
subgraph.attr(rank=actor_id, label=_actor_viz_label(actor))
|
||||
for i, node in enumerate(execution_nodes):
|
||||
optimized_index = node_to_optimized_index.get(node)
|
||||
node_viz_id, node_viz_label = _node_viz_id_and_label(
|
||||
node, i, optimized_index
|
||||
)
|
||||
color = "red" if optimized_index != i else "black"
|
||||
subgraph.node(node_viz_id, node_viz_label, color=color)
|
||||
node_to_viz_id[node] = node_viz_id
|
||||
|
||||
for actor, execution_nodes in actor_to_execution_schedule.items():
|
||||
for i, node in enumerate(execution_nodes):
|
||||
node_viz_id = node_to_viz_id[node]
|
||||
for out_edge, viz_info in node.out_edges.items():
|
||||
label, control_dependency = viz_info
|
||||
out_task_idx, out_op_type = out_edge
|
||||
out_node = graph[out_task_idx][out_op_type]
|
||||
out_node_viz_id = node_to_viz_id[out_node]
|
||||
color = "blue" if label == "accelerator" else "black"
|
||||
style = "dashed" if control_dependency else "solid"
|
||||
dot.edge(
|
||||
node_viz_id, out_node_viz_id, label=label, color=color, style=style
|
||||
)
|
||||
|
||||
# Add legend
|
||||
with dot.subgraph(name="cluster_legend") as legend:
|
||||
legend.attr(label="Legend", labelloc="t", fontsize="20", bgcolor="lightgrey")
|
||||
|
||||
# Single node and its explanation
|
||||
legend.node("example_node", "[0] bwd C 10,10\n")
|
||||
explanation = (
|
||||
'<<TABLE BORDER="0" CELLBORDER="0" CELLSPACING="0">' # noqa
|
||||
'<TR><TD ALIGN="LEFT"><B>Node description format:</B></TD></TR>'
|
||||
'<TR><TD ALIGN="LEFT">[<task_index>] <method_name> <operation> <orig_index>, <overlap_index></TD></TR>' # noqa
|
||||
"<TR><TD></TD></TR>"
|
||||
'<TR><TD ALIGN="LEFT"><B>Node description fields:</B></TD></TR>'
|
||||
'<TR><TD ALIGN="LEFT">operation: is R(READ), C(COMPUTE), or W(WRITE)</TD></TR>' # noqa
|
||||
'<TR><TD ALIGN="LEFT">orig_index: the index in the original execution schedule</TD></TR>' # noqa
|
||||
'<TR><TD ALIGN="LEFT">overlap_index: the index in the overlap-communication optimized execution schedule</TD></TR>' # noqa
|
||||
'<TR><TD ALIGN="LEFT">If this is different from orig_index, the node is highlighted in <FONT COLOR="red">red color</FONT></TD></TR>' # noqa
|
||||
"<TR><TD></TD></TR>"
|
||||
'<TR><TD ALIGN="LEFT"><B>Node grouping:</B></TD></TR>'
|
||||
'<TR><TD ALIGN="LEFT">The nodes belonging to the same actor are grouped in the same rectangle</TD></TR>' # noqa
|
||||
'<TR><TD ALIGN="LEFT">The actor class name and the actor id are shown in the rectangle</TD></TR>' # noqa
|
||||
"<TR><TD></TD></TR>"
|
||||
'<TR><TD ALIGN="LEFT"><B>Edges:</B></TD></TR>'
|
||||
'<TR><TD ALIGN="LEFT">black color (without label): data dependency</TD></TR>' # noqa
|
||||
'<TR><TD ALIGN="LEFT">black color (annotated with "shm"): shared memory channel</TD></TR>' # noqa
|
||||
'<TR><TD ALIGN="LEFT"><FONT COLOR="blue">blue color</FONT> (annotated with "accelerator): accelerator channel</TD></TR>' # noqa
|
||||
'<TR><TD ALIGN="LEFT">dashed edge: control dependency between compute operations</TD></TR>' # noqa
|
||||
"</TABLE>>"
|
||||
)
|
||||
|
||||
legend.node("example_explanation", explanation, shape="plaintext")
|
||||
legend.edge("example_node", "example_explanation", style="invis")
|
||||
|
||||
logger.info(
|
||||
"Writing compiled graph schedule visualization "
|
||||
"to compiled_graph_schedule.png"
|
||||
)
|
||||
dot.render("compiled_graph_schedule", format="png", view=False)
|
||||
|
||||
|
||||
def _generate_actor_to_execution_schedule(
|
||||
graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]],
|
||||
) -> Dict["ray.actor.ActorHandle", List[_DAGOperationGraphNode]]:
|
||||
"""
|
||||
Generate an execution schedule for each actor. The schedule is a list of
|
||||
operation nodes to be executed. The function uses a topological sort
|
||||
algorithm to generate the schedule.
|
||||
|
||||
Args:
|
||||
graph: A graph where each node is a _DAGOperationGraphNode. The key is
|
||||
`task_idx`, the index to retrieve its task from `idx_to_task`, and
|
||||
the value is a dictionary that maps the _DAGNodeOperationType (READ,
|
||||
COMPUTE, or WRITE) to the corresponding _DAGOperationGraphNode. It is
|
||||
generated by `_build_dag_node_operation_graph`.
|
||||
|
||||
Returns:
|
||||
actor_to_execution_schedule: A dictionary that maps an actor handle to
|
||||
the execution schedule which is a list of operation nodes to be
|
||||
executed.
|
||||
"""
|
||||
|
||||
# Mapping from the actor handle to the execution schedule which is a list
|
||||
# of operations to be executed.
|
||||
actor_to_execution_schedule: Dict[
|
||||
"ray.actor.ActorHandle", List[_DAGOperationGraphNode]
|
||||
] = defaultdict(list)
|
||||
|
||||
# A dictionary mapping an actor id to a list of candidate nodes. The list
|
||||
# is maintained as a priority queue, so the head of the queue, i.e.,
|
||||
# `candidates[0]`, is the node with the smallest `bind_index`.
|
||||
actor_to_candidates: Dict[
|
||||
"ray._raylet.ActorID", List[_DAGOperationGraphNode]
|
||||
] = defaultdict(list)
|
||||
for _, node_dict in graph.items():
|
||||
for _, node in node_dict.items():
|
||||
# A node with a zero in-degree edge means all of its dependencies
|
||||
# have been satisfied, including both data and control dependencies.
|
||||
# Therefore, it is a candidate for execution.
|
||||
if node.in_degree == 0:
|
||||
_push_candidate_node_if_ready(actor_to_candidates, graph, node)
|
||||
|
||||
visited_nodes = set()
|
||||
|
||||
# Use topological sort algorithm to generate the execution schedule.
|
||||
while True:
|
||||
# Select a list of nodes to be executed. There are three cases:
|
||||
# 1. If a selected node is not an accelerator operation, only itself is returned.
|
||||
# 2. If a selected node is an accelerator write operation, the corresponding accelerator
|
||||
# read operations are also returned.
|
||||
# 3. If a selected node is an accelerator collective operation, all the nodes in
|
||||
# its collective operation are returned.
|
||||
nodes = _select_next_nodes(actor_to_candidates, graph)
|
||||
if nodes is None:
|
||||
break
|
||||
# Add the selected nodes to the execution schedule.
|
||||
for node in nodes:
|
||||
assert node not in visited_nodes
|
||||
visited_nodes.add(node)
|
||||
actor_to_execution_schedule[node.actor_handle].append(node)
|
||||
# Update the in-degree of the downstream nodes.
|
||||
for node in nodes:
|
||||
for out_node_task_idx, out_node_type in node.out_edges:
|
||||
out_node = graph[out_node_task_idx][out_node_type]
|
||||
if out_node in visited_nodes:
|
||||
# If the downstream node is already visited, it has been added
|
||||
# to the execution schedule. They are the accelerator read nodes in
|
||||
# case 2.
|
||||
continue
|
||||
out_node.in_edges.pop((node.task_idx, node.operation.type))
|
||||
if out_node.in_degree == 0:
|
||||
_push_candidate_node_if_ready(actor_to_candidates, graph, out_node)
|
||||
assert len(visited_nodes) == len(graph) * 3, "Expected all nodes to be visited"
|
||||
|
||||
return actor_to_execution_schedule
|
||||
|
||||
|
||||
def _generate_overlapped_execution_schedule(
|
||||
actor_to_execution_schedule: Dict[
|
||||
"ray.actor.ActorHandle", List[_DAGOperationGraphNode]
|
||||
],
|
||||
) -> Dict["ray.actor.ActorHandle", List[_DAGOperationGraphNode]]:
|
||||
"""
|
||||
From an existing execution schedule, generate a new schedule by overlapping
|
||||
computation and communication.
|
||||
|
||||
Currently, the algorithm generates a new schedule for each actor as follows:
|
||||
For each accelerator read operation (i.e., recv), scan backwards to find the nearest
|
||||
compute node to swap with so that the accelerator read operation can be overlapped
|
||||
with computation.
|
||||
|
||||
Collective operations are not yet supported.
|
||||
|
||||
Args:
|
||||
actor_to_execution_schedule: A dictionary that maps an actor handle to
|
||||
the existing execution schedule for the actor. The schedule is a list
|
||||
is a list of operations to be executed.
|
||||
|
||||
Returns:
|
||||
A dictionary that maps an actor handle to the overlapped execution schedule
|
||||
for the actor.
|
||||
"""
|
||||
|
||||
actor_to_overlapped_schedule: Dict[
|
||||
"ray.actor.ActorHandle", List[_DAGOperationGraphNode]
|
||||
] = copy.deepcopy(actor_to_execution_schedule)
|
||||
for overlapped_schedule in actor_to_overlapped_schedule.values():
|
||||
for i in range(len(overlapped_schedule)):
|
||||
if (
|
||||
overlapped_schedule[i].operation.type == _DAGNodeOperationType.READ
|
||||
and overlapped_schedule[i].requires_accelerator
|
||||
):
|
||||
# For each accelerator read operation (i.e., recv), scan backwards
|
||||
# to find the nearest compute node to swap with so that
|
||||
# the accelerator read operation can be overlapped with computation.
|
||||
for j in range(i - 1, -1, -1):
|
||||
if (
|
||||
overlapped_schedule[j].operation.type
|
||||
== _DAGNodeOperationType.COMPUTE
|
||||
):
|
||||
# Found a desired compute operation, make the swap
|
||||
accelerator_read_op = overlapped_schedule[i]
|
||||
prev_ops = overlapped_schedule[j:i]
|
||||
overlapped_schedule[j + 1 : i + 1] = prev_ops
|
||||
overlapped_schedule[j] = accelerator_read_op
|
||||
break
|
||||
if (
|
||||
overlapped_schedule[j].operation.type
|
||||
== _DAGNodeOperationType.READ
|
||||
or overlapped_schedule[j].operation.type
|
||||
== _DAGNodeOperationType.WRITE
|
||||
) and overlapped_schedule[j].requires_accelerator:
|
||||
# Found an accelerator read/write operation, skip the overlap
|
||||
# optimization to keep relative order of accelerator operations
|
||||
break
|
||||
return actor_to_overlapped_schedule
|
||||
|
||||
|
||||
def _extract_execution_schedule(
|
||||
actor_to_execution_schedule: Dict[
|
||||
"ray.actor.ActorHandle", List[_DAGOperationGraphNode]
|
||||
],
|
||||
) -> Dict["ray.actor.ActorHandle", List[_DAGNodeOperation]]:
|
||||
"""
|
||||
Extract _DAGNodeOperation from _DAGOperationGraphNode in the schedule
|
||||
and discard unnecessary information.
|
||||
"""
|
||||
return {
|
||||
actor: [node.operation for node in nodes]
|
||||
for actor, nodes in actor_to_execution_schedule.items()
|
||||
}
|
||||
Reference in New Issue
Block a user