VLab Issue #58894 - SLURM updates for free-forecast mode

parm/config/config.fv3ic

@@ -8,7 +8,11 @@ echo "BEGIN: config.fv3ic"
 # Task and thread configuration
 export wtime_fv3ic="00:30:00"
 export npe_fv3ic=1
-export npe_node_fv3ic=1 
+export npe_node_fv3ic=1
+if [ $machine = THEIA ]; then
+  export npe_fv3ic=4
+  export npe_node_fv3ic=4
+fi 
 export nth_fv3ic=${NTHREADS_CHGRES:-24}
 
 echo "END: config.fv3ic"

ush/rocoto/setup_workflow.py

@@ -177,7 +177,7 @@ def get_definitions(base):
 
     strings.append('\t<!ENTITY QUEUE      "%s">\n' % base['QUEUE'])
     if base['machine'] == 'THEIA' and wfu.check_slurm():
-        strings.append('\t<!ENTITY QUEUE_ARCH "%s">\n' % base['QUEUE'])
+        strings.append('\t<!ENTITY QUEUE_ARCH "%s">\n' % base['QUEUE_ARCH'])
         strings.append('\t<!ENTITY PARTITION_ARCH "%s">\n' % base['QUEUE_ARCH'])
     else:
         strings.append('\t<!ENTITY QUEUE_ARCH "%s">\n' % base['QUEUE_ARCH'])

ush/rocoto/setup_workflow_fcstonly.py

@@ -27,6 +27,7 @@
 
 import os
 import sys
+import re
 import numpy as np
 from datetime import datetime
 from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
@@ -127,7 +128,13 @@ def get_definitions(base):
     strings.append('\t<!-- Machine related entities -->\n')
     strings.append('\t<!ENTITY ACCOUNT    "%s">\n' % base['ACCOUNT'])
     strings.append('\t<!ENTITY QUEUE      "%s">\n' % base['QUEUE'])
-    strings.append('\t<!ENTITY QUEUE_ARCH "%s">\n' % base['QUEUE_ARCH'])
+
+    if base['machine'] == 'THEIA' and wfu.check_slurm():
+       strings.append('\t<!ENTITY QUEUE_ARCH "%s">\n' % base['QUEUE'])
+       strings.append('\t<!ENTITY PARTITION_ARCH "%s">\n' % base['QUEUE_ARCH'])
+    else:
+       strings.append('\t<!ENTITY QUEUE_ARCH "%s">\n' % base['QUEUE_ARCH'])
+
     strings.append('\t<!ENTITY SCHEDULER  "%s">\n' % wfu.get_scheduler(base['machine']))
     strings.append('\n')
     strings.append('\t<!-- Toggle HPSS archiving -->\n')
@@ -152,7 +159,9 @@ def get_resources(dict_configs, cdump='gdas'):
     strings.append('\t<!-- BEGIN: Resource requirements for the workflow -->\n')
     strings.append('\n')
 
-    machine = dict_configs['base']['machine']
+    base = dict_configs['base']
+    machine = base.get('machine', 'WCOSS_C')
+    #machine = dict_configs['base']['machine']
 
     for task in taskplan:
 
@@ -163,9 +172,14 @@ def get_resources(dict_configs, cdump='gdas'):
         taskstr = '%s_%s' % (task.upper(), cdump.upper())
 
         strings.append('\t<!ENTITY QUEUE_%s     "%s">\n' % (taskstr, queuestr))
+        if base['machine'] == 'THEIA' and wfu.check_slurm() and task == 'arch':
+            strings.append('\t<!ENTITY PARTITION_%s "&PARTITION_ARCH;">\n' % taskstr )
+	elif base['machine'] == 'THEIA' and wfu.check_slurm() and task == 'getic':
+            strings.append('\t<!ENTITY PARTITION_%s "&PARTITION_ARCH;">\n' % taskstr )
         strings.append('\t<!ENTITY WALLTIME_%s  "%s">\n' % (taskstr, wtimestr))
         strings.append('\t<!ENTITY RESOURCES_%s "%s">\n' % (taskstr, resstr))
-        strings.append('\t<!ENTITY MEMORY_%s    "%s">\n' % (taskstr, memstr))
+        if len(memstr) != 0:
+            strings.append('\t<!ENTITY MEMORY_%s    "%s">\n' % (taskstr, memstr))
         strings.append('\t<!ENTITY NATIVE_%s    "%s">\n' % (taskstr, natstr))
 
         strings.append('\n')
@@ -361,6 +375,20 @@ def create_xml(dict_configs):
     resources = get_resources(dict_configs, cdump=base['CDUMP'])
     workflow = get_workflow_body(dict_configs, cdump=base['CDUMP'])
 
+    # Removes <memory>&MEMORY_JOB_DUMP</memory> post mortem from gdas tasks
+    temp_workflow = ''
+    memory_dict = []
+    for each_resource_string in re.split(r'(\s+)', resources):
+        if 'MEMORY' in each_resource_string:
+            memory_dict.append(each_resource_string)
+    for each_line in re.split(r'(\s+)', workflow):
+        if 'MEMORY' not in each_line:
+            temp_workflow += each_line
+        else:
+            if any( substring in each_line for substring in memory_dict):
+                temp_workflow += each_line
+    workflow = temp_workflow
+
     # Start writing the XML file
     fh = open('%s/%s.xml' % (base['EXPDIR'], base['PSLOT']), 'w')
 

ush/rocoto/workflow_utils.py

@@ -207,7 +207,7 @@ def create_wf_task(task, cdump='gdas', cycledef=None, envar=None, dependency=Non
                  'partition' : '&PARTITION_%s_%s;' % (task.upper(),cdump.upper()), \
                  'final': final}
 
-    if task in ['arch','earc'] and check_slurm():
+    if task in ['getic','arch','earc'] and check_slurm():
         task_dict['partition'] = '&PARTITION_%s_%s;' % (task.upper(),cdump.upper())
     else:
         task_dict['partition'] = None
@@ -323,8 +323,8 @@ def get_resources(machine, cfg, task, cdump='gdas'):
         resstr = '<cores>%d</cores>' % tasks
 
     queuestr = '&QUEUE_ARCH;'
-    # Tricky logic added for Thiea arch queues beasue partition
-    # is a subset up queue for service queues (for now)
+    # Tricky logic added for Theia arch queues because partition
+    # is a subset of queue for service queues (for now)
     if task in ['arch', 'earc', 'getic']:
         if machine in ['THEIA'] and check_slurm():
             queuestr = '&QUEUE;'