feat: support omit_in_job attribute on submit args to drop args invalid inside a job

Schedulers like flux give each job an ephemeral nested instance where
some submit args (e.g. -p <partition>) are invalid, so args marked
omit_in_job="true" are dropped when submission happens from inside an
already-running batch job, detected via well-known scheduler
environment variables.

Closes #5018

Author: jasonb5

CIME/XML/env_batch.py

@@ -24,6 +24,22 @@
 
 logger = logging.getLogger(__name__)
 
+# Static mapping of batch system type to the well-known environment
+# variable that indicates the current process is running inside an
+# active job for that scheduler. Batch schedulers change infrequently
+# so this is maintained in code rather than in per-machine config.
+IN_JOB_ENVIRONMENT_VARIABLES = {
+    "flux": "FLUX_JOB_ID",
+    "lsf": "LSB_JOBID",
+    "pbs": "PBS_JOBID",
+    "pbspro": "PBS_JOBID",
+    "moab": "PBS_JOBID",
+    "slurm": "SLURM_JOB_ID",
+    "slurm_single_node": "SLURM_JOB_ID",
+    "cobalt": "COBALT_JOBID",
+    "cobalt_theta": "COBALT_JOBID",
+}
+
 # pragma pylint: disable=attribute-defined-outside-init
 
 
@@ -704,7 +720,37 @@ def _process_args(self, case, submit_arg_nodes, job, resolve=True):
 
         return submitargs
 
+    def is_in_batch_job(self, environ=None):
+        """Checks whether the current process is running inside a batch job.
+
+        Detection is based on the presence of the scheduler specific
+        environment variable for the case's batch system, e.g.
+        ``FLUX_JOB_ID`` for flux or ``SLURM_JOB_ID`` for slurm. This is
+        used to drop submit args marked ``omit_in_job`` which are only
+        valid when submitting from outside a job, e.g. flux nested
+        instances define no partitions so ``-p`` must be omitted when
+        resubmitting from inside a job.
+
+        Args:
+            environ (dict, optional): Environment mapping to check,
+                defaults to ``os.environ``.
+
+        Returns:
+            bool: True if inside an active batch job, otherwise False.
+        """
+        if environ is None:
+            environ = os.environ
+
+        env_var = IN_JOB_ENVIRONMENT_VARIABLES.get(self._batchtype)
+
+        return env_var is not None and env_var in environ
+
     def _get_argument(self, case, arg):
+        omit_in_job = self.get(arg, "omit_in_job", default="false")
+
+        if omit_in_job.lower() in ("true", "1") and self.is_in_batch_job():
+            raise ValueError()
+
         flag = self.get(arg, "flag")
 
         name = self.get(arg, "name")

CIME/data/config/xml_schemas/config_batch.xsd

@@ -132,11 +132,13 @@
           <xs:complexType>
             <xs:attribute name="flag" use="required"/>
             <xs:attribute name="name"/>
+            <xs:attribute name="omit_in_job" type="xs:boolean"/>
           </xs:complexType>
         </xs:element>
         <xs:element name="argument" maxOccurs="unbounded">
           <xs:complexType mixed="true">
             <xs:attribute name="job_queue"/>
+            <xs:attribute name="omit_in_job" type="xs:boolean"/>
           </xs:complexType>
         </xs:element>
       </xs:choice>

CIME/tests/test_unit_xml_env_batch.py

@@ -514,8 +514,7 @@ def test_get_job_deps(self):
 
     def test_get_submit_args_job_queue(self):
         with tempfile.NamedTemporaryFile() as tfile:
-            tfile.write(
-                b"""<?xml version="1.0"?>
+            tfile.write(b"""<?xml version="1.0"?>
 <file id="env_batch.xml" version="2.0">
   <header>
       These variables may be changed anytime during a run, they
@@ -548,8 +547,7 @@ def test_get_submit_args_job_queue(self):
     </queues>
   </batch_system>
 </file>
-"""
-            )
+""")
 
             tfile.seek(0)
 
@@ -571,8 +569,7 @@ def test_get_submit_args_job_queue(self):
     @mock.patch.dict(os.environ, {"TEST": "GOOD"})
     def test_get_submit_args(self):
         with tempfile.NamedTemporaryFile() as tfile:
-            tfile.write(
-                b"""<?xml version="1.0"?>
+            tfile.write(b"""<?xml version="1.0"?>
 <file id="env_batch.xml" version="2.0">
   <header>
       These variables may be changed anytime during a run, they
@@ -629,8 +626,7 @@ def test_get_submit_args(self):
     </queues>
   </batch_system>
 </file>
-"""
-            )
+""")
 
             tfile.seek(0)
 
@@ -1302,5 +1298,107 @@ def run_get_job_overrides(
         return overrides
 
 
+XML_OMIT_IN_JOB = b"""<?xml version="1.0"?>
+<file id="env_batch.xml" version="2.0">
+  <header>
+      These variables may be changed anytime during a run, they
+      control arguments to the batch submit command.
+    </header>
+  <group id="config_batch">
+    <entry id="BATCH_SYSTEM" value="flux">
+      <type>char</type>
+      <valid_values>flux,slurm,pbs,lsf,none</valid_values>
+      <desc>The batch system type to use for this machine.</desc>
+    </entry>
+  </group>
+  <batch_system type="flux">
+    <submit_args>
+      <arg flag="--fixed" name="$PROJECT" omit_in_job="true"/>
+    </submit_args>
+  </batch_system>
+  <batch_system MACH="docker" type="flux">
+    <submit_args>
+      <argument>-o exit-timeout=none</argument>
+      <argument omit_in_job="true">-p pbatch</argument>
+    </submit_args>
+  </batch_system>
+</file>
+"""
+
+
+def _create_omit_in_job_batch(tmp_path):
+    infile = tmp_path / "env_batch.xml"
+
+    infile.write_bytes(XML_OMIT_IN_JOB)
+
+    batch = EnvBatch(infile=str(infile))
+
+    case = mock.MagicMock()
+
+    case.get_value.side_effect = lambda *args, **kwargs: {
+        "BATCH_SPEC_FILE": str(infile),
+        "PROJECT": "CIME",
+        "JOB_QUEUE": "pbatch",
+    }.get(args[0])
+
+    case.get_resolved_value.side_effect = lambda val: val
+
+    return batch, case
+
+
+def test_get_submit_args_omit_in_job_not_in_job(tmp_path, monkeypatch):
+    # Context
+    batch, case = _create_omit_in_job_batch(tmp_path)
+
+    monkeypatch.delenv("FLUX_JOB_ID", raising=False)
+
+    # Act
+    submit_args = batch.get_submit_args(case, ".case.run")
+
+    # Assert
+    assert submit_args == "  --fixed CIME -o exit-timeout=none -p pbatch"
+
+
+def test_get_submit_args_omit_in_job_in_job(tmp_path, monkeypatch):
+    # Context
+    batch, case = _create_omit_in_job_batch(tmp_path)
+
+    monkeypatch.setenv("FLUX_JOB_ID", "fuzzybunny")
+
+    # Act
+    submit_args = batch.get_submit_args(case, ".case.run")
+
+    # Assert
+    assert submit_args == "  -o exit-timeout=none"
+
+
+def test_get_submit_args_omit_in_job_other_scheduler_env(tmp_path, monkeypatch):
+    # Context
+    batch, case = _create_omit_in_job_batch(tmp_path)
+
+    monkeypatch.delenv("FLUX_JOB_ID", raising=False)
+
+    # Only the current batch system's env var is considered
+    monkeypatch.setenv("SLURM_JOB_ID", "1234")
+
+    # Act
+    submit_args = batch.get_submit_args(case, ".case.run")
+
+    # Assert
+    assert submit_args == "  --fixed CIME -o exit-timeout=none -p pbatch"
+
+
+def test_is_in_batch_job_unknown_batch_system(monkeypatch):
+    # Context
+    batch = EnvBatch()
+
+    batch._batchtype = "made_up_scheduler"
+
+    monkeypatch.setenv("SLURM_JOB_ID", "1234")
+
+    # Act/Assert
+    assert not batch.is_in_batch_job()
+
+
 if __name__ == "__main__":
     unittest.main()