nb2slurm 0.0.1.dev1__py3-none-any.whl

nb2slurm/__init__.py

@@ -0,0 +1,37 @@
+"""nb2slurm: scale a single-subject notebook workflow to many subjects on SLURM.
+
+The public surface is intentionally tiny so the whole workflow can be driven
+from a notebook with no command line:
+
+    import nb2slurm
+
+    wf = nb2slurm.Workflow(
+        name="square",
+        notebooks=["notebooks/0_settings.ipynb", "notebooks/1_compute.ipynb"],
+        kernel="python3",
+        varying=["item_id"],
+        resources=dict(nodes=1, cpus=2, time="00:10:00"),
+    )
+    wf.build()                       # render scripts/ into the project
+    wf.submit([1, 2, 3], ssh=cfg)    # sbatch one job per item, over SSH
+    wf.status(ssh=cfg)               # squeue -> list of dicts
+    wf.cancel(ssh=cfg)              # scancel the jobs we submitted
+
+Inside the notebooks themselves, use the Settings helper:
+
+    nb2slurm.Settings.write(outdir, {...})   # first notebook
+    settings = nb2slurm.Settings.load(path)  # later notebooks
+"""
+
+from .workflow import Workflow
+from .environment import Environment
+from .ssh import SSHConfig, generate_key, public_key
+from .settings import Settings
+from .done import Done
+from .structure import Structure
+from .config import save_config, load_config
+from .runtime import on_hpc
+
+__all__ = ["Workflow", "Environment", "SSHConfig", "generate_key", "public_key",
+           "Settings", "Done", "Structure", "save_config", "load_config", "on_hpc"]
+__version__ = "0.0.1.dev1"

nb2slurm/config.py

@@ -0,0 +1,70 @@
+"""Save / load a run's control settings to one JSON file.
+
+The config notebook builds a Workflow (+ SSHConfig) and calls ``save_config``;
+the build / submit / sync notebooks call ``load_config`` to get them back. This
+keeps a single source of truth on disk, so the notebooks never duplicate settings.
+
+    # in 0_config.ipynb
+    nb2slurm.save_config("control_config.json", workflow=wf, ssh=cfg)
+
+    # in 1_build / 2_submit / 3_sync
+    wf, cfg = nb2slurm.load_config("control_config.json")
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict
+from pathlib import Path
+from typing import Optional
+
+from .environment import Environment
+from .ssh import SSHConfig
+from .workflow import Workflow
+
+
+def _workflow_to_dict(wf: Workflow) -> dict:
+    d = asdict(wf)               # recurses Environment, resources, mounts
+    d.pop("submitted_jobs", None)  # runtime state, not configuration
+    return d
+
+
+def _workflow_from_dict(d: dict) -> Workflow:
+    d = dict(d)
+    d.pop("submitted_jobs", None)
+    env = d.pop("environment", None)
+    extras = d.pop("extra_environments", None) or []
+    return Workflow(
+        **d,
+        environment=Environment(**env) if env else None,
+        extra_environments=[Environment(**e) for e in extras],
+    )
+
+
+def _ssh_to_dict(ssh: SSHConfig) -> dict:
+    d = asdict(ssh)
+    d.pop("password", None)      # never persist secrets to disk
+    return d
+
+
+def save_config(path: str | Path, *, workflow: Workflow,
+                ssh: Optional[SSHConfig] = None) -> Path:
+    """Write the workflow (and optional SSH) config to ``path`` as JSON."""
+    data = {"workflow": _workflow_to_dict(workflow)}
+    if ssh is not None:
+        data["ssh"] = _ssh_to_dict(ssh)
+    path = Path(path)
+    path.write_text(json.dumps(data, indent=2), encoding="utf-8")
+    return path
+
+
+def load_config(path: str | Path) -> tuple[Workflow, Optional[SSHConfig]]:
+    """Read a config file back into ``(workflow, ssh)``.
+
+    ``ssh`` is ``None`` if none was saved. Passwords are never stored, so set
+    ``ssh.password`` yourself afterwards if your cluster needs one.
+    """
+    data = json.loads(Path(path).read_text(encoding="utf-8"))
+    wf = _workflow_from_dict(data["workflow"])
+    ssh = SSHConfig(**data["ssh"]) if data.get("ssh") else None
+    return wf, ssh

nb2slurm/done.py

@@ -0,0 +1,65 @@
+"""Idempotent 'already done?' bookkeeping, as an object.
+
+Generalised from cci.py: a single CSV records which subjects have completed, so
+re-submitting the whole set skips finished work. A file lock makes it safe when
+many jobs write concurrently on a shared filesystem.
+
+    done = nb2slurm.Done("done/done.csv")
+    if done.is_done(key):
+        ...
+    done.mark(key)
+"""
+
+from __future__ import annotations
+
+import csv
+from pathlib import Path
+
+from filelock import FileLock
+
+LOCK_TIMEOUT = 60 * 3
+
+
+class Done:
+    """A CSV ledger of completed subjects, with concurrency-safe writes."""
+
+    def __init__(self, csv_file: str | Path):
+        self.csv_file = Path(csv_file)
+
+    @property
+    def _lock(self) -> str:
+        return str(self.csv_file) + ".lock"
+
+    def is_done(self, key: str) -> bool:
+        """Return True if ``key`` is already recorded."""
+        if not self.csv_file.exists():
+            return False
+        with FileLock(self._lock, timeout=LOCK_TIMEOUT):
+            with open(self.csv_file, newline="") as f:
+                reader = csv.reader(f)
+                next(reader, None)  # header
+                for row in reader:
+                    if row and row[0] == str(key):
+                        return True
+        return False
+
+    def mark(self, key: str) -> None:
+        """Record ``key`` as done (no-op if already present)."""
+        self.csv_file.parent.mkdir(parents=True, exist_ok=True)
+        if not self.csv_file.exists():
+            with open(self.csv_file, "w", newline="") as f:
+                csv.writer(f).writerow(["key"])
+        with FileLock(self._lock, timeout=LOCK_TIMEOUT):
+            existing = set()
+            with open(self.csv_file, newline="") as f:
+                reader = csv.reader(f)
+                next(reader, None)
+                for row in reader:
+                    if row:
+                        existing.add(row[0])
+            if str(key) not in existing:
+                with open(self.csv_file, "a", newline="") as f:
+                    csv.writer(f).writerow([key])
+                print(f"Marked {key} as done.")
+            else:
+                print(f"{key} already recorded as done.")

nb2slurm/environment.py

@@ -0,0 +1,140 @@
+"""Create the conda environment + Jupyter kernel the workflow runs in.
+
+Most users of nb2slurm are not Linux/conda experts, but the generated SLURM job
+does ``conda activate <env>`` and papermill needs a *registered Jupyter kernel*
+to execute the notebooks. This module writes an ``environment.yml`` and creates
+the environment + kernel on the cluster (or locally), so the user never touches
+the command line.
+
+    from nb2slurm import Environment
+
+    env = Environment(
+        name="myenv",
+        kernel="myenv",                       # must match Workflow(kernel=...)
+        conda_packages=["xarray", "numpy"],
+        pip_packages=["nb2slurm", "ewatercycle"],
+    )
+    env.write()                # -> environment.yml
+    env.create(ssh=cfg)        # build env + register kernel on the HPC
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional
+
+from .ssh import CommandResult, SSHConfig, run_shell
+
+
+@dataclass
+class Environment:
+    name: str
+    kernel: str
+    python: str = "3.11"
+    channels: list[str] = field(default_factory=lambda: ["conda-forge"])
+    conda_packages: list[str] = field(default_factory=list)
+    pip_packages: list[str] = field(default_factory=lambda: ["nb2slurm"])
+
+    def to_yaml(self) -> str:
+        """Render an ``environment.yml`` for conda/mamba."""
+        lines = [f"name: {self.name}", "channels:"]
+        lines += [f"  - {c}" for c in self.channels]
+        lines.append("dependencies:")
+        lines.append(f"  - python={self.python}")
+        lines.append("  - pip")
+        lines.append("  - ipykernel")  # required so papermill can run the notebooks
+        lines += [f"  - {p}" for p in self.conda_packages]
+        if self.pip_packages:
+            lines.append("  - pip:")
+            lines += [f"      - {p}" for p in self.pip_packages]
+        return "\n".join(lines) + "\n"
+
+    def write(self, project_dir: str | Path = ".", filename: str = "environment.yml") -> Path:
+        """Write the ``environment.yml`` into the project directory."""
+        path = Path(project_dir) / filename
+        path.write_text(self.to_yaml(), encoding="utf-8")
+        return path
+
+    def _exists_test(self) -> str:
+        """A shell test (exit 0 = env exists). Matches the name as a whole path
+        component so ``montecarlo`` doesn't match ``montecarlo2``."""
+        return f'conda env list | grep -qE "[/ ]{self.name}([ /]|$)"'
+
+    def _create_command(self, filename: str = "environment.yml") -> str:
+        # mamba is much faster than conda; use it when available.
+        return (
+            "set -e; "
+            # There's no TTY over ssh, so any interactive prompt would hang the
+            # build forever. Belt and suspenders: set always-yes AND pipe `yes`
+            # into conda (CONDA_ALWAYS_YES alone is ignored by some mamba builds
+            # for the 'Confirm changes? [Y/n]' transaction prompt).
+            "export CONDA_ALWAYS_YES=yes; "
+            "CONDA=conda; command -v mamba >/dev/null 2>&1 && CONDA=mamba; "
+            # update an existing env in place rather than hitting the interactive
+            # 'Found conda-prefix ... Overwrite? [y/N]' prompt.
+            f"if {self._exists_test()}; then "
+            f'  echo "Updating existing environment {self.name} with $CONDA..."; '
+            f"  yes | $CONDA env update -f {filename}; "
+            f"else "
+            f'  echo "Creating environment {self.name} with $CONDA..."; '
+            f"  yes | $CONDA env create -f {filename}; "
+            f"fi; "
+            f"conda run -n {self.name} python -m ipykernel install --user "
+            f'--name {self.kernel} --display-name "{self.kernel}"; '
+            f'echo "Environment {self.name} ready; kernel {self.kernel} registered."'
+        )
+
+    def _remove_command(self) -> str:
+        return (
+            "export CONDA_ALWAYS_YES=yes; "
+            f'echo "Removing environment {self.name} and kernel {self.kernel}..."; '
+            # `|| true`: removing a non-existent env/kernel is not an error here
+            f"conda env remove -n {self.name} || true; "
+            f"jupyter kernelspec remove -f {self.kernel} 2>/dev/null || true; "
+            f'echo "Removed {self.name}."'
+        )
+
+    def exists(
+        self,
+        ssh: Optional[SSHConfig] = None,
+        project_dir: str | Path = ".",
+    ) -> bool:
+        """Return True if the conda env already exists — on the HPC (ssh) or locally."""
+        return run_shell(self._exists_test(), ssh, str(project_dir)).exit_status == 0
+
+    def remove(
+        self,
+        ssh: Optional[SSHConfig] = None,
+        project_dir: str | Path = ".",
+        stream: bool = True,
+    ) -> CommandResult:
+        """Delete the conda env and its Jupyter kernel — on the HPC (ssh) or locally.
+
+        Safe to call when nothing is there yet (a missing env/kernel is ignored).
+        Use it to recover from a half-built env or to force a clean rebuild.
+        """
+        return run_shell(self._remove_command(), ssh, str(project_dir), stream=stream).check()
+
+    def create(
+        self,
+        ssh: Optional[SSHConfig] = None,
+        project_dir: str | Path = ".",
+        filename: str = "environment.yml",
+        stream: bool = True,
+        overwrite: bool = False,
+    ) -> CommandResult:
+        """Create the env and register the kernel — on the HPC (ssh) or locally.
+
+        Writes ``environment.yml`` first if it is missing. If the env already
+        exists it is *updated* in place; pass ``overwrite=True`` to delete and
+        rebuild it from scratch. ``stream=True`` (the default) echoes conda/pip
+        output live, since a solve + downloads can take minutes and would
+        otherwise look like a hang.
+        """
+        if not (Path(project_dir) / filename).exists():
+            self.write(project_dir, filename)
+        if overwrite:
+            self.remove(ssh=ssh, project_dir=project_dir, stream=stream)
+        command = self._create_command(filename)
+        return run_shell(command, ssh, str(project_dir), stream=stream).check()

nb2slurm/render.py

@@ -0,0 +1,48 @@
+"""Render the bundled Jinja2 templates into a project's scripts/ directory."""
+
+from __future__ import annotations
+
+from importlib import resources
+from pathlib import Path
+from typing import Any, Mapping
+
+from jinja2 import Environment
+
+
+def _env() -> Environment:
+    # No autoescaping: these are shell/python scripts, not HTML.
+    # Move the comment delimiters off the default "{# #}" so bash parameter
+    # expansions like ${#arr[@]} don't get mistaken for Jinja comments.
+    return Environment(
+        trim_blocks=True,
+        lstrip_blocks=True,
+        keep_trailing_newline=True,
+        comment_start_string="<#nb2slurm#",
+        comment_end_string="#nb2slurm#>",
+    )
+
+
+def render_template(template_name: str, context: Mapping[str, Any]) -> str:
+    """Render one bundled template (e.g. 'run_workflow.py.j2') to a string."""
+    source = (
+        resources.files("nb2slurm.templates")
+        .joinpath(template_name)
+        .read_text(encoding="utf-8")
+    )
+    return _env().from_string(source).render(**context)
+
+
+def write_rendered(
+    template_name: str,
+    out_path: str | Path,
+    context: Mapping[str, Any],
+    executable: bool = False,
+) -> Path:
+    """Render a template and write it to ``out_path``."""
+    out_path = Path(out_path)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    out_path.write_text(render_template(template_name, context), encoding="utf-8")
+    if executable:
+        mode = out_path.stat().st_mode
+        out_path.chmod(mode | 0o111)
+    return out_path

nb2slurm/runtime.py

@@ -0,0 +1,43 @@
+"""Detect whether the current notebook is running under nb2slurm on a cluster.
+
+Checking ``Path.home()`` for a username is fragile (breaks for other users, can't
+tell a cloud VM from a laptop). Instead we look at environment variables that are
+only present in a batch job:
+
+* the ``SLURM_*`` variables SLURM sets in every job, and
+* the ``NB2SLURM`` sentinel that the generated ``job.slurm`` exports (so this also
+  works on non-SLURM batch setups that run nb2slurm's job script).
+
+Use it in your notebooks for the things that genuinely differ between
+interactive and batch runs — machine-specific data paths, skipping ``!pip install``:
+
+    import nb2slurm
+    if nb2slurm.on_hpc():
+        data_dir = "/project/ewater/Data"
+    else:
+        data_dir = "/data/shared"
+
+It also cleans up importing a helper from ``scripts/``. On the cluster the job
+runs from the project root, so ``from scripts.foo import bar`` just works; run
+interactively from a ``notebooks/`` subfolder it doesn't, so add the project root
+to the path only when running locally:
+
+    import sys
+    from pathlib import Path
+    import nb2slurm
+    if not nb2slurm.on_hpc():
+        sys.path.append(str(Path().resolve().parent))
+    from scripts.montecarlo import estimate_pi
+"""
+
+from __future__ import annotations
+
+import os
+
+# present inside a SLURM job; NB2SLURM is exported by the generated job.slurm
+_BATCH_ENV_VARS = ("NB2SLURM", "SLURM_JOB_ID", "SLURM_JOBID")
+
+
+def on_hpc() -> bool:
+    """Return True if running inside an nb2slurm/SLURM batch job."""
+    return any(var in os.environ for var in _BATCH_ENV_VARS)

nb2slurm/settings.py

@@ -0,0 +1,46 @@
+"""The per-run settings file, as an object.
+
+The paper's notebook structure has the *first* notebook write a settings file
+(JSON) that every later notebook reads. This keeps the varying parameters in one
+place: only notebook 0 is parameterised by papermill, the rest just load it.
+
+    import nb2slurm
+
+    # first notebook (parameterised by nb2slurm):
+    nb2slurm.Settings.write(outdir, {"region_id": region_id, "outdir": outdir})
+
+    # every later notebook:
+    settings = nb2slurm.Settings.load(settings_path)
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, Mapping
+
+
+class Settings:
+    """Read/write the per-run ``settings.json``."""
+
+    filename = "settings.json"
+
+    @staticmethod
+    def write(outdir: str | Path, settings: Mapping[str, Any]) -> Path:
+        """Write ``settings`` to ``<outdir>/settings.json`` and return the path.
+
+        Call this from your first notebook. ``outdir`` is supplied to that
+        notebook by the generated runner, so you never hardcode a path.
+        """
+        outdir = Path(outdir)
+        outdir.mkdir(parents=True, exist_ok=True)
+        path = outdir / Settings.filename
+        with open(path, "w") as f:
+            json.dump(dict(settings), f, indent=2)
+        return path
+
+    @staticmethod
+    def load(settings_path: str | Path) -> dict[str, Any]:
+        """Load a settings file. Call this at the top of every later notebook."""
+        with open(settings_path) as f:
+            return json.load(f)

nb2slurm/ssh.py

@@ -0,0 +1,211 @@
+"""Minimal SSH transport so the managing notebook can drive SLURM with no CLI.
+
+This mimics the command line / ssh that the paper says should be hidden from the
+user: sbatch/squeue/scancel run on the cluster, but the user only writes Python.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional, Tuple
+
+
+@dataclass
+class CommandResult:
+    command: str
+    exit_status: int
+    stdout: str
+    stderr: str
+
+    def check(self) -> "CommandResult":
+        if self.exit_status != 0:
+            raise RuntimeError(
+                f"Remote command failed ({self.exit_status}): {self.command}\n{self.stderr}"
+            )
+        return self
+
+
+@dataclass
+class SSHConfig:
+    """Connection details for the HPC login node.
+
+    Provide either ``key_filename`` or ``password`` (or rely on an agent/known
+    config). ``remote_dir`` is the project directory on the cluster that the
+    generated scripts live in; commands are run from there.
+    """
+
+    host: str
+    user: str
+    remote_dir: str
+    port: int = 22
+    key_filename: Optional[str] = None
+    password: Optional[str] = None
+    extra_connect_kwargs: dict = field(default_factory=dict)
+
+    def key_path(self) -> Optional[str]:
+        """The private key path with ``~`` expanded, or ``None`` if unset.
+
+        paramiko opens ``key_filename`` directly and does **not** expand ``~``,
+        so we resolve it here (e.g. ``~/.ssh/id_rsa`` -> the absolute path).
+        """
+        return os.path.expanduser(self.key_filename) if self.key_filename else None
+
+    def rsync_ssh(self) -> str:
+        """The ``-e`` transport string rsync should use (ssh + port + key)."""
+        parts = ["ssh"]
+        if self.port != 22:
+            parts += ["-p", str(self.port)]
+        if self.key_filename:
+            parts += ["-i", self.key_path()]
+        return " ".join(parts)
+
+    def rsync_target(self, subpath: str = "") -> str:
+        """A ``user@host:remote_dir/<subpath>`` spec for rsync."""
+        base = self.remote_dir.rstrip("/")
+        return f"{self.user}@{self.host}:{base}/{subpath}" if subpath else f"{self.user}@{self.host}:{base}/"
+
+    def run(self, command: str, cwd: Optional[str] = None,
+            stream: bool = False) -> CommandResult:
+        """Run a single command on the cluster and return its result.
+
+        Output is drained continuously while the command runs, so a chatty
+        command (``conda env create``, ``pip install``) can't fill paramiko's
+        channel window and deadlock against ``recv_exit_status``. Pass
+        ``stream=True`` to also echo output live — useful for long-running
+        builds where you'd otherwise see nothing until they finish.
+        """
+        import paramiko  # imported lazily so the package imports without a cluster
+
+        cwd = cwd or self.remote_dir
+        wrapped = f"cd {cwd} && {command}" if cwd else command
+
+        client = paramiko.SSHClient()
+        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        try:
+            client.connect(
+                hostname=self.host,
+                port=self.port,
+                username=self.user,
+                key_filename=self.key_path(),
+                password=self.password,
+                **self.extra_connect_kwargs,
+            )
+            chan = client.get_transport().open_session()
+            chan.exec_command(wrapped)
+
+            out_parts: list[str] = []
+            err_parts: list[str] = []
+
+            def _drain() -> bool:
+                got = False
+                while chan.recv_ready():
+                    chunk = chan.recv(32768).decode("utf-8", "replace")
+                    out_parts.append(chunk)
+                    if stream:
+                        print(chunk, end="", flush=True)
+                    got = True
+                while chan.recv_stderr_ready():
+                    chunk = chan.recv_stderr(32768).decode("utf-8", "replace")
+                    err_parts.append(chunk)
+                    if stream:
+                        print(chunk, end="", flush=True)
+                    got = True
+                return got
+
+            # keep reading so the remote side never blocks on a full window
+            while not chan.exit_status_ready():
+                if not _drain():
+                    time.sleep(0.05)
+            while _drain():  # whatever is left after exit
+                pass
+            status = chan.recv_exit_status()
+        finally:
+            client.close()
+        return CommandResult(wrapped, status, "".join(out_parts), "".join(err_parts))
+
+
+def _pub_path(path: str) -> Path:
+    """The ``.pub`` file for a key path (accepts the private path or the .pub)."""
+    p = Path(os.path.expanduser(path))
+    return p if p.suffix == ".pub" else Path(str(p) + ".pub")
+
+
+def public_key(path: str = "~/.ssh/id_rsa") -> str:
+    """Return the public key line for ``path`` (reads ``<path>.pub``).
+
+    This is the text you paste into your HPC — ``print(nb2slurm.public_key())``
+    then copy it into the cluster's key-upload page (or ``~/.ssh/authorized_keys``
+    on a login node, if your HPC lets you edit it directly).
+    """
+    return _pub_path(path).read_text().strip()
+
+
+def generate_key(path: str = "~/.ssh/id_rsa", bits: int = 4096,
+                 comment: Optional[str] = None, overwrite: bool = False,
+                 show: bool = True) -> Tuple[Path, Path]:
+    """Create an RSA SSH keypair at ``path`` (+ ``<path>.pub``).
+
+    Returns ``(private_path, public_path)``. The private key is written 0600 and
+    the public key in ``authorized_keys`` format. An existing key is left alone
+    unless ``overwrite=True``, so this is safe to call repeatedly. Point your
+    ``SSHConfig(key_filename=...)`` at ``path``.
+
+    nb2slurm can't install the key for you — many HPCs disable password login, so
+    there's no way in. With ``show=True`` (default) the public key is printed so
+    you can copy it into your cluster's key-upload page (or its
+    ``~/.ssh/authorized_keys``); ``nb2slurm.public_key(path)`` reprints it later.
+    """
+    import paramiko  # lazy: keep the package importable without a crypto backend
+
+    priv = Path(os.path.expanduser(path))
+    pub = _pub_path(path)
+    if priv.exists() and not overwrite:
+        if show:
+            print(f"key already exists at {priv}; its public key is:\n\n{public_key(path)}")
+        return priv, pub
+    priv.parent.mkdir(parents=True, exist_ok=True)
+
+    key = paramiko.RSAKey.generate(bits)
+    key.write_private_key_file(str(priv))
+    pub.write_text(f"ssh-rsa {key.get_base64()} {comment or ''}".strip() + "\n")
+    for p, mode in ((priv, 0o600), (pub, 0o644)):
+        try:
+            os.chmod(p, mode)
+        except OSError:
+            pass  # Windows without POSIX perms; OpenSSH there enforces via ACLs
+    if show:
+        print(
+            f"created SSH key: {priv} (private) and {pub} (public)\n\n"
+            "Add the PUBLIC key below to your HPC - via its key-upload page, or by\n"
+            "appending it to ~/.ssh/authorized_keys on a login node. nb2slurm can't\n"
+            "do this step for you (clusters usually disable password login):\n\n"
+            f"{public_key(path)}"
+        )
+    return priv, pub
+
+
+def run_shell(command: str, ssh: Optional[SSHConfig] = None,
+              cwd: str = ".", stream: bool = False) -> CommandResult:
+    """Run a shell command on the cluster (via ``ssh``) or locally (subprocess).
+
+    Shared by Workflow and Environment so the ssh-vs-local branch lives in one
+    place. ``stream=True`` echoes output live (for long-running commands).
+    """
+    if ssh is not None:
+        return ssh.run(command, stream=stream)
+    if stream:
+        proc = subprocess.Popen(command, shell=True, cwd=str(cwd), text=True,
+                                stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+        parts: list[str] = []
+        for line in proc.stdout:  # tee: capture and echo
+            parts.append(line)
+            print(line, end="", flush=True)
+        proc.wait()
+        return CommandResult(command, proc.returncode, "".join(parts), "")
+    proc = subprocess.run(command, shell=True, cwd=str(cwd),
+                          capture_output=True, text=True)
+    return CommandResult(command, proc.returncode, proc.stdout, proc.stderr)