88 lines
4.0 KiB
Python
88 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
import argparse
|
|
import hashlib
|
|
import inspect
|
|
import json
|
|
import os
|
|
import shutil
|
|
from pathlib import Path
|
|
import sys
|
|
|
|
|
|
def main():
|
|
p=argparse.ArgumentParser()
|
|
p.add_argument('--config', required=True); p.add_argument('--job', required=True)
|
|
a=p.parse_args()
|
|
from config import load_config
|
|
cfg=load_config(a.config)
|
|
runtime=Path(cfg['_paths']['runtime'])/'python'
|
|
os.environ['PYTHONPATH']=str(runtime)
|
|
sys.path.insert(0,str(runtime))
|
|
import exoplasim
|
|
from results import prepare_results
|
|
job=json.loads(Path(a.job).read_text())
|
|
spec=job['spec']; run=spec['run']; model=spec['model']
|
|
work=Path(job['directory'])/'runtime'
|
|
instance=exoplasim.Model(workdir=str(work),modelname='job_'+job['id'],resolution=model['resolution'],layers=model['layers'],ncpus=model['ncpus'])
|
|
params=dict(spec.get('parameters',{}))
|
|
map_paths=job.get('map_snapshots',{})
|
|
resource_maps=[]
|
|
steps=run.get('steps')
|
|
params['runsteps']=steps
|
|
if steps is not None:
|
|
cadence=params.get('snapshots',inspect.signature(instance.configure).parameters['snapshots'].default)
|
|
params['snapshots']=min(cadence,steps)
|
|
instance.cfgpostprocessor(times=None,timeaverage=False)
|
|
surface=spec.get('surface')
|
|
if surface:
|
|
from terrain import prepare_surface
|
|
surface_config=dict(cfg); surface_config['parameters']=params
|
|
maps=prepare_surface(surface,model,surface_config,str(work/'surface'))
|
|
params['landmap']=maps['landmap']; params['topomap']=maps['topomap']
|
|
if map_paths:
|
|
from terrain import write_map_sra, _reference_gravity, _RESOLUTIONS
|
|
gravity=params.get('gravity')
|
|
if gravity is None: gravity=_reference_gravity(cfg)
|
|
gravity=float(gravity)
|
|
from mapcatalog import FIELDS
|
|
if surface:
|
|
landmask=map_paths.get('landmask') or maps['landmap']
|
|
else:
|
|
landmask=map_paths.get('landmask') or work/f"N{_RESOLUTIONS[model['resolution']]:03d}_surf_{FIELDS['landmask']['native_code']:04d}.sra"
|
|
native_inputs=Path(job['directory'])/'inputs'/'native'
|
|
native_inputs.mkdir(parents=True,exist_ok=True)
|
|
for field, snapshot in map_paths.items():
|
|
contract=FIELDS[field]
|
|
native_path=write_map_sra(snapshot,model['resolution'],field,native_inputs/f'{field}.sra',gravity,landmask)
|
|
if contract['native_parameter'] == 'resources':
|
|
resource_maps.append((native_path,contract['native_code']))
|
|
else:
|
|
params[contract['native_parameter']]=native_path
|
|
instance.configure(**params)
|
|
# configure removes SRA files when height or land maps are supplied.
|
|
for native_path,code in resource_maps:
|
|
destination=work/f"N{_RESOLUTIONS[model['resolution']]:03d}_surf_{code:04d}.sra"
|
|
shutil.copyfile(native_path,destination)
|
|
instance.resources.append(str(destination))
|
|
if steps is not None:
|
|
for key in ('N_RUN_YEARS','N_RUN_MONTHS','N_RUN_DAYS'):
|
|
instance._edit_namelist('plasim_namelist',key,'0')
|
|
instance.exportcfg()
|
|
jobroot=Path(job['directory'])
|
|
for cycle in range(int(run['cycles'])):
|
|
instance.run(years=1,crashifbroken=True,clean=False)
|
|
output=work/f'MOST.{cycle:05d}.npz'
|
|
if not output.is_file():
|
|
candidates=sorted(work.glob('MOST.*.npz'))
|
|
if not candidates: raise FileNotFoundError('Model produced no MOST NPZ output')
|
|
output=candidates[-1]
|
|
dataset_id=hashlib.sha256(f"{job['id']}:{cycle+1}".encode()).hexdigest()[:32]
|
|
dest=jobroot/'artifacts'/'datasets'/dataset_id
|
|
result=prepare_results(output,dest,dataset_id,config=cfg)
|
|
budget=int(cfg.get('max_result_bytes',1024**3))
|
|
used=sum(p.stat().st_size for p in (jobroot/'artifacts').rglob('*') if p.is_file())
|
|
if used > budget: raise ValueError('Result output exceeds configured storage budget.')
|
|
print(json.dumps({'cycle':cycle+1,'cycles':run['cycles'],'dataset':dataset_id,'output':str(output)}),flush=True)
|
|
|
|
if __name__=='__main__': main()
|