def run_setup(self):
print("worker pid: %d" % self.pid, flush=True)
self.rebuild_buffers()
self.exe_graph = graph.ExecutionGraph(self.mapping)
self.cur_chunk = None
self.out_chunk = None
self.runner = self.model.get_runner(self.exe_graph.get_dataspecs(True),
shared=True)
def run_setup(self):
# import os
# self.pid = os.getpid()
print("reader pid: %d"%self.pid,flush=True)
logger.info("reader pid: %d",self.pid)
self.rebuild_buffers()
self.exe_graph = graph.ExecutionGraph(self.mapping)
def input_map_to_param_df(input_map):
param_map = dict([(k,v) for k,v in input_map.items() if v.node_type == 'parameter'])
eg = graph.ExecutionGraph(param_map)
out_df = pd.DataFrame(columns=['min_val','max_val','value','fixed','description'])
for k,v in eg.input_graph.items():
param = v['exe']
out_df.loc[k] = dict(value = param.value, min_val = param.min_val, \
max_val = param.max_val, fixed = param.fixed, description = param.description)
return out_df
def build_static_graph(self, model_keys):
parameterised, fixed_mapping, fixed_endpoints = graph.split_parameterised_mapping(
self.mapping, model_keys)
efixed = graph.ExecutionGraph(fixed_mapping)
res = efixed.get_data_flat(self.coords, self.masks, multi=True)
fixed_data = dict((k, res[k]) for k in fixed_endpoints)
dspecs = efixed.get_dataspecs()
fixed_mapping = expand_const(fixed_mapping)
static_mapping = dict([
k,
nodes.static(fixed_data[k], dspecs[k], fixed_mapping[k].out_type)
] for k in fixed_endpoints)
static_mapping.update(parameterised)
estatic = graph.ExecutionGraph(static_mapping)
return estatic
def build_io(self, node_settings): #input_settings,required_inputs):
'''
Assume that we have NCD files we can load from - probably there are other sources...
Build shared memory dictionaries for each of our cell_workers
'''
# print("Building inputs...")
input_settings = node_settings['inputs']
self.shm_inputs = {}
self.shm_outputs = {}
self.outputs = {}
# inputs = {}
igraph = graph.ExecutionGraph(input_settings.mapping)
node_settings['input_dataspecs'] = igraph.get_dataspecs(True)
for cid in self.catchments:
ovs = node_settings['output_variables']
self.shm_outputs[cid] = create_shm_dict(
ovs, (self.timesteps, self.catchments[cid].cell_count))
self.outputs[cid] = shm_to_nd_dict(**self.shm_outputs[cid])
coords = gen_coordset(self.run_period, self.catchments[cid])
input_build = igraph.get_data_flat(coords,
self.catchments[cid].mask)
# self.shm_inputs[cid] = {}
shapes = {}
for n in igraph.input_graph:
if not type(igraph.input_graph[n]['exe']) == ConstNode:
try:
shapes[n] = input_build[n].shape
except AttributeError:
shapes[n] = None
self.igraph = igraph
self.shm_inputs[cid] = create_shm_dict_inputs(shapes)
_inputs_np = shm_to_nd_dict_inputs(**self.shm_inputs[cid])
for n in igraph.input_graph:
if not type(igraph.input_graph[n]['exe']) == ConstNode:
# inputs[cid][n] = input_build[n]
if shapes[n] is None or len(shapes[n]) == 0:
_inputs_np[n][0] = input_build[n]
else:
_inputs_np[n][...] = input_build[n][...]
def __init__(self,model,imapping,omapping=None):
self.input_runner = graph.ExecutionGraph(imapping)
self.model_runner = model.get_runner(self.input_runner.get_dataspecs(True))
self.outputs = None
if omapping is not None:
output_vars = []
for v in self.model_runner.template['OUTPUTS_AVG'] + self.model_runner.template['OUTPUTS_CELL']:
output_vars.append(v)
for v in self.model_runner.template['OUTPUTS_HRU']:
output_vars.extend([v+'_sr',v+'_dr'])
for v in output_vars:
omapping[v] = nodes.model_output(v)
self.outputs = graph.OutputGraph(omapping)
def _init_graphs(self):
out_keys = self.get_required_mapping_outputs()
if self._sep_subtask_graph:
self.graphs = {}
for i in range(self.nsubtasks):
if self._statify:
cur_dom = self.local_dom.realise(subtask=i, task=0)
dyn_keys = [
k for k, v in self.mapping.items()
if 'dynamic' in v.properties
]
cur_graph = graph.ExecutionGraph(
graph.build_fixed_mapping(
self.mapping, out_keys, dyn_keys,
cur_dom.coords['time'].mapping.value_obj,
cur_dom.coords['latlon'].mapping.value_obj))
else:
cur_graph = graph.ExecutionGraph(self.mapping)
self.graphs[i] = cur_graph
else:
self.graphs = {}
if self._statify:
cur_dom = self.local_dom.realise(subtask=0, task=0)
dyn_keys = [
k for k, v in self.mapping.items()
if 'dynamic' in v.properties
]
cur_graph = graph.ExecutionGraph(
graph.build_fixed_mapping(
self.mapping, out_keys, dyn_keys,
cur_dom.coords['time'].mapping.value_obj,
cur_dom.coords['latlon'].mapping.value_obj))
else:
cur_graph = graph.ExecutionGraph(self.mapping)
self.graphs[0] = cur_graph
self.cur_graph = self.graphs[0]
def __init__(self, model, imapping, omapping=None, extent=None):
if extent is None:
extent = extents.get_default_extent()
imapping = graph.get_input_tree(model.get_input_keys(), imapping)
#+++
#Document the use of this manually, don't just change the graph behind the scenes...
#imapping = graph.map_rescaling_nodes(imapping,extent)
self.input_runner = graph.ExecutionGraph(imapping)
self.model_runner = model.get_runner(
self.input_runner.get_dataspecs(True))
self.outputs = None
if omapping is not None:
self.outputs = graph.OutputGraph(omapping)
def __init__(self, len_period, catchments, inputs, outputs):
self.catchments = catchments
self.inputs = inputs
self.outputs = outputs
self.params = None
self.len_period = len_period
from awrams.models import awral
imap = awral.get_default_mapping()
dspec = {}
for k in 'tmin_f', 'tmax_f', 'precip_f', 'solar_f', 'fday', 'radcskyt':
dspec[k] = ['time', 'latitude', 'longitude']
sp = [k for k in imap.mapping if k.endswith('_grid')]
for k in sp:
dspec[k] = ['latitude', 'longitude']
dspec['height'] = ['hypsometric_percentile', 'latitude', 'longitude']
dspec['hypsperc_f'] = ['hypsometric_percentile']
data_map = {k: {} for k in dspec}
for cid in self.catchments:
input = shm_to_nd_dict_inputs(**(self.inputs[cid]))
for k in dspec:
data_map[k][cid] = input[k]
for k in dspec:
imap.mapping[k] = forcing_from_dict(data_map[k], k, dims=dspec[k])
self.igraph = graph.ExecutionGraph(imap.mapping)
self.sim = awral.get_runner(self.igraph.get_dataspecs(True))
def run(self,input_map,output_map,period,extent): #periods,chunks):
'''
Should be the basis for new-style sim server
Currently no file output, but runs inputgraph/model quite happily...
'''
import time
start = time.time()
chunks = extents.subdivide_extent(extent,self.spatial_chunk)
periods = dt.split_period(period,'a')
self.logger.info("Getting I/O dataspecs...")
#+++ Document rescaling separately, don't just change the graph behind the scenes...
#mapping = graph.map_rescaling_nodes(input_map.mapping,extent)
mapping = input_map
filtered = graph.get_input_tree(self.model.get_input_keys(),mapping)
input_nodes = {}
worker_nodes = {}
output_nodes = {}
for k,v in filtered.items():
if 'io' in v.properties:
input_nodes[k] = v
worker_nodes[k] = nodes.const(None)
else:
worker_nodes[k] = v
for k,v in output_map.items():
try:
if v.properties['io'] == 'from_model':
output_nodes[k] = v
except: # AttributeError:
pass
# print("EXCEPTION",k,v)
igraph = graph.ExecutionGraph(input_nodes)
self._set_max_dims(igraph)
input_dspecs = igraph.get_dataspecs(True)
#+++ No guarantee this will close files. Put in separate function?
del igraph
model_dspecs = graph.ExecutionGraph(mapping).get_dataspecs(True)
output_dspecs = graph.OutputGraph(output_nodes).get_dataspecs(True)
self.model.init_shared(model_dspecs)
### initialise output ncfiles
self.logger.info("Initialising output files...")
outgraph = graph.OutputGraph(output_map)
outgraph.initialise(period,extent)
#+++ Can we guarantee that statespecs will be 64bit for recycling?
# NWORKERS = 2
# READ_AHEAD = 1
sspec = DataSpec('array',['cell'],np.float64)
state_specs = {}
for k in self.model.get_state_keys():
init_k = 'init_' + k
input_dspecs[init_k] = sspec
state_specs[k] = sspec
self.logger.info("Building buffers...")
input_bufs = create_managed_buffergroups(input_dspecs,self.max_dims,self.num_workers+self.read_ahead)
state_bufs = create_managed_buffergroups(state_specs,self.max_dims,self.num_workers*2+self.read_ahead)
output_bufs = create_managed_buffergroups(output_dspecs,self.max_dims,self.num_workers+self.read_ahead)
all_buffers = dict(inputs=input_bufs,states=state_bufs,outputs=output_bufs)
smc = SharedMemClient(all_buffers,False)
control_master = mp.Queue()
control_status = mp.Queue()
state_returnq =mp.Queue()
chunkq = mp.Queue()
chunkoutq = mp.Queue()
reader_inq = dict(control=mp.Queue(),state_return=state_returnq)
reader_outq = dict(control=control_master,chunks=chunkq)
writer_inq = dict(control=mp.Queue(),chunks=chunkoutq)
writer_outq = dict(control=control_master,log=mp.Queue()) #,chunks=chunkq)
child_control_qs = [reader_inq['control'],writer_inq['control'],writer_outq['log']]
self.logger.info("Running simulation...")
workers = []
for w in range(self.num_workers):
worker_inq = dict(control=mp.Queue(),chunks=chunkq)
worker_outq = dict(control=control_master,state_return=state_returnq,chunks=chunkoutq)
worker_p = mg.ModelGraphRunner(worker_inq,worker_outq,all_buffers,chunks,periods,worker_nodes,self.model)
worker_p.start()
workers.append(worker_p)
child_control_qs.append(worker_inq['control'])
control = ControlMaster(control_master, control_status, child_control_qs)
control.start()
reader_p = input_reader.InputGraphRunner(reader_inq,reader_outq,all_buffers,chunks,periods,input_nodes,self.model.get_state_keys())
reader_p.start()
writer_p = writer.OutputGraphRunner(writer_inq,writer_outq,all_buffers,chunks,periods,output_map)
writer_p.start()
log = True
while log:
msg = writer_outq['log'].get()
if msg['subject'] == 'terminate':
log = False
else:
self.logger.info(msg['subject'])
writer_p.join()
for w in workers:
w.qin['control'].put(message('terminate'))
# control_master.get_nowait()
w.join()
reader_inq['control'].put(message('terminate'))
control_master.put(message('finished'))
problem = False
msg = control_status.get()
if msg['subject'] == 'exception_raised':
problem = True
control.join()
reader_p.join()
if problem:
raise Exception("Problem detected")
self.logger.info("elapsed time: %.2f",time.time() - start)
def run_setup(self):
self.rebuild_buffers()
self.exe_graph = graph.ExecutionGraph(self.mapping)