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)