labnotes 0.0.125__tar.gz → 0.0.127__tar.gz

{labnotes-0.0.125 → labnotes-0.0.127}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: labnotes
-Version: 0.0.125
+Version: 0.0.127
 Summary: Browse experiment designs, notes, and run data from a folder.
 Project-URL: Repository, https://github.com/thekevinscott/labnotes
 Author-email: Kevin Scott <me@thekevinscott.com>

{labnotes-0.0.125 → labnotes-0.0.127}/src/labnotes/context.md

@@ -89,6 +89,11 @@ command that prints one JSON object, such as the git commit. labnotes runs it
 once from the root when `run` starts and copies it into every run's `run.json`
 as `meta`. If it fails or prints anything else, no run starts. Queued runs keep
 the value from the start, even if the code changed while they waited.
+`setup` and `teardown` (same places; `--setup`/`--teardown` override) are shell
+commands run once per `run` from the root: setup before the first run starts,
+teardown after the last one ends, even if runs failed, were stopped, or setup
+failed (so teardown must cope with a partial setup). A failing setup starts no
+run; a failing teardown fails the invocation.
 A run counts when it exited 0 and its `run.json` doesn't hold
 `"excluded": true`, so excluding a run by hand means adding that key; the next
 `run` replaces it. Other keys you add to `run.json` are kept and show in `run`.
@@ -125,7 +130,10 @@ Runs can take a long time. Start `run` in the background, then check on it with
 
 While a run goes, labnotes appends a memory sample to `memory.jsonl` in its run
 folder every 5 seconds: one JSON line, `{"time": <unix seconds>, "rss": <bytes>,
-"processes": <count>}`, for the command and everything it started. It keeps
+"processes": <count>}`, for the command's process tree and its labeled Docker
+containers. `LABNOTES_RUN` is the absolute run folder path; scripts label containers
+with `docker run --label "labnotes.run=$LABNOTES_RUN" ...` to count their memory.
+`processes` counts host processes in the run's tree, not containers. It keeps
 every sample, so a run killed for memory still shows the climb.
 `memory_interval` in `experiment.toml` (or `[experiment]` in `.labnotes.toml`)
 sets the seconds between samples; 0 turns it off.

{labnotes-0.0.125 → labnotes-0.0.127}/src/labnotes/daemon/daemon_app.py

@@ -80,6 +80,8 @@ def daemon_app(config: LabnotesConfig) -> Starlette:
                 body.get("concurrency"),  # type: ignore[arg-type]
                 body.get("only"),  # type: ignore[arg-type]
                 body.get("args"),  # type: ignore[arg-type]
+                body.get("setup"),  # type: ignore[arg-type]
+                body.get("teardown"),  # type: ignore[arg-type]
             )
         except ValueError as error:
             return refuse(400, str(error))

{labnotes-0.0.125 → labnotes-0.0.127}/src/labnotes/daemon/remote_start.py

@@ -18,6 +18,8 @@ def remote_start(
     timeout: float | None = None,
     only: dict[str, str] | None = None,
     args: list[str] | None = None,
+    setup: str | None = None,
+    teardown: str | None = None,
 ) -> Start:
     async def start(report: Report, listen: Listen) -> list[dict[str, object]]:
         try:
@@ -50,6 +52,8 @@ def remote_start(
                     "timeout": timeout,
                     "only": only,
                     "args": args,
+                    "setup": setup,
+                    "teardown": teardown,
                 },
             )
             if reply.status_code != 202:

{labnotes-0.0.125 → labnotes-0.0.127}/src/labnotes/experiments.py

@@ -66,6 +66,18 @@ class PassthroughCommand(click.Command):
     help="With run: stop any run still going after this many seconds."
     " Unset, runs have no time limit.",
 )
+@click.option(
+    "--setup",
+    metavar="COMMAND",
+    help="With run: shell command to run once before any run starts."
+    " Overrides experiment.toml and .labnotes.toml.",
+)
+@click.option(
+    "--teardown",
+    metavar="COMMAND",
+    help="With run: shell command to run once after every run ends."
+    " Overrides experiment.toml and .labnotes.toml.",
+)
 @click.pass_context
 def experiments(
     ctx: click.Context,
@@ -75,6 +87,8 @@ def experiments(
     concurrency: int | None,
     experiment_config: Path | None,
     timeout: int | None,
+    setup: str | None,
+    teardown: str | None,
 ) -> None:
     passthrough: list[str] | None = ctx.meta.get(PASSTHROUGH)
     config = load_config(ctx)
@@ -103,6 +117,8 @@ def experiments(
                 timeout,
                 only,
                 passthrough or None,
+                setup,
+                teardown,
             )
             try:
                 if sys.stdout.isatty() and sys.stderr.isatty():

{labnotes-0.0.125 → labnotes-0.0.127}/src/labnotes/log_path.py

@@ -1,8 +1,8 @@
 import os
 from pathlib import Path
 
-# The one env contract between runner and script; config is otherwise never env.
 LOG_ENV = "LABNOTES_LOG"
+RUN_ENV = "LABNOTES_RUN"
 
 
 def log_path() -> Path | None:

{labnotes-0.0.125 → labnotes-0.0.127}/src/labnotes/models.py

@@ -64,3 +64,5 @@ class ExperimentConfig:
     retries: int = DEFAULT_RETRIES
     meta: str | None = None
     memory_interval: float = DEFAULT_MEMORY_INTERVAL
+    setup: str | None = None
+    teardown: str | None = None

labnotes-0.0.127/src/labnotes/runner/cgroup_memory.py

@@ -0,0 +1,17 @@
+from pathlib import Path
+
+
+def cgroup_memory(path: Path) -> int:
+    current = int((path / "memory.current").read_text())
+    try:
+        stats = (path / "memory.stat").read_text().splitlines()
+    except OSError:
+        return current
+    for line in stats:
+        name, _, value = line.partition(" ")
+        if name == "inactive_file":
+            try:
+                return max(0, current - int(value))
+            except ValueError:
+                return current
+    return current

labnotes-0.0.127/src/labnotes/runner/container_cgroup.py

@@ -0,0 +1,15 @@
+from pathlib import Path
+
+CGROUP_ROOT = Path("/sys/fs/cgroup")
+PROC_ROOT = Path("/proc")
+
+
+def container_cgroup(pid: int) -> Path | None:
+    try:
+        lines = (PROC_ROOT / str(pid) / "cgroup").read_text().splitlines()
+    except OSError:
+        return None
+    for line in lines:
+        if line.startswith("0::"):
+            return CGROUP_ROOT / line.removeprefix("0::").lstrip("/")
+    return None

labnotes-0.0.127/src/labnotes/runner/container_memory.py

@@ -0,0 +1,100 @@
+import asyncio
+import shutil
+import time
+from pathlib import Path
+
+from labnotes.runner.cgroup_memory import cgroup_memory
+from labnotes.runner.container_cgroup import container_cgroup
+from labnotes.runner.docker_output import docker_output
+
+LABEL = "labnotes.run"
+
+
+class ContainerMemory:
+    """Track this run's Docker cgroups without starting a CLI process per sample."""
+
+    def __init__(self, run_id: str) -> None:
+        self.run_id = run_id
+        self.cgroups: dict[str, Path] = {}
+        self.events: asyncio.subprocess.Process | None = None
+        self.watcher: asyncio.Task[None] | None = None
+
+    async def start(self) -> None:
+        if shutil.which("docker") is None:
+            return
+        since = str(int(time.time()))
+        ids = await asyncio.to_thread(
+            docker_output,
+            "ps",
+            "--quiet",
+            "--no-trunc",
+            "--filter",
+            f"label={LABEL}={self.run_id}",
+        )
+        if ids is None:
+            return
+        for container_id in ids.splitlines():
+            await self._add(container_id)
+        try:
+            self.events = await asyncio.create_subprocess_exec(
+                "docker",
+                "events",
+                "--since",
+                since,
+                "--filter",
+                "type=container",
+                "--filter",
+                f"label={LABEL}={self.run_id}",
+                "--format",
+                "{{.Actor.ID}} {{.Action}}",
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.DEVNULL,
+            )
+        except OSError:
+            return
+        self.watcher = asyncio.create_task(self._watch())
+
+    async def _add(self, container_id: str) -> None:
+        raw_pid = await asyncio.to_thread(
+            docker_output, "inspect", "--format", "{{.State.Pid}}", container_id
+        )
+        if raw_pid is None:
+            return
+        try:
+            pid = int(raw_pid.strip())
+        except ValueError:
+            return
+        if pid > 0 and (path := container_cgroup(pid)) is not None:
+            self.cgroups[container_id] = path
+
+    async def _watch(self) -> None:
+        assert self.events is not None
+        assert self.events.stdout is not None
+        while line := await self.events.stdout.readline():
+            container_id, _, action = line.decode().strip().partition(" ")
+            if action == "start":
+                await self._add(container_id)
+            elif action in {"die", "destroy", "stop"}:
+                self.cgroups.pop(container_id, None)
+
+    def usage(self) -> int:
+        total = 0
+        for path in self.cgroups.values():
+            try:
+                total += cgroup_memory(path)
+            except (OSError, ValueError):
+                pass
+        return total
+
+    async def close(self) -> None:
+        if self.watcher is not None:
+            self.watcher.cancel()
+            await asyncio.gather(self.watcher, return_exceptions=True)
+        if self.events is not None:
+            if self.events.returncode is None:
+                self.events.terminate()
+            try:
+                await asyncio.wait_for(self.events.wait(), timeout=1)
+            except TimeoutError:
+                self.events.kill()
+                await self.events.wait()

{labnotes-0.0.125 → labnotes-0.0.127}/src/labnotes/runner/describe_event.py

@@ -33,6 +33,12 @@ def describe_event(event: dict[str, object]) -> str:
             return "No runs left."
         case "concurrency":
             return f"{event['concurrency']} at a time."
+        case "setting up":
+            return "Setting up."
+        case "tearing down":
+            return "Tearing down."
+        case "teardown failed":
+            return str(event["error"])
         case "stopping":
             return "Stopping runs. Ctrl-C again to kill them."
         case _:

labnotes-0.0.127/src/labnotes/runner/docker_output.py

@@ -0,0 +1,11 @@
+import subprocess
+
+
+def docker_output(*args: str) -> str | None:
+    try:
+        result = subprocess.run(
+            ["docker", *args], capture_output=True, text=True, timeout=3, check=False
+        )
+    except (OSError, subprocess.TimeoutExpired):
+        return None
+    return result.stdout if result.returncode == 0 else None

{labnotes-0.0.125 → labnotes-0.0.127}/src/labnotes/runner/execute_run.py

@@ -7,7 +7,7 @@ from pathlib import Path
 
 import psutil
 
-from labnotes.log_path import LOG_ENV
+from labnotes.log_path import LOG_ENV, RUN_ENV
 from labnotes.runner.exit_signal import exit_signal
 from labnotes.runner.record_memory import record_memory
 from labnotes.runner.shell_command import shell_command
@@ -66,6 +66,7 @@ async def execute_run(
                     "CLICOLOR_FORCE": "1",
                     **os.environ,
                     LOG_ENV: str(folder / "log.jsonl"),
+                    RUN_ENV: str(folder.resolve()),
                 },
                 stdout=stdout,
                 stderr=stderr,

labnotes-0.0.127/src/labnotes/runner/invocation_hook.py

@@ -0,0 +1,16 @@
+import asyncio
+from pathlib import Path
+
+
+async def invocation_hook(root: Path, hook: str, command: str) -> None:
+    process = await asyncio.create_subprocess_shell(
+        command,
+        cwd=root,
+        stdout=asyncio.subprocess.DEVNULL,
+        stderr=asyncio.subprocess.PIPE,
+    )
+    _, stderr = await process.communicate()
+    if process.returncode != 0:
+        message = f"{hook} command {command} exited {process.returncode}"
+        detail = stderr.decode(errors="replace").strip()
+        raise ValueError(f"{message}: {detail}" if detail else message)

{labnotes-0.0.125 → labnotes-0.0.127}/src/labnotes/runner/prepare_experiment.py

@@ -1,3 +1,4 @@
+import dataclasses
 from pathlib import Path
 
 from labnotes.models import ExperimentConfig, LabnotesConfig
@@ -14,6 +15,8 @@ async def prepare_experiment(
     concurrency: int | None,
     only: dict[str, str] | None = None,
     args: list[str] | None = None,
+    setup: str | None = None,
+    teardown: str | None = None,
 ) -> tuple[ExperimentConfig, int]:
     if config.experiment not in (None, name):
         raise ValueError(
@@ -33,6 +36,10 @@ async def prepare_experiment(
         config.retries,
         config.defaults,
     )
+    if setup is not None:
+        experiment = dataclasses.replace(experiment, setup=setup)
+    if teardown is not None:
+        experiment = dataclasses.replace(experiment, teardown=teardown)
     if only:
         check_subset(experiment.params, only)
     if args:

{labnotes-0.0.125 → labnotes-0.0.127}/src/labnotes/runner/read_experiment_config.py

@@ -24,7 +24,10 @@ SETTINGS = {
     "meta",
     "memory_interval",
     "views",
+    "setup",
+    "teardown",
 }
+HOOKS = ("setup", "teardown")
 VIEWS = {"runs"}
 
 
@@ -81,6 +84,10 @@ def read_experiment_config(
     memory_interval = values.get("memory_interval", DEFAULT_MEMORY_INTERVAL)
     if not is_non_negative_number(memory_interval):
         raise ValueError(f"{path}: memory_interval must be 0 or more seconds")
+    for hook in HOOKS:
+        run = values.get(hook)
+        if run is not None and (not isinstance(run, str) or not run.strip()):
+            raise ValueError(f"{path}: {hook} must be a string")
     views = values.get("views", {})
     if not isinstance(views, dict):
         raise ValueError(f"{path}: views must be a table, like [views]")
@@ -101,4 +108,6 @@ def read_experiment_config(
         retries,
         meta,
         memory_interval,
+        values.get("setup"),
+        values.get("teardown"),
     )

labnotes-0.0.127/src/labnotes/runner/record_memory.py

@@ -0,0 +1,28 @@
+import asyncio
+import json
+import time
+from pathlib import Path
+
+from labnotes.runner.container_memory import ContainerMemory
+from labnotes.utils.tree_usage import tree_usage
+
+
+async def record_memory(path: Path, pid: int, created: float, interval: float) -> None:
+    containers = ContainerMemory(str(path.parent.resolve()))
+    try:
+        await containers.start()
+        with path.open("a", encoding="utf-8") as samples:
+            while True:
+                usage = tree_usage(pid, created)
+                if usage is not None:
+                    rss, processes = usage
+                    sample = {
+                        "time": time.time(),
+                        "rss": rss + containers.usage(),
+                        "processes": processes,
+                    }
+                    samples.write(json.dumps(sample) + "\n")
+                    samples.flush()
+                await asyncio.sleep(interval)
+    finally:
+        await containers.close()

{labnotes-0.0.125 → labnotes-0.0.127}/src/labnotes/runner/run_experiment.py

@@ -14,6 +14,7 @@ from labnotes.models import (
 from labnotes.retries import RETRY_EXIT_CODE, RETRY_SECONDS
 from labnotes.runner.execute_run import execute_run
 from labnotes.runner.in_subset import in_subset
+from labnotes.runner.invocation_hook import invocation_hook
 from labnotes.runner.invocation_meta import invocation_meta
 from labnotes.runner.new_param_run_id import new_param_run_id
 from labnotes.runner.new_stamped_run_id import new_stamped_run_id
@@ -233,8 +234,25 @@ async def run_experiment(
         ):
             retry_later(run)
 
+    async def teardown() -> None:
+        if experiment.teardown is not None:
+            report({"event": "tearing down"})
+            await invocation_hook(config.root, "teardown", experiment.teardown)
+
     listen(Controls(interrupt, set_concurrency, kill, finish))
-    fill()
-    await done.wait()
-    await asyncio.gather(*started)
+    # Like vitest's globalSetup: a failed setup still gets its teardown.
+    try:
+        if experiment.setup is not None:
+            report({"event": "setting up"})
+            await invocation_hook(config.root, "setup", experiment.setup)
+        fill()
+        await done.wait()
+        await asyncio.gather(*started)
+    except BaseException:
+        try:
+            await teardown()
+        except Exception as failure:  # noqa: BLE001
+            report({"event": "teardown failed", "error": str(failure)})
+        raise
+    await teardown()
     return results

labnotes-0.0.125/src/labnotes/runner/record_memory.py

@@ -1,18 +0,0 @@
-import asyncio
-import json
-import time
-from pathlib import Path
-
-from labnotes.utils.tree_usage import tree_usage
-
-
-async def record_memory(path: Path, pid: int, created: float, interval: float) -> None:
-    with path.open("a", encoding="utf-8") as samples:
-        while True:
-            usage = tree_usage(pid, created)
-            if usage is not None:
-                rss, processes = usage
-                sample = {"time": time.time(), "rss": rss, "processes": processes}
-                samples.write(json.dumps(sample) + "\n")
-                samples.flush()
-            await asyncio.sleep(interval)