# coding: utf-8 import os import sys from typing import Dict, List, Tuple import pytest from ray.actor import ActorHandle from ray.dag import ClassMethodNode, InputNode, MultiOutputNode from ray.dag.compiled_dag_node import CompiledTask from ray.dag.dag_node_operation import ( _add_edge, _build_dag_node_operation_graph, _DAGNodeOperation, _DAGNodeOperationType, _DAGOperationGraphNode, _extract_execution_schedule, _generate_actor_to_execution_schedule, _select_next_nodes, ) from ray.tests.conftest import * # noqa if sys.platform != "linux" and sys.platform != "darwin": pytest.skip("Skipping, requires Linux or Mac.", allow_module_level=True) def mock_actor_handle_init(self, actor_id: str): self._ray_actor_id = actor_id self._ray_weak_ref = False def mock_class_method_call_init(self): self._is_class_method_output = False def mock_init(self): pass def generate_dag_graph_nodes( exec_task_idx, task_idx, actor_handle, requires_nccl_read=False, requires_nccl_compute=False, requires_nccl_write=False, ): graph_nodes = {} for op_type in _DAGNodeOperationType: requires_nccl = ( (op_type == _DAGNodeOperationType.READ and requires_nccl_read) or (op_type == _DAGNodeOperationType.COMPUTE and requires_nccl_compute) or (op_type == _DAGNodeOperationType.WRITE and requires_nccl_write) ) graph_nodes[op_type] = _DAGOperationGraphNode( _DAGNodeOperation(exec_task_idx, op_type), task_idx, actor_handle, requires_nccl, ) return graph_nodes def set_sync_idxs_p2p( graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]], write_idx: int, read_idx: int, ) -> None: write_node = graph[write_idx][_DAGNodeOperationType.WRITE] read_node = graph[read_idx][_DAGNodeOperationType.READ] p2p_idxs = { (write_idx, _DAGNodeOperationType.WRITE), (read_idx, _DAGNodeOperationType.READ), } for node in [write_node, read_node]: node.sync_idxs.update(p2p_idxs) node.pending_sync_idxs.update(p2p_idxs) def set_sync_idxs_collective( graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]], task_idxs: List[int], ) -> None: collective_idxs = { (task_idx, _DAGNodeOperationType.COMPUTE) for task_idx in task_idxs } for task_idx in task_idxs: node = graph[task_idx][_DAGNodeOperationType.COMPUTE] node.sync_idxs.update(collective_idxs) node.pending_sync_idxs.update(collective_idxs) def _generate_and_extract_execution_schedule(graph): return _extract_execution_schedule(_generate_actor_to_execution_schedule(graph)) class TestSelectNextNodes: """ Test whether `_select_next_nodes` function selects the next nodes for topological sort to generate execution schedule correctly. task_idx: Each DAG node has a unique global index. exec_task_idx: The DAG node's index in the actor's `executable_tasks` list. """ def test_two_candidates_on_same_actor(self, monkeypatch): """ Simulate the case where there are two candidates on the same actor. The candidate with the smaller index in the `executable_tasks` list should be selected. driver -> fake_actor.op -> fake_actor.op -> driver In the example above, both READ operations on the fake_actor have zero in-degree. The operation with the smaller index in the executable_tasks list should be selected first; therefore, the one on the left side will be selected first. """ monkeypatch.setattr(ActorHandle, "__init__", mock_actor_handle_init) fake_actor = ActorHandle("fake_actor") # The DAG node has a global index of 1, and its index in the # actor's `executable_tasks` list is 0. task_idx_1 = 1 dag_node_1 = _DAGOperationGraphNode( _DAGNodeOperation(0, _DAGNodeOperationType.READ), task_idx_1, fake_actor, False, ) # The DAG node has a global index of 2, and its index in the # actor's `executable_tasks` list is 1. task_idx_2 = 2 dag_node_2 = _DAGOperationGraphNode( _DAGNodeOperation(1, _DAGNodeOperationType.READ), task_idx_2, fake_actor, False, ) mock_actor_to_candidates = { fake_actor._actor_id: [ dag_node_1, dag_node_2, ], } next_nodes = _select_next_nodes(mock_actor_to_candidates, None) assert len(next_nodes) == 1 assert next_nodes[0] == dag_node_1 def test_only_one_nccl_write(self, monkeypatch): """ Simulate the case where there is only one candidate which is a NCCL WRITE operation. In this case, `_select_next_nodes` should return both the NCCL WRITE operation and the corresponding READ operation. driver -> fake_actor_1.op -> fake_actor_2.op -> driver In the example above, communication between fake_actor_1 and fake_actor_2 is done using NCCL. The following test case simulates a scenario where the READ and COMPUTE operations on fake_actor_1 have already been added to the execution schedule. """ monkeypatch.setattr(ActorHandle, "__init__", mock_actor_handle_init) fake_actor_1, task_idx_1, exec_task_idx_1 = ActorHandle("fake_actor_1"), 1, 0 fake_actor_2, task_idx_2, exec_task_idx_2 = ActorHandle("fake_actor_2"), 2, 0 mock_graph = { task_idx_1: generate_dag_graph_nodes( exec_task_idx_1, task_idx_1, fake_actor_1, requires_nccl_write=True, ), task_idx_2: generate_dag_graph_nodes( exec_task_idx_2, task_idx_2, fake_actor_2, requires_nccl_read=True, ), } del mock_graph[task_idx_1][_DAGNodeOperationType.READ] del mock_graph[task_idx_1][_DAGNodeOperationType.COMPUTE] _add_edge( mock_graph[task_idx_1][_DAGNodeOperationType.WRITE], mock_graph[task_idx_2][_DAGNodeOperationType.READ], ) _add_edge( mock_graph[task_idx_2][_DAGNodeOperationType.READ], mock_graph[task_idx_2][_DAGNodeOperationType.COMPUTE], ) _add_edge( mock_graph[task_idx_2][_DAGNodeOperationType.COMPUTE], mock_graph[task_idx_2][_DAGNodeOperationType.WRITE], ) set_sync_idxs_p2p(mock_graph, task_idx_1, task_idx_2) mock_actor_to_candidates = { fake_actor_1._actor_id: [ mock_graph[task_idx_1][_DAGNodeOperationType.WRITE] ], fake_actor_2._actor_id: [ mock_graph[task_idx_2][_DAGNodeOperationType.READ] ], } next_nodes = _select_next_nodes(mock_actor_to_candidates, mock_graph) assert next_nodes == [ mock_graph[task_idx_1][_DAGNodeOperationType.WRITE], mock_graph[task_idx_2][_DAGNodeOperationType.READ], ] def test_two_nccl_writes(self, monkeypatch): """ Simulate a scenario where there are two candidates that are NCCL WRITE operations. In this case, _select_next_nodes can choose either of the two NCCL WRITE operations and their corresponding READ operations. driver -> fake_actor_1.op -> fake_actor_2.op -> driver | | -> fake_actor_2.op -> fake_actor_1.op - In the example above, communication between fake_actor_1 and fake_actor_2 is done using NCCL. The following test case simulates a scenario where the READ and COMPUTE operations on both the DAG nodes with smaller bind_index on fake_actor_1 and fake_actor_2 have already been added to the execution schedule. """ monkeypatch.setattr(ActorHandle, "__init__", mock_actor_handle_init) fake_actor_1 = ActorHandle("fake_actor_1") task_idx_1_0, exec_task_idx_1_0 = 1, 0 task_idx_1_1, exec_task_idx_1_1 = 3, 1 fake_actor_2 = ActorHandle("fake_actor_2") task_idx_2_0, exec_task_idx_2_0 = 2, 0 task_idx_2_1, exec_task_idx_2_1 = 4, 1 # Run the test 10 times to ensure that the result of `_select_next_nodes` # is deterministic. for _ in range(20): mock_graph = { task_idx_1_0: generate_dag_graph_nodes( exec_task_idx_1_0, task_idx_1_0, fake_actor_1, requires_nccl_write=True, ), task_idx_1_1: generate_dag_graph_nodes( exec_task_idx_1_1, task_idx_1_1, fake_actor_1, requires_nccl_read=True, ), task_idx_2_0: generate_dag_graph_nodes( exec_task_idx_2_0, task_idx_2_0, fake_actor_2, requires_nccl_write=True, ), task_idx_2_1: generate_dag_graph_nodes( exec_task_idx_2_1, task_idx_2_1, fake_actor_2, requires_nccl_read=True, ), } del mock_graph[task_idx_1_0][_DAGNodeOperationType.READ] del mock_graph[task_idx_1_0][_DAGNodeOperationType.COMPUTE] del mock_graph[task_idx_2_0][_DAGNodeOperationType.READ] del mock_graph[task_idx_2_0][_DAGNodeOperationType.COMPUTE] _add_edge( mock_graph[task_idx_1_0][_DAGNodeOperationType.WRITE], mock_graph[task_idx_2_1][_DAGNodeOperationType.READ], ) _add_edge( mock_graph[task_idx_2_0][_DAGNodeOperationType.WRITE], mock_graph[task_idx_1_1][_DAGNodeOperationType.READ], ) _add_edge( mock_graph[task_idx_2_1][_DAGNodeOperationType.READ], mock_graph[task_idx_2_1][_DAGNodeOperationType.COMPUTE], ) _add_edge( mock_graph[task_idx_2_1][_DAGNodeOperationType.COMPUTE], mock_graph[task_idx_2_1][_DAGNodeOperationType.WRITE], ) _add_edge( mock_graph[task_idx_1_1][_DAGNodeOperationType.READ], mock_graph[task_idx_1_1][_DAGNodeOperationType.COMPUTE], ) _add_edge( mock_graph[task_idx_1_1][_DAGNodeOperationType.COMPUTE], mock_graph[task_idx_1_1][_DAGNodeOperationType.WRITE], ) set_sync_idxs_p2p(mock_graph, task_idx_1_0, task_idx_2_1) set_sync_idxs_p2p(mock_graph, task_idx_2_0, task_idx_1_1) mock_actor_to_candidates = { fake_actor_1._actor_id: [ mock_graph[task_idx_1_0][_DAGNodeOperationType.WRITE], mock_graph[task_idx_1_1][_DAGNodeOperationType.READ], ], fake_actor_2._actor_id: [ mock_graph[task_idx_2_0][_DAGNodeOperationType.WRITE], mock_graph[task_idx_2_1][_DAGNodeOperationType.READ], ], } next_nodes = _select_next_nodes(mock_actor_to_candidates, mock_graph) assert next_nodes == [ mock_graph[task_idx_1_0][_DAGNodeOperationType.WRITE], mock_graph[task_idx_2_1][_DAGNodeOperationType.READ], ] def test_only_one_nccl_collective(self, monkeypatch): """ Simulate the case where there is only one candidate which is a NCCL collective operation. In this case, `_select_next_nodes` should return all the NCCL collective nodes. driver -> fake_actor_1.allreduce_1 -> driver | | -> fake_actor_2.allreduce_1 -> """ monkeypatch.setattr(ActorHandle, "__init__", mock_actor_handle_init) fake_actor_1, dag_idx_1, local_idx_1 = ActorHandle("fake_actor_1"), 1, 0 fake_actor_2, dag_idx_2, local_idx_2 = ActorHandle("fake_actor_2"), 2, 0 mock_graph = { dag_idx_1: generate_dag_graph_nodes( local_idx_1, dag_idx_1, fake_actor_1, requires_nccl_compute=True, ), dag_idx_2: generate_dag_graph_nodes( local_idx_2, dag_idx_2, fake_actor_2, requires_nccl_compute=True, ), } set_sync_idxs_collective(mock_graph, [dag_idx_1, dag_idx_2]) mock_actor_to_candidates = { fake_actor_1._actor_id: [ mock_graph[dag_idx_1][_DAGNodeOperationType.COMPUTE] ], fake_actor_2._actor_id: [ mock_graph[dag_idx_2][_DAGNodeOperationType.COMPUTE] ], } next_nodes = _select_next_nodes(mock_actor_to_candidates, mock_graph) assert set(next_nodes) == { mock_graph[dag_idx_1][_DAGNodeOperationType.COMPUTE], mock_graph[dag_idx_2][_DAGNodeOperationType.COMPUTE], } def test_two_nccl_collectives(self, monkeypatch): """ Simulate the case where there are two candidates that are NCCL collective operations. In this case, `_select_next_nodes` should return all the NCCL collective nodes that are bond earlier. driver -> fake_actor_1.allreduce_1 -> driver | | -> fake_actor_2.allreduce_1 -> | | -> fake_actor_3.allreduce_2 -> | | -> fake_actor_4.allreduce_2 -> """ monkeypatch.setattr(ActorHandle, "__init__", mock_actor_handle_init) fake_actor_1, dag_idx_1, local_idx_1 = ActorHandle("fake_actor_1"), 1, 0 fake_actor_2, dag_idx_2, local_idx_2 = ActorHandle("fake_actor_2"), 2, 0 fake_actor_3, dag_idx_3, local_idx_3 = ActorHandle("fake_actor_3"), 3, 0 fake_actor_4, dag_idx_4, local_idx_4 = ActorHandle("fake_actor_4"), 4, 0 mock_graph = { dag_idx_1: generate_dag_graph_nodes( local_idx_1, dag_idx_1, fake_actor_1, requires_nccl_compute=True, ), dag_idx_2: generate_dag_graph_nodes( local_idx_2, dag_idx_2, fake_actor_2, requires_nccl_compute=True, ), dag_idx_3: generate_dag_graph_nodes( local_idx_3, dag_idx_3, fake_actor_3, requires_nccl_compute=True, ), dag_idx_4: generate_dag_graph_nodes( local_idx_4, dag_idx_4, fake_actor_4, requires_nccl_compute=True, ), } set_sync_idxs_collective(mock_graph, [dag_idx_1, dag_idx_2]) set_sync_idxs_collective(mock_graph, [dag_idx_3, dag_idx_4]) mock_actor_to_candidates = { fake_actor_1._actor_id: [ mock_graph[dag_idx_1][_DAGNodeOperationType.COMPUTE] ], fake_actor_2._actor_id: [ mock_graph[dag_idx_2][_DAGNodeOperationType.COMPUTE] ], fake_actor_3._actor_id: [ mock_graph[dag_idx_3][_DAGNodeOperationType.COMPUTE] ], fake_actor_4._actor_id: [ mock_graph[dag_idx_4][_DAGNodeOperationType.COMPUTE] ], } next_nodes = _select_next_nodes(mock_actor_to_candidates, mock_graph) assert set(next_nodes) == { mock_graph[dag_idx_1][_DAGNodeOperationType.COMPUTE], mock_graph[dag_idx_2][_DAGNodeOperationType.COMPUTE], } next_nodes = _select_next_nodes(mock_actor_to_candidates, mock_graph) assert set(next_nodes) == { mock_graph[dag_idx_3][_DAGNodeOperationType.COMPUTE], mock_graph[dag_idx_4][_DAGNodeOperationType.COMPUTE], } class TestBuildDAGNodeOperationGraph: """ Test whether `_build_dag_node_operation_graph` function adds the correct edges between the nodes in the operation graph base on the 3 rules mentioned in the doc string of `_build_dag_node_operation_graph`. """ def check_edges_between_read_compute_write( self, graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]], task_idx: int, expected_num_edges: List[Tuple[int, int]], ): """ Check whether edges from READ to COMPUTE, and from COMPUTE to WRITE, belonging to the same task are added. Args: graph: The operation graph generated by `_build_dag_node_operation_graph`. task_idx: The global index of the task used to access the task in `idx_to_task`. expected_num_edges: A list of tuples where each tuple contains the expected number of in-edges and out-edges for READ, COMPUTE, and WRITE operations. """ assert len(expected_num_edges) == 3 assert len(graph[task_idx]) == 3 read_node = graph[task_idx][_DAGNodeOperationType.READ] compute_node = graph[task_idx][_DAGNodeOperationType.COMPUTE] write_node = graph[task_idx][_DAGNodeOperationType.WRITE] for idx, node in enumerate([read_node, compute_node, write_node]): assert node.in_degree == expected_num_edges[idx][0] assert len(node.out_edges) == expected_num_edges[idx][1] assert (task_idx, _DAGNodeOperationType.COMPUTE) in read_node.out_edges assert (task_idx, _DAGNodeOperationType.READ) in compute_node.in_edges assert (task_idx, _DAGNodeOperationType.WRITE) in compute_node.out_edges assert (task_idx, _DAGNodeOperationType.COMPUTE) in write_node.in_edges def check_edge_between_writer_and_reader( self, graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]], writer_task_idx: int, reader_task_idx: int, ): """ Check whether the edge from writer's WRITE to reader's READ operation is added. Args: graph: The operation graph generated by `_build_dag_node_operation_graph`. writer_task_idx: The index of the task used to access the task that the writer belongs to in `idx_to_task`. reader_task_idx: The index of the task used to access the task that the reader belongs to in `idx_to_task`. """ write_node = graph[writer_task_idx][_DAGNodeOperationType.WRITE] read_node = graph[reader_task_idx][_DAGNodeOperationType.READ] assert (reader_task_idx, _DAGNodeOperationType.READ) in write_node.out_edges assert (writer_task_idx, _DAGNodeOperationType.WRITE) in read_node.in_edges def check_edge_between_compute_nodes( self, graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]], task_idx_1: int, task_idx_2: int, ): """ Check whether the edge from COMPUTE with `bind_index` i to COMPUTE with `bind_index` i+1 if they belong to the same actor. Args: graph: The operation graph generated by `_build_dag_node_operation_graph`. task_idx_1: The index of the task used to access the task in `idx_to_task`. task_idx_2: The index of the task used to access the task in `idx_to_task`. Note that both tasks belong to the same actor, and the `bind_index` of the second task is equal to the `bind_index` of the first task plus one. """ compute_node_1 = graph[task_idx_1][_DAGNodeOperationType.COMPUTE] compute_node_2 = graph[task_idx_2][_DAGNodeOperationType.COMPUTE] assert (task_idx_2, _DAGNodeOperationType.COMPUTE) in compute_node_1.out_edges assert (task_idx_1, _DAGNodeOperationType.COMPUTE) in compute_node_2.in_edges def test_edges_between_read_compute_write(self, monkeypatch): """ driver -> fake_actor.op -> driver This test case aims to verify whether the function correctly adds edges between READ/COMPUTE and COMPUTE/WRITE operations on the same actor. """ monkeypatch.setattr(ClassMethodNode, "__init__", mock_class_method_call_init) monkeypatch.setattr(MultiOutputNode, "__init__", mock_init) idx_to_task = { 0: CompiledTask(0, InputNode()), 1: CompiledTask(1, ClassMethodNode()), 2: CompiledTask(2, MultiOutputNode()), } fake_actor = "fake_actor" task_idx = 1 actor_to_operation_nodes = { fake_actor: [ list(generate_dag_graph_nodes(0, task_idx, fake_actor).values()) ] } graph = _build_dag_node_operation_graph(idx_to_task, actor_to_operation_nodes) assert len(graph) == 1 self.check_edges_between_read_compute_write( graph, task_idx, [(0, 1), (1, 1), (1, 0)] ) def test_edge_between_writer_and_reader(self, monkeypatch): """ driver -> fake_actor_1.op -> fake_actor_2.op -> driver This test case aims to verify whether the function correctly adds an edge from the writer's WRITE operation to the reader's READ operation. """ monkeypatch.setattr(ClassMethodNode, "__init__", mock_class_method_call_init) monkeypatch.setattr(MultiOutputNode, "__init__", mock_init) fake_actor_1, task_idx_1 = "fake_actor_1", 1 fake_actor_2, task_idx_2 = "fake_actor_2", 2 idx_to_task = { 0: CompiledTask(0, InputNode()), 1: CompiledTask(1, ClassMethodNode()), 2: CompiledTask(2, ClassMethodNode()), 3: CompiledTask(3, MultiOutputNode()), } idx_to_task[1].downstream_task_idxs = {2: fake_actor_2} actor_to_operation_nodes = { fake_actor_1: [ list(generate_dag_graph_nodes(0, task_idx_1, fake_actor_1).values()) ], fake_actor_2: [ list(generate_dag_graph_nodes(0, task_idx_2, fake_actor_2).values()) ], } graph = _build_dag_node_operation_graph(idx_to_task, actor_to_operation_nodes) assert len(graph) == 2 self.check_edges_between_read_compute_write( graph, task_idx_1, [(0, 1), (1, 1), (1, 1)] ) self.check_edges_between_read_compute_write( graph, task_idx_2, [(1, 1), (1, 1), (1, 0)] ) self.check_edge_between_writer_and_reader(graph, task_idx_1, task_idx_2) def test_edge_between_compute_nodes(self, monkeypatch): """ driver -> fake_actor.op -> fake_actor.op -> driver This test case aims to verify whether the function correctly adds an edge from the COMPUTE operation with `bind_index` i to the COMPUTE operation with `bind_index` i+1 if they belong to the same actor. """ monkeypatch.setattr(ClassMethodNode, "__init__", mock_class_method_call_init) monkeypatch.setattr(MultiOutputNode, "__init__", mock_init) fake_actor = "fake_actor" task_idx_1, task_idx_2 = 1, 2 idx_to_task = { 0: CompiledTask(0, InputNode()), task_idx_1: CompiledTask(task_idx_1, ClassMethodNode()), task_idx_2: CompiledTask(task_idx_2, ClassMethodNode()), 3: CompiledTask(3, MultiOutputNode()), } idx_to_task[task_idx_1].downstream_task_idxs = {task_idx_2: fake_actor} actor_to_operation_nodes = { fake_actor: [ list(generate_dag_graph_nodes(0, task_idx_1, fake_actor).values()), list(generate_dag_graph_nodes(1, task_idx_2, fake_actor).values()), ], } graph = _build_dag_node_operation_graph(idx_to_task, actor_to_operation_nodes) assert len(graph) == 2 self.check_edges_between_read_compute_write( graph, task_idx_1, [(0, 1), (1, 2), (1, 1)] ) self.check_edges_between_read_compute_write( graph, task_idx_2, [(1, 1), (2, 1), (1, 0)] ) self.check_edge_between_writer_and_reader(graph, task_idx_1, task_idx_2) self.check_edge_between_compute_nodes(graph, task_idx_1, task_idx_2) def test_two_actors(self, monkeypatch): """ driver -> fake_actor_1.op -> fake_actor_2.op -> driver | | -> fake_actor_2.op -> fake_actor_1.op - This test includes two actors, each with two tasks. The test case covers all three rules for adding edges between operation nodes in the operation graph. """ monkeypatch.setattr(ClassMethodNode, "__init__", mock_class_method_call_init) monkeypatch.setattr(MultiOutputNode, "__init__", mock_init) fake_actor_1, task_idx_1, task_idx_3 = "fake_actor_1", 1, 3 fake_actor_2, task_idx_2, task_idx_4 = "fake_actor_2", 2, 4 idx_to_task = { 0: CompiledTask(0, InputNode()), task_idx_1: CompiledTask(task_idx_1, ClassMethodNode()), task_idx_2: CompiledTask(task_idx_2, ClassMethodNode()), task_idx_3: CompiledTask(task_idx_3, ClassMethodNode()), task_idx_4: CompiledTask(task_idx_4, ClassMethodNode()), 5: CompiledTask(5, MultiOutputNode()), } idx_to_task[task_idx_1].downstream_task_idxs = {task_idx_4: fake_actor_2} idx_to_task[task_idx_2].downstream_task_idxs = {task_idx_3: fake_actor_1} actor_to_operation_nodes = { fake_actor_1: [ list(generate_dag_graph_nodes(0, task_idx_1, fake_actor_1).values()), list(generate_dag_graph_nodes(1, task_idx_3, fake_actor_1).values()), ], fake_actor_2: [ list(generate_dag_graph_nodes(0, task_idx_2, fake_actor_2).values()), list(generate_dag_graph_nodes(1, task_idx_4, fake_actor_2).values()), ], } graph = _build_dag_node_operation_graph(idx_to_task, actor_to_operation_nodes) assert len(graph) == 4 self.check_edges_between_read_compute_write( graph, task_idx_1, [(0, 1), (1, 2), (1, 1)] ) self.check_edges_between_read_compute_write( graph, task_idx_2, [(0, 1), (1, 2), (1, 1)] ) self.check_edges_between_read_compute_write( graph, task_idx_3, [(1, 1), (2, 1), (1, 0)] ) self.check_edges_between_read_compute_write( graph, task_idx_4, [(1, 1), (2, 1), (1, 0)] ) self.check_edge_between_writer_and_reader(graph, task_idx_1, task_idx_4) self.check_edge_between_writer_and_reader(graph, task_idx_2, task_idx_3) class TestGenerateActorToExecutionSchedule: """ Test whether `_generate_actor_to_execution_schedule` function generates the correct execution schedule for each actor. """ def add_edge_between_read_compute_write( self, operations: Dict[_DAGNodeOperationType, _DAGOperationGraphNode] ): """ Add edges between READ and COMPUTE, and between COMPUTE and WRITE operations on the same actor. Args: operations: A dictionary where the key is the operation type and the value is the operation node. """ assert len(operations) == 3 _add_edge( operations[_DAGNodeOperationType.READ], operations[_DAGNodeOperationType.COMPUTE], ) _add_edge( operations[_DAGNodeOperationType.COMPUTE], operations[_DAGNodeOperationType.WRITE], ) def add_data_dependeny( self, writer_operations: Dict[_DAGNodeOperationType, _DAGOperationGraphNode], reader_operations: Dict[_DAGNodeOperationType, _DAGOperationGraphNode], ): """ Add a data dependency between the WRITE operation of the writer and the READ operation of the reader. Args: writer_operations: A dictionary where the key is the operation type and the value is the operation node of the writer. reader_operations: A dictionary where the key is the operation type and the value is the operation node of the reader. """ _add_edge( writer_operations[_DAGNodeOperationType.WRITE], reader_operations[_DAGNodeOperationType.READ], ) def add_control_dependency( self, operations_1: Dict[_DAGNodeOperationType, _DAGOperationGraphNode], operations_2: Dict[_DAGNodeOperationType, _DAGOperationGraphNode], ): """ Add a control dependency between the COMPUTE operation of the task with bind_index i and the COMPUTE operation of the task with bind_index i+1 on the same actor. Args: operations_1: A dictionary where the key is the operation type and the value is the operation node of the task with bind_index i. operations_2: A dictionary where the key is the operation type and the value is the operation node of the task with bind_index i+1. """ _add_edge( operations_1[_DAGNodeOperationType.COMPUTE], operations_2[_DAGNodeOperationType.COMPUTE], ) def test_single_actor_1(self, monkeypatch): """ driver -> fake_actor.op (task_idx_1) -> fake_actor.op (task_idx_2) -> driver Test the case where there is only one actor and no NCCL operations. Because there is no NCCL operation, all operations with smaller `bind_index` should be executed before the operations with larger `bind_index` on the same actor. """ monkeypatch.setattr(ActorHandle, "__init__", mock_actor_handle_init) fake_actor = ActorHandle("fake_actor") task_idx_1, exec_task_idx_1 = 1, 0 task_idx_2, exec_task_idx_2 = 2, 1 graph = { task_idx_1: generate_dag_graph_nodes( exec_task_idx_1, task_idx_1, fake_actor ), task_idx_2: generate_dag_graph_nodes( exec_task_idx_2, task_idx_2, fake_actor ), } self.add_edge_between_read_compute_write(graph[task_idx_1]) self.add_edge_between_read_compute_write(graph[task_idx_2]) self.add_data_dependeny(graph[task_idx_1], graph[task_idx_2]) self.add_control_dependency(graph[task_idx_1], graph[task_idx_2]) actor_to_execution_schedule = _generate_and_extract_execution_schedule(graph) assert len(actor_to_execution_schedule) == 1 assert len(actor_to_execution_schedule[fake_actor]) == 6 assert actor_to_execution_schedule[fake_actor] == [ graph[task_idx_1][_DAGNodeOperationType.READ].operation, graph[task_idx_1][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_1][_DAGNodeOperationType.WRITE].operation, graph[task_idx_2][_DAGNodeOperationType.READ].operation, graph[task_idx_2][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_2][_DAGNodeOperationType.WRITE].operation, ] def test_single_actor_2(self, monkeypatch): """ driver -> fake_actor.op (task_idx_1) -> fake_actor.op (task_idx_2) -> driver | | -> fake_actor.op (task_idx_3) - When the `dad_idx_1.WRITE` operation is picked, both `task_idx_2.READ` and `task_idx_3.READ` operations should be zero in-degree. In this case, the one with the smaller `bind_index` should be selected first. That is, `task_idx_2.READ` should be selected first. """ monkeypatch.setattr(ActorHandle, "__init__", mock_actor_handle_init) fake_actor = ActorHandle("fake_actor") task_idx_1, exec_task_idx_1 = 1, 0 task_idx_2, exec_task_idx_2 = 2, 1 task_idx_3, exec_task_idx_3 = 3, 2 graph = { task_idx_1: generate_dag_graph_nodes( exec_task_idx_1, task_idx_1, fake_actor ), task_idx_2: generate_dag_graph_nodes( exec_task_idx_2, task_idx_2, fake_actor ), task_idx_3: generate_dag_graph_nodes( exec_task_idx_3, task_idx_3, fake_actor ), } self.add_edge_between_read_compute_write(graph[task_idx_1]) self.add_edge_between_read_compute_write(graph[task_idx_2]) self.add_edge_between_read_compute_write(graph[task_idx_3]) self.add_data_dependeny(graph[task_idx_1], graph[task_idx_2]) self.add_data_dependeny(graph[task_idx_1], graph[task_idx_3]) self.add_control_dependency(graph[task_idx_1], graph[task_idx_2]) self.add_control_dependency(graph[task_idx_2], graph[task_idx_3]) actor_to_execution_schedule = _generate_and_extract_execution_schedule(graph) assert len(actor_to_execution_schedule) == 1 assert len(actor_to_execution_schedule[fake_actor]) == 9 assert actor_to_execution_schedule[fake_actor] == [ graph[task_idx_1][_DAGNodeOperationType.READ].operation, graph[task_idx_1][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_1][_DAGNodeOperationType.WRITE].operation, graph[task_idx_2][_DAGNodeOperationType.READ].operation, graph[task_idx_2][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_2][_DAGNodeOperationType.WRITE].operation, graph[task_idx_3][_DAGNodeOperationType.READ].operation, graph[task_idx_3][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_3][_DAGNodeOperationType.WRITE].operation, ] def test_two_actors_no_nccl(self, monkeypatch): """ driver -> actor_1.op (task_idx_1_1) -> actor_2.op (task_idx_2_2) -> driver | | -> actor_2.op (task_idx_2_1) -> actor_1.op (task_idx_1_2) - Test the case where there are two actors and no NCCL operations. Because there is no NCCL operation, all operations with smaller `bind_index` should be executed before the operations with larger `bind_index` on the same actor. """ monkeypatch.setattr(ActorHandle, "__init__", mock_actor_handle_init) fake_actor_1 = ActorHandle("fake_actor_1") task_idx_1_1, exec_task_idx_1_1 = 1, 0 task_idx_1_2, exec_task_idx_1_2 = 4, 1 fake_actor_2 = ActorHandle("fake_actor_2") task_idx_2_1, exec_task_idx_2_1 = 2, 0 task_idx_2_2, exec_task_idx_2_2 = 3, 1 graph = { task_idx_1_1: generate_dag_graph_nodes( exec_task_idx_1_1, task_idx_1_1, fake_actor_1 ), task_idx_2_1: generate_dag_graph_nodes( exec_task_idx_2_1, task_idx_2_1, fake_actor_2 ), task_idx_2_2: generate_dag_graph_nodes( exec_task_idx_2_2, task_idx_2_2, fake_actor_2 ), task_idx_1_2: generate_dag_graph_nodes( exec_task_idx_1_2, task_idx_1_2, fake_actor_1 ), } self.add_edge_between_read_compute_write(graph[task_idx_1_1]) self.add_edge_between_read_compute_write(graph[task_idx_1_2]) self.add_edge_between_read_compute_write(graph[task_idx_2_1]) self.add_edge_between_read_compute_write(graph[task_idx_2_2]) self.add_data_dependeny(graph[task_idx_1_1], graph[task_idx_2_2]) self.add_data_dependeny(graph[task_idx_2_1], graph[task_idx_1_2]) self.add_control_dependency(graph[task_idx_1_1], graph[task_idx_1_2]) self.add_control_dependency(graph[task_idx_2_1], graph[task_idx_2_2]) actor_to_execution_schedule = _generate_and_extract_execution_schedule(graph) assert len(actor_to_execution_schedule) == 2 assert len(actor_to_execution_schedule[fake_actor_1]) == 6 assert len(actor_to_execution_schedule[fake_actor_2]) == 6 assert actor_to_execution_schedule[fake_actor_1] == [ graph[task_idx_1_1][_DAGNodeOperationType.READ].operation, graph[task_idx_1_1][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_1_1][_DAGNodeOperationType.WRITE].operation, graph[task_idx_1_2][_DAGNodeOperationType.READ].operation, graph[task_idx_1_2][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_1_2][_DAGNodeOperationType.WRITE].operation, ] assert actor_to_execution_schedule[fake_actor_2] == [ graph[task_idx_2_1][_DAGNodeOperationType.READ].operation, graph[task_idx_2_1][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_2_1][_DAGNodeOperationType.WRITE].operation, graph[task_idx_2_2][_DAGNodeOperationType.READ].operation, graph[task_idx_2_2][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_2_2][_DAGNodeOperationType.WRITE].operation, ] def test_two_actors_with_nccl(self, monkeypatch): """ driver -> actor_1.op (task_idx_1_1) -> actor_2.op (task_idx_2_2) -> driver | | -> actor_2.op (task_idx_2_1) -> actor_1.op (task_idx_1_2) - In this test, the communication between fake_actor_1 and fake_actor_2 is done using NCCL. When the task_idx_1.WRITE operation is picked, the task_idx_2.READ operation is also added to the execution schedule because of the NCCL operation. """ monkeypatch.setattr(ActorHandle, "__init__", mock_actor_handle_init) fake_actor_1 = ActorHandle("fake_actor_1") task_idx_1_1, exec_task_idx_1_1 = 1, 0 task_idx_1_2, exec_task_idx_1_2 = 4, 1 fake_actor_2 = ActorHandle("fake_actor_2") task_idx_2_1, exec_task_idx_2_1 = 2, 0 task_idx_2_2, exec_task_idx_2_2 = 3, 1 graph = { task_idx_1_1: generate_dag_graph_nodes( exec_task_idx_1_1, task_idx_1_1, fake_actor_1, requires_nccl_write=True, ), task_idx_2_1: generate_dag_graph_nodes( exec_task_idx_2_1, task_idx_2_1, fake_actor_2, requires_nccl_write=True, ), task_idx_2_2: generate_dag_graph_nodes( exec_task_idx_2_2, task_idx_2_2, fake_actor_2, requires_nccl_read=True, ), task_idx_1_2: generate_dag_graph_nodes( exec_task_idx_1_2, task_idx_1_2, fake_actor_1, requires_nccl_read=True, ), } set_sync_idxs_p2p(graph, task_idx_1_1, task_idx_2_2) set_sync_idxs_p2p(graph, task_idx_2_1, task_idx_1_2) self.add_edge_between_read_compute_write(graph[task_idx_1_1]) self.add_edge_between_read_compute_write(graph[task_idx_1_2]) self.add_edge_between_read_compute_write(graph[task_idx_2_1]) self.add_edge_between_read_compute_write(graph[task_idx_2_2]) self.add_data_dependeny(graph[task_idx_1_1], graph[task_idx_2_2]) self.add_data_dependeny(graph[task_idx_2_1], graph[task_idx_1_2]) self.add_control_dependency(graph[task_idx_1_1], graph[task_idx_1_2]) self.add_control_dependency(graph[task_idx_2_1], graph[task_idx_2_2]) actor_to_execution_schedule = _generate_and_extract_execution_schedule(graph) assert len(actor_to_execution_schedule) == 2 assert len(actor_to_execution_schedule[fake_actor_1]) == 6 assert len(actor_to_execution_schedule[fake_actor_2]) == 6 assert actor_to_execution_schedule[fake_actor_1] == [ graph[task_idx_1_1][_DAGNodeOperationType.READ].operation, graph[task_idx_1_1][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_1_1][_DAGNodeOperationType.WRITE].operation, graph[task_idx_1_2][_DAGNodeOperationType.READ].operation, graph[task_idx_1_2][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_1_2][_DAGNodeOperationType.WRITE].operation, ] assert actor_to_execution_schedule[fake_actor_2] == [ # `actor_2.task_idx_2_2.READ` (P2P recv) is scheduled together with # `actor_1.task_idx_1_1.WRITE` (P2P send). graph[task_idx_2_2][_DAGNodeOperationType.READ].operation, graph[task_idx_2_1][_DAGNodeOperationType.READ].operation, graph[task_idx_2_1][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_2_1][_DAGNodeOperationType.WRITE].operation, graph[task_idx_2_2][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_2_2][_DAGNodeOperationType.WRITE].operation, ] def test_simulate_pp_2workers_2batches_1f1b_with_nccl(self, monkeypatch): """ This test simulates a simple 1F1B pipeline parallelism for training with 2 workers and 2 batches. w1: fwd_b1 fwd_b2 bwd_b1 bwd_b2 w2: fwd_b1 bwd_b1 fwd_b2 bwd_b2 The communication between workers is done using NCCL. The communication within the worker actor is done using IntraProcessChannel. """ monkeypatch.setattr(ActorHandle, "__init__", mock_actor_handle_init) worker_1 = ActorHandle("worker_1") task_idx_1_1, exec_task_idx_1_1 = 1, 0 task_idx_1_2, exec_task_idx_1_2 = 2, 1 task_idx_1_3, exec_task_idx_1_3 = 3, 2 task_idx_1_4, exec_task_idx_1_4 = 4, 3 worker_2 = ActorHandle("worker_2") task_idx_2_1, exec_task_idx_2_1 = 5, 0 task_idx_2_2, exec_task_idx_2_2 = 6, 1 task_idx_2_3, exec_task_idx_2_3 = 7, 2 task_idx_2_4, exec_task_idx_2_4 = 8, 3 graph = { task_idx_1_1: generate_dag_graph_nodes( exec_task_idx_1_1, task_idx_1_1, worker_1, requires_nccl_write=True, ), task_idx_1_2: generate_dag_graph_nodes( exec_task_idx_1_2, task_idx_1_2, worker_1, requires_nccl_write=True, ), task_idx_1_3: generate_dag_graph_nodes( exec_task_idx_1_3, task_idx_1_3, worker_1, requires_nccl_read=True, ), task_idx_1_4: generate_dag_graph_nodes( exec_task_idx_1_4, task_idx_1_4, worker_1, requires_nccl_read=True, ), task_idx_2_1: generate_dag_graph_nodes( exec_task_idx_2_1, task_idx_2_1, worker_2, requires_nccl_read=True, ), task_idx_2_2: generate_dag_graph_nodes( exec_task_idx_2_2, task_idx_2_2, worker_2, requires_nccl_write=True, ), task_idx_2_3: generate_dag_graph_nodes( exec_task_idx_2_3, task_idx_2_3, worker_2, requires_nccl_read=True, ), task_idx_2_4: generate_dag_graph_nodes( exec_task_idx_2_4, task_idx_2_4, worker_2, requires_nccl_write=True, ), } set_sync_idxs_p2p(graph, task_idx_1_1, task_idx_2_1) set_sync_idxs_p2p(graph, task_idx_1_2, task_idx_2_3) set_sync_idxs_p2p(graph, task_idx_2_2, task_idx_1_3) set_sync_idxs_p2p(graph, task_idx_2_4, task_idx_1_4) self.add_edge_between_read_compute_write(graph[task_idx_1_1]) self.add_edge_between_read_compute_write(graph[task_idx_1_2]) self.add_edge_between_read_compute_write(graph[task_idx_1_3]) self.add_edge_between_read_compute_write(graph[task_idx_1_4]) self.add_edge_between_read_compute_write(graph[task_idx_2_1]) self.add_edge_between_read_compute_write(graph[task_idx_2_2]) self.add_edge_between_read_compute_write(graph[task_idx_2_3]) self.add_edge_between_read_compute_write(graph[task_idx_2_4]) self.add_data_dependeny(graph[task_idx_1_1], graph[task_idx_2_1]) self.add_data_dependeny(graph[task_idx_2_1], graph[task_idx_2_2]) self.add_data_dependeny(graph[task_idx_2_2], graph[task_idx_1_3]) self.add_data_dependeny(graph[task_idx_1_2], graph[task_idx_2_3]) self.add_data_dependeny(graph[task_idx_2_3], graph[task_idx_2_4]) self.add_data_dependeny(graph[task_idx_2_4], graph[task_idx_1_4]) self.add_control_dependency(graph[task_idx_1_1], graph[task_idx_1_2]) self.add_control_dependency(graph[task_idx_1_2], graph[task_idx_1_3]) self.add_control_dependency(graph[task_idx_1_3], graph[task_idx_1_4]) self.add_control_dependency(graph[task_idx_2_1], graph[task_idx_2_2]) self.add_control_dependency(graph[task_idx_2_2], graph[task_idx_2_3]) self.add_control_dependency(graph[task_idx_2_3], graph[task_idx_2_4]) actor_to_execution_schedule = _generate_and_extract_execution_schedule(graph) assert len(actor_to_execution_schedule) == 2 assert len(actor_to_execution_schedule[worker_1]) == 12 assert len(actor_to_execution_schedule[worker_2]) == 12 assert actor_to_execution_schedule[worker_1] == [ graph[task_idx_1_1][_DAGNodeOperationType.READ].operation, graph[task_idx_1_1][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_1_1][_DAGNodeOperationType.WRITE].operation, graph[task_idx_1_2][_DAGNodeOperationType.READ].operation, graph[task_idx_1_2][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_1_2][_DAGNodeOperationType.WRITE].operation, graph[task_idx_1_3][_DAGNodeOperationType.READ].operation, graph[task_idx_1_3][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_1_3][_DAGNodeOperationType.WRITE].operation, graph[task_idx_1_4][_DAGNodeOperationType.READ].operation, graph[task_idx_1_4][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_1_4][_DAGNodeOperationType.WRITE].operation, ] assert actor_to_execution_schedule[worker_2] == [ graph[task_idx_2_1][_DAGNodeOperationType.READ].operation, graph[task_idx_2_1][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_2_1][_DAGNodeOperationType.WRITE].operation, # `actor_2.task_idx_2_3.READ` (P2P recv) is scheduled together with # `actor_1.task_idx_1_2.WRITE` (P2P send). graph[task_idx_2_3][_DAGNodeOperationType.READ].operation, graph[task_idx_2_2][_DAGNodeOperationType.READ].operation, graph[task_idx_2_2][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_2_2][_DAGNodeOperationType.WRITE].operation, graph[task_idx_2_3][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_2_3][_DAGNodeOperationType.WRITE].operation, graph[task_idx_2_4][_DAGNodeOperationType.READ].operation, graph[task_idx_2_4][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_2_4][_DAGNodeOperationType.WRITE].operation, ] def test_simulate_pp_2workers_2batches_1f1b_no_nccl(self, monkeypatch): """ This test simulates a simple 1F1B pipeline parallelism for training with 2 workers and 2 batches. w1: fwd_b1 fwd_b2 bwd_b1 bwd_b2 w2: fwd_b1 bwd_b1 fwd_b2 bwd_b2 Because there is no NCCL operation, all operations with smaller `bind_index` should be executed before the operations with larger `bind_index` on the same actor. """ monkeypatch.setattr(ActorHandle, "__init__", mock_actor_handle_init) worker_1 = ActorHandle("worker_1") task_idx_1_1, exec_task_idx_1_1 = 1, 0 task_idx_1_2, exec_task_idx_1_2 = 2, 1 task_idx_1_3, exec_task_idx_1_3 = 3, 2 task_idx_1_4, exec_task_idx_1_4 = 4, 3 worker_2 = ActorHandle("worker_2") task_idx_2_1, exec_task_idx_2_1 = 5, 0 task_idx_2_2, exec_task_idx_2_2 = 6, 1 task_idx_2_3, exec_task_idx_2_3 = 7, 2 task_idx_2_4, exec_task_idx_2_4 = 8, 3 # No NCCL operation. graph = { task_idx_1_1: generate_dag_graph_nodes( exec_task_idx_1_1, task_idx_1_1, worker_1 ), task_idx_1_2: generate_dag_graph_nodes( exec_task_idx_1_2, task_idx_1_2, worker_1 ), task_idx_1_3: generate_dag_graph_nodes( exec_task_idx_1_3, task_idx_1_3, worker_1 ), task_idx_1_4: generate_dag_graph_nodes( exec_task_idx_1_4, task_idx_1_4, worker_1 ), task_idx_2_1: generate_dag_graph_nodes( exec_task_idx_2_1, task_idx_2_1, worker_2 ), task_idx_2_2: generate_dag_graph_nodes( exec_task_idx_2_2, task_idx_2_2, worker_2 ), task_idx_2_3: generate_dag_graph_nodes( exec_task_idx_2_3, task_idx_2_3, worker_2 ), task_idx_2_4: generate_dag_graph_nodes( exec_task_idx_2_4, task_idx_2_4, worker_2 ), } self.add_edge_between_read_compute_write(graph[task_idx_1_1]) self.add_edge_between_read_compute_write(graph[task_idx_1_2]) self.add_edge_between_read_compute_write(graph[task_idx_1_3]) self.add_edge_between_read_compute_write(graph[task_idx_1_4]) self.add_edge_between_read_compute_write(graph[task_idx_2_1]) self.add_edge_between_read_compute_write(graph[task_idx_2_2]) self.add_edge_between_read_compute_write(graph[task_idx_2_3]) self.add_edge_between_read_compute_write(graph[task_idx_2_4]) self.add_data_dependeny(graph[task_idx_1_1], graph[task_idx_2_1]) self.add_data_dependeny(graph[task_idx_2_1], graph[task_idx_2_2]) self.add_data_dependeny(graph[task_idx_2_2], graph[task_idx_1_3]) self.add_data_dependeny(graph[task_idx_1_2], graph[task_idx_2_3]) self.add_data_dependeny(graph[task_idx_2_3], graph[task_idx_2_4]) self.add_data_dependeny(graph[task_idx_2_4], graph[task_idx_1_4]) self.add_control_dependency(graph[task_idx_1_1], graph[task_idx_1_2]) self.add_control_dependency(graph[task_idx_1_2], graph[task_idx_1_3]) self.add_control_dependency(graph[task_idx_1_3], graph[task_idx_1_4]) self.add_control_dependency(graph[task_idx_2_1], graph[task_idx_2_2]) self.add_control_dependency(graph[task_idx_2_2], graph[task_idx_2_3]) self.add_control_dependency(graph[task_idx_2_3], graph[task_idx_2_4]) actor_to_execution_schedule = _generate_and_extract_execution_schedule(graph) assert len(actor_to_execution_schedule) == 2 assert len(actor_to_execution_schedule[worker_1]) == 12 assert len(actor_to_execution_schedule[worker_2]) == 12 assert actor_to_execution_schedule[worker_1] == [ graph[task_idx_1_1][_DAGNodeOperationType.READ].operation, graph[task_idx_1_1][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_1_1][_DAGNodeOperationType.WRITE].operation, graph[task_idx_1_2][_DAGNodeOperationType.READ].operation, graph[task_idx_1_2][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_1_2][_DAGNodeOperationType.WRITE].operation, graph[task_idx_1_3][_DAGNodeOperationType.READ].operation, graph[task_idx_1_3][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_1_3][_DAGNodeOperationType.WRITE].operation, graph[task_idx_1_4][_DAGNodeOperationType.READ].operation, graph[task_idx_1_4][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_1_4][_DAGNodeOperationType.WRITE].operation, ] assert actor_to_execution_schedule[worker_2] == [ graph[task_idx_2_1][_DAGNodeOperationType.READ].operation, graph[task_idx_2_1][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_2_1][_DAGNodeOperationType.WRITE].operation, graph[task_idx_2_2][_DAGNodeOperationType.READ].operation, graph[task_idx_2_2][_DAGNodeOperationType.COMPUTE].operation, # The order of `task_idx_2_3.READ` and `task_idx_2_2.WRITE` is important. # It is different from the case where there is an NCCL operation. graph[task_idx_2_2][_DAGNodeOperationType.WRITE].operation, graph[task_idx_2_3][_DAGNodeOperationType.READ].operation, graph[task_idx_2_3][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_2_3][_DAGNodeOperationType.WRITE].operation, graph[task_idx_2_4][_DAGNodeOperationType.READ].operation, graph[task_idx_2_4][_DAGNodeOperationType.COMPUTE].operation, graph[task_idx_2_4][_DAGNodeOperationType.WRITE].operation, ] if __name__ == "__main__": if os.environ.get("PARALLEL_CI"): sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__])) else: sys.exit(pytest.main(["-sv", __file__]))