Commit fba88ae8 authored by Yin, Junqi's avatar Yin, Junqi

save neighbors and graph matrix

parent 3ecdf289
......@@ -23,7 +23,7 @@ import pcode.utils.logging as logging
from pcode.utils.timer import Timer
import platform
def print_neighbors(conf):
def print_neighbors(conf, save=True):
neighbors_info = conf.graph.get_neighborhood()
neighbor_ranks = [
neighbor_rank
......@@ -32,6 +32,12 @@ def print_neighbors(conf):
]
print("NEIGHBOR_INFO: rank %d is on node %s with neighbors: %s"
% (conf.graph.rank, platform.node(), tuple(neighbor_ranks)))
if save:
conf.logger.log(f"rank: {conf.graph.rank}")
conf.logger.log(f"node: {platform.node()}")
conf.logger.log(f"neighbors: {neighbor_ranks}")
conf.logger.log(f"matrix: {conf.graph.matrix}")
def setup_ddp(backend):
""""Initialize Apex DDP"""
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment