diff --git a/README.md b/README.md index c4387e0..7718e44 100644 --- a/README.md +++ b/README.md @@ -165,6 +165,42 @@ 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("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("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. +- 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. + ## Sweeps Loop with `override()`. Runs are grouped in W&B for easy comparison: @@ -205,6 +241,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/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/__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..fe931be 100644 --- a/src/lite_runner/params.py +++ b/src/lite_runner/params.py @@ -3,8 +3,9 @@ from __future__ import annotations 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 @@ -316,3 +317,39 @@ 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. + + ``$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.""" + 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..9fd165d 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,26 @@ 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 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: + display = _interp_aux_command(c, output_dir) + if flags.dry_run: + logger.info("[pre:%s] (dry-run) %s", c.name, shlex.join(display)) + continue + rc = _run_aux_command(c, output_dir, "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 +642,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 +668,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: + display = _interp_aux_command(c, output_dir) + if flags.dry_run: + logger.info("[post:%s] (dry-run) %s", c.name, shlex.join(display)) + continue + try: + 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 + _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 +847,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 +905,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 +965,122 @@ 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, + 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 = _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() + + 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..21ef7d7 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,212 @@ 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_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") + 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