Skip to content
This repository has been archived by the owner on Nov 25, 2024. It is now read-only.

Commit

Permalink
Add reset_communicators to clean the communicators
Browse files Browse the repository at this point in the history
It is necessary when users use wholegraph back-to-back, i.e., use
wholegraph in pytest and needs to run it multiple times (wg.init ->
wg.finalize -> wg.init -> wg.finalize...)
  • Loading branch information
chang-l committed Dec 6, 2023
1 parent 0c61783 commit 54ef09c
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 1 deletion.
11 changes: 11 additions & 0 deletions python/pylibwholegraph/pylibwholegraph/torch/comm.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,17 @@
all_comm_local_rank = 0
all_comm_local_size = 1

def reset_communicators():
global all_comm_world_rank, all_comm_world_size, all_comm_local_rank, all_comm_local_size
global global_communicators, local_node_communicator, local_device_communicator
global_communicators = {}
local_node_communicator = None
local_device_communicator = None

all_comm_world_rank = 0
all_comm_world_size = 1
all_comm_local_rank = 0
all_comm_local_size = 1

def set_world_info(world_rank: int, world_size: int, local_rank: int, local_size: int):
"""
Expand Down
3 changes: 2 additions & 1 deletion python/pylibwholegraph/pylibwholegraph/torch/initialize.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
import torch
import torch.utils.dlpack
import pylibwholegraph.binding.wholememory_binding as wmb
from .comm import set_world_info, get_global_communicator, get_local_node_communicator
from .comm import set_world_info, get_global_communicator, get_local_node_communicator, reset_communicators


def init(world_rank: int, world_size: int, local_rank: int, local_size: int):
Expand Down Expand Up @@ -73,3 +73,4 @@ def finalize():
:return: None
"""
wmb.finalize()
reset_communicators()

0 comments on commit 54ef09c

Please sign in to comment.