def run_setup(self):
# import os
# self.pid = os.getpid()
# print("writer pid: %d"%self.pid,flush=True)
#logger.info("writer pid: %d",self.pid)
self.rebuild_buffers()
self.exe_graph = graph.OutputGraph(self.mapping)
def build_output_graph():
from awrams.utils.nodegraph import nodes,graph
from awrams.models import awral
from awrams.models.awral import ffi_wrapper as fw
from awrams.models.awral.template import DEFAULT_TEMPLATE
output_map = awral.get_output_nodes(DEFAULT_TEMPLATE)
print(output_map)
outpath = '/data/cwd_awra_data/awra_test_outputs/sbaronha/sim_test_outputs/'
output_map.mapping.update({
's0_avg': nodes.transform(nodes.average,['s0_dr','s0_sr']),
's0_avg_save': nodes.write_to_annual_ncfile(outpath,'s0_avg'),
'ss_avg': nodes.transform(nodes.average,['ss_dr','ss_sr']),
'ss_avg_save': nodes.write_to_annual_ncfile(outpath,'ss_avg'),
# 'sd_avg': nodes.transform(nodes.average,['s0_dr','s0_sr']),
# 'sd_avg_save': nodes.write_to_annual_ncfile('./','s0_avg'),
#
# 'qtot_avg_save': nodes.write_to_annual_ncfile('./','qtot'),
# 'etot_avg_save': nodes.write_to_annual_ncfile('./','etot')
})
outputs = graph.OutputGraph(output_map.mapping)
# print(outputs.get_dataspecs())
# print(outputs.get_dataspecs(flat=True))
return outputs
def setup():
from os.path import join, dirname
from awrams.utils import datetools as dt
from awrams.utils.nodegraph import nodes, graph
from awrams.models import awral
from awrams.models.awral.template import DEFAULT_TEMPLATE
from awrams.utils.mapping_types import period_to_tc
global period
period = dt.dates('dec 2010 - jan 2011')
global input_map
input_map = awral.get_default_mapping()
change_path_to_forcing(input_map)
global output_map
output_map = awral.get_output_nodes(DEFAULT_TEMPLATE)
global outpath
outpath = join(dirname(__file__), '..', '..', 'test_data', 'simulation',
'outputs')
output_map.mapping['s0_ncsave'] = nodes.write_to_annual_ncfile(
outpath, 's0')
outgraph = graph.OutputGraph(output_map.mapping)
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__(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 build_output_graph():
from awrams.utils.nodegraph import nodes, graph
from awrams.models import awral
from awrams.models.awral import ffi_wrapper as fw
from awrams.models.awral.template import DEFAULT_TEMPLATE
output_map = awral.get_output_nodes(DEFAULT_TEMPLATE)
print(output_map)
output_map.mapping.update({
's0_avg':
nodes.transform(nodes.average, ['s0_dr', 's0_sr']),
's0_avg_save':
nodes.write_to_annual_ncfile('./', 's0_avg')
})
outputs = graph.OutputGraph(output_map.mapping)
print(outputs.get_dataspecs())
print(outputs.get_dataspecs(flat=True))
return outputs
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 _init_graphs(self):
self.graph = graph.OutputGraph(self.mapping)
self.graph.initialise(self.period, self.extent)