Addressing more issues....

GEOS-ESM · Sep 7, 2023 · 65e38ff · 65e38ff
1 parent a3ea245
commit 65e38ff
Show file tree

Hide file tree

Showing 8 changed files with 152 additions and 53 deletions.
diff --git a/pre/remap_restart/remap_catchANDcn.py b/pre/remap_restart/remap_catchANDcn.py
@@ -87,6 +87,13 @@ def remap(self):
         NPE = 120
      else:
         NPE = 160
+
+     QOS = "#SBATCH --qos="+config['slurm']['qos']
+     TIME ="#SBATCH --time=1:00:00"
+     if NPE >= 160:
+       assert config['slurm']['qos'] != 'debug', "qos should be allnccs"
+       TIME = "#SBATCH --time=12:00:00"
+     PARTITION = "#SBATCH --partition=" + config['slurm']['partition']
 
      account    = config['slurm']['account']
      # even the input is binary, the output would be nc4
@@ -120,10 +127,11 @@ def remap(self):
      mk_catch_j_template = """#!/bin/csh -f
 #SBATCH --account={account}
 #SBATCH --ntasks={NPE}
-#SBATCH --time=1:00:00
 #SBATCH --job-name=mk_catchANDcn
-#SBATCH --qos=debug
 #SBATCH --output={log_name}
+{TIME}
+{QOS}
+{PARTITION}
 #
 
 source {Bin}/g5_modules
@@ -143,7 +151,7 @@ def remap(self):
      catch1script =  mk_catch_j_template.format(Bin = bindir, account = account, out_bcs = out_bcsdir, \
                   model = model, out_dir = out_dir, surflay = surflay, log_name = log_name, NPE = NPE,  \
                   in_wemin   = in_wemin, out_wemin = out_wemin, out_tilefile = out_tilefile, in_tilefile = in_tilefile, \
-                  in_rstfile = in_rstfile, out_rstfile = out_rstfile, time = yyyymmddhh_ )
+                  in_rstfile = in_rstfile, out_rstfile = out_rstfile, time = yyyymmddhh_, TIME = TIME, PARTITION = PARTITION, QOS=QOS )
 
      script_name = './mk_catchANDcn.j'
 

diff --git a/pre/remap_restart/remap_command.py b/pre/remap_restart/remap_command.py
@@ -72,7 +72,7 @@ def parse_args(program_description):
     p_command.add_argument('-qos',    default = "debug", help= "queue of slurm job", choices=['debug', 'allnccs'])
     account = get_account()
     p_command.add_argument('-account', default = account,  help= "account of slurm job")
-    p_command.add_argument('-constraint', default= 'sky', help= "machine of slurm job")
+    p_command.add_argument('-partition', default= 'compute', help= "partition of slurm job")
     p_command.add_argument('-rs', default= '3', help='flag indicating which restarts to regrid: 1 (upper air); 2 (surface) 3 (both)', choices=['1','2','3'])
 
     # Parse using parse_known_args so we can pass the rest to the remap scripts
@@ -148,7 +148,7 @@ def get_answers_from_command_line(cml):
 
    answers["slurm:account"]    = cml.account
    answers["slurm:qos"]        = cml.qos
-   answers["slurm:constraint"] = cml.constraint
+   answers["slurm:partition"]  = cml.partition
 
    return answers
 

diff --git a/pre/remap_restart/remap_params.py b/pre/remap_restart/remap_params.py
@@ -71,7 +71,7 @@ def params_for_air(self, config_tpl):
   def options_for_slurm(self, config_tpl):
     config_tpl['slurm']['account'] = self.slurm_options['account']
     config_tpl['slurm']['qos'] = self.slurm_options['qos']
-    config_tpl['slurm']['constraint'] = self.slurm_options['constraint']
+    config_tpl['slurm']['partition'] = self.slurm_options['partition']
     return config_tpl
 
   def params_for_surface(self, config_tpl):

diff --git a/pre/remap_restart/remap_params.tpl b/pre/remap_restart/remap_params.tpl
@@ -55,4 +55,4 @@ output:
 slurm:
   account:
   qos:
-  constraint:
+  partition:
diff --git a/pre/remap_restart/remap_questions.py b/pre/remap_restart/remap_questions.py
@@ -45,32 +45,38 @@ def echo_bcs(x,opt):
     print("\nUsers can change the paths in the generated remap_params.yaml file later on")
   return False
 
+def default_partition(x):
+   if x['slurm:qos'] == 'debug':
+      x['slurm:partition'] = 'compute'
+      return False
+   return True
+
 def ask_questions():
 
    questions = [
         {
             "type": "confirm",
             "name": "input:shared:MERRA-2",
-            "message": "Would you like to remap archived MERRA-2 restarts?",
+            "message": "Would you like to remap archived MERRA-2 restarts?\n",
             "default": False,
         },
         {
             "type": "path",
             "name": "input:shared:rst_dir",
-            "message": "Enter the input directory containing restart files to be remapped:",
+            "message": "Enter the input directory containing restart files to be remapped:\n",
             "when": lambda x: not x['input:shared:MERRA-2'],
         },
         {
             "type": "text",
             "name": "input:shared:yyyymmddhh",
-            "message": "From what restart date/time would you like to remap? (must be 10 digits: yyyymmddhh)",
+            "message": "From what restart date/time would you like to remap? (must be 10 digits: yyyymmddhh)\n",
             "validate": lambda text: len(text)==10 ,
-            "when": lambda x: not x['input:shared:MERRA-2'],
+            "when": lambda x: not x['input:shared:MERRA-2'] and fvcore_time(x),
         },
         {
             "type": "text",
             "name": "input:shared:yyyymmddhh",
-            "message": "From what restart date would you like to remap? (must be 8 digits: yyyymmdd, hour=21z)",
+            "message": "From what restart date would you like to remap? (must be 8 digits: yyyymmdd, hour=21z)\n",
             "validate": lambda text: len(text)==8 ,
             "when": lambda x: x['input:shared:MERRA-2'],
         },
@@ -103,7 +109,7 @@ def ask_questions():
         {
             "type": "select",
             "name": "input:shared:model",
-            "message": "Select ocean model that matches input restarts:",
+            "message": "Select ocean model that matches input restarts:\n",
             "choices": ["data", "MOM5", "MOM6"],
             "default": "data",
             "when": lambda x: not x['input:shared:MERRA-2']
@@ -112,7 +118,7 @@ def ask_questions():
         {
             "type": "select",
             "name": "input:shared:ogrid",
-            "message": "Select data ocean grid that matches input restarts:",
+            "message": "Select data ocean grid that matches input restarts:\n",
             "choices": ['360x180   (Reynolds)','1440x720  (MERRA-2)','2880x1440 (OSTIA)','CS  (same as atmosphere OSTIA cubed-sphere grid)'],
             "default": lambda x: data_ocean_default(x.get('input:shared:agrid')),
             "when": lambda x: x.get('input:shared:model') == 'data' and not x['input:shared:MERRA-2'],
@@ -130,7 +136,7 @@ def ask_questions():
         {
             "type": "select",
             "name": "input:shared:ogrid",
-            "message": "Select coupled (MOM5, MOM6) ocean grid that matches input restarts:",
+            "message": "Select coupled (MOM5, MOM6) ocean grid that matches input restarts:\n",
             "choices": ['72x36','360x200','720x410','1440x1080'],
             "when": lambda x: x.get('input:shared:model') == 'MOM5' or x.get('input:shared:model')== 'MOM6'
         },
@@ -150,14 +156,14 @@ def ask_questions():
         {
             "type": "select",
             "name": "output:shared:model",
-            "message": "Select ocean model for new restarts:",
+            "message": "Select ocean model for new restarts:\n",
             "choices": ["data", "MOM5", "MOM6"],
             "default": "data",
         },
         {
             "type": "select",
             "name": "output:shared:ogrid",
-            "message": "Select data ocean grid for new restarts:",
+            "message": "Select data ocean grid for new restarts:\n",
             "choices": ['360x180   (Reynolds)','1440x720  (MERRA-2)','2880x1440 (OSTIA)','CS  (same as atmosphere OSTIA cubed-sphere grid)'],
             "default": lambda x: data_ocean_default(x.get('output:shared:agrid')),
             "when": lambda x: x['output:shared:model'] == 'data',
@@ -173,59 +179,85 @@ def ask_questions():
         {
             "type": "select",
             "name": "output:shared:ogrid",
-            "message": "Select couple ocean grid for new restarts:",
+            "message": "Select couple ocean grid for new restarts:\n",
             "choices": ['72x36','360x200','720x410','1440x1080'],
             "when": lambda x: x['output:shared:model'] != 'data',
         },
 
         {
             "type": "text",
             "name": "output:air:nlevel",
-            "message": "Enter new atmospheric levels: (71 72 91 127 132 137 144 181)",
+            "message": "Enter new atmospheric levels: (71 72 91 127 132 137 144 181)\n",
             "default": "72",
         },
 
         # to show the message, we ask output first
         {
             "type": "select",
-            "name": "output:shared:bc_version",
-            "message": f'''Enter BC version for new restarts:
+            "name": "input:shared:bc_version",
+            "message": f'''\nEnter BC version that matches input restarts:
     BC version                 ADAS tags                   GCM tags
     ----------              ---------------          ------------------
     GM4: Ganymed-4_0    5_12_2 ... 5_16_5         Ganymed-4_0      ... Heracles-5_4_p3
     ICA: Icarus         5_17_0 ... 5_24_0_p1      Icarus, Jason    ...  10.18   
     NL3: Icarus-NLv3    5_25_1 ... present        Icarus_NL, 10.19 ... present
-
-    Development:            Notes
-                        -------------------
-        v06:     NL3 + JPL veg height + PEATMAP + MODIS snow alb
-              \n ''',
+    Development: more choices
+              \n\n ''',
             "choices": ['NL3', 'ICA','GM4','Development'],
-            "default": "NL3",
+            "when": lambda x: not x["input:shared:MERRA-2"],
         },
 
         {
             "type": "select",
-            "name": "output:shared:bc_version",
-            "message": "Enter BC version for new restarts; this option is in testing stage:",
+            "name": "input:shared:bc_version",
+            "message": f'''\nMore choices of BC versions that macthes input restarts (this option is in testing stage):
+
+             v06:     NL3 + JPL veg height + PEATMAP + MODIS snow alb\n\n''',
             "choices": ['v06'],
-            "when": lambda x:  x["output:shared:bc_version"] == 'Development',
+            "when": lambda x: x["input:shared:bc_version"] == 'Development',
         },
 
         {
             "type": "select",
-            "name": "input:shared:bc_version",
-            "message": "Enter BC version that matches input restarts:", 
+            "name": "output:shared:bc_version",
+            "message": f'''\nEnter BC version for new restarts:
+    BC version                 ADAS tags                   GCM tags
+    ----------              ---------------          ------------------
+    GM4: Ganymed-4_0    5_12_2 ... 5_16_5         Ganymed-4_0      ... Heracles-5_4_p3
+    ICA: Icarus         5_17_0 ... 5_24_0_p1      Icarus, Jason    ...  10.18   
+    NL3: Icarus-NLv3    5_25_1 ... present        Icarus_NL, 10.19 ... present
+    Development: more choices
+              \n\n''',
+            "choices": ['NL3', 'ICA','GM4','Development'],
+            "default": "NL3",
+            "when": lambda x: x["input:shared:MERRA-2"],
+        },
+
+        {
+            "type": "select",
+            "name": "output:shared:bc_version",
+            "message": "Enter BC version for new restarts:\n",
             "choices": ['NL3', 'ICA','GM4','Development'],
+            "default": "NL3",
             "when": lambda x: not x["input:shared:MERRA-2"],
         },
 
         {
             "type": "select",
-            "name": "input:shared:bc_version",
-            "message": "Enter BC version for new restarts; this option is in testing stage:",
+            "name": "output:shared:bc_version",
+            "message": f'''\nMore choices of BC versions for new restarts (this option is in testing stage):
+
+             v06:     NL3 + JPL veg height + PEATMAP + MODIS snow alb\n\n''',
             "choices": ['v06'],
-            "when": lambda x: not x["input:shared:MERRA-2"] and x["input:shared:bc_version"] == 'Development',
+            "when": lambda x:  x["output:shared:bc_version"] == 'Development' and x["input:shared:bc_version"] not in ['v06'],
+        },
+
+        {
+            "type": "select",
+            "name": "output:shared:bc_version",
+            "message": "\nMore choices of BC version for new restarts (this option is in testing stage):\n",
+            "choices": ['v06'],
+            "when": lambda x:  x["output:shared:bc_version"] == 'Development' and x["input:shared:bc_version"] in ['v06'],
         },
 
         {
@@ -270,25 +302,26 @@ def ask_questions():
         {
             "type": "text",
             "name": "input:surface:wemin",
-            "message": "What is value of wemin (minimum snow water equivalent parameter) for surface inputs?",
-            "default": lambda x: wemin_default(x.get('input:shared:bc_version'))
+            "message": "What is value of wemin (minimum snow water equivalent parameter) for surface inputs?\n",
+            "default": lambda x: wemin_default(x.get('input:shared:bc_version')),
+            "when": lambda x: show_wemin_default(x),
         },
         {
             "type": "text",
             "name": "output:surface:wemin",
-            "message": "What is value of wemin (minimum snow water equivalent parameter) for new surface restarts?",
+            "message": "What is value of wemin (minimum snow water equivalent parameter) for new surface restarts?\n",
             "default": lambda x: wemin_default(x.get('output:shared:bc_version'))
         },
         {
             "type": "text",
             "name": "input:surface:zoom",
-            "message": "What is value of zoom (parameter of radius search, smaller value means larger radius) for surface inputs [1-8]?",
+            "message": "What is value of zoom (parameter of radius search, smaller value means larger radius) for surface inputs [1-8]?\n",
             "default": lambda x: zoom_default(x)
         },
         {
             "type": "text",
             "name": "output:shared:expid",
-            "message": "Enter new restarts expid:",
+            "message": "Enter new restarts expid (prefix of the output restart names): \n",
             "default": "",
         },
         {
@@ -312,10 +345,11 @@ def ask_questions():
             "default": get_account(),
         },
         {
-            "type": "select",
-            "name": "slurm:constraint",
-            "message": "constraint?",
-            "choices": ['sky', 'cas'],
+            "type": "text",
+            "name": "slurm:partition",
+            "message": "partition?",
+            "default": "compute",
+            "when": lambda x : default_partition(x),
         },
    ]
    answers = questionary.prompt(questions)

diff --git a/pre/remap_restart/remap_restarts.py b/pre/remap_restart/remap_restarts.py
@@ -75,11 +75,12 @@ def main():
       raw_config = get_config_from_answers(answers)
       cmd = get_command_line_from_answers(answers)
       write_cmd(answers['output:shared:out_dir'], cmd)
-      with open("raw_answers.yaml", "w") as f:
-        yaml.dump(raw_config, f)
+      # just for debugging
+      # with open("raw_answers.yaml", "w") as f:
+      #   yaml.dump(raw_config, f)
       params = remap_params(raw_config)
       config = params.config
-      config_yaml = 'remap_params.yaml'
+      config_yaml = answers['output:shared:out_dir']+'/remap_params.yaml'
 
   print('\n')
   if config:
@@ -101,6 +102,14 @@ def main():
 
        # write config to yaml file
      config_to_yaml(config, config_yaml,noprompt = noprompt)
+
+     if not noprompt :
+       submit = questionary.confirm("Submit the jobs now ?" , default=True).ask()
+       if not submit :
+         print("\nYou can submit the jobs by the command later on: \n")
+         print("./remap_restarts.py config_file -c " + config_yaml + "\n")
+         sys.exit(0)
+
   print(config_yaml)
   # upper air
   upper = upperair(params_file=config_yaml)

diff --git a/pre/remap_restart/remap_upper.py b/pre/remap_restart/remap_upper.py
@@ -141,7 +141,7 @@ def remap(self):
        assert config['slurm']['qos'] != 'debug', "qos should be allnccs"
        TIME = "#SBATCH --time=12:00:00"
 
-     CONSTR = "#SBATCH --constraint=" + config['slurm']['constraint']
+     PARTITION = "#SBATCH --partition=" + config['slurm']['partition']
 
      log_name = out_dir+'/remap_upper_log'
 
@@ -190,7 +190,7 @@ def remap(self):
 #SBATCH --output={log_name}
 {TIME}
 {QOS}
-{CONSTR}
+{PARTITION}
 
 unlimit
 
@@ -254,7 +254,7 @@ def remap(self):
      remap_upper_script = remap_template.format(Bin=bindir, account = account, \
              out_dir = out_dir, log_name = log_name, drymassFLG = drymassFLG, \
              imout = imout, nwrit = nwrit, NPE = NPE, \
-             QOS = QOS, TIME = TIME, CONSTR = CONSTR, nlevel = nlevel, hydrostatic = hydrostatic,
+             QOS = QOS, TIME = TIME, PARTITION = PARTITION, nlevel = nlevel, hydrostatic = hydrostatic,
              stretch_str = stretch_str)
 
      script_name = './remap_upper.j'