From 786467e32ecb1dfe9ef961875a6360b0fe0525bf Mon Sep 17 00:00:00 2001 From: Jason Boutte Date: Tue, 28 Jul 2026 14:52:49 -0700 Subject: [PATCH 1/2] 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 ) 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 --- CIME/XML/env_batch.py | 46 +++++++ CIME/data/config/xml_schemas/config_batch.xsd | 2 + CIME/tests/test_unit_xml_env_batch.py | 114 ++++++++++++++++-- 3 files changed, 154 insertions(+), 8 deletions(-) diff --git a/CIME/XML/env_batch.py b/CIME/XML/env_batch.py index 00147b86b3c..cc75e99b2ab 100644 --- a/CIME/XML/env_batch.py +++ b/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") diff --git a/CIME/data/config/xml_schemas/config_batch.xsd b/CIME/data/config/xml_schemas/config_batch.xsd index 10b8368538d..55bb2da1a6e 100644 --- a/CIME/data/config/xml_schemas/config_batch.xsd +++ b/CIME/data/config/xml_schemas/config_batch.xsd @@ -132,11 +132,13 @@ + + diff --git a/CIME/tests/test_unit_xml_env_batch.py b/CIME/tests/test_unit_xml_env_batch.py index 5b1295c4a68..b1e2dc71f10 100755 --- a/CIME/tests/test_unit_xml_env_batch.py +++ b/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""" + tfile.write(b"""
These variables may be changed anytime during a run, they @@ -548,8 +547,7 @@ def test_get_submit_args_job_queue(self): -""" - ) +""") 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""" + tfile.write(b"""
These variables may be changed anytime during a run, they @@ -629,8 +626,7 @@ def test_get_submit_args(self): -""" - ) +""") tfile.seek(0) @@ -1302,5 +1298,107 @@ def run_get_job_overrides( return overrides +XML_OMIT_IN_JOB = b""" + +
+ These variables may be changed anytime during a run, they + control arguments to the batch submit command. +
+ + + char + flux,slurm,pbs,lsf,none + The batch system type to use for this machine. + + + + + + + + + + -o exit-timeout=none + -p pbatch + + + +""" + + +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() From f5d5c3f8c6ed37c6b96e01602e4beae15f63c559 Mon Sep 17 00:00:00 2001 From: Jason Boutte Date: Tue, 28 Jul 2026 15:02:54 -0700 Subject: [PATCH 2/2] fix: black formatting --- CIME/tests/test_unit_xml_env_batch.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/CIME/tests/test_unit_xml_env_batch.py b/CIME/tests/test_unit_xml_env_batch.py index b1e2dc71f10..659cf41d8c4 100755 --- a/CIME/tests/test_unit_xml_env_batch.py +++ b/CIME/tests/test_unit_xml_env_batch.py @@ -514,7 +514,8 @@ def test_get_job_deps(self): def test_get_submit_args_job_queue(self): with tempfile.NamedTemporaryFile() as tfile: - tfile.write(b""" + tfile.write( + b"""
These variables may be changed anytime during a run, they @@ -547,7 +548,8 @@ def test_get_submit_args_job_queue(self): -""") +""" + ) tfile.seek(0) @@ -569,7 +571,8 @@ 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""" + tfile.write( + b"""
These variables may be changed anytime during a run, they @@ -626,7 +629,8 @@ def test_get_submit_args(self): -""") +""" + ) tfile.seek(0)