Source code for triqs.utility.mpi_mpi4py

# Copyright (c) 2013 Commissariat à l'énergie atomique et aux énergies alternatives (CEA)
# Copyright (c) 2013 Centre national de la recherche scientifique (CNRS)
# Copyright (c) 2019-2023 Simons Foundation
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You may obtain a copy of the License at
#     https:#www.gnu.org/licenses/gpl-3.0.txt
#
# Authors: John Bonini, Alexander Hampel, Olivier Parcollet, Nils Wentzell


import os,sys,datetime, time
myprint_err = lambda x : sys.stderr.write("%s\n"%x)
myprint_out = lambda x : sys.stdout.write("%s\n"%x)

from mpi4py import MPI

world = MPI.COMM_WORLD
rank = world.Get_rank()
size = world.Get_size()

if rank==0 :
    myprint_err ("Starting run with %s MPI rank(s) at : %s"%(size,str(datetime.datetime.now())))

master =0

All_Nodes_report = False

[docs] def bcast(x, root = 0): return world.bcast(x, root = root)
[docs] def send(val, dest): world.send(val, dest = dest)
#print "node ", rank, " has sent ", val, " to node ", dest
[docs] def recv(source = 0): r = world.recv(source = source) #print "node ", rank, " receives ", r, " from node ", source return r
[docs] def barrier(poll_msec=1): """ Use asynchronous synchronization, otherwise mpi.barrier uses up all the CPU time during the run of subprocess. if poll_msec is 0.0 or None the default mpi barrier of mpi4py is used. Parameters ---------- poll_msec: float, time step for pinging the status of the sleeping ranks in msec """ if not poll_msec: world.barrier() else: req = world.Ibarrier() while not req.Test(): time.sleep(poll_msec / 1000)
[docs] def all_reduce(x, comm=world, op=MPI.SUM): try: return comm.allreduce(x, op=op) except (AttributeError, TypeError) as e: if isinstance(x, MPI.Intracomm): report( "WARNING: the signature for all_reduce is now 'all_reduce(x, comm=world, op=MPI.SUM)'\n\tattempting compatibility with old signature" ) return world.allreduce(comm) raise e
Verbosity_Level_Report_Max = 1
[docs] def report(*x,**opt): """ print on the master, do nothing on the nodes """ try : if opt['Verbosity_Level'] > Verbosity_Level_Report_Max : return except : pass if All_Nodes_report and rank!=0 : for y in x: open('report_node%s'%rank,'a').write(str(y) + '\n') # print y if rank==0 : myprint,myflush = (myprint_err,sys.stderr.flush) if 'stderr' in opt and opt['stderr'] else (myprint_out,sys.stdout.flush) for y in x: myprint( y) # open('report','a').write(str(y) + '\n') # print y myflush() # be sure to flush the buffer!
[docs] def slice_inf(imin,imax) : j=(imax - imin + 1)//size i= imax - imin + 1 - size*j return imin + rank*(j+1) if rank<=i-1 else imin + rank*j + i
[docs] def slice_sup(imin,imax) : j=(imax - imin + 1)//size i= imax - imin + 1 - size*j; return imin + (rank+1)*(j+1) -1 if rank<=i-1 else imin + (rank+1)*j + i - 1
[docs] def slice_array(A) : """Given an array A, it returns a VIEW of a slice over the first dim on the node""" imax = A.shape[0]-1 return A[slice_inf(0,imax):slice_sup(0,imax)+1] # +1 due to the slice convention
[docs] def is_master_node(): return rank == master
HostNames = {}
[docs] def master_gets_host_names(): """ find the host name of all the nodes """ from socket import gethostname global HostNames HostNames[rank] = gethostname() if is_master_node() : for proc in range (1,size) : HostNames[proc],status = recv(proc) else: send(gethostname(),master) barrier() assert len(HostNames)==size," internal pb MPI module" if is_master_node() : print("Hostnames : ") for u,host in list(HostNames.items()) : print("Node %d on machine %s"%(u,host))