From f2143d1f9f7c5c85a073831b3c3401674e7d0abf Mon Sep 17 00:00:00 2001 From: Tsvika Shapira Date: Tue, 28 Apr 2026 17:01:22 +0300 Subject: [PATCH 1/7] feat(runner): record meta/env in run config Extra env vars passed via Runner.env were logged to stderr but never written to run_info.json (or W&B config), so reproducing a run from the JSON alone was incomplete. Add meta/env as a dict; None values (vars explicitly unset for the subprocess) survive as JSON null. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/lite_runner/runner.py | 1 + tests/test_runner.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/src/lite_runner/runner.py b/src/lite_runner/runner.py index c96c0f7..a1901fa 100644 --- a/src/lite_runner/runner.py +++ b/src/lite_runner/runner.py @@ -495,6 +495,7 @@ def run( config["meta/hostname"] = os.uname().nodename config["meta/datetime"] = timestamp.isoformat() config["meta/command"] = shlex.join(r.command) + config["meta/env"] = dict(r.env) # Init WandbBackend first (needs to happen early to get run_name) backend_classes: list[type[WandbBackend | JsonBackend | DryRunBackend]] = [] diff --git a/tests/test_runner.py b/tests/test_runner.py index 130f2b4..ca68081 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -1133,6 +1133,7 @@ def test_full_run_no_wandb(tmp_path: Path) -> None: metrics=[Metric("val", pattern=r"x=([\d.]+)")], tags=["v1"], run_group="test-group", + env={"MY_FLAG": "1", "DROP_ME": None}, ) with ( @@ -1159,6 +1160,7 @@ def test_full_run_no_wandb(tmp_path: Path) -> None: # Check config assert run_info["config"]["git/repo"] == "test-repo" assert run_info["config"]["meta/output_dir"] == str(output_dir) + assert run_info["config"]["meta/env"] == {"MY_FLAG": "1", "DROP_ME": None} # Check metrics extracted assert run_info["metrics"]["val"] == 42.0 From 3430d717681e51e1d9e93a1636a1b83bb6394c1c Mon Sep 17 00:00:00 2001 From: Tsvika Shapira Date: Tue, 28 Apr 2026 17:31:24 +0300 Subject: [PATCH 2/7] feat(runner): record meta/cwd, meta/user, param_source/* in run config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reproducibility additions to the recorded run config: - meta/cwd, meta/user — capture working directory and invoking user - param_source/ — mirrors Runner.param_sources so the JSON shows whether each param value came from cli/default/fixed/prompt/override Co-Authored-By: Claude Opus 4.7 (1M context) --- src/lite_runner/runner.py | 5 +++++ tests/test_runner.py | 2 ++ 2 files changed, 7 insertions(+) diff --git a/src/lite_runner/runner.py b/src/lite_runner/runner.py index a1901fa..1223554 100644 --- a/src/lite_runner/runner.py +++ b/src/lite_runner/runner.py @@ -5,6 +5,7 @@ import argparse import copy import datetime +import getpass import hashlib import importlib.metadata import logging @@ -489,10 +490,14 @@ def run( config: dict[str, object] = {} for k, v in r.param_values.items(): config[f"param/{k}"] = "" if _contains_unset(v) else v + for k, v in r.param_sources.items(): + config[f"param_source/{k}"] = v for k, v in git_info.items(): config[f"git/{k}"] = v timestamp = datetime.datetime.now(tz=datetime.timezone.utc) config["meta/hostname"] = os.uname().nodename + config["meta/user"] = getpass.getuser() + config["meta/cwd"] = os.getcwd() # noqa: PTH109 config["meta/datetime"] = timestamp.isoformat() config["meta/command"] = shlex.join(r.command) config["meta/env"] = dict(r.env) diff --git a/tests/test_runner.py b/tests/test_runner.py index ca68081..d49fe85 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -1161,6 +1161,8 @@ def test_full_run_no_wandb(tmp_path: Path) -> None: assert run_info["config"]["git/repo"] == "test-repo" assert run_info["config"]["meta/output_dir"] == str(output_dir) assert run_info["config"]["meta/env"] == {"MY_FLAG": "1", "DROP_ME": None} + assert run_info["config"]["meta/cwd"] + assert run_info["config"]["meta/user"] # Check metrics extracted assert run_info["metrics"]["val"] == 42.0 From 695a37b78356bd4094b62795663f0295290fd6d1 Mon Sep 17 00:00:00 2001 From: Tsvika Shapira Date: Tue, 28 Apr 2026 17:33:54 +0300 Subject: [PATCH 3/7] feat(runner): add Runner.secret_env for redacted environment variables For env vars (tokens, API keys) the subprocess needs but that should never appear in logs or recorded config: secret_env values pass through to the subprocess unchanged, but render as *** in the Env: log line and are stored as "***" in meta/env so the key list is still recorded. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/lite_runner/runner.py | 7 ++++++- tests/test_runner.py | 15 +++++++++++++-- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/src/lite_runner/runner.py b/src/lite_runner/runner.py index 1223554..39500bf 100644 --- a/src/lite_runner/runner.py +++ b/src/lite_runner/runner.py @@ -177,6 +177,8 @@ class Runner: metrics: Regex patterns to extract from stdout via :class:`Metric`. tags: run tags. env: Extra environment variables for the subprocess. + secret_env: Like ``env``, but values are redacted as ``***`` in logs + and recorded config (subprocess still receives the real value). project: project name (default: git repo name). run_group: run group for sweeps. """ @@ -186,6 +188,7 @@ class Runner: outputs: list[Output] = field(default_factory=list) metrics: list[Metric] = field(default_factory=list) env: dict[str, str | None] = field(default_factory=dict) + secret_env: dict[str, str] = field(default_factory=dict) project: str | None = None run_group: str | None = None tags: list[str] = field(default_factory=list) @@ -500,7 +503,7 @@ def run( config["meta/cwd"] = os.getcwd() # noqa: PTH109 config["meta/datetime"] = timestamp.isoformat() config["meta/command"] = shlex.join(r.command) - config["meta/env"] = dict(r.env) + config["meta/env"] = {**dict(r.env), **dict.fromkeys(r.secret_env, "***")} # Init WandbBackend first (needs to happen early to get run_name) backend_classes: list[type[WandbBackend | JsonBackend | DryRunBackend]] = [] @@ -599,6 +602,7 @@ def run( for b in backend_list: b.update_config({"meta/full_command": shlex.join(cmd)}) set_env = {k: v for k, v in r.env.items() if v is not None} + set_env.update(dict.fromkeys(r.secret_env, "***")) unset_env = [k for k, v in r.env.items() if v is None] if set_env or unset_env: parts = [] @@ -837,6 +841,7 @@ def stream_pipe( run_env.pop(k, None) else: run_env[k] = v + run_env.update(self.secret_env) if "COLUMNS" not in run_env: with suppress(OSError): run_env["COLUMNS"] = str(os.get_terminal_size().columns) diff --git a/tests/test_runner.py b/tests/test_runner.py index d49fe85..1e653dd 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -1127,13 +1127,16 @@ def test_full_run_no_wandb(tmp_path: Path) -> None: """--no-wandb skips W&B but still runs command and writes run_info.json.""" runner = Runner( command=( - f"{sys.executable} -c \"import sys; print('hello'); print('x=42.0')\"" + f'{sys.executable} -c "import os, sys; ' + "print('hello'); print('x=42.0'); " + "print('TOKEN=' + os.environ['MY_TOKEN'])\"" ), params=[], metrics=[Metric("val", pattern=r"x=([\d.]+)")], tags=["v1"], run_group="test-group", env={"MY_FLAG": "1", "DROP_ME": None}, + secret_env={"MY_TOKEN": "supersecret"}, ) with ( @@ -1160,10 +1163,18 @@ def test_full_run_no_wandb(tmp_path: Path) -> None: # Check config assert run_info["config"]["git/repo"] == "test-repo" assert run_info["config"]["meta/output_dir"] == str(output_dir) - assert run_info["config"]["meta/env"] == {"MY_FLAG": "1", "DROP_ME": None} + assert run_info["config"]["meta/env"] == { + "MY_FLAG": "1", + "DROP_ME": None, + "MY_TOKEN": "***", + } assert run_info["config"]["meta/cwd"] assert run_info["config"]["meta/user"] + # secret reaches the subprocess, but does not leak into recorded JSON + assert "TOKEN=supersecret" in (output_dir / "stdout.log").read_text() + assert "supersecret" not in run_info_path.read_text() + # Check metrics extracted assert run_info["metrics"]["val"] == 42.0 From 71d9359d9ab6b823cbecff669320b855e6027047 Mon Sep 17 00:00:00 2001 From: Tsvika Shapira Date: Tue, 28 Apr 2026 17:34:12 +0300 Subject: [PATCH 4/7] docs(readme): document Runner.secret_env Co-Authored-By: Claude Opus 4.7 (1M context) --- README.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/README.md b/README.md index 8538085..0c15aa9 100644 --- a/README.md +++ b/README.md @@ -202,6 +202,9 @@ Runner( "CUDA_VISIBLE_DEVICES": "0", "NOISY_VAR": None, }, # set or unset env vars + secret_env={ + "HF_TOKEN": "hf_xxx", + }, # like env, but redacted in logs / recorded config project="my-project", # default: git repo name run_group="my-sweep", # W&B run group for sweeps (None = no grouping) ) From 7f97b2c0306091b5151e9b37375ad9b6dd482360 Mon Sep 17 00:00:00 2001 From: Tsvika Shapira Date: Tue, 28 Apr 2026 17:35:03 +0300 Subject: [PATCH 5/7] docs(readme): list new meta/* keys and param_source/* in W&B logging table Co-Authored-By: Claude Opus 4.7 (1M context) --- README.md | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 0c15aa9..c4387e0 100644 --- a/README.md +++ b/README.md @@ -254,14 +254,15 @@ Methods: ## What gets logged to W&B -| Location | Content | -| ----------------------- | --------------------------------------------------- | -| `run.config["param/*"]` | All param values | -| `run.config["git/*"]` | commit, branch, repo, dirty | -| `run.config["meta/*"]` | hostname, datetime, command | -| `run.summary` | exit_code, duration_seconds, status, metrics | -| Artifacts | Log files, code snapshot, artifact-type outputs | -| Media | Videos and images from `path-*` type params/outputs | +| Location | Content | +| ------------------------------ | ---------------------------------------------------------------------------------------- | +| `run.config["param/*"]` | All param values | +| `run.config["param_source/*"]` | Where each param value came from (cli, default, fixed, override, prompt) | +| `run.config["git/*"]` | commit, branch, repo, dirty | +| `run.config["meta/*"]` | hostname, user, cwd, datetime, command, full_command, output_dir, env (secrets as `***`) | +| `run.summary` | exit_code, duration_seconds, status, metrics | +| Artifacts | Log files, code snapshot, artifact-type outputs | +| Media | Videos and images from `path-*` type params/outputs | ## Using with Claude Code From 389f7aad07bb97747bebfc98abfb6b3b375eae76 Mon Sep 17 00:00:00 2001 From: Tsvika Shapira Date: Tue, 28 Apr 2026 17:44:22 +0300 Subject: [PATCH 6/7] feat(runner): add pre_commands and post_commands New Command dataclass for auxiliary subprocesses run alongside the main command. Output is streamed to the terminal and saved to per-command log files (_{,_stdout,_stderr}.log), which are also uploaded to backends as text artifacts so the recorded run shows what setup or teardown actually did. - pre_commands run after the code snapshot but before the main subprocess. A non-zero exit aborts the run; main and post commands are skipped, but the backends are still finalized cleanly with the failing exit code. - post_commands run after the main subprocess. Failures are warned but never change the run outcome (consistent with the never-fail post-run philosophy). Skipped only if a pre-command failed. Also extracts _build_run_env() so aux commands inherit the same env (including secret_env) as the main subprocess. Co-Authored-By: Claude Opus 4.7 (1M context) --- README.md | 34 +++++++ src/lite_runner/__init__.py | 3 +- src/lite_runner/params.py | 30 +++++++ src/lite_runner/runner.py | 174 +++++++++++++++++++++++++++++++++--- tests/test_runner.py | 125 +++++++++++++++++++++++++- 5 files changed, 352 insertions(+), 14 deletions(-) diff --git a/README.md b/README.md index c4387e0..437c1ce 100644 --- a/README.md +++ b/README.md @@ -165,6 +165,38 @@ Last match wins. Patterns are matched against both stdout and stderr. Stored in Supported types: `"float"` (default), `"int"`, `"str"`, `"timedelta"`. +## Pre/Post commands + +Run auxiliary commands before or after the main subprocess. Useful for capturing +build steps, environment snapshots, or post-run summaries: + + + +```python +from lite_runner import Command + +runner = Runner( + command="python train.py", + pre_commands=[ + Command("build", "make"), # build before training + Command("pip-list", "uv pip list"), # snapshot the env + ], + post_commands=[ + Command("disk-usage", "du -sh ./output"), # capture output size + ], +) +``` + + + +Each command's stdout/stderr is streamed to the terminal and saved to +`_.log` (combined), `__stdout.log`, and +`__stderr.log` in the run's output directory; the log files are +also uploaded to W&B. + +- A non-zero exit from a **pre-command** aborts the run (main and post are skipped). +- A non-zero exit from a **post-command** is logged but does not change the run outcome. + ## Sweeps Loop with `override()`. Runs are grouped in W&B for easy comparison: @@ -205,6 +237,8 @@ Runner( secret_env={ "HF_TOKEN": "hf_xxx", }, # like env, but redacted in logs / recorded config + pre_commands=[Command("build", "make")], # run before main; failure aborts + post_commands=[Command("du", "du -sh ./")], # run after main; failure logged project="my-project", # default: git repo name run_group="my-sweep", # W&B run group for sweeps (None = no grouping) ) diff --git a/src/lite_runner/__init__.py b/src/lite_runner/__init__.py index e4cc7aa..89941a1 100644 --- a/src/lite_runner/__init__.py +++ b/src/lite_runner/__init__.py @@ -7,7 +7,7 @@ from ._version import version as _version from .backends import JsonBackend, LogBackend, WandbBackend -from .params import UNSET, Metric, Output, Param, ParamType +from .params import UNSET, Command, Metric, Output, Param, ParamType from .runner import Runner, RunResult __version__ = _version @@ -15,6 +15,7 @@ __all__ = [ "UNSET", + "Command", "JsonBackend", "LogBackend", "Metric", diff --git a/src/lite_runner/params.py b/src/lite_runner/params.py index e56497c..47e0b86 100644 --- a/src/lite_runner/params.py +++ b/src/lite_runner/params.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +import shlex from collections.abc import Sequence from dataclasses import dataclass from typing import Any, Literal, TypeGuard @@ -316,3 +317,32 @@ class Metric: name: str pattern: str type: str = "float" + + +@dataclass +class Command: + """An auxiliary command to run before or after the main subprocess. + + Useful for setup/teardown work whose output should be captured for + reproducibility (e.g. ``uv pip list`` to snapshot the environment, ``make`` + to build before running). Output is streamed to the terminal and saved to + ``_{,_stdout,_stderr}.log`` in the run's output directory. + + Args: + name: Identifier used in log filenames; must be filesystem-safe. + command: Shell command (str is split via shlex; list is used as-is). + """ + + name: str + command: str | list[str] + + def __post_init__(self) -> None: + """Split string command into a list and validate name.""" + if isinstance(self.command, str): + self.command = shlex.split(self.command) + if not self.name or "/" in self.name or "\\" in self.name: + msg = ( + f"Command name {self.name!r} must be non-empty and not contain " + "'/' or '\\'" + ) + raise ValueError(msg) diff --git a/src/lite_runner/runner.py b/src/lite_runner/runner.py index 39500bf..23b3b8b 100644 --- a/src/lite_runner/runner.py +++ b/src/lite_runner/runner.py @@ -46,6 +46,7 @@ _PARAM_TYPE_MAP, _SKIP_INPUT, UNSET, + Command, Metric, Output, Param, @@ -179,6 +180,10 @@ class Runner: env: Extra environment variables for the subprocess. secret_env: Like ``env``, but values are redacted as ``***`` in logs and recorded config (subprocess still receives the real value). + pre_commands: Auxiliary commands run before the main subprocess. + A non-zero exit aborts the run; output is captured to log files. + post_commands: Auxiliary commands run after the main subprocess. + Failures are logged but never abort the run. project: project name (default: git repo name). run_group: run group for sweeps. """ @@ -189,6 +194,8 @@ class Runner: metrics: list[Metric] = field(default_factory=list) env: dict[str, str | None] = field(default_factory=dict) secret_env: dict[str, str] = field(default_factory=dict) + pre_commands: list[Command] = field(default_factory=list) + post_commands: list[Command] = field(default_factory=list) project: str | None = None run_group: str | None = None tags: list[str] = field(default_factory=list) @@ -597,6 +604,25 @@ def run( except Exception as e: # noqa: BLE001, PERF203 logger.warning("%s pre-run logging failed: %s", type(b).__name__, e) + # Pre-commands: abort the run if any fails (non-zero exit) + run_env = r._build_run_env() # noqa: SLF001 + pre_failed = False + pre_failure_exit_code = 0 + for c in r.pre_commands: + cmd_list = c.command if isinstance(c.command, list) else [c.command] + if flags.dry_run: + logger.info("[pre:%s] (dry-run) %s", c.name, shlex.join(cmd_list)) + continue + rc = _run_aux_command(c, output_dir, run_env, "pre") + _log_aux_files(backend_list, output_dir, c.name, "pre") + if rc != 0: + logger.error( + "Pre-command %r failed (exit %d); aborting run", c.name, rc + ) + pre_failed = True + pre_failure_exit_code = rc + break + # Build command cmd = r.build_command(interpolated_params) for b in backend_list: @@ -615,14 +641,20 @@ def run( logger.info("Env: %s", " ; ".join(parts)) logger.info("Command:\n%s", shlex.join(cmd)) - # Execute + # Execute (skipped if a pre-command failed) logger.info( "Run started at %s", datetime.datetime.now(tz=datetime.timezone.utc) .astimezone() .strftime("%H:%M:%S %Z (%z)"), ) - if not flags.dry_run: + if pre_failed: + exit_code = pre_failure_exit_code + duration = 0.0 + stdout_text = "" + stderr_text = "" + aborted = False + elif not flags.dry_run: exit_code, duration, stdout_text, stderr_text, aborted = r.execute( cmd, output_dir ) @@ -635,6 +667,25 @@ def run( logger.info("Run finished") + # Post-commands: best-effort, never abort the run. + # Skipped only if a pre-command failed (main never ran). + if not pre_failed: + for c in r.post_commands: + cmd_list = c.command if isinstance(c.command, list) else [c.command] + if flags.dry_run: + logger.info("[post:%s] (dry-run) %s", c.name, shlex.join(cmd_list)) + continue + try: + rc = _run_aux_command(c, output_dir, run_env, "post") + except Exception as e: # noqa: BLE001 + logger.warning("Post-command %r raised: %s", c.name, e) + rc = -1 + _log_aux_files(backend_list, output_dir, c.name, "post") + if rc != 0: + logger.warning( + "Post-command %r failed (exit %d); continuing", c.name, rc + ) + # Post-run: never raise, always try to finish backends r.post_run( backend_list, @@ -795,6 +846,24 @@ def build_command(self, param_values: dict[str, object]) -> list[str]: # Subprocess execution # ----------------------------------------------------------------------- + def _build_run_env(self) -> dict[str, str]: + """Build the environment dict for subprocess execution. + + Inherits ``os.environ``, applies ``self.env`` (None unsets a key), + layers ``self.secret_env`` on top, and ensures ``COLUMNS`` is set. + """ + run_env: dict[str, str] = {**os.environ} + for k, v in self.env.items(): + if v is None: + run_env.pop(k, None) + else: + run_env[k] = v + run_env.update(self.secret_env) + if "COLUMNS" not in run_env: + with suppress(OSError): + run_env["COLUMNS"] = str(os.get_terminal_size().columns) + return run_env + def execute( self, cmd: list[str], output_dir: Path ) -> tuple[int, float, str, str, bool]: @@ -835,16 +904,7 @@ def stream_pipe( log_stdout = stack.enter_context((output_dir / "stdout.log").open("w")) log_stderr = stack.enter_context((output_dir / "stderr.log").open("w")) - run_env: dict[str, str] = {**os.environ} - for k, v in self.env.items(): - if v is None: - run_env.pop(k, None) - else: - run_env[k] = v - run_env.update(self.secret_env) - if "COLUMNS" not in run_env: - with suppress(OSError): - run_env["COLUMNS"] = str(os.get_terminal_size().columns) + run_env = self._build_run_env() proc = subprocess.Popen( # noqa: S603 cmd, @@ -904,6 +964,96 @@ def stream_pipe( # --------------------------------------------------------------------------- +def _run_aux_command( + cmd: Command, + output_dir: Path, + run_env: dict[str, str], + phase: str, +) -> int: + """Run a pre/post command, streaming to terminal and writing log files. + + Writes ``_.log`` (combined) plus ``__stdout.log`` + and ``__stderr.log``. Returns the exit code; never raises + on a non-zero exit. + """ + base = f"{phase}_{cmd.name}" + combined_path = output_dir / f"{base}.log" + stdout_path = output_dir / f"{base}_stdout.log" + stderr_path = output_dir / f"{base}_stderr.log" + cmd_list = cmd.command if isinstance(cmd.command, list) else [cmd.command] + logger.info("[%s:%s] %s", phase, cmd.name, shlex.join(cmd_list)) + + lock = threading.Lock() + + with ExitStack() as stack: + combined = stack.enter_context(combined_path.open("w")) + log_stdout = stack.enter_context(stdout_path.open("w")) + log_stderr = stack.enter_context(stderr_path.open("w")) + + def stream( + pipe: IO[bytes], + sys_stream: TextIO, + file_log: TextIO, + *, + prefix: str = "", + ) -> None: + while True: + chunk = pipe.read1(8192) if hasattr(pipe, "read1") else pipe.read(8192) + if not chunk: + break + text = chunk.decode("utf-8", errors="replace") + sys_stream.write(text) + sys_stream.flush() + file_log.write(text) + file_log.flush() + with lock: + combined.write(prefix + text if prefix else text) + combined.flush() + pipe.close() + + proc = subprocess.Popen( # noqa: S603 + cmd_list, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=run_env, + ) + assert proc.stdout is not None # noqa: S101 + assert proc.stderr is not None # noqa: S101 + t_out = threading.Thread( + target=stream, args=(proc.stdout, sys.stdout, log_stdout) + ) + t_err = threading.Thread( + target=stream, + args=(proc.stderr, sys.stderr, log_stderr), + kwargs={"prefix": "[stderr] "}, + ) + t_out.start() + t_err.start() + proc.wait() + t_out.join() + t_err.join() + return proc.returncode + + +def _log_aux_files( + backends: Sequence[LogBackend], + output_dir: Path, + cmd_name: str, + phase: str, +) -> None: + """Upload aux-command log files to all backends as text artifacts.""" + base = f"{phase}_{cmd_name}" + for suffix in ("", "_stdout", "_stderr"): + path = output_dir / f"{base}{suffix}.log" + if not path.exists(): + continue + for b in backends: + try: + b.log_file(path, "text", f"{phase}/{cmd_name}{suffix}") + except Exception as e: # noqa: BLE001, PERF203 + logger.warning("%s log %s failed: %s", type(b).__name__, path.name, e) + + def _subst_output(v: object, out: str) -> object: if isinstance(v, str): v = v.replace("$output", out) diff --git a/tests/test_runner.py b/tests/test_runner.py index 1e653dd..0def5d2 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -11,7 +11,7 @@ import pytest from conftest import _FAKE_GIT_INFO, _make_runner, _mock_wb_run -from lite_runner import UNSET, Metric, Param, Runner +from lite_runner import UNSET, Command, Metric, Param, Runner from lite_runner.backends import DryRunBackend, JsonBackend from lite_runner.runner import ( _collect_git_info, @@ -1195,6 +1195,129 @@ def test_full_run_no_wandb(tmp_path: Path) -> None: assert run_info["config"]["wandb/url"] == "(no wandb)" +def _run_with_aux( + tmp_path: Path, + *, + pre_commands: list[Command] | None = None, + post_commands: list[Command] | None = None, + main_exit: int = 0, +) -> Path: + """Run a Runner with given aux commands and return its output_dir.""" + main_cmd = ( + f"{sys.executable} -c \"import sys; print('main'); sys.exit({main_exit})\"" + ) + runner = Runner( + command=main_cmd, + params=[], + pre_commands=pre_commands or [], + post_commands=post_commands or [], + ) + with ( + patch("lite_runner.runner._collect_git_info", return_value=_FAKE_GIT_INFO), + patch("lite_runner.backends.create_repo_archive", return_value=None), + patch("lite_runner.backends.create_repo_diff", return_value=None), + patch("lite_runner.runner.RUNS_DIR", tmp_path / "lite_runs"), + pytest.MonkeyPatch.context() as mp, + ): + mp.setattr(sys, "exit", lambda _code=0: None) + runner.run(no_wandb=True, no_interactive=True) + project_dir = tmp_path / "lite_runs" / "test-repo" + return next(project_dir.iterdir()) + + +def test_pre_command_runs_and_writes_logs(tmp_path: Path) -> None: + """Pre-command runs before main, output captured to pre_*.log files.""" + output_dir = _run_with_aux( + tmp_path, + pre_commands=[ + Command( + name="snap", + command=f"{sys.executable} -c \"print('pre-out'); " + "import sys; print('pre-err', file=sys.stderr)\"", + ), + ], + ) + assert (output_dir / "pre_snap.log").exists() + assert "pre-out" in (output_dir / "pre_snap_stdout.log").read_text() + assert "pre-err" in (output_dir / "pre_snap_stderr.log").read_text() + # main also ran + assert "main" in (output_dir / "stdout.log").read_text() + + +def test_pre_command_failure_skips_main(tmp_path: Path) -> None: + """A non-zero pre-command aborts the run; main is not executed.""" + output_dir = _run_with_aux( + tmp_path, + pre_commands=[ + Command( + name="bad", + command=f'{sys.executable} -c "import sys; sys.exit(7)"', + ), + ], + ) + # pre log written + assert (output_dir / "pre_bad_stdout.log").exists() + # main never ran (its log file is created only inside execute()) + assert not (output_dir / "stdout.log").exists() + # exit code recorded as the pre-command failure + run_info = json.loads((output_dir / "run_info.json").read_text()) + assert run_info["summary"]["exit_code"] == 7 + assert run_info["summary"]["status"] == "failed" + + +def test_post_command_runs_after_main(tmp_path: Path) -> None: + """Post-command runs after main, output captured to post_*.log.""" + output_dir = _run_with_aux( + tmp_path, + post_commands=[ + Command(name="cleanup", command=f"{sys.executable} -c \"print('after')\""), + ], + ) + assert "after" in (output_dir / "post_cleanup_stdout.log").read_text() + + +def test_post_command_failure_does_not_abort(tmp_path: Path) -> None: + """A failing post-command logs a warning but doesn't change run outcome.""" + output_dir = _run_with_aux( + tmp_path, + post_commands=[ + Command( + name="bad", + command=f'{sys.executable} -c "import sys; sys.exit(3)"', + ), + ], + ) + run_info = json.loads((output_dir / "run_info.json").read_text()) + # main exit was 0; post failure shouldn't change that + assert run_info["summary"]["exit_code"] == 0 + assert run_info["summary"]["status"] == "success" + assert (output_dir / "post_bad_stdout.log").exists() + + +def test_post_command_skipped_when_pre_fails(tmp_path: Path) -> None: + """If a pre-command fails, post-commands don't run.""" + output_dir = _run_with_aux( + tmp_path, + pre_commands=[ + Command( + name="bad", + command=f'{sys.executable} -c "import sys; sys.exit(1)"', + ), + ], + post_commands=[ + Command(name="never", command=f"{sys.executable} -c \"print('nope')\""), + ], + ) + assert not (output_dir / "post_never_stdout.log").exists() + + +def test_command_name_validation() -> None: + with pytest.raises(ValueError, match="must be non-empty"): + Command(name="", command="echo hi") + with pytest.raises(ValueError, match="must be non-empty"): + Command(name="bad/name", command="echo hi") + + def test_input_files_copied_to_output_dir(tmp_path: Path) -> None: """Input files (log_when=before) are copied to output_dir/input/.""" # Create a fake input image From 1bdb571c5d7ec62c6e10924d801f34ce64ec5e8e Mon Sep 17 00:00:00 2001 From: Tsvika Shapira Date: Tue, 28 Apr 2026 18:03:34 +0300 Subject: [PATCH 7/7] feat(runner): aux commands support \$output and per-command env - Command gains an env field (dict[str, str | None]); aux commands now build their own env from os.environ + cmd.env. They no longer inherit Runner.env / Runner.secret_env, so e.g. CUDA_VISIBLE_DEVICES set for the main subprocess doesn't leak into uv-pip-list-style snapshots. - \$output is interpolated in command tokens and env values (same syntax as Param value=). - Pre/post commands run in declared list order (already true; documented). - Updated examples/run_example.py with pre 'uv sync -v' and post 'uv pip list' to demonstrate the canonical use case. Co-Authored-By: Claude Opus 4.7 (1M context) --- README.md | 20 ++++++---- examples/run_example.py | 10 ++++- src/lite_runner/params.py | 9 ++++- src/lite_runner/runner.py | 47 +++++++++++++++++----- tests/test_runner.py | 83 +++++++++++++++++++++++++++++++++++++++ 5 files changed, 149 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index 437c1ce..7718e44 100644 --- a/README.md +++ b/README.md @@ -178,22 +178,26 @@ from lite_runner import Command runner = Runner( command="python train.py", pre_commands=[ - Command("build", "make"), # build before training - Command("pip-list", "uv pip list"), # snapshot the env + Command("uv-sync", "uv sync -v"), # sync env before run + Command("ls-out", "ls -la $output", env={"LC_ALL": "C"}), # $output is interpolated ], post_commands=[ - Command("disk-usage", "du -sh ./output"), # capture output size + Command("uv-pip-list", "uv pip list"), # snapshot resolved versions ], ) ``` -Each command's stdout/stderr is streamed to the terminal and saved to -`_.log` (combined), `__stdout.log`, and -`__stderr.log` in the run's output directory; the log files are -also uploaded to W&B. - +- Each command's stdout/stderr is streamed to the terminal and saved to + `_.log` (combined), `__stdout.log`, and + `__stderr.log` in the run's output directory; the log files are + also uploaded to W&B. +- Commands run sequentially in the order they appear in the list. +- `$output` is interpolated in both command tokens and `env` values. +- Each `Command` has its own optional `env` (None to unset). Aux commands do + **not** inherit `Runner.env` or `Runner.secret_env` — they get a clean env + built from `os.environ` plus per-command overrides. - A non-zero exit from a **pre-command** aborts the run (main and post are skipped). - A non-zero exit from a **post-command** is logged but does not change the run outcome. diff --git a/examples/run_example.py b/examples/run_example.py index cdd7de3..5936ff3 100644 --- a/examples/run_example.py +++ b/examples/run_example.py @@ -1,7 +1,7 @@ # Run with: uv run python examples/run_example.py """Example run config for the fake model.""" -from lite_runner import Metric, Output, Param, Runner +from lite_runner import Command, Metric, Output, Param, Runner runner = Runner( command="python examples/fake_model.py", @@ -39,6 +39,14 @@ ], tags=["example"], env={"FAKE_MODEL_DEBUG": "1"}, + pre_commands=[ + # Sync the env before the run; abort if it fails. + Command("uv-sync", "uv sync -v"), + ], + post_commands=[ + # Snapshot the resolved package versions for reproducibility. + Command("uv-pip-list", "uv pip list"), + ], ) if __name__ == "__main__": diff --git a/src/lite_runner/params.py b/src/lite_runner/params.py index 47e0b86..fe931be 100644 --- a/src/lite_runner/params.py +++ b/src/lite_runner/params.py @@ -5,7 +5,7 @@ import logging import shlex from collections.abc import Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any, Literal, TypeGuard import questionary @@ -328,13 +328,20 @@ class Command: to build before running). Output is streamed to the terminal and saved to ``_{,_stdout,_stderr}.log`` in the run's output directory. + ``$output`` is interpolated to the run's output directory in both the + command tokens and ``env`` values. + Args: name: Identifier used in log filenames; must be filesystem-safe. command: Shell command (str is split via shlex; list is used as-is). + env: Per-command environment variables (None unsets a key). Layered + on top of ``os.environ``; does NOT inherit ``Runner.env`` or + ``Runner.secret_env`` — aux commands get their own clean env. """ name: str command: str | list[str] + env: dict[str, str | None] = field(default_factory=dict) def __post_init__(self) -> None: """Split string command into a list and validate name.""" diff --git a/src/lite_runner/runner.py b/src/lite_runner/runner.py index 23b3b8b..9fd165d 100644 --- a/src/lite_runner/runner.py +++ b/src/lite_runner/runner.py @@ -604,16 +604,17 @@ def run( except Exception as e: # noqa: BLE001, PERF203 logger.warning("%s pre-run logging failed: %s", type(b).__name__, e) - # Pre-commands: abort the run if any fails (non-zero exit) - run_env = r._build_run_env() # noqa: SLF001 + # Pre-commands: abort the run if any fails (non-zero exit). + # Run in declared list order; aux commands use their own env (do NOT + # inherit Runner.env / secret_env). pre_failed = False pre_failure_exit_code = 0 for c in r.pre_commands: - cmd_list = c.command if isinstance(c.command, list) else [c.command] + display = _interp_aux_command(c, output_dir) if flags.dry_run: - logger.info("[pre:%s] (dry-run) %s", c.name, shlex.join(cmd_list)) + logger.info("[pre:%s] (dry-run) %s", c.name, shlex.join(display)) continue - rc = _run_aux_command(c, output_dir, run_env, "pre") + rc = _run_aux_command(c, output_dir, "pre") _log_aux_files(backend_list, output_dir, c.name, "pre") if rc != 0: logger.error( @@ -671,12 +672,12 @@ def run( # Skipped only if a pre-command failed (main never ran). if not pre_failed: for c in r.post_commands: - cmd_list = c.command if isinstance(c.command, list) else [c.command] + display = _interp_aux_command(c, output_dir) if flags.dry_run: - logger.info("[post:%s] (dry-run) %s", c.name, shlex.join(cmd_list)) + logger.info("[post:%s] (dry-run) %s", c.name, shlex.join(display)) continue try: - rc = _run_aux_command(c, output_dir, run_env, "post") + rc = _run_aux_command(c, output_dir, "post") except Exception as e: # noqa: BLE001 logger.warning("Post-command %r raised: %s", c.name, e) rc = -1 @@ -964,10 +965,35 @@ def stream_pipe( # --------------------------------------------------------------------------- +def _interp_aux_command(cmd: Command, output_dir: Path) -> list[str]: + """Interpolate ``$output`` in *cmd*'s tokens and return as a list.""" + out = str(output_dir) + raw = cmd.command if isinstance(cmd.command, list) else [cmd.command] + return [str(_subst_output(t, out)) for t in raw] + + +def _build_aux_env(cmd: Command, output_dir: Path) -> dict[str, str]: + """Build env for an aux command: ``os.environ`` + ``cmd.env`` (None unsets). + + Does NOT inherit ``Runner.env`` / ``secret_env``. ``$output`` is + interpolated in env values. ``COLUMNS`` is set if not already present. + """ + out = str(output_dir) + run_env: dict[str, str] = {**os.environ} + for k, v in cmd.env.items(): + if v is None: + run_env.pop(k, None) + else: + run_env[k] = str(_subst_output(v, out)) + if "COLUMNS" not in run_env: + with suppress(OSError): + run_env["COLUMNS"] = str(os.get_terminal_size().columns) + return run_env + + def _run_aux_command( cmd: Command, output_dir: Path, - run_env: dict[str, str], phase: str, ) -> int: """Run a pre/post command, streaming to terminal and writing log files. @@ -980,7 +1006,8 @@ def _run_aux_command( combined_path = output_dir / f"{base}.log" stdout_path = output_dir / f"{base}_stdout.log" stderr_path = output_dir / f"{base}_stderr.log" - cmd_list = cmd.command if isinstance(cmd.command, list) else [cmd.command] + cmd_list = _interp_aux_command(cmd, output_dir) + run_env = _build_aux_env(cmd, output_dir) logger.info("[%s:%s] %s", phase, cmd.name, shlex.join(cmd_list)) lock = threading.Lock() diff --git a/tests/test_runner.py b/tests/test_runner.py index 0def5d2..21ef7d7 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -1311,6 +1311,89 @@ def test_post_command_skipped_when_pre_fails(tmp_path: Path) -> None: assert not (output_dir / "post_never_stdout.log").exists() +def test_aux_command_interpolates_output(tmp_path: Path) -> None: + """$output in a Command's tokens and env values is interpolated to output_dir.""" + output_dir = _run_with_aux( + tmp_path, + pre_commands=[ + Command( + name="show", + command=( + f'{sys.executable} -c "' + "import os; print('ARG=' + __import__('sys').argv[1]); " + "print('ENV=' + os.environ['MYDIR'])\" $output" + ), + env={"MYDIR": "$output"}, + ), + ], + ) + log = (output_dir / "pre_show_stdout.log").read_text() + assert f"ARG={output_dir}" in log + assert f"ENV={output_dir}" in log + + +def test_aux_command_uses_own_env_not_runner_env(tmp_path: Path) -> None: + """Aux commands do NOT inherit Runner.env / secret_env.""" + main_cmd = ( + f'{sys.executable} -c "import os, sys; ' + "print('MAIN_FLAG=' + os.environ.get('MAIN_FLAG', 'unset'))\"" + ) + runner = Runner( + command=main_cmd, + params=[], + env={"MAIN_FLAG": "from-runner"}, + secret_env={"MAIN_TOKEN": "secret"}, + pre_commands=[ + Command( + name="check", + command=f'{sys.executable} -c "' + "import os; " + "print('PRE_FLAG=' + os.environ.get('MAIN_FLAG', 'unset')); " + "print('PRE_TOKEN=' + os.environ.get('MAIN_TOKEN', 'unset')); " + "print('PRE_OWN=' + os.environ.get('PRE_OWN', 'unset'))\"", + env={"PRE_OWN": "from-cmd"}, + ), + ], + ) + with ( + patch("lite_runner.runner._collect_git_info", return_value=_FAKE_GIT_INFO), + patch("lite_runner.backends.create_repo_archive", return_value=None), + patch("lite_runner.backends.create_repo_diff", return_value=None), + patch("lite_runner.runner.RUNS_DIR", tmp_path / "lite_runs"), + pytest.MonkeyPatch.context() as mp, + ): + mp.setattr(sys, "exit", lambda _code=0: None) + runner.run(no_wandb=True, no_interactive=True) + + output_dir = next((tmp_path / "lite_runs" / "test-repo").iterdir()) + pre_log = (output_dir / "pre_check_stdout.log").read_text() + main_log = (output_dir / "stdout.log").read_text() + # Pre command sees its own env but NOT main's env or secret_env + assert "PRE_OWN=from-cmd" in pre_log + assert "PRE_FLAG=unset" in pre_log + assert "PRE_TOKEN=unset" in pre_log + # Main command still sees Runner.env + assert "MAIN_FLAG=from-runner" in main_log + + +def test_aux_commands_run_in_list_order(tmp_path: Path) -> None: + """pre_commands and post_commands execute in the order they were declared.""" + output_dir = _run_with_aux( + tmp_path, + pre_commands=[ + Command(name="first", command=f"{sys.executable} -c \"print('1')\""), + Command(name="second", command=f"{sys.executable} -c \"print('2')\""), + Command(name="third", command=f"{sys.executable} -c \"print('3')\""), + ], + ) + # Files are created sequentially; check by mtime + times = [ + (output_dir / f"pre_{n}.log").stat().st_mtime_ns + for n in ("first", "second", "third") + ] + assert times == sorted(times) + + def test_command_name_validation() -> None: with pytest.raises(ValueError, match="must be non-empty"): Command(name="", command="echo hi")