def _init_on_workers(module, name, args, kwargs):
import importlib
from pyscf.gto import mole
from pyscf.pbc.gto import cell
from mpi4pyscf.tools import mpi
if args is None and kwargs is None: # Not to call __init__ function on workers
if module is None: # master proccess
obj = name
if hasattr(obj, 'cell'):
mol_str = obj.cell.dumps()
elif hasattr(obj, 'mol'):
mol_str = obj.mol.dumps()
else:
mol_str = None
mpi.comm.bcast((mol_str, obj.pack()))
else:
cls = getattr(importlib.import_module(module), name)
obj = cls.__new__(cls)
mol_str, obj_attr = mpi.comm.bcast(None)
obj.unpack_(obj_attr)
if mol_str is not None:
if '"ke_cutoff"' in mol_str:
obj.cell = cell.loads(mol_str)
elif '_bas' in mol_str:
obj.mol = mole.loads(mol_str)
elif module is None: # master proccess
obj = name
else:
# Guess whether the args[0] is the serialized mole or cell objects
if isinstance(args[0], str) and args[0][0] == '{':
if '"ke_cutoff"' in args[0]:
c = cell.loads(args[0])
args = (c, ) + args[1:]
elif '_bas' in args[0]:
m = mole.loads(args[0])
args = (m, ) + args[1:]
cls = getattr(importlib.import_module(module), name)
obj = cls(*args, **kwargs)
key = id(obj)
mpi._registry[key] = obj
regs = mpi.comm.gather(key)
return regs
def _init_on_workers(module, name, args, kwargs):
import importlib
from pyscf.gto import mole
from pyscf.pbc.gto import cell
from mpi4pyscf.tools import mpi
if args is None and kwargs is None: # Not to call __init__ function on workers
if module is None: # master proccess
obj = name
if hasattr(obj, 'cell'):
mol_str = obj.cell.dumps()
elif hasattr(obj, 'mol'):
mol_str = obj.mol.dumps()
else:
mol_str = None
mpi.comm.bcast((mol_str, obj.pack()))
else:
cls = getattr(importlib.import_module(module), name)
obj = cls.__new__(cls)
mol_str, obj_attr = mpi.comm.bcast(None)
obj.unpack_(obj_attr)
if mol_str is not None:
if '"ke_cutoff"' in mol_str:
obj.cell = cell.loads(mol_str)
elif '_bas' in mol_str:
obj.mol = mole.loads(mol_str)
elif module is None: # master proccess
obj = name
else:
# Guess whether the args[0] is the serialized mole or cell objects
if isinstance(args[0], str) and args[0][0] == '{':
if '"ke_cutoff"' in args[0]:
c = cell.loads(args[0])
args = (c,) + args[1:]
elif '_bas' in args[0]:
m = mole.loads(args[0])
args = (m,) + args[1:]
cls = getattr(importlib.import_module(module), name)
obj = cls(*args, **kwargs)
key = id(obj)
mpi._registry[key] = obj
regs = mpi.comm.gather(key)
return regs
def _distribute_call(module, name, reg_procs, args, kwargs):
import importlib
dev = reg_procs
if module is None: # Master process
fn = name
else:
fn = getattr(importlib.import_module(module), name)
if dev is None:
pass
elif isinstance(dev, str) and dev[0] == '{':
# Guess whether dev is Mole or Cell, then deserialize dev
from pyscf.gto import mole
from pyscf.pbc.gto import cell
if '"ke_cutoff"' in dev:
dev = cell.loads(dev)
elif '_bas' in dev:
dev = mole.loads(dev)
else:
from mpi4pyscf.tools import mpi
dev = mpi._registry[reg_procs[mpi.rank]]
return fn(dev, *args, **kwargs)
def _distribute_call(module, name, reg_procs, args, kwargs):
import importlib
dev = reg_procs
if module is None: # Master process
fn = name
else:
fn = getattr(importlib.import_module(module), name)
if dev is None:
pass
elif isinstance(dev, str) and dev[0] == '{':
# Guess whether dev is Mole or Cell, then deserialize dev
from pyscf.gto import mole
from pyscf.pbc.gto import cell
if '"ke_cutoff"' in dev:
dev = cell.loads(dev)
elif '_bas' in dev:
dev = mole.loads(dev)
else:
from mpi4pyscf.tools import mpi
dev = mpi._registry[reg_procs[mpi.rank]]
return fn(dev, *args, **kwargs)
