diff --git a/payu/schedulers/pbs.py b/payu/schedulers/pbs.py index 4d5410bf..adcee9b5 100644 --- a/payu/schedulers/pbs.py +++ b/payu/schedulers/pbs.py @@ -6,6 +6,7 @@ # Standard library import os +from pathlib import Path import re import sys import shlex @@ -129,6 +130,17 @@ def submit(self, pbs_script, pbs_config, pbs_vars=None, python_exe=None): envmod.setup() envmod.module('load', 'pbs') + # Check for custom container launcher script environment variable + launcher_script = os.environ.get('ENV_LAUNCHER_SCRIPT_PATH') + if ( + launcher_script + and Path(launcher_script).is_file() + and os.access(launcher_script, os.X_OK) + ): + # Prepend the container launcher script to the python command + # so the python executable is accessible in the container + python_exe = f'{launcher_script} {python_exe}' + # Construct job submission command cmd = 'qsub {flags} -- {python} {script}'.format( flags=' '.join(pbs_flags), diff --git a/test/test_pbs.py b/test/test_pbs.py index 7389c17f..2637b20e 100644 --- a/test/test_pbs.py +++ b/test/test_pbs.py @@ -5,6 +5,7 @@ from pathlib import Path import shutil import sys +from unittest.mock import patch import pdb import pytest @@ -236,6 +237,47 @@ def test_run(): assert(args.remaining[-1].endswith(payu_cmd)) +@patch("payu.schedulers.pbs.pbs_env_init", return_value=True) +@patch("payu.schedulers.pbs.check_exe_path", side_effect=lambda x, y: y) +@pytest.mark.parametrize( + "env_exists,file_exists,file_exe,expected_cmd", + [ + # Test backwards compatibility with no launcher script + (False, False, False, "/path/to/python payu-run"), + # With only launcher script env set + (True, False, False, "/path/to/python payu-run"), + # With launch script env and file exists + (True, True, False, "/path/to/python payu-run"), + # With launch script env, file exists, and file is executable + (True, True, True, "{tmp_path}/launcher.sh /path/to/python payu-run") + ], +) +def test_submit_launcher_script_setting( + mock_pbs_env_init, mock_check_exe_path, + env_exists, file_exists, file_exe, expected_cmd, tmp_path, monkeypatch +): + config = { + "control_path": "/path/to/experiment" + } + + # Setup based on test parameters + if env_exists: + monkeypatch.setenv("ENV_LAUNCHER_SCRIPT_PATH", + f"{tmp_path}/launcher.sh") + if file_exists: + launcher_script_path = tmp_path / "launcher.sh" + launcher_script_path.write_text("#!/bin/bash\necho 'Running...'\n") + if file_exe: + launcher_script_path.chmod(0o755) + + # Generate the qsub command + pbs_cmd = pbs.PBS().submit("payu-run", config, + python_exe="/path/to/python") + + _, cmd = pbs_cmd.split("--") + assert cmd.strip() == expected_cmd.format(tmp_path=tmp_path) + + def test_tenacity(): # This should fail and do nothing