<s> import xml <EOL> import joblib <EOL> import subprocess <EOL> import xml . etree . ElementTree as et <EOL> from config import PYTHON_CMD , CLUSTER_NODES , SSH_CMD , CLUSTER_DIR , CPU_FREE_TOL , RAM_FREE_TOL , NUM_CPUS , GPU_FREE_TOL <EOL> def is_node_free ( node_name ) : <EOL> print node_name <EOL> try : <EOL> cpu_usage = float ( subprocess . Popen ( SSH_CMD . split ( '<STR_LIT:U+0020>' ) + [ node_name , <EOL> PYTHON_CMD , '<STR_LIT>' % CLUSTER_DIR ] , <EOL> stdout = subprocess . PIPE ) . communicate ( ) [ <NUM_LIT:0> ] ) <EOL> cpu_usage /= <NUM_LIT> <EOL> mem_usage = float ( subprocess . Popen ( SSH_CMD . split ( '<STR_LIT:U+0020>' ) + [ node_name , <EOL> PYTHON_CMD , '<STR_LIT>' % CLUSTER_DIR ] , <EOL> stdout = subprocess . PIPE ) . communicate ( ) [ <NUM_LIT:0> ] ) <EOL> except Exception as e : <EOL> print e <EOL> return False <EOL> print cpu_usage , mem_usage <EOL> return cpu_usage < CPU_FREE_TOL and mem_usage < RAM_FREE_TOL <EOL> def get_free_nodes ( cluster , parallel = True ) : <EOL> nodes = [ cluster + str ( node ) for node in CLUSTER_NODES [ cluster ] ] <EOL> if parallel : <EOL> is_free = joblib . Parallel ( n_jobs = NUM_CPUS ) ( <EOL> joblib . delayed ( is_node_free ) ( node ) for node in nodes ) <EOL> else : <EOL> is_free = list ( ) <EOL> for node in nodes : <EOL> is_free . append ( is_node_free ( node ) ) <EOL> free_nodes = [ nodes [ k ] for k in range ( len ( nodes ) ) if is_free [ k ] ] <EOL> return free_nodes <EOL> def run_cpu_job ( node , cmd , stdout = subprocess . PIPE , blocking = True ) : <EOL> cmd = '<STR_LIT>' % ( SSH_CMD , node , cmd ) <EOL> print '<STR_LIT>' % node <EOL> print cmd <EOL> if not blocking : <EOL> subprocess . Popen ( cmd , stdout = stdout , shell = True ) <EOL> else : <EOL> output = subprocess . Popen ( cmd , <EOL> stdout = stdout , shell = True ) . communicate ( ) [ <NUM_LIT:0> ] <EOL> return output <EOL> def get_free_gpus ( node_name ) : <EOL> '''<STR_LIT>''' <EOL> output = subprocess . Popen ( SSH_CMD . split ( '<STR_LIT:U+0020>' ) + [ node_name , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] , <EOL> stdout = subprocess . PIPE ) . communicate ( ) [ <NUM_LIT:0> ] <EOL> if not output : <EOL> print '<STR_LIT>' % node_name <EOL> return [ ] <EOL> try : <EOL> tree = et . fromstring ( output . strip ( ) ) <EOL> except xml . parsers . expat . ExpatError : <EOL> print '<STR_LIT>' , output . strip ( ) <EOL> return [ ] <EOL> gpus = tree . findall ( '<STR_LIT>' ) <EOL> print '<STR_LIT>' % ( len ( gpus ) , node_name ) <EOL> free_gpus = [ ] <EOL> for i , gpu in enumerate ( gpus ) : <EOL> mem = gpu . findall ( '<STR_LIT>' ) <EOL> if len ( mem ) == <NUM_LIT:0> : <EOL> mem = gpu . findall ( '<STR_LIT>' ) <EOL> if len ( mem ) == <NUM_LIT:0> : <EOL> print '<STR_LIT>' % node_name <EOL> return [ ] <EOL> mem = mem [ <NUM_LIT:0> ] <EOL> tot = int ( mem . findall ( '<STR_LIT>' ) [ <NUM_LIT:0> ] . text . split ( ) [ <NUM_LIT:0> ] ) <EOL> used = int ( mem . findall ( '<STR_LIT>' ) [ <NUM_LIT:0> ] . text . split ( ) [ <NUM_LIT:0> ] ) <EOL> print used , '<STR_LIT:/>' , tot <EOL> if float ( used ) / tot < GPU_FREE_TOL : <EOL> free_gpus . append ( i ) <EOL> return free_gpus <EOL> def get_all_free_gpus ( cluster , parallel = True ) : <EOL> nodes = [ cluster + str ( node ) for node in CLUSTER_NODES [ cluster ] ] <EOL> if parallel : <EOL> free_gpus = joblib . Parallel ( n_jobs = NUM_CPUS ) ( <EOL> joblib . delayed ( get_free_gpus ) ( node ) for node in nodes ) <EOL> else : <EOL> free_gpus = list ( ) <EOL> for node in nodes : <EOL> free_gpus . append ( get_free_gpus ( node ) ) <EOL> free_gpus_map = dict ( ) <EOL> for fg , n in zip ( free_gpus , nodes ) : <EOL> if len ( fg ) > <NUM_LIT:0> : <EOL> free_gpus_map [ n ] = fg <EOL> return free_gpus_map <EOL> def get_next_free_gpu_sequential ( all_free_gpus , used_gpus , flagged_gpus ) : <EOL> keys = sorted ( all_free_gpus . keys ( ) ) <EOL> for key in keys : <EOL> for gpu in all_free_gpus [ key ] : <EOL> gpu_name = key + '<STR_LIT:_>' + str ( gpu ) <EOL> if gpu_name not in used_gpus and gpu_name not in flagged_gpus : <EOL> return key , gpu <EOL> return None , None <EOL> def run_gpu_job ( node , gpu , cmd , blocking = True , stdout = None ) : <EOL> print '<STR_LIT>' % ( node , gpu ) <EOL> if not blocking : <EOL> cmd = '<STR_LIT>' % ( SSH_CMD , node , gpu , cmd ) <EOL> print cmd <EOL> subprocess . Popen ( cmd , <EOL> stdout = stdout , stderr = subprocess . STDOUT , shell = True ) <EOL> else : <EOL> cmd = '<STR_LIT>' % ( SSH_CMD , node , gpu , cmd ) <EOL> print cmd <EOL> if not stdout : <EOL> stdout = subprocess . PIPE <EOL> output = subprocess . Popen ( cmd , <EOL> stdout = stdout , shell = True ) . communicate ( ) [ <NUM_LIT:0> ] <EOL> return output </s>
