update batch script generator

desihub · Nov 2, 2023 · ca1911d · ca1911d
1 parent 8ee4fd5
commit ca1911d
Show file tree

Hide file tree

Showing 3 changed files with 86 additions and 29 deletions.
diff --git a/py/specprodDB/batch.py b/py/specprodDB/batch.py
@@ -9,24 +9,25 @@
 import os
 from sys import argv
 from argparse import ArgumentParser
+from . import __version__ as specprod_db_version
 
 
 template = """#!/bin/{shell}
 #SBATCH --qos={qos}
 #SBATCH --constraint={constraint}
 #SBATCH --nodes=1
 #SBATCH --time={time}
-#SBATCH --job-name=load_specprod_db_{schema}_{stage}
+#SBATCH --job-name=load_specprod_db_{script_schema}_{stage}
 #SBATCH --output={job_dir}/%x-%j.log
 #SBATCH --licenses=SCRATCH,cfs
 #SBATCH --account=desi
 #SBATCH --mail-type=end,fail
 #SBATCH --mail-user={email}
-module {load} specprod-db/main
+module {load} specprod-db/{load_version}
 {export_root}
 {export_specprod}
 srun --ntasks=1 load_specprod_db {overwrite} \\
-    --load {stage} --schema ${{SPECPROD}} ${{DESI_ROOT}}
+    --load {stage} --schema {schema} ${{DESI_ROOT}}
 """
 
 times = {'exposures': '00:10:00',
@@ -63,8 +64,8 @@ def get_options():
     #                   type=int, default=50000, metavar='N',
     #                   help="Load N rows at a time (default %(default)s).")
     prsr.add_argument('-s', '--schema', action='store', dest='schema',
-                      metavar='SCHEMA',
-                      help='Set the schema name in the PostgreSQL database.')
+                      metavar='SCHEMA', default='${SPECPROD}',
+                      help='Set the schema name in the PostgreSQL database (default "%(default)s").')
     prsr.add_argument('-S', '--swap', action='store_true', dest='swap',
                       help='Perform "module swap" instead of "module load".')
     # prsr.add_argument('-t', '--tiles-path', action='store', dest='tilespath', metavar='PATH',
@@ -74,10 +75,11 @@ def get_options():
     # prsr.add_argument('-U', '--username', action='store', dest='username',
     #                   metavar='USERNAME', default='desi_admin',
     #                   help='If specified, connect to a PostgreSQL database with USERNAME (default "%(default)s").')
-    # prsr.add_argument('-v', '--verbose', action='store_true', dest='verbose',
-    #                   help='Print extra information.')
+    prsr.add_argument('-v', '--specprod-version', metavar='VERSION', dest='specprod_version',
+                      help='Set the specprod-db version to VERSION.')
     prsr.add_argument('email', metavar='EMAIL', help='Send batch messages to EMAIL.')
     prsr.add_argument('root', metavar='DIR', help='Load the data in DIR.')
+    prsr.add_argument('specprod', metavar='SPECPROD', help='Set the value of SPECPROD, which may be different from the schema name.')
     options = prsr.parse_args()
     return options
 
@@ -99,22 +101,30 @@ def prepare_template(options):
         extension = 'csh'
         shell = 'tcsh'
         export_root = f'setenv DESI_ROOT {options.root}'
-        export_specprod = f'setenv SPECPROD {options.schema}'
+        export_specprod = f'setenv SPECPROD {options.specprod}'
     else:
         extension = 'sh'
         shell = 'bash'
         export_root = f'export DESI_ROOT={options.root}'
-        export_specprod = f'export SPECPROD={options.schema}'
+        export_specprod = f'export SPECPROD={options.specprod}'
     scripts = dict()
     load = 'load'
     if options.swap:
         load = 'swap'
+    if options.specprod_version is None:
+        load_version = specprod_db_version
+    else:
+        load_version = options.specprod_version
     for stage in ('exposures', 'photometry', 'targetphot', 'target', 'redshift', 'fiberassign'):
         if stage == 'exposures':
             overwrite = '--overwrite'
         else:
             overwrite = ''
-        script_name = 'load_specprod_db_{schema}_{stage}.{extension}'.format(schema=options.schema, stage=stage, extension=extension)
+        if options.schema == '${SPECPROD}':
+            script_schema = options.specprod
+        else:
+            script_schema = options.schema
+        script_name = 'load_specprod_db_{schema}_{stage}.{extension}'.format(schema=script_schema, stage=stage, extension=extension)
         try:
             wall_time = times[stage]
         except KeyError:
@@ -123,11 +133,13 @@ def prepare_template(options):
              'qos': options.qos,
              'constraint': options.constraint,
              'time': wall_time,
+             'script_schema': script_schema,
              'load': load,
              'schema': options.schema,
              'stage': stage,
              'job_dir': options.job_dir,
              'email': options.email,
+             'load_version': load_version,
              'export_root': export_root,
              'export_specprod': export_specprod,
              'overwrite': overwrite}

diff --git a/py/specprodDB/load.py b/py/specprodDB/load.py
@@ -1750,4 +1750,11 @@ def main():
             return 1
         log.info("Finished target bitmask corrections for %s zpix table.",
                  specprod)
+    if options.load == 'fiberassign':
+        #
+        # Automatically VACUUM.
+        #
+        log.info("Issuing VACUUM command.")
+        dbSession.execute("VACUUM FULL VERBOSE ANALYZE;")
+        log.info("Finished with VACUUM command.")
     return 0
diff --git a/py/specprodDB/test/test_batch.py b/py/specprodDB/test/test_batch.py
@@ -6,14 +6,15 @@
 import unittest
 from unittest.mock import patch, mock_open, call
 from ..batch import get_options, prepare_template, write_scripts
+from .. import __version__ as specprod_db_version
 
 
 class TestBatch(unittest.TestCase):
     """Test specprodDB.batch.
     """
     @classmethod
     def setUpClass(cls):
-        pass
+        cls.maxDiff = None
 
     @classmethod
     def tearDownClass(cls):
@@ -25,38 +26,39 @@ def setUp(self):
     def tearDown(self):
         pass
 
-    @patch('sys.argv', ['prepare_batch_specprod_db', '--csh', '[email protected]', '/global/cfs/cdirs/desi'])
+    @patch('sys.argv', ['prepare_batch_specprod_db', '--csh', '[email protected]', '/global/cfs/cdirs/desi', 'fuji'])
     def test_get_options(self):
         """Test option parser.
         """
         options = get_options()
         self.assertTrue(options.csh)
         self.assertEqual(options.root, '/global/cfs/cdirs/desi')
+        self.assertEqual(options.specprod, 'fuji')
 
-    @patch('sys.argv', ['prepare_batch_specprod_db', '--csh', '--schema', 'fuji', '[email protected]', '/global/cfs/cdirs/desi'])
+    @patch('sys.argv', ['prepare_batch_specprod_db', '--csh', '--schema', 'fuji_test', '[email protected]', '/global/cfs/cdirs/desi', 'fuji'])
     def test_prepare_template_csh(self):
         """Test conversion of options to scripts with csh.
         """
         with patch.dict('os.environ', {'DESI_ROOT': '/global/cfs/cdirs/desi', 'DESI_TARGET': '/global/cfs/cdirs/desi/target'}):
             options = get_options()
             scripts = prepare_template(options)
-        self.assertIn('load_specprod_db_fuji_exposures.csh', scripts)
+        self.assertIn('load_specprod_db_fuji_test_exposures.csh', scripts)
 
-    @patch('sys.argv', ['prepare_batch_specprod_db', '--swap', '--schema', 'fuji', '[email protected]', '/global/cfs/cdirs/desi'])
+    @patch('sys.argv', ['prepare_batch_specprod_db', '--swap', '--schema', 'fuji_test', '[email protected]', '/global/cfs/cdirs/desi', 'fuji'])
     def test_prepare_template_bash(self):
         """Test conversion of options to scripts with bash.
         """
         with patch.dict('os.environ', {'DESI_ROOT': '/global/cfs/cdirs/desi', 'DESI_TARGET': '/global/cfs/cdirs/desi/target'}):
             options = get_options()
             scripts = prepare_template(options)
-        self.assertIn('load_specprod_db_fuji_exposures.sh', scripts)
-        self.assertIn('module swap', scripts['load_specprod_db_fuji_exposures.sh'])
+        self.assertIn('load_specprod_db_fuji_test_exposures.sh', scripts)
+        self.assertIn('module swap', scripts['load_specprod_db_fuji_test_exposures.sh'])
 
-    @patch('sys.argv', ['prepare_batch_specprod_db', '--qos', 'bigmem', '--constraint', 'haswell', '--schema', 'fuji', '[email protected]', '/global/cfs/cdirs/desi'])
+    @patch('sys.argv', ['prepare_batch_specprod_db', '--qos', 'bigmem', '--constraint', 'haswell', '[email protected]', '/global/cfs/cdirs/desi', 'fuji'])
     def test_prepare_template_bash_qos(self):
         """Test conversion of options to scripts with bash and alternate qos/constraint.
         """
-        exposures = """#!/bin/bash
+        expected = f"""#!/bin/bash
 #SBATCH --qos=bigmem
 #SBATCH --constraint=haswell
 #SBATCH --nodes=1
@@ -67,15 +69,52 @@ def test_prepare_template_bash_qos(self):
 #SBATCH --account=desi
 #SBATCH --mail-type=end,fail
 #SBATCH [email protected]
-module load specprod-db/main
+module load specprod-db/{specprod_db_version}
 export DESI_ROOT=/global/cfs/cdirs/desi
 export SPECPROD=fuji
 srun --ntasks=1 load_specprod_db --overwrite \\
-    --load exposures --schema ${SPECPROD} ${DESI_ROOT}
+    --load exposures --schema ${{SPECPROD}} ${{DESI_ROOT}}
 """
-        photometry = """#!/bin/bash
-#SBATCH --qos=bigmem
-#SBATCH --constraint=haswell
+        with patch('os.environ', {'HOME': '/home/test'}):
+            options = get_options()
+            scripts = prepare_template(options)
+        self.assertIn('load_specprod_db_fuji_exposures.sh', scripts)
+        self.assertEqual(scripts['load_specprod_db_fuji_exposures.sh'], expected)
+
+    @patch('sys.argv', ['prepare_batch_specprod_db', '--schema', 'fuji_test', '[email protected]', '/global/cfs/cdirs/desi', 'fuji'])
+    def test_prepare_template_bash_schema(self):
+        """Test conversion of options to scripts with bash and alternate schema name.
+        """
+        expected = f"""#!/bin/bash
+#SBATCH --qos=regular
+#SBATCH --constraint=cpu
+#SBATCH --nodes=1
+#SBATCH --time=04:00:00
+#SBATCH --job-name=load_specprod_db_fuji_test_photometry
+#SBATCH --output=/home/test/Documents/Jobs/%x-%j.log
+#SBATCH --licenses=SCRATCH,cfs
+#SBATCH --account=desi
+#SBATCH --mail-type=end,fail
+#SBATCH [email protected]
+module load specprod-db/{specprod_db_version}
+export DESI_ROOT=/global/cfs/cdirs/desi
+export SPECPROD=fuji
+srun --ntasks=1 load_specprod_db  \\
+    --load photometry --schema fuji_test ${{DESI_ROOT}}
+"""
+        with patch('os.environ', {'HOME': '/home/test'}):
+            options = get_options()
+            scripts = prepare_template(options)
+        self.assertIn('load_specprod_db_fuji_test_photometry.sh', scripts)
+        self.assertEqual(scripts['load_specprod_db_fuji_test_photometry.sh'], expected)
+
+    @patch('sys.argv', ['prepare_batch_specprod_db', '--specprod-version', 'main', '[email protected]', '/global/cfs/cdirs/desi', 'fuji'])
+    def test_prepare_template_bash_version(self):
+        """Test conversion of options to scripts with bash and alternate version name.
+        """
+        expected = f"""#!/bin/bash
+#SBATCH --qos=regular
+#SBATCH --constraint=cpu
 #SBATCH --nodes=1
 #SBATCH --time=04:00:00
 #SBATCH --job-name=load_specprod_db_fuji_photometry
@@ -88,16 +127,15 @@ def test_prepare_template_bash_qos(self):
 export DESI_ROOT=/global/cfs/cdirs/desi
 export SPECPROD=fuji
 srun --ntasks=1 load_specprod_db  \\
-    --load photometry --schema ${SPECPROD} ${DESI_ROOT}
+    --load photometry --schema ${{SPECPROD}} ${{DESI_ROOT}}
 """
         with patch('os.environ', {'HOME': '/home/test'}):
             options = get_options()
             scripts = prepare_template(options)
-        self.assertIn('load_specprod_db_fuji_exposures.sh', scripts)
-        self.assertEqual(scripts['load_specprod_db_fuji_exposures.sh'], exposures)
-        self.assertEqual(scripts['load_specprod_db_fuji_photometry.sh'], photometry)
+        self.assertIn('load_specprod_db_fuji_photometry.sh', scripts)
+        self.assertEqual(scripts['load_specprod_db_fuji_photometry.sh'], expected)
 
-    @patch('sys.argv', ['prepare_batch_specprod_db', '--schema', 'fuji', '[email protected]', '/global/cfs/cdirs/desi'])
+    @patch('sys.argv', ['prepare_batch_specprod_db', '--schema', 'fuji', '[email protected]', '/global/cfs/cdirs/desi', 'fuji'])
     def test_write_scripts(self):
         """Test conversion of options to scripts with bash.
         """