Land inherit_env/login_shell on main; bump to 0.7.0

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:

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,15 +125,19 @@ 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)}
             stdout_dest, stderr_dest = self._open_output_files(
                 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,

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
 

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" }

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(

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 <normal>.",
+        ) 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 <normal>.",
+        ) 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(