From 48863dfe99e27ffb76376cbf01382e4a2152ea26 Mon Sep 17 00:00:00 2001 From: Konrad Rokicki Date: Wed, 8 Jul 2026 19:57:38 -0400 Subject: [PATCH 1/2] Add inherit_env and login_shell submit options By default jobs inherit the submitting process's full environment (Popen os.environ merge locally; bsub's environment copy on LSF), which leaks whatever the submitter has set. The new kwargs let callers construct the job environment instead: - inherit_env=False: local jobs get a minimal baseline (identity vars + fixed PATH) plus explicit env; LSF submits with `-env none`, and any explicit env vars move into the script as exports since bsub will no longer copy them. - login_shell=True: run under a login shell (`bash -l` locally, `bsub -L /bin/bash` on LSF) so the target user's own profile builds PATH/modules/conda, as an SSH session would. Defaults preserve existing behavior. Co-Authored-By: Claude Fable 5 --- cluster_api/core.py | 23 +++++++++++++++- cluster_api/executors/local.py | 36 ++++++++++++++++++++++--- cluster_api/executors/lsf.py | 36 +++++++++++++++++++++++++ tests/test_local.py | 48 ++++++++++++++++++++++++++++++++++ tests/test_lsf.py | 40 ++++++++++++++++++++++++++++ 5 files changed, 178 insertions(+), 5 deletions(-) diff --git a/cluster_api/core.py b/cluster_api/core.py index d9eaba1..e66fe87 100644 --- a/cluster_api/core.py +++ b/cluster_api/core.py @@ -82,14 +82,27 @@ async def submit( epilogue: list[str] | None = None, env: dict[str, str] | None = None, metadata: dict[str, Any] | None = None, + inherit_env: bool = True, + login_shell: bool = False, ) -> JobRecord: - """Submit a job to the scheduler.""" + """Submit a job to the scheduler. + + Args: + inherit_env: When True (default), the job inherits the submitting + process's environment. When False, the job starts from a + minimal environment: only *env* plus scheduler-provided + variables reach the job. + login_shell: When True, the job runs under a login shell so the + target user's own profile builds the environment (PATH, + modules, conda), as if they had SSH'd to the node. + """ resources = resources or ResourceSpec() full_name = _sanitize_job_name(f"{self._prefix}-{name}" if self._prefix else name) job_id, script_path = await self._submit_job( command, full_name, resources, prologue, epilogue, env, cwd=resources.work_dir, + inherit_env=inherit_env, login_shell=login_shell, ) record = JobRecord( @@ -117,6 +130,8 @@ async def submit_array( env: dict[str, str] | None = None, metadata: dict[str, Any] | None = None, max_concurrent: int | None = None, + inherit_env: bool = True, + login_shell: bool = False, ) -> JobRecord: """Submit a job array to the scheduler.""" resources = resources or ResourceSpec() @@ -125,6 +140,7 @@ async def submit_array( job_id, script_path = await self._submit_array_job( command, full_name, array_range, resources, prologue, epilogue, env, max_concurrent, cwd=resources.work_dir, + inherit_env=inherit_env, login_shell=login_shell, ) meta = {**(metadata or {}), "array_range": array_range} @@ -159,6 +175,8 @@ async def _submit_job( env: dict[str, str] | None = None, *, cwd: str | None = None, + inherit_env: bool = True, + login_shell: bool = False, ) -> tuple[str, str | None]: """Submit a single job. @@ -179,10 +197,13 @@ async def _submit_array_job( max_concurrent: int | None = None, *, cwd: str | None = None, + inherit_env: bool = True, + login_shell: bool = False, ) -> tuple[str, str | None]: """Submit an array job. Override in subclasses.""" return await self._submit_job( command, name, resources, prologue, epilogue, env, cwd=cwd, + inherit_env=inherit_env, login_shell=login_shell, ) def _job_id_from_submit_output(self, out: str) -> str: diff --git a/cluster_api/executors/local.py b/cluster_api/executors/local.py index 69eaa1a..d4c43ec 100644 --- a/cluster_api/executors/local.py +++ b/cluster_api/executors/local.py @@ -17,6 +17,22 @@ logger = logging.getLogger(__name__) +# PATH for jobs that don't inherit the submitting process's environment; +# mirrors what cron/sshd hand a fresh process. A login shell's profile will +# prepend the user's own entries on top of this. +_BASELINE_PATH = "/usr/local/bin:/usr/bin:/bin" + +# Identity and locale variables that describe the *user*, not the submitting +# process, so they are safe to carry into a non-inheriting job. +_BASELINE_KEYS = ("HOME", "USER", "LOGNAME", "SHELL", "LANG", "TERM") + + +def _baseline_env(env: dict[str, str] | None) -> dict[str, str]: + """Minimal job environment for inherit_env=False submissions.""" + base = {k: os.environ[k] for k in _BASELINE_KEYS if k in os.environ} + base["PATH"] = _BASELINE_PATH + return {**base, **(env or {})} + class LocalExecutor(Executor): """Runs jobs as local bash subprocesses. Useful for testing.""" @@ -49,6 +65,8 @@ async def _submit_job( env: dict[str, str] | None = None, *, cwd: str | None = None, + inherit_env: bool = True, + login_shell: bool = False, ) -> tuple[str, str | None]: """Render script, write to disk, run as a background subprocess.""" header = self.build_header(name, resources) @@ -58,13 +76,17 @@ async def _submit_job( job_id = str(self._next_id) self._next_id += 1 - full_env = {**os.environ, **(env or {})} + if inherit_env: + full_env = {**os.environ, **(env or {})} + else: + full_env = _baseline_env(env) + bash_cmd = ["bash", "-l", script_path] if login_shell else ["bash", script_path] # Write stdout/stderr directly to log files for real-time access stdout_dest, stderr_dest = self._open_output_files(resources, job_id=job_id) proc = await asyncio.create_subprocess_exec( - "bash", script_path, + *bash_cmd, stdout=stdout_dest, stderr=stderr_dest, env=full_env, @@ -86,6 +108,8 @@ async def _submit_array_job( max_concurrent: int | None = None, *, cwd: str | None = None, + inherit_env: bool = True, + login_shell: bool = False, ) -> tuple[str, str | None]: """Spawn one subprocess per array element with ARRAY_INDEX env var.""" if max_concurrent is not None: @@ -101,7 +125,11 @@ async def _submit_array_job( job_id = str(self._next_id) self._next_id += 1 - full_env = {**os.environ, **(env or {})} + if inherit_env: + full_env = {**os.environ, **(env or {})} + else: + full_env = _baseline_env(env) + bash_cmd = ["bash", "-l", script_path] if login_shell else ["bash", script_path] for index in range(array_range[0], array_range[1] + 1): element_env = {**full_env, "ARRAY_INDEX": str(index)} @@ -109,7 +137,7 @@ async def _submit_array_job( resources, job_id=job_id, element_index=index, ) proc = await asyncio.create_subprocess_exec( - "bash", script_path, + *bash_cmd, stdout=stdout_dest, stderr=stderr_dest, env=element_env, diff --git a/cluster_api/executors/lsf.py b/cluster_api/executors/lsf.py index 497ec33..42ed4fc 100644 --- a/cluster_api/executors/lsf.py +++ b/cluster_api/executors/lsf.py @@ -8,6 +8,7 @@ import logging import math import re +import shlex from datetime import datetime, timezone from typing import Any @@ -134,6 +135,33 @@ def _collect_extra_args(self, resources: ResourceSpec | None = None) -> list[str args.extend(resources.extra_args) return args + @staticmethod + def _env_control_args(inherit_env: bool, login_shell: bool) -> list[str]: + """bsub args controlling the job's starting environment.""" + args: list[str] = [] + if not inherit_env: + # Don't copy the submitting process's environment into the job; + # it starts from scheduler-provided variables only. + args.extend(["-env", "none"]) + if login_shell: + # Initialize the execution environment through the user's login + # shell, as an SSH session to the node would. + args.extend(["-L", "/bin/bash"]) + return args + + @staticmethod + def _env_prologue(env: dict[str, str] | None, inherit_env: bool, + prologue: list[str] | None) -> list[str] | None: + """Fold explicit env vars into the script when bsub won't copy them. + + With ``-env none`` the submission environment never reaches the job, + so variables passed via *env* must be exported by the script itself. + """ + if inherit_env or not env: + return prologue + exports = [f"export {k}={shlex.quote(v)}" for k, v in env.items()] + return [*exports, *(prologue or [])] + async def _bsub( self, script_path: str, env: dict[str, str] | None, extra_args: list[str] | None = None, @@ -158,13 +186,17 @@ async def _submit_job( env: dict[str, str] | None = None, *, cwd: str | None = None, + inherit_env: bool = True, + login_shell: bool = False, ) -> tuple[str, str | None]: """Render script, write to disk, submit via bsub.""" header = self.build_header(name, resources) + prologue = self._env_prologue(env, inherit_env, prologue) script = render_script(self.config, command, header, prologue, epilogue) script_path = write_script(resources.work_dir, script, name, next(self._script_counter)) extra_args = self._collect_extra_args(resources) + extra_args.extend(self._env_control_args(inherit_env, login_shell)) out = await self._bsub(script_path, env, extra_args) return self._job_id_from_submit_output(out), script_path @@ -180,9 +212,12 @@ async def _submit_array_job( max_concurrent: int | None = None, *, cwd: str | None = None, + inherit_env: bool = True, + login_shell: bool = False, ) -> tuple[str, str | None]: """Render script, rewrite for array syntax, submit via bsub.""" header = self.build_header(name, resources) + prologue = self._env_prologue(env, inherit_env, prologue) script = render_script(self.config, command, header, prologue, epilogue) script_path = write_script(resources.work_dir, script, name, next(self._script_counter)) @@ -206,6 +241,7 @@ async def _submit_array_job( f.write(content) extra_args = self._collect_extra_args(resources) + extra_args.extend(self._env_control_args(inherit_env, login_shell)) out = await self._bsub(script_path, env, extra_args) return self._job_id_from_submit_output(out), script_path diff --git a/tests/test_local.py b/tests/test_local.py index f0dfa7f..1c3c9cd 100644 --- a/tests/test_local.py +++ b/tests/test_local.py @@ -43,6 +43,54 @@ async def test_submit_and_poll_failure(self, default_config, work_dir): assert job.exit_code == 1 + async def test_submit_clean_env(self, default_config, work_dir, monkeypatch): + """inherit_env=False keeps the submitter's environment out of the job.""" + monkeypatch.setenv("PIXI_PROJECT_MANIFEST", "/tmp/poison.toml") + monkeypatch.setenv("SECRET_TOKEN", "leaked") + executor = LocalExecutor(default_config) + out_file = Path(work_dir) / "env.txt" + job = await executor.submit( + command=f"printenv > {out_file}", name="cleanenv-test", + resources=ResourceSpec(work_dir=work_dir), + env={"EXPLICIT_VAR": "yes"}, + inherit_env=False, + ) + proc = executor._processes[job.job_id] + await proc.wait() + await executor.poll() + assert job.status == JobStatus.DONE + + content = out_file.read_text() + assert "PIXI_PROJECT_MANIFEST" not in content + assert "SECRET_TOKEN" not in content + assert "EXPLICIT_VAR=yes" in content + assert "HOME=" in content + assert "PATH=/usr/local/bin:/usr/bin:/bin" in content + + async def test_submit_login_shell_invocation(self, default_config, work_dir, monkeypatch): + """login_shell=True runs the script under bash -l.""" + import cluster_api.executors.local as local_mod + + recorded: list[tuple[str, ...]] = [] + real_exec = asyncio.create_subprocess_exec + + async def recording_exec(*args, **kwargs): + recorded.append(args) + return await real_exec(*args, **kwargs) + + monkeypatch.setattr(local_mod.asyncio, "create_subprocess_exec", recording_exec) + executor = LocalExecutor(default_config) + job = await executor.submit( + command="echo hello", name="login-test", + resources=ResourceSpec(work_dir=work_dir), + inherit_env=False, login_shell=True, + ) + proc = executor._processes[job.job_id] + await proc.wait() + + assert recorded[0][:2] == ("bash", "-l") + assert recorded[0][2].endswith(".sh") + async def test_running_job(self, default_config, work_dir): executor = LocalExecutor(default_config) job = await executor.submit( diff --git a/tests/test_lsf.py b/tests/test_lsf.py index 454376a..fe00e33 100644 --- a/tests/test_lsf.py +++ b/tests/test_lsf.py @@ -364,6 +364,46 @@ async def test_submit_email_suppression(self, lsf_config, work_dir): assert env.get("LSB_JOB_REPORT_MAIL") == "N" + async def test_submit_clean_env_login_shell(self, lsf_config, work_dir): + """inherit_env=False adds -env none; login_shell adds -L /bin/bash; + explicit env vars move into the script since bsub won't copy them.""" + executor = LSFExecutor(lsf_config) + with patch.object( + executor, "_call", + new_callable=AsyncMock, + return_value="Job <12345> is submitted to queue .", + ) as mock_call: + job = await executor.submit( + command="echo hello", name="my-job", + resources=ResourceSpec(work_dir=work_dir), + env={"FOO": "bar baz"}, + inherit_env=False, login_shell=True, + ) + cmd = mock_call.call_args[0][0] + assert cmd[cmd.index("-env") + 1] == "none" + assert cmd[cmd.index("-L") + 1] == "/bin/bash" + with open(job.script_path) as f: + script = f.read() + assert "export FOO='bar baz'" in script + # Exports must precede the command + assert script.index("export FOO") < script.index("echo hello") + + async def test_submit_default_env_args_absent(self, lsf_config, work_dir): + """Default submission keeps bsub's normal environment handling.""" + executor = LSFExecutor(lsf_config) + with patch.object( + executor, "_call", + new_callable=AsyncMock, + return_value="Job <12345> is submitted to queue .", + ) as mock_call: + await executor.submit( + command="echo hello", name="my-job", + resources=ResourceSpec(work_dir=work_dir), + ) + cmd = mock_call.call_args[0][0] + assert "-env" not in cmd + assert "-L" not in cmd + async def test_submit_array(self, lsf_config, work_dir): executor = LSFExecutor(lsf_config) with patch.object( From 2d39d976caa88ac19e9629efd27f156dc837c5a0 Mon Sep 17 00:00:00 2001 From: Konrad Rokicki Date: Wed, 8 Jul 2026 22:32:15 -0400 Subject: [PATCH 2/2] updated version to 0.7.0 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 30d032f..5ef0513 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "py-cluster-api" -version = "0.6.1" +version = "0.7.0" description = "Generic Python library for running jobs on HPC clusters" readme = "README.md" license = { file = "LICENSE" }