from multiprocessing import Pool
import subprocess
import tqdm
import warnings
[docs]
class multiprocess_handler():
def __init__(self,verbosity=1,name='MultiprocessHandler'):
self.clear()
self.set_verbosity(verbosity)
self.set_name(name)
[docs]
def run(self):
if self.get_commands() is None:
warnings.warn("No commands specified, please set commands.")
return None
pool=Pool(processes=len(self.get_commands()))
iterator = pool.imap_unordered(self.inner_work,self.get_commands())
if self.verbosity==0:
self.returns = self.verb0(iterator)
elif self.verbosity == 1:
self.returns = self.verb1(iterator)
elif self.verbosity == 2:
self.returns=self.verb2(iterator)
pool.close()
[docs]
def verb0(self,iterator):
return [ i for i in iterator]
[docs]
def verb1(self,iterator):
return [ i for i in tqdm.tqdm(iterator, total=len(self.get_commands()),leave=False,desc=self.get_name()+" running")]
[docs]
def verb2(self,iterator):
res = []
for i in tqdm.tqdm(iterator, total=len(self.get_commands()),leave=False,desc=self.get_name()+" running"):
res.append(i)
tqdm.tqdm.write(f"Completed {i['Command']['Index']}({i['Process']})")
return res
[docs]
def clear(self):
self.commands=None
self.returns=None
self.idx=None
[docs]
def set_verbosity(self,lvl:int=1):
self.verbosity=lvl
[docs]
def set_commands(self,commands:list):
self.commands = [(i,cc) for i,cc in enumerate(commands)]
[docs]
def set_name(self,name:str):
self.name = name
[docs]
def get_commands(self):
return self.commands
[docs]
def get_return(self):
return self.returns
[docs]
def get_verbosity(self):
return self.verbosity
[docs]
def get_name(self):
return self.name
[docs]
def work(self,command):
self.set_index(command[0])
return command[1]
[docs]
def get_index(self):
return self.idx
[docs]
def set_index(self,arg):
self.idx=arg
[docs]
def inner_work(self,command):
res = self.work(command)
index = self.get_index()
return {'Process':index,"Command":{'Index':command[0],'Command':command[1]},'Return':res}
def __str__(self):
printstr = []
if self.get_commands() is not None:
printstr.append("\n Commands:")
for cc in self.get_commands():
printstr.append( " "+f"{cc[0]}) {cc[1]}")
if self.get_return() is not None:
printstr.append("\n Return Values:")
for cc in self.get_return():
printstr.append( " "+f"{cc['Command']['Index']}) Process:{cc['Process']} Return: {cc['Return']}")
first_line = self.get_name()
sep_line = '-'*len(first_line)
if len(printstr)>0:
max_len = max([len(l) for l in printstr])
max_len =min((2*len(first_line),max_len))
if max_len > len(first_line):
first_len = (max_len - len(first_line))//2
first_line =" "*first_len + first_line + " "*(len(first_line)-first_len)
sep_line = '-'*max_len
if len(printstr)>0:
printstr.insert(0,sep_line)
printstr.insert(0,first_line)
char=''
if len(printstr)>0:
char='\n'
return char.join(printstr)
[docs]
class multiprocess_bash(multiprocess_handler):
def __init__(self,verbosity=1,name='MultiprocessBash'):
super().__init__(verbosity=verbosity,name=name)
[docs]
def work(self,command):
command = super().work(command) # to assign index as default
child = subprocess.Popen(command, shell=True)
self.set_index(child.pid) #to overwrite default setting of index
return child.wait()
[docs]
class multiprocess_diana(multiprocess_bash):
def __init__(self,verbosity=1,name='MultiprocessDiana'):
super().__init__(verbosity=verbosity,name=name)
[docs]
def get_return(self):
ret = super().get_return()
if ret is not None:
idx = [rr['Command']['Index'] for rr in ret]
ret = [rr for _,rr in sorted(zip(idx,ret))]
return ret
[docs]
def set_commands(self,commands:list):
not_valid = [cc for cc in commands if 'diana' not in cc]
if len(not_valid)>0:
warnings.warn("Some commands are not valid, no calls to diana")
return None
super().set_commands(commands)
[docs]
def get_failed(self):
ret=self.get_return()
if ret is None:
return ret
failed=[]
for rr in ret:
if rr['Return'] !=0:
failed.append(rr['Command']['Index'])
return failed