import os
import posixpath
import pickle
from typing import Sequence, Tuple, Optional
import logging
import json
from mako.template import Template
import comsdk.misc as aux
from comsdk.communication import CommunicationError
from comsdk.graph import Func, State
dummy_predicate = Func(func=lambda d: True)
dummy_morphism = Func()
job_finished_predicate = Func(func=lambda d: d['job_finished'])
job_unfinished_predicate = Func(func=lambda d: not d['job_finished'])
class InOutMapping(object):
def __init__(self,
keys_mapping={},
relative_keys=(),
default_relative_key=(),
):
self._default_relative_key = default_relative_key if aux.is_sequence(default_relative_key) else (default_relative_key,)
self._relative_keys = relative_keys if aux.is_sequence(relative_keys) else (relative_keys,)
self._keys_mapping = keys_mapping
def __str__(self):
return 'Default relative key: {}\n' \
'Relative keys:\n{}\n' \
'Keys mapping:\n\tLocal -> Global\n\t----------------\n' \
'{}'.format('.'.join(self._default_relative_key),
'\n'.join(['\t' + '.'.join(k) for k in self._relative_keys]),
'\n'.join(['\t' + loc + ' -> ' + '.'.join(glo) for loc, glo in self._keys_mapping]))
def build_proxy_data(self, data, dynamic_keys_mapping={}):
if self._default_relative_key == () and self._relative_keys == () and self._keys_mapping == {} and dynamic_keys_mapping == {}:
return data
else:
#print('\t{}\n\t{}\n\t{}'.format(self._relative_keys, self._keys_mapping, dynamic_keys_mapping))
return aux.ProxyDict(data, self._relative_keys, dict(self._keys_mapping, **dynamic_keys_mapping), self._default_relative_key)
class Edge:
__slots__ = [
'pred_f',
'morph_f',
'_io_mapping',
'preprocess',
'postprocess',
'order',
'comment',
'mandatory_keys',
'use_proxy_data_for_pre_post_processing'
]
def __init__(self, predicate, morphism,
io_mapping=InOutMapping(),
order=0,
comment="",
mandatory_keys=(),
):
self.pred_f = predicate
self.morph_f = morphism
self._io_mapping = io_mapping
self.preprocess = lambda pd: None
self.postprocess = lambda pd: None
self.order = int(0 if order is None else order)
self.comment = comment
self.mandatory_keys = mandatory_keys
self.use_proxy_data_for_pre_post_processing=False
def predicate(self, data, dynamic_keys_mapping={}):
proxy_data = self._io_mapping.build_proxy_data(data, dynamic_keys_mapping)
return self.pred_f.func(proxy_data)
def morph(self, data, dynamic_keys_mapping={}):
#print(self.pred_name, self.morph_name, self.order)
proxy_data = self._io_mapping.build_proxy_data(data, dynamic_keys_mapping)
# print(proxy_data)
if (self.use_proxy_data_for_pre_post_processing):
self.preprocess(proxy_data)
else:
self.preprocess(data)
self._throw_if_not_set(proxy_data, self.mandatory_keys)
self.morph_f.func(proxy_data)
if (self.use_proxy_data_for_pre_post_processing):
self.postprocess(proxy_data)
else:
self.postprocess(data)
def _throw_if_not_set(self, data, mandatory_keys: Sequence[str]):
for k in mandatory_keys:
if k not in data:
logging.exception('EDGE {}: key "{}" is not set whilst being mandatory.\nIOMapping:\n'
'{}'.format(type(self).__name__, k, str(self._io_mapping)))
raise KeyError()
# raise KeyError('EDGE {}: key "{}" is not set whilst being mandatory.\nIOMapping:\n'
# '{}'.format(type(self).__name__, k, str(self._io_mapping)))
class ExecutableProgramEdge(Edge):
'''
Class implementing the edge which uses an external program to morph data.
The program is lauchned via so-called communication which, among others, sets where the program is located and it can be launched.
Environment can be used to launch program on remote resources.
# DESCRIPTION OF KEYS MAPPINGS #
Since data structure is hierarchical, we introduced keys mappings. The edge needs to use some variables
from data which may be located in different (nested) keys of data (we will call these keys "global").
However, it is very convenient to implement the edge imagining that there is no nested structures
and all keys are available in the top-level of data (we will call these keys "local").
To link global and local keys, we introduce keys mapping, which are either dictionaries (local key string -> sequence) or sequences.
If the keys mapping is sequence, we treat it as a relative "path" to all needed keys.
Therefore, we have keys mappings for input and output keys.
# END OF DESCRIPTION OF KEYS MAPPINGS #
We expect that necessary input files are already on remote.
Programs may require three types of arguments:
1) keyword arguments (-somearg something)
2) flags (-someflag)
3) trailing arguments
Local keys determining the corresponding values are located in keyword_names, flag_names and trailing_args_keys.
Finally, data must be somehow updated after finishing. This will be done by updating data according to output_dict (it is just added)
'''
def __init__(self, program_name, comm,
predicate=dummy_predicate,
io_mapping=InOutMapping(),
output_dict={}, # output dict which will be added to the main dictionary (w.r.t. output_keys_mapping)
keyword_names=(), # "local keys" where keyword args are stored
flag_names=(), # "local keys" where flags are stored
trailing_args_keys=(), # "local keys" where trailing args are stored
remote=False,
stdout_processor=None,
chaining_command_at_start=lambda d: '',
chaining_command_at_end=lambda d: '',
):
#predicate = predicate if predicate is not None else dummy_predicate
self._output_dict = output_dict
self._comm = comm
self._program_name = program_name
self._keyword_names = keyword_names
self._flag_names = flag_names
self._trailing_args_keys = trailing_args_keys
self._working_dir_key = '__REMOTE_WORKING_DIR__' if remote else '__WORKING_DIR__'
mandatory_keys = [self._working_dir_key]
self._stdout_processor = stdout_processor
self.chaining_command_at_start = chaining_command_at_start
self.chaining_command_at_end = chaining_command_at_end
super().__init__(predicate, Func(func=self.execute), io_mapping, mandatory_keys=mandatory_keys)
def execute(self, data):
args_str = build_args_line(data, self._keyword_names, self._flag_names, self._trailing_args_keys)
working_dir = data[self._working_dir_key]
stdout_lines, stderr_lines = self._comm.execute_program(self._program_name, args_str, working_dir,
self.chaining_command_at_start(data),
self.chaining_command_at_end(data))
output_data = self._output_dict
if self._stdout_processor:
stdout_data = self._stdout_processor(data, stdout_lines)
data.update(stdout_data)
data.update(output_data)
class QsubScriptEdge(Edge):
'''
Class implementing the edge which builds up the sh-script for qsub.
The script is created via communication.
# DESCRIPTION OF KEYS MAPPINGS #
Since data structure is hierarchical, we introduced keys mappings. The edge needs to use some variables
from data which may be located in different (nested) keys of data (we will call these keys "global").
However, it is very convenient to implement the edge imagining that there is no nested structures
and all keys are available in the top-level of data (we will call these keys "local").
To link global and local keys, we introduce keys mapping, which are either dictionaries (local key string -> sequence) or sequences.
If the keys mapping is sequence, we treat it as a relative "path" to all needed keys.
Therefore, we have keys mappings for input and output keys.
# END OF DESCRIPTION OF KEYS MAPPINGS #
Data will be augmented by 'qsub_script' pointing to the local file.
'''
def __init__(self, program_name, local_comm, remote_comm,
predicate=dummy_predicate,
io_mapping=InOutMapping(),
keyword_names=(), # "local keys" where keyword args are stored
flag_names=(), # "local keys" where flags are stored
trailing_args_keys=(), # "local keys" where trailing args are stored
):
# predicate = predicate if predicate is not None else dummy_predicate
self._local_comm = local_comm
self._remote_comm = remote_comm
self._program_name = program_name
self._keyword_names = keyword_names
self._flag_names = flag_names
self._trailing_args_keys = trailing_args_keys
mandatory_keys = ['__WORKING_DIR__', 'qsub_script_name', 'time_required', 'cores_required']
super().__init__(predicate, Func(func=self.execute), io_mapping, mandatory_keys=mandatory_keys)
def execute(self, data):
if isinstance(data, aux.ProxyDict):
print('QsubScriptEdge -> {}: {}'.format('qsub_script_name', data._keys_mappings['qsub_script_name']))
qsub_script_path = os.path.join(data['__WORKING_DIR__'], data['qsub_script_name'])
args_str = build_args_line(data, self._keyword_names, self._flag_names, self._trailing_args_keys)
program_launch_path = self._remote_comm.host.get_program_launch_path(self._program_name)
command_line = '{} {}'.format(program_launch_path, args_str)
render_sge_template(self._remote_comm.host.sge_template_name, qsub_script_path,
data['cores_required'], data['time_required'], (command_line,))
data.update({'qsub_script': qsub_script_path})
class UploadOnRemoteEdge(Edge):
'''
Class implementing the edge which uploads the data to the remote computer.
It is done via environment which must provide the interface for that.
# DESCRIPTION OF KEYS MAPPINGS #
Since data structure is hierarchical, we introduced keys mappings. The edge needs to use some variables
from data which may be located in different (nested) keys of data (we will call these keys "global").
However, it is very convenient to implement the edge imagining that there is no nested structures
and all keys are available in the top-level of data (we will call these keys "local").
To link global and local keys, we introduce keys mapping, which are either dictionaries (local key string -> sequence) or sequences.
If the keys mapping is sequence, we treat it as a relative "path" to all needed keys.
Therefore, we have keys mappings for input and output keys.
# END OF DESCRIPTION OF KEYS MAPPINGS #
Files for uploading must be found in input_files_keys which is a list of local data keys corresponding to these files.
They will be uploaded in remote working dir which must be in data['__REMOTE_WORKING_DIR__'].
After edge execution, data is going to be updated such that local paths will be replaced by remote ones.
'''
def __init__(self, comm,
predicate=dummy_predicate,
io_mapping=InOutMapping(),
local_paths_keys=(), # "local keys", needed to build a copy list
update_paths=True,
already_remote_path_key=None,
):
# predicate = predicate if predicate is not None else dummy_predicate
self._local_paths_keys = local_paths_keys
self._comm = comm
self._update_paths = update_paths
self._already_remote_path_key = already_remote_path_key
mandatory_keys = list(self._local_paths_keys) + ['__WORKING_DIR__', '__REMOTE_WORKING_DIR__']
if self._already_remote_path_key is not None:
mandatory_keys.append(self._already_remote_path_key)
super().__init__(predicate, Func(func=self.execute), io_mapping, mandatory_keys=mandatory_keys)
def execute(self, data):
if self._already_remote_path_key is not None:
if data[self._already_remote_path_key]:
return
remote_working_dir = data['__REMOTE_WORKING_DIR__']
for key in self._local_paths_keys:
try:
# try data[key] as an absolute path
data[key] = self._comm.copy(data[key], remote_working_dir, mode='from_local')
except CommunicationError as e:
# try data[key] as a relative path
working_dir = data['__WORKING_DIR__']
if isinstance(data, aux.ProxyDict):
print('UploadOnRemoteEdge -> {}: {}'.format(key, data._keys_mappings[key]))
remote_path = self._comm.copy(os.path.join(working_dir, data[key]), remote_working_dir,
mode='from_local')
if self._update_paths:
data[key] = remote_path
class DownloadFromRemoteEdge(Edge):
'''
Class implementing the edge which downloads the data from the remote computer.
It is done via environment which must provide the interface for that.
# DESCRIPTION OF KEYS MAPPINGS #
Since data structure is hierarchical, we introduced keys mappings. The edge needs to use some variables
from data which may be located in different (nested) keys of data (we will call these keys "global").
However, it is very convenient to implement the edge imagining that there is no nested structures
and all keys are available in the top-level of data (we will call these keys "local").
To link global and local keys, we introduce keys mapping, which are either dictionaries (local key string -> sequence) or sequences.
If the keys mapping is sequence, we treat it as a relative "path" to all needed keys.
Therefore, we have keys mappings for input and output keys.
# END OF DESCRIPTION OF KEYS MAPPINGS #
Files for downloading must be found in output_files_keys which is a list of local data keys corresponding to these files.
All these files are relative to the remote working dir and will be downloaded into local working dir
Local working dir must be in data['__LOCAL_WORKING_DIR__'].
Remote working dir must be in data['__REMOTE_WORKING_DIR__'].
After edge execution, data is going to be updated such that remote/relative paths will be replaced by local ones.
'''
def __init__(self, comm,
predicate=dummy_predicate,
io_mapping=InOutMapping(),
remote_paths_keys=(), # "local keys", needed to build a list for downloading
update_paths=True,
show_msg=False,
):
# predicate = predicate if predicate is not None else dummy_predicate
self._remote_paths_keys = remote_paths_keys
self._comm = comm
self._update_paths = update_paths
self._show_msg = show_msg
mandatory_keys = list(self._remote_paths_keys) + ['__WORKING_DIR__', '__REMOTE_WORKING_DIR__']
super().__init__(predicate, Func(func=self.execute), io_mapping, mandatory_keys=mandatory_keys)
def execute(self, data):
working_dir = data['__WORKING_DIR__']
remote_working_dir = data['__REMOTE_WORKING_DIR__']
for key in self._remote_paths_keys:
output_file_or_dir = data[key]
if output_file_or_dir is None:
continue
local_path = None
if output_file_or_dir == '*':
aux.print_msg_if_allowed('\tAll possible output files will be downloaded', allow=self._show_msg)
paths = self._comm.listdir(remote_working_dir)
local_full_paths = ['/'.join([working_dir, file_or_dir]) for file_or_dir in paths]
remote_full_paths = ['/'.join([remote_working_dir, file_or_dir]) for file_or_dir in paths]
for file_or_dir in remote_full_paths:
aux.print_msg_if_allowed('\tAm going to download "{}" to "{}"'.format(file_or_dir, working_dir),
allow=self._show_msg)
self._comm.copy(file_or_dir, working_dir, mode='from_remote', show_msg=self._show_msg)
local_path = local_full_paths
else:
output_file_or_dir_as_list = []
if isinstance(output_file_or_dir, list):
output_file_or_dir_as_list = output_file_or_dir
else:
output_file_or_dir_as_list = [output_file_or_dir]
for f in output_file_or_dir_as_list:
file_or_dir = '/'.join([remote_working_dir, f])
aux.print_msg_if_allowed('\tAm going to download "{}" to "{}"'.format(file_or_dir, working_dir),
allow=self._show_msg)
local_path = self._comm.copy(file_or_dir, working_dir,
mode='from_remote', show_msg=self._show_msg)
if self._update_paths:
data[key] = local_path
def make_cd(key_path):
def _cd(d):
if key_path == '..':
d['__WORKING_DIR__'] = os.path.dirname(d['__WORKING_DIR__'])
if '__REMOTE_WORKING_DIR__' in d:
d['__REMOTE_WORKING_DIR__'] = posixpath.dirname(d['__REMOTE_WORKING_DIR__'])
else:
subdir = aux.recursive_get(d, key_path)
d['__WORKING_DIR__'] = os.path.join(d['__WORKING_DIR__'], subdir)
if '__REMOTE_WORKING_DIR__' in d:
d['__REMOTE_WORKING_DIR__'] = posixpath.join(d['__REMOTE_WORKING_DIR__'], subdir)
return _cd
def make_mkdir(key_path, remote_comm=None):
def _mkdir(d):
remote = '__REMOTE_WORKING_DIR__' in d
dir = os.path.join(d['__WORKING_DIR__'],
d[key_path])
os.mkdir(dir)
if '__REMOTE_WORKING_DIR__' in d:
dir = os.path.join(d['__REMOTE_WORKING_DIR__'],
d[key_path])
remote_comm._mkdirp(dir)
return _mkdir
def make_dump(dump_name_format, format_keys=(), omit=None, method='pickle'):
def _dump(d):
format_params = [aux.recursive_get(d, key) for key in format_keys]
dump_path = os.path.join(d['__WORKING_DIR__'], dump_name_format.format(*format_params))
if omit is None:
dumped_d = d
else:
if (isinstance(d, aux.ProxyDict)):
dumped_d = {key: val for key, val in d._data.items() if not key in omit}
else:
dumped_d = {key: val for key, val in d.items() if not key in omit}
if method == 'pickle':
with open(dump_path, 'wb') as f:
pickle.dump(dumped_d, f)
elif method == 'json':
with open(dump_path, 'w') as f:
json.dump(dumped_d, f)
else:
raise ValueError(f'Method "{method}" is not supported in dumping')
return _dump
def make_composite_func(*funcs):
def _composite(d):
res = None
for func in funcs:
f_res = func(d)
# this trick allows us to combine returning
# and non-returning functions
if f_res is not None:
res = f_res
return res
return _composite
def make_composite_predicate(*preds):
def _composite(d):
for pred in preds:
if not pred(d):
return False
return True
return _composite
def create_local_data_from_global_data(global_data, keys_mapping):
if keys_mapping is None:
return global_data
elif aux.is_sequence(keys_mapping):
return aux.recursive_get(global_data, keys_mapping)
else:
return {local_key: aux.recursive_get(global_data, global_key) for local_key, global_key in keys_mapping.items()}
def update_global_data_according_to_local_data(local_data, global_data, keys_mapping):
if keys_mapping is None:
global_data.update(local_data)
elif aux.is_sequence(keys_mapping):
relative_data = aux.recursive_get(global_data, keys_mapping)
relative_data.update(local_data)
else:
for local_key, global_key in keys_mapping.items():
recursive_set(global_data, global_key, local_data[local_key])
def build_args_line(data, keyword_names, flag_names, trailing_args_keys):
args_str = ''
for keyword in keyword_names:
if keyword in data:
args_str += '-{} {} '.format(keyword, data[keyword])
for flag in flag_names:
if flag in data and data[flag]:
args_str += '-{} '.format(flag)
for place_i, trailing_arg_key in enumerate(trailing_args_keys):
# if we have a sequence under the key, we expand it
if trailing_arg_key in data:
trailing_arg = data[trailing_arg_key]
args_str += ' '.join(map(str, trailing_arg)) if aux.is_sequence(trailing_arg) else trailing_arg
args_str += ' '
return args_str
def render_sge_template(sge_template_name, sge_script_path, cores, time, commands):
with open(os.path.expanduser('~/.comsdk/config_research.json'), 'r') as f:
conf = json.load(f)
sge_templ_path = os.path.join(conf['TEMPLATES_PATH'], sge_template_name)
if not os.path.exists(sge_templ_path): # by default, templates are in templates/, but here we let the user put any path
sge_templ_path = sge_template_name
f = open(sge_templ_path, 'r')
rendered_data = Template(f.read()).render(cores=cores, time=time, commands=commands)
sge_script_file = aux.create_file_mkdir(sge_script_path)
sge_script_file.write(rendered_data)
def connect_branches(branches: Sequence[Tuple[State, State]], edges: Optional[Sequence[Edge]] = None):
if edges is None:
edges = [dummy_edge for _ in range(len(branches) - 1)]
for i, edge in zip(range(1, len(branches)), edges):
_, prev_branch_end = branches[i - 1]
next_branch_start, _ = branches[i]
prev_branch_end.connect_to(next_branch_start, edge=edge)
dummy_edge = Edge(dummy_predicate, Func())