From dd2b25228dba67efd0d2948fde67422a59b8e193 Mon Sep 17 00:00:00 2001 From: Cary Palmer Date: Thu, 22 Oct 2026 23:45:13 -0400 Subject: [PATCH 1/3] flow: opt-in in-process map walks With defaults.inprocess or PUPPETMASTER_FLOW_INPROCESS=1, a map walks its items as threads of the parent walker, not as new walker processes. Run locks, markers, stop and resume do not change. A thread walk that raises leaves the run interrupted, and the map starts it again. --- docs/FLOWS.md | 13 ++++++ puppetmaster/flow.py | 44 ++++++++++++++++- tests/test_flow_inprocess.py | 91 ++++++++++++++++++++++++++++++++++++ 3 files changed, 147 insertions(+), 1 deletion(-) create mode 100644 tests/test_flow_inprocess.py diff --git a/docs/FLOWS.md b/docs/FLOWS.md index e208c6e..3bb3726 100644 --- a/docs/FLOWS.md +++ b/docs/FLOWS.md @@ -444,6 +444,19 @@ a resumed walker kills the command the dead walker left running (when the pid still names that command) and runs it again. A pid that now names another process is left alone. +### Map items in the walker process + +By default a map starts one walker process for each item, and a nested map +does the same again for its own items. Set `"defaults": {"inprocess": true}`, +or `PUPPETMASTER_FLOW_INPROCESS=1` for the walker, to walk the items as +threads of the parent walker instead. A graph value wins. + +- Run locks, markers, stop, cut, resume and reuse work as with processes. +- A thread walk that raises leaves its run interrupted, and the map starts + it again (at most 4 times), as for a walker process that exits. +- It saves one process start for each item. That is about 0.2 s on a plain + machine, but more where each process start is checked or guarded. + ### One shared Codex host A Codex worker usually starts its own `codex exec` process: the launcher, the diff --git a/puppetmaster/flow.py b/puppetmaster/flow.py index 228d49d..cdc1806 100644 --- a/puppetmaster/flow.py +++ b/puppetmaster/flow.py @@ -50,7 +50,9 @@ import shutil import subprocess import sys +import threading import time +import traceback import unicodedata import uuid from dataclasses import dataclass, field @@ -137,6 +139,8 @@ def validate_graph(graph: Any, *, child: bool = False, depth: int = 0, problems.extend(_effort_problems(own_defaults, "defaults")) if "codex_host" in own_defaults and not isinstance(own_defaults["codex_host"], bool): problems.append("defaults.codex_host must be true or false") + if "inprocess" in own_defaults and not isinstance(own_defaults["inprocess"], bool): + problems.append("defaults.inprocess must be true or false") lanes = own_defaults.get("lanes") if lanes is not None: if not isinstance(lanes, dict) or not set(lanes) <= set(LANES): @@ -1558,6 +1562,42 @@ def spawn_background_walk(state_dir: Path, run_id: str, *, answer: Optional[str] return process.pid +INPROCESS_ENV = "PUPPETMASTER_FLOW_INPROCESS" + + +def inprocess_requested(graph: dict, env: Optional[dict] = None) -> bool: + """Whether map items walk as threads of the parent walker, not as new processes.""" + env = os.environ if env is None else env + defaults = graph.get("defaults") if isinstance(graph.get("defaults"), dict) else {} + if "inprocess" in defaults: + return defaults.get("inprocess") is True + return str(env.get(INPROCESS_ENV, "")).strip().lower() in ("1", "true", "on", "yes") + + +def walk_in_thread(state_dir: Path, run_id: str) -> threading.Thread: + """Walk ``run_id`` on a thread of this process. + + The run lock, markers, stop and resume work as for a walker process. Each + walker process pays its own start (interpreter, imports, guards); a thread + does not. A walk that raises leaves the run ``running`` with a free lock, + so the map sees it as interrupted and starts it again. + """ + def target() -> None: + try: + walk(state_dir, run_id) + except FlowError: + pass # another walker holds the run + except Exception: + try: + with open(run_path(state_dir, run_id).with_name("walker.log"), "a", encoding="utf-8") as log: + log.write(traceback.format_exc()) + except OSError: + pass + thread = threading.Thread(target=target, name=f"flow-walk-{run_id}", daemon=True) + thread.start() + return thread + + def _fail_unwalked(state_dir: Path, run_id: str, reason: str) -> None: """Fail a run only if no walker holds it; never race a slow-starting walker's writes.""" lock = _lock(state_dir, run_id, timeout=0) @@ -1614,6 +1654,7 @@ def __init__(self, state_dir: Path, *, backend: str = "sqlite", worker_mode: str self.backend = backend self.worker_mode = worker_mode self.popen = popen or subprocess.Popen + self.custom_spawn = spawn is not None self.spawn = spawn or (lambda state_dir, run_id: spawn_background_walk( state_dir, run_id, popen=self.popen)) self.poll_seconds = poll_seconds @@ -2041,6 +2082,7 @@ def _run_map(self, node: dict, run: FlowRun, prev: str) -> NodeOutcome: def _drive_children(self, run: FlowRun, node: dict, children: dict[str, str]) -> None: concurrency = int(node.get("concurrency", DEFAULT_MAP_CONCURRENCY)) + spawn = self.spawn if self.custom_spawn or not inprocess_requested(run.graph) else walk_in_thread started: dict[str, float] = {} spawns: dict[str, int] = {} finished: dict[str, str] = {} @@ -2073,7 +2115,7 @@ def _drive_children(self, run: FlowRun, node: dict, children: dict[str, str]) -> continue spawns[child_id] = spawns.get(child_id, 0) + 1 started[child_id] = now - self.spawn(self.state_dir, child_id) + spawn(self.state_dir, child_id) live.add(child_id) time.sleep(self.poll_seconds) diff --git a/tests/test_flow_inprocess.py b/tests/test_flow_inprocess.py new file mode 100644 index 0000000..97f1f7e --- /dev/null +++ b/tests/test_flow_inprocess.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +import os +import sys + +sys.path.insert(0, os.path.dirname(__file__)) +import hermetic_env # noqa: F401 # process-wide host-env isolation + +import json +import shutil +import tempfile +import threading +import unittest +from pathlib import Path +from unittest.mock import patch + +from puppetmaster import flow + + +def nested_graph(inprocess=True): + check = f'"{sys.executable}" -c "print(\'VERDICT: PASS - {{{{item}}}} ok\')"' + region = {"entry": "check", "nodes": [{"id": "check", "kind": "shell", "command": check}, + {"id": "done", "kind": "end", "status": "pass"}], + "edges": [{"from": "check", "to": "done", "when": "ok"}]} + group = {"entry": "regions", "nodes": [{"id": "regions", "kind": "map", "items": "{{item.regions}}", + "concurrency": 2, "graph": region}, + {"id": "done", "kind": "end", "status": "pass"}], + "edges": [{"from": "regions", "to": "done", "when": "ok"}]} + return {"id": "inproc", "entry": "groups", "defaults": {"inprocess": inprocess}, + "nodes": [{"id": "groups", "kind": "map", "concurrency": 2, "graph": group, + "items": [{"id": "g0", "regions": ["r0", "r1"]}, {"id": "g1", "regions": ["r2", "r3"]}]}, + {"id": "done", "kind": "end", "status": "pass"}], + "edges": [{"from": "groups", "to": "done", "when": "ok"}]} + + +class InprocessWalkTests(unittest.TestCase): + def setUp(self): + self.folder = Path(tempfile.mkdtemp()) + self.addCleanup(shutil.rmtree, self.folder, True) + self.state = self.folder / "state" + + def run_graph(self, graph): + run = flow.new_run(self.state, graph, cwd=str(self.folder)) + return flow.walk(self.state, run.run_id, execute=flow.JobNodeExecutor(self.state, poll_seconds=0.05)) + + def test_nested_map_walks_without_a_new_process(self): + def no_process(*_a, **_k): + raise AssertionError("an in-process walk started a walker process") + threads = [] + real = flow.walk_in_thread + with patch.object(flow, "spawn_background_walk", no_process), \ + patch.object(flow, "walk_in_thread", lambda s, r: threads.append(real(s, r)) or threads[-1]): + run = self.run_graph(nested_graph()) + self.assertEqual((run.status, run.reason), ("done", "pass"), run.steps) + self.assertEqual(len(threads), 6) # 2 groups and 4 regions + for thread in threads: + thread.join(5) + self.assertFalse(thread.is_alive()) + + def test_graph_false_wins_over_the_operator(self): + self.assertFalse(flow.inprocess_requested({"defaults": {"inprocess": False}}, + env={flow.INPROCESS_ENV: "1"})) + self.assertTrue(flow.inprocess_requested({}, env={flow.INPROCESS_ENV: "1"})) + self.assertFalse(flow.inprocess_requested({}, env={})) + graph = nested_graph() + graph["defaults"]["inprocess"] = "yes" + self.assertIn("defaults.inprocess must be true or false", flow.validate_graph(graph)) + + def test_a_walk_that_raises_is_started_again(self): + calls = {"n": 0} + real_walk = flow.walk + lock = threading.Lock() + + def flaky(state_dir, run_id, **kwargs): + run = flow.load_run(state_dir, run_id) + with lock: + first = run.parent is not None and calls["n"] == 0 + if first: + calls["n"] += 1 + if first: + raise RuntimeError("walker crashed") + return real_walk(state_dir, run_id, **kwargs) + with patch.object(flow, "walk", flaky): + run = flow.new_run(self.state, nested_graph(), cwd=str(self.folder)) + run = real_walk(self.state, run.run_id, execute=flow.JobNodeExecutor(self.state, poll_seconds=0.05)) + self.assertEqual((run.status, run.reason), ("done", "pass"), run.steps) + self.assertEqual(calls["n"], 1) + + +if __name__ == "__main__": + unittest.main() From caa0f0182eb9c1f2737c541a7cd7329ea8cd1aaa Mon Sep 17 00:00:00 2001 From: Cary Palmer Date: Sat, 24 Oct 2026 11:34:37 -0400 Subject: [PATCH 2/3] state: inline runs bind their store per context, not in os.environ _run_inline_workers set PUPPETMASTER_STATE_DIR in the process env for the whole run. Two inline runs on different stores in one process (Marionette runs swarms inline in its backend) then saw each other's state dir: a worker of one run wrote deltas, stdout sidecars and tool offloads into the other run's store, and the restore order could leave the wrong value. - state.bind_state_dir sets a ContextVar. state_dir_from_env reads it first, then the env. Every in-process reader uses it. - The inline role pool and the agentic parallel tool pool copy the context into their threads. - stamp_worker_env and the agentic shell env export the bound dir, so child processes see the same value as before. - A subprocess worker still sets the env in its own process. Repro: two concurrent inline runs held inside their workers with a barrier. Before, the plain run's worker resolved the other store. After, each resolves its own, and os.environ does not change. --- puppetmaster/adapters/_delta_stream.py | 9 +- puppetmaster/adapters/_dir_manifest.py | 4 +- puppetmaster/adapters/_prompts.py | 6 +- puppetmaster/adapters/_streaming.py | 8 +- puppetmaster/adapters/agentic.py | 8 +- puppetmaster/job_brief.py | 4 +- puppetmaster/orchestrator.py | 20 +--- puppetmaster/state.py | 33 +++++- puppetmaster/worker_fence.py | 3 + tests/test_puppetmaster.py | 7 +- tests/test_state_dir_binding.py | 150 +++++++++++++++++++++++++ 11 files changed, 215 insertions(+), 37 deletions(-) create mode 100644 tests/test_state_dir_binding.py diff --git a/puppetmaster/adapters/_delta_stream.py b/puppetmaster/adapters/_delta_stream.py index c481a23..568613a 100644 --- a/puppetmaster/adapters/_delta_stream.py +++ b/puppetmaster/adapters/_delta_stream.py @@ -28,16 +28,17 @@ from puppetmaster.fs_permissions import mkdir_private, open_private from puppetmaster.models import Task from puppetmaster.redaction import redact_secrets +from puppetmaster.state import state_dir_from_env _DELTA_FILE = "agentic_deltas.ndjson" def _resolve_state_dir() -> Optional[Path]: - """The active state dir from ``PUPPETMASTER_STATE_DIR`` (exported by - ``worker_runtime`` after resolving --state-dir), or ``None`` when no state - dir is in scope -- mirroring :func:`_streaming._resolve_sidecar_state_dir`. + """The active state dir (bound by an inline run, or ``PUPPETMASTER_STATE_DIR`` + exported by ``worker_runtime``), or ``None`` when no state dir is in scope -- + mirroring :func:`_streaming._resolve_sidecar_state_dir`. """ - raw = os.environ.get("PUPPETMASTER_STATE_DIR") + raw = state_dir_from_env() if not raw: return None try: diff --git a/puppetmaster/adapters/_dir_manifest.py b/puppetmaster/adapters/_dir_manifest.py index 2a3d528..bb74144 100644 --- a/puppetmaster/adapters/_dir_manifest.py +++ b/puppetmaster/adapters/_dir_manifest.py @@ -11,6 +11,8 @@ from pathlib import Path from typing import Optional +from puppetmaster.state import state_dir_from_env + # Caches and tool state that no task owns. The gate also skips bytecode. SKIP_DIRS = frozenset({ ".git", ".hg", ".svn", "node_modules", "__pycache__", ".venv", "venv", @@ -23,7 +25,7 @@ def _excluded_roots(root: Path) -> list[Path]: """The live Puppetmaster state dir changes during every run; it is not the worker's.""" - raw = os.environ.get("PUPPETMASTER_STATE_DIR") + raw = state_dir_from_env() if not raw: return [] try: diff --git a/puppetmaster/adapters/_prompts.py b/puppetmaster/adapters/_prompts.py index 00d81da..fccd138 100644 --- a/puppetmaster/adapters/_prompts.py +++ b/puppetmaster/adapters/_prompts.py @@ -609,19 +609,17 @@ def _open_store_for_task(task: Task, *, discover: bool = True): so worker subprocesses see the same store the orchestrator wrote into. Returns ``None`` on any miss or error. """ - import os - job_id = getattr(task, "job_id", None) or "" if not job_id: return None from puppetmaster.adapters._streaming import _resolve_sidecar_state_dir - from puppetmaster.state import STATE_DIR_ENV, find_state_dir_for_job, resolve_state_dir + from puppetmaster.state import find_state_dir_for_job, resolve_state_dir, state_dir_from_env from puppetmaster.store_factory import create_store state_dir = _resolve_sidecar_state_dir() if state_dir is None and discover: state_dir = find_state_dir_for_job(job_id) - if state_dir is None and os.environ.get(STATE_DIR_ENV): + if state_dir is None and state_dir_from_env(): try: state_dir = resolve_state_dir() except Exception: diff --git a/puppetmaster/adapters/_streaming.py b/puppetmaster/adapters/_streaming.py index f6c3527..6798759 100644 --- a/puppetmaster/adapters/_streaming.py +++ b/puppetmaster/adapters/_streaming.py @@ -12,6 +12,7 @@ from puppetmaster.invocation import current_attempt from puppetmaster.models import Task, new_id from puppetmaster.redaction import redact_secrets +from puppetmaster.state import state_dir_from_env from puppetmaster.worker_fence import stamp_worker_env _STDOUT_HEAD_CHARS = 1000 @@ -46,10 +47,11 @@ def _resolve_sidecar_state_dir() -> Optional[Path]: Returns ``None`` if no state dir is in scope (e.g. direct adapter unit tests). Falling back to the default state dir would write logs into a workspace-hashed path that may not own the job, so we only honor an - explicit ``PUPPETMASTER_STATE_DIR`` env var (which ``worker_runtime`` - exports after resolving its --state-dir flag). + explicit state dir: the one an inline run bound, or the + ``PUPPETMASTER_STATE_DIR`` env var (which ``worker_runtime`` exports after + resolving its --state-dir flag). """ - raw = os.environ.get("PUPPETMASTER_STATE_DIR") + raw = state_dir_from_env() if not raw: return None try: diff --git a/puppetmaster/adapters/agentic.py b/puppetmaster/adapters/agentic.py index 1937fad..8304d3f 100644 --- a/puppetmaster/adapters/agentic.py +++ b/puppetmaster/adapters/agentic.py @@ -48,6 +48,7 @@ import subprocess import tempfile import time +from contextvars import copy_context from pathlib import Path from typing import Any, Callable, Optional @@ -78,7 +79,7 @@ from puppetmaster.redaction import redact_secrets from puppetmaster.usage import selected_token_usage from puppetmaster.swarm_reasoning import DEFAULT_SWARM_REASONING_EFFORT -from puppetmaster.state import resolve_state_dir +from puppetmaster.state import export_state_dir, resolve_state_dir from puppetmaster.hashline import ( SnapshotStore, apply_patch, @@ -2260,7 +2261,8 @@ def _execute_one(call: dict) -> tuple[dict, str]: # Segmentation / barrier rules are unchanged; only pool size is bounded. worker_count = parallel_executor_max_workers(len(calls)) with concurrent.futures.ThreadPoolExecutor(max_workers=worker_count) as executor: - futures = {executor.submit(_execute_one, call): idx for idx, call in enumerate(calls)} + futures = {executor.submit(copy_context().run, _execute_one, call): idx + for idx, call in enumerate(calls)} results = [None] * len(calls) for future in concurrent.futures.as_completed(futures): idx = futures[future] @@ -2731,7 +2733,7 @@ def _run_owned_shell( owner = new_id("process") env = { - **os.environ, + **export_state_dir(dict(os.environ)), "GIT_PAGER": "cat", "PAGER": "cat", "GIT_TERMINAL_PROMPT": "0", diff --git a/puppetmaster/job_brief.py b/puppetmaster/job_brief.py index ecee70c..4d7c57c 100644 --- a/puppetmaster/job_brief.py +++ b/puppetmaster/job_brief.py @@ -171,12 +171,12 @@ def resolve_job_brief_for_task(task) -> str: if isinstance(inline, str) and inline.strip(): return inline.strip() + ("\n" if not inline.endswith("\n") else "") from puppetmaster.adapters._streaming import _resolve_sidecar_state_dir - from puppetmaster.state import STATE_DIR_ENV, find_state_dir_for_job, resolve_state_dir + from puppetmaster.state import find_state_dir_for_job, resolve_state_dir, state_dir_from_env state_dir = _resolve_sidecar_state_dir() if state_dir is None: state_dir = find_state_dir_for_job(job_id) - if state_dir is None and os.environ.get(STATE_DIR_ENV): + if state_dir is None and state_dir_from_env(): # Explicit env pin — resolve_state_dir short-circuits before git. try: state_dir = resolve_state_dir() diff --git a/puppetmaster/orchestrator.py b/puppetmaster/orchestrator.py index c04a786..70c5190 100644 --- a/puppetmaster/orchestrator.py +++ b/puppetmaster/orchestrator.py @@ -1,11 +1,11 @@ from __future__ import annotations -import contextlib import os import subprocess import sys import time from concurrent.futures import ThreadPoolExecutor +from contextvars import copy_context from dataclasses import dataclass, replace from datetime import datetime, timezone from pathlib import Path @@ -25,6 +25,7 @@ now_iso, ) from puppetmaster.stitcher import Stitcher +from puppetmaster.state import bind_state_dir from puppetmaster.store import SwarmStore from puppetmaster.swarm_reasoning import NO_EFFORT_ADAPTERS, apply_swarm_reasoning from puppetmaster.worker_fence import stamp_worker_env @@ -184,19 +185,6 @@ def _env_int(name: str, default: int) -> int: return default -@contextlib.contextmanager -def _temporary_env_var(name: str, value: str): - previous = os.environ.get(name) - os.environ[name] = value - try: - yield - finally: - if previous is None: - os.environ.pop(name, None) - else: - os.environ[name] = previous - - def merge_routing_payload( payload: dict, decision, extra_fields: Optional[dict] = None, *, registry=None, previous_adapter: Optional[str] = None, @@ -2818,7 +2806,7 @@ def _run_inline_workers( allowed_task_ids: Optional[set[str]] = None, ) -> None: allowed_task_ids = allowed_task_ids or {task.id for task in tasks} - with _temporary_env_var("PUPPETMASTER_STATE_DIR", str(self.store.root)): + with bind_state_dir(self.store.root): while True: record_orchestrator_heartbeat(self.store, job.id) self.store.recover_stale_tasks(job.id) @@ -2882,7 +2870,7 @@ def _run_role(role: str) -> int: pool_size = min(len(roles), _INLINE_ROLE_WORKER_POOL_CAP) with ThreadPoolExecutor(max_workers=pool_size) as executor: - futures = [executor.submit(_run_role, role) for role in roles] + futures = [executor.submit(copy_context().run, _run_role, role) for role in roles] completed = 0 errors: list[BaseException] = [] for future in futures: diff --git a/puppetmaster/state.py b/puppetmaster/state.py index b1a8c0c..2321133 100644 --- a/puppetmaster/state.py +++ b/puppetmaster/state.py @@ -2,6 +2,8 @@ import hashlib from collections import OrderedDict +from contextlib import contextmanager +from contextvars import ContextVar import threading import os import re @@ -9,13 +11,40 @@ import subprocess import sys from pathlib import Path -from typing import Optional, Union +from typing import Iterator, Optional, Union from puppetmaster.readonly import ReadUnavailable STATE_DIR_ENV = "PUPPETMASTER_STATE_DIR" +# An in-process run binds its store here, not in os.environ: two inline runs +# on different stores in one process must not see each other's state dir. +_BOUND_STATE_DIR: ContextVar[Optional[str]] = ContextVar("puppetmaster_state_dir", default=None) + + +@contextmanager +def bind_state_dir(path: Union[Path, str]) -> Iterator[None]: + """Make ``path`` the state dir of this context and of the threads it copies to.""" + token = _BOUND_STATE_DIR.set(str(path)) + try: + yield + finally: + _BOUND_STATE_DIR.reset(token) + + +def state_dir_from_env() -> Optional[str]: + """The state dir bound by an in-process run, else ``PUPPETMASTER_STATE_DIR``.""" + return _BOUND_STATE_DIR.get() or os.environ.get(STATE_DIR_ENV) or None + + +def export_state_dir(env: dict) -> dict: + """Give a child process the bound state dir. Mutates and returns ``env``.""" + bound = _BOUND_STATE_DIR.get() + if bound: + env[STATE_DIR_ENV] = bound + return env + def state_identity(path: Union[Path, str]) -> str: """Return a stable opaque identity for a resolved state directory.""" @@ -28,7 +57,7 @@ def resolve_state_dir(value: Optional[Union[Path, str]] = None, cwd: Optional[Pa root = cwd or Path.cwd() if value: return _resolve_user_path(value, root) - env_value = os.environ.get(STATE_DIR_ENV) + env_value = state_dir_from_env() if env_value: return _resolve_user_path(env_value, root) return default_state_dir(root) diff --git a/puppetmaster/worker_fence.py b/puppetmaster/worker_fence.py index 75bc235..08b89dd 100644 --- a/puppetmaster/worker_fence.py +++ b/puppetmaster/worker_fence.py @@ -11,6 +11,8 @@ import os from typing import Mapping, Optional +from puppetmaster.state import export_state_dir + WORKER_ENV = "PUPPETMASTER_WORKER" WORKER_VALUE = "1" ALLOW_NESTED_ENV = "PUPPETMASTER_ALLOW_NESTED" @@ -71,6 +73,7 @@ def stamp_worker_env( role: str = "", ) -> dict: """Mark ``env`` as a Puppetmaster worker. Mutates and returns ``env``.""" + export_state_dir(env) env[WORKER_ENV] = WORKER_VALUE if job_id: env[JOB_ID_ENV] = str(job_id) diff --git a/tests/test_puppetmaster.py b/tests/test_puppetmaster.py index e6ff881..0cf53a0 100644 --- a/tests/test_puppetmaster.py +++ b/tests/test_puppetmaster.py @@ -28159,7 +28159,9 @@ def test_inline_workers_scope_state_dir_env_to_store_root(self) -> None: seen = [] def complete_task(runtime): - seen.append(os.environ.get("PUPPETMASTER_STATE_DIR")) + from puppetmaster.state import state_dir_from_env + + seen.append((state_dir_from_env(), os.environ.get("PUPPETMASTER_STATE_DIR"))) stored = store.get_task_by_id(task.id) store.update_task_status(stored, TaskStatus.COMPLETE) return 1 @@ -28173,7 +28175,8 @@ def complete_task(runtime): Orchestrator(store)._run_inline_workers(job, [task]) self.assertEqual(os.environ["PUPPETMASTER_STATE_DIR"], "parent-state") - self.assertEqual(seen, [str(store.root)]) + # The run binds its store; the process env keeps the parent value. + self.assertEqual(seen, [(str(store.root), "parent-state")]) def test_redact_secrets_scrubs_argument_supplied_keys(self) -> None: from puppetmaster.redaction import ( diff --git a/tests/test_state_dir_binding.py b/tests/test_state_dir_binding.py new file mode 100644 index 0000000..fa29bf9 --- /dev/null +++ b/tests/test_state_dir_binding.py @@ -0,0 +1,150 @@ +"""An inline run binds its store per context, never in os.environ. + +Two inline runs on different stores in one process (Marionette runs swarms +inline in its backend) must each resolve their own state dir. Before the +binding, ``_run_inline_workers`` set ``PUPPETMASTER_STATE_DIR`` process-wide, +so worker output of one run could land in the other run's store. +""" +from __future__ import annotations + +import os +import threading +import unittest +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest.mock import patch + +from puppetmaster.adapters import _delta_stream, _streaming +from puppetmaster.adapters.agentic import AgenticAdapter +from puppetmaster.models import Task +from puppetmaster.orchestrator import Orchestrator +from puppetmaster.state import STATE_DIR_ENV, bind_state_dir, resolve_state_dir, state_dir_from_env +from puppetmaster.store import SwarmStore +from puppetmaster.worker_fence import stamp_worker_env +from puppetmaster.workers import WorkerSpec + + +def _seen() -> dict: + return { + "bound": state_dir_from_env(), + "sidecar": str(_streaming._resolve_sidecar_state_dir()), + "delta": str(_delta_stream._resolve_state_dir()), + "resolved": str(resolve_state_dir()), + "env": os.environ.get(STATE_DIR_ENV), + } + + +class StateDirBindingTests(unittest.TestCase): + def setUp(self) -> None: + patcher = patch.dict(os.environ) + patcher.start() + self.addCleanup(patcher.stop) + os.environ.pop(STATE_DIR_ENV, None) + + def test_binding_is_local_to_each_thread_context(self) -> None: + inside = threading.Barrier(2, timeout=5) + seen: dict = {} + + def run(name: str, path: str) -> None: + with bind_state_dir(path): + inside.wait() + seen[name] = _seen() + inside.wait() + seen[name + "_after"] = state_dir_from_env() + + with TemporaryDirectory() as a, TemporaryDirectory() as b: + threads = [threading.Thread(target=run, args=("a", a)), threading.Thread(target=run, args=("b", b))] + for t in threads: + t.start() + for t in threads: + t.join() + for name, path in (("a", a), ("b", b)): + self.assertEqual(seen[name]["bound"], path) + self.assertEqual(seen[name]["sidecar"], path) + self.assertEqual(seen[name]["delta"], path) + self.assertEqual(Path(seen[name]["resolved"]).resolve(), Path(path).resolve()) + self.assertIsNone(seen[name]["env"]) + self.assertIsNone(seen[name + "_after"]) + + def test_binding_wins_over_the_process_env(self) -> None: + os.environ[STATE_DIR_ENV] = "/env/state" + self.assertEqual(state_dir_from_env(), "/env/state") + with bind_state_dir("/bound/state"): + self.assertEqual(state_dir_from_env(), "/bound/state") + self.assertEqual(state_dir_from_env(), "/env/state") + + def test_child_env_gets_the_bound_state_dir(self) -> None: + self.assertNotIn(STATE_DIR_ENV, stamp_worker_env({})) + self.assertEqual(stamp_worker_env({STATE_DIR_ENV: "/kept"})[STATE_DIR_ENV], "/kept") + with bind_state_dir("/bound/state"): + self.assertEqual(stamp_worker_env({STATE_DIR_ENV: "/old"})[STATE_DIR_ENV], "/bound/state") + + def test_concurrent_inline_runs_each_see_their_own_store(self) -> None: + from puppetmaster.worker_runtime import WorkerRuntime + + original = WorkerRuntime.run_until_idle + inside = threading.Barrier(2, timeout=10) + seen: dict = {} + lock = threading.Lock() + + def run_until_idle(runtime: WorkerRuntime) -> int: + first = False + with lock: + if runtime.store.root not in seen: + seen[runtime.store.root] = None + first = True + if first: + inside.wait() + observed = _seen() + inside.wait() + with lock: + seen[runtime.store.root] = observed + return original(runtime) + + with TemporaryDirectory() as tmp: + stores = [SwarmStore(Path(tmp) / name) for name in ("plain", "sealed")] + errors: list = [] + + def run(store: SwarmStore) -> None: + try: + Orchestrator(store).run( + "binding probe", specs=[WorkerSpec(role="explore", instruction="look")], + worker_mode="inline", + ) + except BaseException as exc: # surfaced below + errors.append(exc) + + with patch("puppetmaster.orchestrator.WorkerRuntime.run_until_idle", run_until_idle): + threads = [threading.Thread(target=run, args=(s,)) for s in stores] + for t in threads: + t.start() + for t in threads: + t.join() + self.assertEqual(errors, []) + for store in stores: + observed = seen[store.root] + self.assertEqual(observed["bound"], str(store.root)) + self.assertEqual(observed["sidecar"], str(store.root)) + self.assertEqual(observed["delta"], str(store.root)) + self.assertIsNone(observed["env"]) + self.assertNotIn(STATE_DIR_ENV, os.environ) + + def test_agentic_parallel_tool_calls_see_the_binding(self) -> None: + adapter = AgenticAdapter() + seen: list = [] + + def execute_tool(name, args, cwd, implement, task): + seen.append(state_dir_from_env()) + return "ok" + + adapter._execute_tool = execute_tool + task = Task(job_id="job_x", role="explore", instruction="look") + calls = [{"name": "read_file", "arguments": {"path": str(i)}} for i in range(4)] + with bind_state_dir("/bound/state"): + results = adapter._execute_tool_segment_parallel(calls, Path("."), False, task) + self.assertEqual([out for _, out in results], ["ok"] * 4) + self.assertEqual(seen, ["/bound/state"] * 4) + + +if __name__ == "__main__": + unittest.main() From 097d6aa3bcb4641e3924dfa964b16876cce399ac Mon Sep 17 00:00:00 2001 From: Cary Palmer Date: Sat, 24 Oct 2026 11:53:53 -0400 Subject: [PATCH 3/3] Release v1.48.0: inline runs bind their store per context; opt-in in-process map walks --- docs/CHANGELOG.md | 27 +++++++++++++++++++++++++++ puppetmaster/__init__.py | 2 +- pyproject.toml | 2 +- 3 files changed, 29 insertions(+), 2 deletions(-) diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 02ec88f..92a49ad 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -1,3 +1,30 @@ +## v1.48.0 — 2026-10-24 + +**An inline run binds its store per context, not in the process env.** +`_run_inline_workers` set `PUPPETMASTER_STATE_DIR` in `os.environ` for the +whole run. Two inline runs on different stores in one process (Marionette +runs swarms inline in its backend) then saw each other's state dir. A worker +of one run wrote its deltas, stdout sidecars and tool offloads into the +other run's store, and the restore order could leave the wrong value. + +- `state.bind_state_dir` sets a context variable. `state.state_dir_from_env` + reads it first, then the env. Every in-process reader uses it. +- The inline role pool and the agentic parallel tool pool copy the context + into their threads. +- `stamp_worker_env` and the agentic shell env give child processes the + bound dir, so children see the same value as before. +- A subprocess worker still sets the env in its own process. +- Repro: two concurrent inline runs, held inside their workers with a + barrier. Before, one run's worker resolved the other store. Now each run + resolves its own, and `os.environ` does not change. + +**Opt-in: map items as threads of the walker.** Set +`"defaults": {"inprocess": true}` or `PUPPETMASTER_FLOW_INPROCESS=1` to walk +map items as threads of the parent walker, not as new walker processes. Run +locks, markers, stop, cut, resume and reuse do not change. A thread walk that +raises leaves its run interrupted, and the map starts it again. See +`docs/FLOWS.md`. + ## v1.47.0 — 2026-10-21 **Opt-in: one shared Codex host for a flow walk.** A Codex worker starts its diff --git a/puppetmaster/__init__.py b/puppetmaster/__init__.py index fcef00a..983e5a9 100644 --- a/puppetmaster/__init__.py +++ b/puppetmaster/__init__.py @@ -1,4 +1,4 @@ """Puppetmaster: distributed agent workers with stitched shared memory.""" __all__ = ["__version__"] -__version__ = "1.47.0" +__version__ = "1.48.0" diff --git a/pyproject.toml b/pyproject.toml index c8dd9e4..86d25d1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "puppetmaster-ai" -version = "1.47.0" +version = "1.48.0" description = "Gunicorn-style agent swarms with Redis-like shared memory and artifact stitching. (Imported as `puppetmaster`; published as `puppetmaster-ai` because the bare PyPI name is held by an abandoned 2019 project, name-reassignment pending.)" readme = "README.md" requires-python = ">=3.9"