Merge pull request #2235 from desihub/multi-reservation

add reservation parsing for cpu vs. gpu

py/desispec/test/test_workflow_batch.py

@@ -0,0 +1,72 @@
+"""
+Test desispec.workflow.batch
+"""
+
+import os
+import unittest
+
+import numpy as np
+from astropy.table import Table, vstack
+
+from desispec.workflow import batch
+
+class TestWorkflowBatch(unittest.TestCase):
+    """Test desispec.workflow.calibration_selection
+    """
+
+    #- Tests might alter NERSC host so reset after each test
+    @classmethod
+    def setUpClass(cls):
+        cls._cached_nersc_host = os.getenv('NERSC_HOST')  # None if not set
+
+    def tearDown(self):
+        if self._cached_nersc_host is None:
+            if 'NERSC_HOST' in os.environ:
+                del os.environ['NERSC_HOST']
+            else:
+                pass
+        else:
+            os.environ['NERSC_HOST'] = self._cached_nersc_host
+
+    def test_parse_reservation(self):
+        """Test parse_reservation at NERSC"""
+        os.environ['NERSC_HOST'] = 'perlmutter'
+
+        # single reservation -> same name regardless of CPU or GPU
+        self.assertEqual(batch.parse_reservation('blat', 'arc'),       'blat')
+        self.assertEqual(batch.parse_reservation('blat', 'ccdcalib'),  'blat')
+        self.assertEqual(batch.parse_reservation('blat', 'flat'),      'blat')
+        self.assertEqual(batch.parse_reservation('blat', 'tilenight'), 'blat')
+        self.assertEqual(batch.parse_reservation('blat', 'ztile'),     'blat')
+
+        # blat,foo -> cpu_reservation, gpu_reservtion
+        self.assertEqual(batch.parse_reservation('blat,foo', 'arc'),       'blat')
+        self.assertEqual(batch.parse_reservation('blat,foo', 'ccdcalib'),  'blat')
+        self.assertEqual(batch.parse_reservation('blat,foo', 'flat'),      'foo')
+        self.assertEqual(batch.parse_reservation('blat,foo', 'tilenight'), 'foo')
+        self.assertEqual(batch.parse_reservation('blat,foo', 'ztile'),     'foo')
+
+        # "none" special cases to no reservation
+        self.assertEqual(batch.parse_reservation('blat,none', 'arc'),       'blat')
+        self.assertEqual(batch.parse_reservation('blat,none', 'ccdcalib'),  'blat')
+        self.assertEqual(batch.parse_reservation('blat,none', 'flat'),      None)
+        self.assertEqual(batch.parse_reservation('blat,none', 'tilenight'), None)
+        self.assertEqual(batch.parse_reservation('blat,none', 'ztile'),     None)
+
+        self.assertEqual(batch.parse_reservation('none,foo', 'arc'),       None)
+        self.assertEqual(batch.parse_reservation('none,foo', 'ccdcalib'),  None)
+        self.assertEqual(batch.parse_reservation('none,foo', 'flat'),      'foo')
+        self.assertEqual(batch.parse_reservation('none,foo', 'tilenight'), 'foo')
+        self.assertEqual(batch.parse_reservation('none,foo', 'ztile'),     'foo')
+
+        # And None -> None
+        self.assertEqual(batch.parse_reservation(None, 'arc'),       None)
+        self.assertEqual(batch.parse_reservation(None, 'ccdcalib'),  None)
+        self.assertEqual(batch.parse_reservation(None, 'flat'),      None)
+        self.assertEqual(batch.parse_reservation(None, 'tilenight'), None)
+        self.assertEqual(batch.parse_reservation(None, 'ztile'),     None)
+
+        # test errors
+        with self.assertRaises(ValueError):
+            batch.parse_reservation('blat,foo,bar', 'arc')
+

py/desispec/test/test_workflow_queue.py

@@ -1,5 +1,5 @@
 """
-Test desispec.workflow.calibration_selection
+Test desispec.workflow.queue
 """
 
 import os

py/desispec/workflow/batch.py

@@ -11,6 +11,7 @@
 
 from desiutil.log import get_logger
 
+_config_cache = dict()
 def get_config(name):
     """
     Return configuration dictionary for system `name`
@@ -30,13 +31,20 @@ def get_config(name):
     if name is None:
         name = default_system()
 
+    global _config_cache
+    if name in _config_cache:
+        return _config_cache[name]
+
     configfile = resources.files('desispec').joinpath('data/batch_config.yaml')
     with open(configfile) as fx:
         config = yaml.safe_load(fx)
 
     #- Add the name for reference, in case it was default selected
     config['name'] = name
 
+    #- Add to cache so that we don't have to re-read batch_config.yaml every time
+    _config_cache[name] = config[name]
+
     return config[name]
 
 def default_system(jobdesc=None, no_gpu=False):
@@ -76,6 +84,44 @@ def default_system(jobdesc=None, no_gpu=False):
 
     return name
 
+def parse_reservation(reservation, jobdesc):
+    """
+    Parse reservation name into cpu/gpu reservation based upon jobdesc
+
+    Args:
+        reservation (str): resvname or resvname_cpu,resvname_gpu or None
+        jobdesc (str): job description string e.g. 'arc', 'flat', 'tilenight'
+
+    Returns:
+        cpu_reservation_name, gpu_reservation_name
+
+    If a single reservation name is provided, return both cpu/gpu as the same.
+    If either is 'none' (case-insensitive), return None for that reservation
+    """
+    if reservation is None:
+        return reservation
+
+    tmp = reservation.split(',')
+    if len(tmp) == 1:
+        reservation_cpu = reservation_gpu = reservation
+    elif len(tmp) == 2:
+        reservation_cpu, reservation_gpu = tmp
+    else:
+        raise ValueError(f'Unable to parse {reservation} as rescpu,resgpu')
+
+    if reservation_cpu.lower() == 'none':
+        reservation_cpu = None
+
+    if reservation_gpu.lower() == 'none':
+        reservation_gpu = None
+
+    system_name = default_system(jobdesc)
+    config = get_config(system_name)
+
+    if 'gpus_per_node' not in config or config['gpus_per_node'] == 0:
+        return reservation_cpu
+    else:
+        return reservation_gpu
 
 
 

py/desispec/workflow/processing.py

@@ -26,6 +26,7 @@
     get_desi_proc_batch_file_path, \
     get_desi_proc_tilenight_batch_file_pathname, \
     create_desi_proc_tilenight_batch_script, create_linkcal_batch_script
+from desispec.workflow.batch import parse_reservation
 from desispec.workflow.utils import pathjoin, sleep_and_report, \
     load_override_file
 from desispec.workflow.tableio import write_table, load_table
@@ -680,8 +681,11 @@ def submit_batch_script(prow, dry_run=0, reservation=None, strictly_successful=F
     batch_params = ['sbatch', '--parsable']
     if dep_str != '':
         batch_params.append(f'{dep_str}')
+
+    reservation = parse_reservation(reservation, prow['JOBDESC'])
     if reservation is not None:
         batch_params.append(f'--reservation={reservation}')
+
     batch_params.append(f'{script_path}')
 
     if dry_run: