Add Execution Strategies: Client-Server and Shared Memory implementations (#120)

agentlightning/execution/base.py

@@ -0,0 +1,37 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+from typing import Protocol
+
+from agentlightning.store.base import LightningStore
+
+from .events import Event
+
+logger = logging.getLogger(__name__)
+
+
+class AlgorithmBundle(Protocol):
+    async def __call__(self, store: LightningStore, event: Event) -> None:
+        """Initalization and execution logic."""
+
+
+class RunnerBundle(Protocol):
+    async def __call__(self, store: LightningStore, worker_id: int, event: Event) -> None:
+        """Initalization and execution logic."""
+
+
+class ExecutionStrategy:
+    """When trainer has created the executable of algorithm and runner in two bundles,
+    the execution strategy defines how to run them together, and how many parallel runners to run.
+
+    The store is the centric place for the two bundles to communicate.
+
+    The algorithm and runner's behavior (whether runner should perform one step or run forever,
+    whether the algo would send out the tasks or not) are defined inside the bundle,
+    and does not belong to the execution strategy.
+
+    The execute should support Ctrl+C to exit gracefully.
+    """
+
+    def execute(self, algorithm: AlgorithmBundle, runner: RunnerBundle, store: LightningStore) -> None:
+        raise NotImplementedError()

agentlightning/execution/client_server.py

@@ -0,0 +1,379 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import logging
+import multiprocessing
+import os
+import signal
+import time
+from multiprocessing.context import BaseContext
+from typing import Callable, Iterable, Literal, cast
+
+from agentlightning.store.base import LightningStore
+from agentlightning.store.client_server import LightningStoreClient, LightningStoreServer
+
+from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle
+from .events import Event, MultiprocessingEvent
+
+logger = logging.getLogger(__name__)
+
+
+class ClientServerExecutionStrategy(ExecutionStrategy):
+    """Run algorithm (server) and runners (clients) as separate processes over HTTP.
+
+    **Execution Roles:**
+
+    - "algorithm": Start the HTTP server (`LightningStoreServer`) in-process and run the
+      algorithm bundle against it.
+    - "runner": Connect to an already running server via `LightningStoreClient` and
+      execute runner bundles (optionally in multiple processes).
+    - "both": Spawn the runner processes first, then launch the algorithm/server
+      bundle on the main process. This mode orchestrates the full loop locally.
+
+    When role == "both", you may choose which side runs on the main process via
+    `main_process` (debug helper). Running the runner bundle on the main process
+    is only supported with `n_runners == 1`.
+
+    Important: When `main_process == "runner"`, the algorithm runs in a subprocess
+    with the LightningStore server. This means any state modifications made during
+    execution remain in that subprocess and are NOT reflected in the original store
+    object passed to `execute()`. The main process runner accesses the store only
+    through the HTTP client interface.
+
+    **Abort / Stop Model (four-step escalation):**
+
+    1. Cooperative stop:
+       A shared :class:`~agentlightning.execution.events.MultiprocessingEvent`
+       (`stop_evt`) is passed to *all* bundles. Bundles should check it to exit.
+       Any crash (algorithm or runner) sets `stop_evt` so the other side can
+       stop cooperatively. Ctrl+C on the main process also flips the event.
+    2. KeyboardInterrupt synth:
+       Remaining subprocesses receive `SIGINT` to trigger `KeyboardInterrupt`
+       handlers.
+    3. Termination:
+       Stubborn subprocesses get `terminate()` (SIGTERM on POSIX).
+    4. Kill:
+        As a last resort we call `kill()` (SIGKILL on POSIX).
+
+    Notes:
+        This mirrors the semantics implemented in :mod:`shared_memory`, but adapted
+        to multiple processes and the HTTP client/server boundary.
+    """
+
+    alias: str = "cs"
+
+    def __init__(
+        self,
+        role: Literal["algorithm", "runner", "both"],
+        server_host: str = "localhost",
+        server_port: int = 4747,
+        n_runners: int = 1,
+        graceful_timeout: float = 5.0,
+        terminate_timeout: float = 5.0,
+        main_process: Literal["algorithm", "runner"] = "algorithm",
+    ) -> None:
+        """Configure the strategy.
+
+        Args:
+            role: Which side(s) to run in this process.
+            server_host: Interface the HTTP server binds to when running the
+                algorithm bundle locally.
+            server_port: Port for the HTTP server in "algorithm"/"both" modes.
+            n_runners: Number of runner processes to spawn in "runner"/"both".
+            graceful_timeout: How long to wait (seconds) after setting the stop
+                event before escalating to signals.
+            terminate_timeout: How long to wait between escalation steps beyond
+                the cooperative phase (re-used for SIGINT, terminate, and kill).
+            main_process: Which bundle runs on the main process when
+                `role == "both"`. `"runner"` requires `n_runners == 1` and
+                is primarily intended for debugging.
+        """
+        self.role = role
+        self.n_runners = n_runners
+        self.server_host = server_host
+        self.server_port = server_port
+        self.graceful_timeout = graceful_timeout
+        self.terminate_timeout = terminate_timeout
+        if main_process not in ("algorithm", "runner"):
+            raise ValueError("main_process must be 'algorithm' or 'runner'")
+        if main_process == "runner":
+            if role != "both":
+                raise ValueError("main_process='runner' is only supported when role='both'")
+            if n_runners != 1:
+                raise ValueError("main_process='runner' requires n_runners to be 1")
+        self.main_process = main_process
+
+    async def _execute_algorithm(self, algorithm: AlgorithmBundle, store: LightningStore, stop_evt: Event) -> None:
+        logger.info("Starting LightningStore server on %s:%s", self.server_host, self.server_port)
+        server_store = LightningStoreServer(store, host=self.server_host, port=self.server_port)
+        server_started = False
+
+        try:
+            await server_store.start()
+            server_started = True
+            logger.debug("Algorithm bundle starting against endpoint %s", server_store.endpoint)
+            await algorithm(server_store, stop_evt)
+            logger.debug("Algorithm bundle completed successfully")
+        except KeyboardInterrupt:
+            logger.warning("Algorithm received KeyboardInterrupt; signaling stop event")
+            stop_evt.set()
+            raise
+        except BaseException:
+            logger.exception("Algorithm bundle crashed; signaling stop event")
+            stop_evt.set()
+            raise
+        finally:
+            if server_started:
+                try:
+                    await server_store.stop()
+                except Exception:
+                    logger.exception("Error stopping LightningStore server")
+                else:
+                    logger.debug("LightningStore server shutdown completed")
+
+    async def _execute_runner(self, runner: RunnerBundle, worker_id: int, stop_evt: Event) -> None:
+        client_store = LightningStoreClient(f"http://{self.server_host}:{self.server_port}")
+        try:
+            logger.debug("Runner %s connecting to server at %s:%s", worker_id, self.server_host, self.server_port)
+            await runner(client_store, worker_id, stop_evt)
+            logger.debug("Runner %s completed successfully", worker_id)
+        except KeyboardInterrupt:
+            logger.warning("Runner %s received KeyboardInterrupt; signaling stop event", worker_id)
+            stop_evt.set()
+            raise
+        except BaseException:
+            logger.exception("Runner %s crashed; signaling stop event", worker_id)
+            stop_evt.set()
+            raise
+        finally:
+            try:
+                await client_store.close()
+            except Exception:
+                logger.exception("Error closing LightningStore client for runner %s", worker_id)
+            else:
+                logger.debug("Runner %s closed LightningStore client", worker_id)
+
+    def _spawn_runners(
+        self,
+        runner: RunnerBundle,
+        stop_evt: Event,
+        *,
+        ctx: BaseContext,
+    ) -> list[multiprocessing.Process]:
+        """Used when `role == "runner"` or `role == "both"` and `n_runners > 1`."""
+        processes: list[multiprocessing.Process] = []
+
+        def _runner_sync(runner: RunnerBundle, worker_id: int, stop_evt: Event) -> None:
+            # Runners are executed in child processes; each process owns its own
+            # event loop to keep the asyncio scheduler isolated.
+            asyncio.run(self._execute_runner(runner, worker_id, stop_evt))
+
+        for i in range(self.n_runners):
+            process = cast(
+                multiprocessing.Process,
+                ctx.Process(target=_runner_sync, args=(runner, i, stop_evt), name=f"runner-{i}"),  # type: ignore
+            )
+            process.start()
+            logger.debug("Spawned runner process %s (pid=%s)", process.name, process.pid)
+            processes.append(process)
+
+        return processes
+
+    def _spawn_algorithm_process(
+        self,
+        algorithm: AlgorithmBundle,
+        store: LightningStore,
+        stop_evt: Event,
+        *,
+        ctx: BaseContext,
+    ) -> multiprocessing.Process:
+        """Used when `main_process == "runner"`."""
+
+        def _algorithm_sync(algorithm: AlgorithmBundle, store: LightningStore, stop_evt: Event) -> None:
+            asyncio.run(self._execute_algorithm(algorithm, store, stop_evt))
+
+        process = cast(
+            multiprocessing.Process,
+            ctx.Process(target=_algorithm_sync, args=(algorithm, store, stop_evt), name="algorithm"),  # type: ignore
+        )
+        process.start()
+        logger.debug("Spawned algorithm process %s (pid=%s)", process.name, process.pid)
+        return process
+
+    def _join_until_deadline(
+        self,
+        processes: Iterable[multiprocessing.Process],
+        timeout: float,
+    ) -> list[multiprocessing.Process]:
+        """Join ``processes`` until ``timeout`` elapses, returning those still alive."""
+        deadline = time.monotonic() + timeout
+        still_alive: list[multiprocessing.Process] = []
+        for process in processes:
+            remaining = deadline - time.monotonic()
+            if remaining > 0:
+                process.join(remaining)
+            else:
+                process.join(0)
+            if process.is_alive():
+                still_alive.append(process)
+        return still_alive
+
+    def _signal_processes(
+        self,
+        processes: Iterable[multiprocessing.Process],
+        action: Callable[[multiprocessing.Process], None],
+    ) -> None:
+        """Invoke ``action`` on each process while suppressing individual failures."""
+        for process in processes:
+            try:
+                action(process)
+            except Exception:
+                logger.exception("Error signaling process %s (pid=%s)", process.name, process.pid)
+
+    def _shutdown_processes(
+        self,
+        processes: list[multiprocessing.Process],
+        stop_evt: Event,
+    ) -> None:
+        """4-step escalation shutdown of ``processes``."""
+        if not processes:
+            logger.debug("No subprocesses to shutdown")
+            return
+
+        if not stop_evt.is_set():
+            logger.debug("Sending cooperative stop signal to subprocesses")
+            stop_evt.set()
+        else:
+            logger.debug("Stop event already set; waiting for subprocesses to exit")
+
+        alive = self._join_until_deadline(processes, self.graceful_timeout)
+        if not alive:
+            return
+
+        logger.warning(
+            "Subprocesses still alive after cooperative wait; sending SIGINT to %s",
+            ", ".join(p.name or str(p.pid) for p in alive),
+        )
+        # SIGINT is not reliable on Windows, but we do not consider such case yet.
+        self._signal_processes(alive, lambda p: os.kill(cast(int, p.pid), signal.SIGINT))
+        alive = self._join_until_deadline(alive, self.terminate_timeout)
+        if not alive:
+            return
+
+        logger.warning(
+            "Subprocesses still alive after SIGINT wait; sending terminate() to %s",
+            ", ".join(p.name or str(p.pid) for p in alive),
+        )
+        self._signal_processes(alive, lambda p: p.terminate())
+
+        alive = self._join_until_deadline(alive, self.terminate_timeout)
+        if not alive:
+            return
+
+        logger.error(
+            "Subprocesses still alive after terminate(); sending kill() to %s",
+            ", ".join(p.name or str(p.pid) for p in alive),
+        )
+        self._signal_processes(alive, lambda p: p.kill())
+        alive = self._join_until_deadline(alive, self.terminate_timeout)
+
+        if alive:
+            logger.error(
+                "Subprocesses failed to exit even after kill(): %s", ", ".join(p.name or str(p.pid) for p in alive)
+            )
+
+    def _check_process_exitcodes(self, processes: Iterable[multiprocessing.Process]) -> None:
+        """Raise an error if any managed process exited with a non-zero status."""
+        failed = [p for p in processes if p.exitcode not in (0, None)]
+        if failed:
+            formatted = ", ".join(f"{p.name or p.pid} (exitcode={p.exitcode})" for p in failed)
+            raise RuntimeError(f"Subprocesses failed: {formatted}")
+
+    def execute(self, algorithm: AlgorithmBundle, runner: RunnerBundle, store: LightningStore) -> None:
+        logger.info(
+            "Starting client-server execution with %d runner(s) [role=%s, main_process=%s]",
+            self.n_runners,
+            self.role,
+            self.main_process,
+        )
+
+        # Re-use the active multiprocessing context so the event and processes
+        # agree on the start method (fork/spawn/forkserver).
+        ctx = multiprocessing.get_context()
+        stop_evt = MultiprocessingEvent(ctx=ctx)
+        # Track spawned processes so we can enforce termination ordering and
+        # surface non-zero exit codes back to the caller.
+        processes: list[multiprocessing.Process] = []
+
+        exception: BaseException | None = None
+        keyboard_interrupt = False
+
+        try:
+            if self.role == "algorithm":
+                logger.info("Running algorithm solely...")
+                asyncio.run(self._execute_algorithm(algorithm, store, stop_evt))
+            elif self.role == "runner":
+                if self.n_runners == 1:
+                    logger.info("Running runner solely...")
+                    asyncio.run(self._execute_runner(runner, 0, stop_evt))
+                else:
+                    logger.info("Spawning runner processes...")
+                    processes = self._spawn_runners(runner, stop_evt, ctx=ctx)
+                    # Wait for the processes to finish naturally.
+                    for process in processes:
+                        process.join()
+                    self._check_process_exitcodes(processes)
+            elif self.role == "both":
+                if self.main_process == "algorithm":
+                    logger.info("Spawning runner processes...")
+                    processes = self._spawn_runners(runner, stop_evt, ctx=ctx)
+                    try:
+                        logger.info("Running algorithm...")
+                        asyncio.run(self._execute_algorithm(algorithm, store, stop_evt))
+                    finally:
+                        # Always request the runner side to unwind once the
+                        # algorithm/server portion finishes (successfully or not).
+                        stop_evt.set()
+                else:  # main_process == "runner"
+                    if self.n_runners > 1:
+                        raise ValueError("main_process='runner' requires n_runners to be 1")
+
+                    logger.info("Spawning algorithm process...")
+                    algorithm_process = self._spawn_algorithm_process(algorithm, store, stop_evt, ctx=ctx)
+                    processes = [algorithm_process]
+
+                    # Run the lone runner cooperatively in-process so users can
+                    # attach a debugger. The algorithm + HTTP server live in
+                    # the background process spawned above (the provided
+                    # store must therefore be picklable when using spawn).
+                    logger.info("Running runner...")
+                    asyncio.run(self._execute_runner(runner, 0, stop_evt))
+
+                    # Wait for the algorithm process to finish.
+                    algorithm_process.join()
+            else:
+                raise ValueError(f"Unknown role: {self.role}")
+        except KeyboardInterrupt:
+            logger.warning("KeyboardInterrupt received; initiating shutdown")
+            stop_evt.set()
+            keyboard_interrupt = True
+        except BaseException as exc:
+            logger.exception("Unhandled exception in execute method")
+            stop_evt.set()
+            # Preserve the original exception so we can avoid masking it during
+            # the cleanup phase.
+            exception = exc
+            raise
+        finally:
+            logger.info("Shutting down subprocesses")
+            self._shutdown_processes(processes, stop_evt)
+            if processes:
+                try:
+                    self._check_process_exitcodes(processes)
+                except RuntimeError as err:
+                    if exception is not None or keyboard_interrupt:
+                        # We already propagate/handled a different failure, so
+                        # emit a warning instead of raising a secondary error.
+                        logger.warning("Subprocesses ended abnormally during shutdown: %s", err)
+                    else:
+                        raise

agentlightning/execution/events.py

@@ -0,0 +1,75 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import multiprocessing as mp
+import threading
+from multiprocessing.context import BaseContext
+from typing import Optional, Protocol
+
+
+class Event(Protocol):
+    """
+    A minimal protocol similar to threading.Event.
+
+    Methods:
+        set(): Signal event like a cancellation (idempotent).
+        clear(): Reset to the non-set state.
+        is_set() -> bool: True if event has been signaled.
+        wait(timeout: Optional[float] = None) -> bool:
+            Block until event is set or timeout. Returns True if event has signaled.
+    """
+
+    def set(self) -> None: ...
+    def clear(self) -> None: ...
+    def is_set(self) -> bool: ...
+    def wait(self, timeout: Optional[float] = None) -> bool: ...
+
+
+class ThreadingEvent:
+    """
+    An Event implementation using threading.Event.
+
+    Provides a thread-safe event object for signaling between threads.
+    """
+
+    __slots__ = ("_evt",)
+
+    def __init__(self) -> None:
+        self._evt = threading.Event()
+
+    def set(self) -> None:
+        self._evt.set()
+
+    def clear(self) -> None:
+        self._evt.clear()
+
+    def is_set(self) -> bool:
+        return self._evt.is_set()
+
+    def wait(self, timeout: Optional[float] = None) -> bool:
+        return self._evt.wait(timeout)
+
+
+class MultiprocessingEvent:
+    """
+    An Event implementation using multiprocessing.Event.
+
+    Provides a process-safe event object for signaling between processes.
+    Optionally accepts a multiprocessing context for custom process start methods.
+    """
+
+    __slots__ = ("_evt",)
+
+    def __init__(self, *, ctx: Optional[BaseContext] = None) -> None:
+        self._evt = (ctx or mp).Event()
+
+    def set(self) -> None:
+        self._evt.set()
+
+    def clear(self) -> None:
+        self._evt.clear()
+
+    def is_set(self) -> bool:
+        return self._evt.is_set()
+
+    def wait(self, timeout: Optional[float] = None) -> bool:
+        return self._evt.wait(timeout)

agentlightning/execution/inter_process.py

@@ -0,0 +1,10 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .base import ExecutionStrategy
+
+
+class InterProcessExecutionStrategy(ExecutionStrategy):
+
+    alias: str = "ipc"
+
+    # TODO: to be implemented

agentlightning/execution/shared_memory.py

@@ -0,0 +1,264 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import logging
+import threading
+from contextlib import suppress
+from queue import SimpleQueue
+from typing import Any, Awaitable, Callable, List, Literal, Optional, Tuple
+
+from agentlightning.store.base import LightningStore
+from agentlightning.store.threading import LightningStoreThreaded
+
+from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle
+from .events import Event, ThreadingEvent
+
+logger = logging.getLogger(__name__)
+
+
+class SharedMemoryExecutionStrategy(ExecutionStrategy):
+    """Run algorithm and runners in a single process with threads sharing memory.
+
+    Termination & abort model:
+
+    - One shared ThreadingEvent (`stop_evt`) is passed to *all* bundles.
+    - The main thread (only) receives KeyboardInterrupt on Ctrl+C; we set `stop_evt` there.
+    - If any bundle raises, we set `stop_evt` from that thread to stop the rest.
+    - After the main-thread bundle finishes normally:
+      - If main_thread is "algorithm", we also set `stop_evt` to stop the runners.
+      - If main_thread is "runner", we do not set `stop_evt` to stop the algorithm.
+        We instead wait for the algorithm to finish naturally.
+    - Background threads are daemons; we join briefly and log any stragglers.
+
+    Notes: Signals other than SIGINT (e.g., SIGTERM) are not intercepted; we respect
+    Python's default behavior for them.
+    """
+
+    alias: str = "shm"
+
+    def __init__(
+        self,
+        n_runners: int = 1,
+        main_thread: Literal["algorithm", "runner"] = "runner",
+        join_timeout: float = 15.0,
+        graceful_delay: float = 5.0,
+        poll_interval: float = 0.05,
+    ) -> None:
+        if main_thread not in ("algorithm", "runner"):
+            raise ValueError("main_thread must be 'algorithm' or 'runner'")
+        if main_thread == "runner" and n_runners != 1:
+            raise ValueError("When main_thread is 'runner', n_runners must be 1")
+        self.n_runners = n_runners
+        self.main_thread = main_thread
+        self.join_timeout = join_timeout
+        self.graceful_delay = graceful_delay
+        self.poll_interval = poll_interval
+
+    async def _run_until_completed_or_canceled(self, coro: Awaitable[Any], stop_evt: Event) -> Any:
+        """Run `coro` until it finishes or a cooperative stop is requested.
+
+        Control flow:
+          1) Start the bundle coroutine as `task`.
+          2) Start a watcher task that waits for `stop_evt` *without blocking* the loop
+             by periodically polling the threading event.
+          3) When the stop event flips:
+               a) Give the bundle *graceful_delay* seconds to finish on its own,
+                  because well-behaved bundles will check the event and return.
+               b) If still running after the grace period, cancel the bundle task.
+          4) Ensure both tasks are awaited; swallow `CancelledError` where appropriate.
+
+        This is a *backup* mechanism for bundles that might not poll the event
+        frequently; cooperative shutdown (checking `stop_evt` yourself) is still preferred.
+        """
+        task: asyncio.Task[Any] = asyncio.create_task(coro)  # type: ignore
+        task_exception: Optional[BaseException] = None
+
+        async def watcher() -> None:
+            # Poll the threading event without blocking the event loop. Using a
+            # background thread via ``asyncio.to_thread`` makes cancellation
+            # difficult because ``ThreadingEvent.wait`` is not interruptible.
+            # Instead we cooperatively check the flag from the loop so the
+            # watcher task stays cancellable and tests don't hang when the
+            # bundle finishes naturally before the stop event is set.
+            while not stop_evt.is_set():
+                await asyncio.sleep(self.poll_interval)
+
+            # Grace period: let a cooperative bundle exit on its own.
+            try:
+                # At this point of waiting, the main task should already see the stop event.
+                await asyncio.wait_for(asyncio.shield(task), timeout=self.graceful_delay)  # type: ignore
+                logger.debug("Bundle finished by itself during grace period.")
+                return  # bundle finished by itself during grace period
+            except asyncio.TimeoutError:
+                # Still running after the grace window.
+                pass
+            except asyncio.CancelledError:
+                # If someone else canceled the task already, we're done.
+                logger.debug("Bundle already canceled by someone else; exiting watcher.")
+                return
+
+            # Still running after the grace window: cancel it.
+            if not task.done():
+                logger.debug("Graceful delay elapsed; canceling bundle task...")
+                task.cancel()
+
+        watcher_task = asyncio.create_task(watcher())
+        result: Any = None
+
+        try:
+            # We don't wait on FIRST_COMPLETED here, because we want the watcher
+            # to be able to grant a grace window after stop_evt flips.
+            await asyncio.wait(
+                {task, watcher_task}, return_when=asyncio.FIRST_COMPLETED
+            )  # pyright: ignore[reportUnknownArgumentType]
+        finally:
+            # If the main task hasn't completed yet (e.g., watcher scheduled cancel),
+            # finish the cancellation handshake.
+            if not task.done():
+                try:
+                    await asyncio.wait_for(task, timeout=self.graceful_delay)  # second chance
+                except asyncio.TimeoutError:
+                    logger.error(
+                        "Bundle task did not stop after cancellation; abandoning task."
+                        "This thread could live until the process exits."
+                    )
+                    # We return without awaiting it. asyncio.run will still try to cancel
+                    # pending tasks on loop close; if the task ignores cancellation, this
+                    # thread may still stick. It's the best we can do in Python.
+                    # We don't raise an exception here, but the thread could be a zombie.
+                    return result
+            else:
+                # Task completed naturally; retrieve result.
+                try:
+                    result = await task  # type: ignore
+                except asyncio.CancelledError:
+                    pass
+                except BaseException as exc:
+                    task_exception = exc
+
+            watcher_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await watcher_task
+
+        if task_exception is not None:
+            raise task_exception
+
+        return result  # type: ignore
+
+    def _run_algorithm(
+        self,
+        algorithm: AlgorithmBundle,
+        store: LightningStore,
+        stop_evt: Event,
+        thread_exceptions: Optional[SimpleQueue[BaseException]],
+    ) -> None:
+        try:
+            asyncio.run(self._run_until_completed_or_canceled(algorithm(store, stop_evt), stop_evt))
+        except asyncio.CancelledError:
+            logger.info("Algorithm bundle canceled due to stop signal.")
+        except BaseException as exc:
+            logger.exception("Algorithm bundle crashed; signaling stop to others.")
+            if thread_exceptions is not None:
+                thread_exceptions.put(exc)
+            stop_evt.set()
+            raise
+
+    def _run_runner(
+        self,
+        runner: RunnerBundle,
+        store: LightningStore,
+        worker_id: int,
+        stop_evt: Event,
+        thread_exceptions: Optional[SimpleQueue[BaseException]],
+    ) -> None:
+        try:
+            asyncio.run(self._run_until_completed_or_canceled(runner(store, worker_id, stop_evt), stop_evt))
+        except asyncio.CancelledError:
+            logger.info("Runner bundle (worker_id=%s) canceled due to stop signal.", worker_id)
+        except BaseException as exc:
+            logger.exception("Runner bundle crashed (worker_id=%s); signaling stop to others.", worker_id)
+            if thread_exceptions is not None:
+                thread_exceptions.put(exc)
+            stop_evt.set()
+            raise
+
+    def execute(self, algorithm: AlgorithmBundle, runner: RunnerBundle, store: LightningStore) -> None:
+        logger.info(
+            "Starting shm execution with %d runner(s); main thread runs '%s'",
+            self.n_runners,
+            self.main_thread,
+        )
+
+        # Create stop event and thread-safe store.
+        stop_evt = ThreadingEvent()
+        thread_safe_store = LightningStoreThreaded(store)
+
+        thread_exceptions: SimpleQueue[BaseException] = SimpleQueue()
+        raised_from_thread: Optional[BaseException] = None
+
+        def make_thread(name: str, target: Callable[..., Any], args: Tuple[Any, ...]) -> threading.Thread:
+            t = threading.Thread(name=name, target=target, args=args, daemon=True)
+            t.start()
+            return t
+
+        threads: List[threading.Thread] = []
+
+        try:
+            if self.main_thread == "algorithm":
+                # Start runner threads; algorithm runs on main thread.
+                for i in range(self.n_runners):
+                    thread = make_thread(
+                        name=f"runner-{i}",
+                        target=self._run_runner,
+                        args=(runner, thread_safe_store, i, stop_evt, thread_exceptions),
+                    )
+                    threads.append(thread)
+
+                # Ctrl+C here raises KeyboardInterrupt on this stack.
+                # Main thread doesn't need to collect exceptions.
+                self._run_algorithm(algorithm, thread_safe_store, stop_evt, None)
+
+                # If algo finishes naturally, request runners to stop.
+                stop_evt.set()
+
+            else:  # main_thread == "runner"
+                # Start algorithm in background; runner runs on main thread.
+                thread = make_thread(
+                    name="algorithm",
+                    target=self._run_algorithm,
+                    args=(algorithm, thread_safe_store, stop_evt, thread_exceptions),
+                )
+                threads.append(thread)
+
+                # Ctrl+C here raises KeyboardInterrupt on this stack.
+                # Main thread doesn't need to collect exceptions.
+                self._run_runner(runner, thread_safe_store, 0, stop_evt, None)
+
+                # If runner finishes naturally, WAIT FOR ALGORITHM TO FINISH.
+                thread.join()
+
+            if not thread_exceptions.empty():
+                raised_from_thread = thread_exceptions.get()
+
+        except KeyboardInterrupt:
+            logger.warning("KeyboardInterrupt received on main thread; initiating cooperative shutdown...")
+            stop_evt.set()
+        finally:
+            # Attempt a clean join; if some threads don't comply, log and move on.
+            for t in threads:
+                logger.debug("Joining thread %s...", t.name)
+                t.join(timeout=self.join_timeout)
+
+            alive = [t.name for t in threads if t.is_alive()]
+            if alive:
+                logger.error(
+                    "Threads still alive after %.1fs: %s. They are daemons; continuing shutdown.",
+                    self.join_timeout,
+                    ", ".join(alive),
+                )
+
+            if raised_from_thread is None and not thread_exceptions.empty():
+                raised_from_thread = thread_exceptions.get()
+
+        if raised_from_thread is not None:
+            raise raised_from_thread

agentlightning/store/client_server.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import asyncio
 import logging
+import os
 from typing import Any, Dict, List, Literal, Optional, Sequence, Union
 
 import aiohttp
@@ -87,31 +88,96 @@ class LightningStoreServer(LightningStore):
         self.store = store
         self.host = host
         self.port = port
-        self.app = FastAPI(title="LightningStore Server")
+        self.app: FastAPI | None = FastAPI(title="LightningStore Server")
         self._setup_routes()
-        self._uvicorn_config = uvicorn.Config(self.app, host=self.host, port=self.port, log_level="info")
-        self._uvicorn_server = uvicorn.Server(self._uvicorn_config)
+        self._uvicorn_config: uvicorn.Config | None = uvicorn.Config(
+            self.app, host=self.host, port=self.port, log_level="info"
+        )
+        self._uvicorn_server: uvicorn.Server | None = uvicorn.Server(self._uvicorn_config)
+
+        # Process-awareness:
+        # LightningStoreServer holds a plain Python object (self.store) in one process
+        # (the process that runs uvicorn/FastAPI).
+        # When you multiprocessing.Process(...) and call methods on a different LightningStore instance
+        # (or on a copy inherited via fork), you’re mutating another process’s memory, not the server’s memory.
+        # So we need to track the owner process (whoever creates the server),
+        # and only mutate the store in that process.
+        self._owner_pid = os.getpid()
+        self._client: Optional[LightningStoreClient] = None
+
+    def __getstate__(self):
+        """
+        Control pickling to prevent server state from being sent to subprocesses.
+
+        When LightningStoreServer is pickled (e.g., passed to a subprocess), we only
+        serialize the underlying store and connection details. The FastAPI app and
+        uvicorn server are excluded as they should not be transferred between processes.
+
+        The subprocess should create its own server instance if needed.
+        """
+        return {
+            "store": self.store,
+            "host": self.host,
+            "port": self.port,
+            "_owner_pid": self._owner_pid,
+        }
+
+    def __setstate__(self, state: Dict[str, Any]):
+        """
+        Restore from pickle by reconstructing only the essential attributes.
+
+        Note: This creates a new server instance without FastAPI/uvicorn initialized.
+        Call __init__() pattern or create a new LightningStoreServer if you need
+        a fully functional server in the subprocess.
+        """
+        self.store = state["store"]
+        self.host = state["host"]
+        self.port = state["port"]
+        self._owner_pid = state["_owner_pid"]
+        # Do NOT reconstruct app, _uvicorn_config, _uvicorn_server
+        # to avoid transferring server state to subprocess
 
     @property
     def endpoint(self) -> str:
         return f"http://{self.host}:{self.port}"
 
     async def start(self):
-        """Starts the FastAPI server in the background."""
+        """Starts the FastAPI server in the background.
+
+        You need to call this method in the same process as the server was created in.
+        """
+        assert self._uvicorn_server is not None
         logger.info(f"Starting server at {self.endpoint}")
         asyncio.create_task(self._uvicorn_server.serve())
         await asyncio.sleep(1)  # Allow time for server to start up.
 
     async def stop(self):
-        """Gracefully stops the running FastAPI server."""
+        """Gracefully stops the running FastAPI server.
+
+        You need to call this method in the same process as the server was created in.
+        """
+        assert self._uvicorn_server is not None
         if self._uvicorn_server.started:
             logger.info("Stopping server...")
             self._uvicorn_server.should_exit = True
             await asyncio.sleep(1)  # Allow time for graceful shutdown.
             logger.info("Server stopped.")
 
+    def _backend(self) -> LightningStore:
+        """Returns the object to delegate to in *this* process.
+
+        - In the owner process: delegate to the in-process store.
+        - In a different process: delegate to a HTTP client talking to the server.
+        """
+        if os.getpid() == self._owner_pid:
+            return self.store
+        if self._client is None:
+            self._client = LightningStoreClient(self.endpoint)
+        return self._client
+
     def _setup_routes(self):
         """Set up FastAPI routes for all store operations."""
+        assert self.app is not None
 
         @self.app.post("/start_rollout", response_model=AttemptedRollout)
         async def start_rollout(request: RolloutRequest):  # pyright: ignore[reportUnusedFunction]
@@ -210,7 +276,7 @@ class LightningStoreServer(LightningStore):
                 metadata=request.metadata if not isinstance(request.metadata, PydanticUnset) else UNSET,
             )
 
-    # Delegate methods -------------------------------------------------
+    # Delegate methods
     async def start_rollout(
         self,
         input: TaskInput,
@@ -218,7 +284,7 @@ class LightningStoreServer(LightningStore):
         resources_id: str | None = None,
         metadata: Dict[str, Any] | None = None,
     ) -> AttemptedRollout:
-        return await self.store.start_rollout(input, mode, resources_id, metadata)
+        return await self._backend().start_rollout(input, mode, resources_id, metadata)
 
     async def enqueue_rollout(
         self,
@@ -227,42 +293,42 @@ class LightningStoreServer(LightningStore):
         resources_id: str | None = None,
         metadata: Dict[str, Any] | None = None,
     ) -> RolloutV2:
-        return await self.store.enqueue_rollout(input, mode, resources_id, metadata)
+        return await self._backend().enqueue_rollout(input, mode, resources_id, metadata)
 
     async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
-        return await self.store.dequeue_rollout()
+        return await self._backend().dequeue_rollout()
 
     async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
-        return await self.store.start_attempt(rollout_id)
+        return await self._backend().start_attempt(rollout_id)
 
     async def query_rollouts(
         self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
     ) -> List[RolloutV2]:
-        return await self.store.query_rollouts(status=status, rollout_ids=rollout_ids)
+        return await self._backend().query_rollouts(status=status, rollout_ids=rollout_ids)
 
     async def query_attempts(self, rollout_id: str) -> List[Attempt]:
-        return await self.store.query_attempts(rollout_id)
+        return await self._backend().query_attempts(rollout_id)
 
     async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
-        return await self.store.get_latest_attempt(rollout_id)
+        return await self._backend().get_latest_attempt(rollout_id)
 
     async def get_rollout_by_id(self, rollout_id: str) -> Optional[RolloutV2]:
-        return await self.store.get_rollout_by_id(rollout_id)
+        return await self._backend().get_rollout_by_id(rollout_id)
 
     async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
-        return await self.store.update_resources(resources_id, resources)
+        return await self._backend().update_resources(resources_id, resources)
 
     async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
-        return await self.store.get_resources_by_id(resources_id)
+        return await self._backend().get_resources_by_id(resources_id)
 
     async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
-        return await self.store.get_latest_resources()
+        return await self._backend().get_latest_resources()
 
     async def add_span(self, span: Span) -> Span:
-        return await self.store.add_span(span)
+        return await self._backend().add_span(span)
 
     async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
-        return await self.store.get_next_span_sequence_id(rollout_id, attempt_id)
+        return await self._backend().get_next_span_sequence_id(rollout_id, attempt_id)
 
     async def add_otel_span(
         self,
@@ -271,17 +337,17 @@ class LightningStoreServer(LightningStore):
         readable_span: ReadableSpan,
         sequence_id: int | None = None,
     ) -> Span:
-        return await self.store.add_otel_span(rollout_id, attempt_id, readable_span, sequence_id)
+        return await self._backend().add_otel_span(rollout_id, attempt_id, readable_span, sequence_id)
 
     async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[RolloutV2]:
-        return await self.store.wait_for_rollouts(rollout_ids=rollout_ids, timeout=timeout)
+        return await self._backend().wait_for_rollouts(rollout_ids=rollout_ids, timeout=timeout)
 
     async def query_spans(
         self,
         rollout_id: str,
         attempt_id: str | Literal["latest"] | None = None,
     ) -> List[Span]:
-        return await self.store.query_spans(rollout_id, attempt_id)
+        return await self._backend().query_spans(rollout_id, attempt_id)
 
     async def update_rollout(
         self,
@@ -293,7 +359,7 @@ class LightningStoreServer(LightningStore):
         config: RolloutConfig | Unset = UNSET,
         metadata: Optional[Dict[str, Any]] | Unset = UNSET,
     ) -> RolloutV2:
-        return await self.store.update_rollout(
+        return await self._backend().update_rollout(
             rollout_id=rollout_id,
             input=input,
             mode=mode,
@@ -312,7 +378,7 @@ class LightningStoreServer(LightningStore):
         last_heartbeat_time: float | Unset = UNSET,
         metadata: Optional[Dict[str, Any]] | Unset = UNSET,
     ) -> Attempt:
-        return await self.store.update_attempt(
+        return await self._backend().update_attempt(
             rollout_id=rollout_id,
             attempt_id=attempt_id,
             status=status,

tests/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Microsoft. All rights reserved.

tests/execution/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Microsoft. All rights reserved.

tests/execution/test_shared_memory.py

@@ -0,0 +1,357 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import logging
+import time
+from typing import Any, List, Optional
+
+import pytest
+from _pytest.logging import LogCaptureFixture
+
+from agentlightning.execution.events import Event, ThreadingEvent
+from agentlightning.execution.shared_memory import SharedMemoryExecutionStrategy
+from agentlightning.store.base import LightningStore
+
+from ..store.dummy_store import DummyLightningStore, minimal_dummy_store
+
+
+@pytest.fixture
+def store() -> DummyLightningStore:
+    return minimal_dummy_store()
+
+
+def tiny_sleep(seconds: float) -> float:
+    """Sleep a tiny bit and return the elapsed time (monotonic)."""
+    start = time.monotonic()
+    time.sleep(seconds)
+    return time.monotonic() - start
+
+
+# Helper bundles for tests
+
+
+def make_cooperative_algorithm(
+    started: List[str],
+    finished: List[str],
+    poll_delay: float = 0.005,
+):
+    async def algo(store: LightningStore, event: Event) -> None:
+        started.append("algo")
+        # cooperatively exit when asked
+        while not event.is_set():
+            await asyncio.sleep(poll_delay)
+        finished.append("algo")
+
+    return algo
+
+
+def make_cooperative_runner(
+    started: List[int],
+    finished: List[int],
+    poll_delay: float = 0.005,
+):
+    async def runner(store: LightningStore, worker_id: int, event: Event) -> None:
+        started.append(worker_id)
+        while not event.is_set():
+            await asyncio.sleep(poll_delay)
+        finished.append(worker_id)
+
+    return runner
+
+
+def make_sleepy_coro(delay: float, result: Any):
+    async def _coro():
+        await asyncio.sleep(delay)
+        return result
+
+    return _coro()
+
+
+def make_never_polls_coro(total_time: float):
+    """A coro that just sleeps in chunks; cancellation lands at await points only."""
+
+    async def _coro():
+        end = time.monotonic() + total_time
+        while time.monotonic() < end:
+            await asyncio.sleep(0.05)
+        return "done"
+
+    return _coro()
+
+
+def make_slow_cleanup_on_cancel_coro(cleanup_time: float):
+    """On cancellation, do slow cleanup before exiting (tests second-chance timeout)."""
+
+    async def _coro():
+        try:
+            # Idle until canceled
+            await asyncio.Event().wait()
+        except asyncio.CancelledError:
+            await asyncio.sleep(cleanup_time)
+            raise
+
+    return _coro()
+
+
+# Unit tests: _run_until_completed_or_canceled
+
+
+def test_run_until_completes_naturally(caplog: LogCaptureFixture):
+    caplog.set_level(logging.DEBUG)
+    strat = SharedMemoryExecutionStrategy(n_runners=1, main_thread="runner", graceful_delay=0.05, join_timeout=0.2)
+
+    evt = ThreadingEvent()
+
+    result = asyncio.run(
+        strat._run_until_completed_or_canceled(make_sleepy_coro(0.02, "ok"), evt)  # pyright: ignore[reportPrivateUsage]
+    )
+    assert result == "ok"
+
+
+def test_run_until_stops_gracefully_with_event(caplog: LogCaptureFixture):
+    caplog.set_level(logging.DEBUG)
+    strat = SharedMemoryExecutionStrategy(n_runners=1, main_thread="runner", graceful_delay=0.05, join_timeout=0.2)
+
+    evt = ThreadingEvent()
+
+    async def cooperative(event: Event):
+        # Stop promptly after event flips
+        while not event.is_set():
+            await asyncio.sleep(0.005)
+        return "bye"
+
+    # Set the stop event before we start to exercise the "finish during grace" path
+    evt.set()
+    result = asyncio.run(
+        strat._run_until_completed_or_canceled(cooperative(evt), evt)  # pyright: ignore[reportPrivateUsage]
+    )
+    assert result == "bye"
+    # Should have logged that the bundle finished during grace period
+    assert any("finished by itself during grace period" in rec.message for rec in caplog.records)
+
+
+def test_run_until_cancels_after_grace_if_not_quick_to_stop(caplog: LogCaptureFixture):
+    caplog.set_level(logging.DEBUG)
+    # Grace: 30ms; the task won't stop within that, so cancel fires; then second-chance wait succeeds.
+    graceful = 0.03
+    strat = SharedMemoryExecutionStrategy(n_runners=1, main_thread="runner", graceful_delay=graceful, join_timeout=0.2)
+
+    evt = ThreadingEvent()
+
+    async def slow_to_notice():
+        # Doesn't poll stop_evt; cancellation must interrupt sleeps
+        await asyncio.sleep(1.0)
+
+    async def runner():
+        # flip stop after a short time
+        asyncio.get_running_loop().call_later(0.01, evt.set)
+        return await strat._run_until_completed_or_canceled(  # pyright: ignore[reportPrivateUsage]
+            slow_to_notice(), evt
+        )
+
+    t0 = time.monotonic()
+    result = asyncio.run(runner())
+    elapsed = time.monotonic() - t0
+    # We don't get a "result" because the task was canceled
+    assert result is None
+    # Should be roughly >= grace window, but way less than full 1.0s sleep
+    assert elapsed >= graceful * 0.9
+    assert elapsed < 0.5
+    assert any("Graceful delay elapsed; canceling bundle task..." in rec.message for rec in caplog.records)
+
+
+def test_run_until_second_chance_timeout_logs_and_returns(caplog: LogCaptureFixture):
+    caplog.set_level(logging.DEBUG)
+    # The task takes longer to clean up than graceful_delay; we hit the second-chance timeout branch.
+    graceful = 0.03
+    cleanup = 0.08
+    strat = SharedMemoryExecutionStrategy(n_runners=1, main_thread="runner", graceful_delay=graceful, join_timeout=0.2)
+
+    evt = ThreadingEvent()
+
+    async def runner():
+        # Schedule stop, then run a task that sleeps extra on cancel
+        asyncio.get_running_loop().call_later(0.005, evt.set)
+        return await strat._run_until_completed_or_canceled(  # pyright: ignore[reportPrivateUsage]
+            make_slow_cleanup_on_cancel_coro(cleanup), evt
+        )
+
+    t0 = time.monotonic()
+    result = asyncio.run(runner())
+    elapsed = time.monotonic() - t0
+
+    # No result expected because the task gets canceled; we abandon await after timeout.
+    assert result is None
+    assert elapsed >= graceful  # we at least burned the grace + a bit
+    assert any("did not stop after cancellation; abandoning task" in rec.message for rec in caplog.records)
+
+
+# Unit tests: _run_algorithm and _run_runner
+
+
+def test_run_algorithm_sets_stop_on_exception(store: DummyLightningStore):
+    strat = SharedMemoryExecutionStrategy()
+
+    evt = ThreadingEvent()
+
+    async def boom(store: LightningStore, event: Event) -> None:
+        await asyncio.sleep(0.005)
+        raise ValueError("algo crash")
+
+    with pytest.raises(ValueError):
+        strat._run_algorithm(boom, store, evt, None)  # pyright: ignore[reportPrivateUsage]
+
+    assert evt.is_set(), "stop_evt must be set when algorithm raises"
+
+
+def test_run_runner_sets_stop_on_exception(store: DummyLightningStore):
+    strat = SharedMemoryExecutionStrategy()
+
+    evt = ThreadingEvent()
+
+    async def boom(store: LightningStore, worker_id: int, event: Event) -> None:
+        await asyncio.sleep(0.005)
+        raise RuntimeError("runner crash")
+
+    with pytest.raises(RuntimeError):
+        strat._run_runner(boom, store, 0, evt, None)  # pyright: ignore[reportPrivateUsage]
+
+    assert evt.is_set(), "stop_evt must be set when a runner raises"
+
+
+# Integration tests: execute(...)
+
+
+def test_execute_main_algorithm_normal_stop_sets_event(store: DummyLightningStore):
+    started_r: List[int] = []
+    finished_r: List[int] = []
+    started_a: List[str] = []
+    finished_a: List[str] = []
+
+    strat = SharedMemoryExecutionStrategy(n_runners=2, main_thread="algorithm", graceful_delay=0.02, join_timeout=0.2)
+
+    runner = make_cooperative_runner(started_r, finished_r, poll_delay=0.005)
+
+    # Algorithm finishes quickly; then execute() should set stop_evt for runners.
+    async def algo(store: LightningStore, event: Event) -> None:
+        started_a.append("algo")
+        await asyncio.sleep(0.02)
+        finished_a.append("algo")
+
+    strat.execute(algo, runner, store)
+
+    # Runners should have started and finished due to stop_evt set by execute()
+    assert sorted(started_r) == [0, 1]
+    assert sorted(finished_r) == [0, 1]
+    assert started_a == ["algo"]
+    assert finished_a == ["algo"]
+
+
+def test_execute_main_runner_waits_for_algorithm_natural_finish(store: DummyLightningStore):
+    # Policy: when main_thread='runner' and runner finishes, do NOT set stop; wait for algo to finish.
+    started_r: List[int] = []
+    finished_r: List[int] = []
+    started_a: List[str] = []
+    finished_a: List[str] = []
+    captured_evt: List[Optional[Any]] = [None]
+
+    strat = SharedMemoryExecutionStrategy(n_runners=1, main_thread="runner", graceful_delay=0.02, join_timeout=0.2)
+
+    async def algo(store: LightningStore, event: Event) -> None:
+        captured_evt[0] = event
+        started_a.append("algo")
+        await asyncio.sleep(0.05)  # natural finish
+        finished_a.append("algo")
+
+    async def runner(store: LightningStore, worker_id: int, event: Event) -> None:
+        started_r.append(worker_id)
+        await asyncio.sleep(0.01)  # finish quickly
+        finished_r.append(worker_id)
+
+    strat.execute(algo, runner, store)
+
+    assert started_r == [0]
+    assert finished_r == [0]
+    assert started_a == ["algo"]
+    assert finished_a == ["algo"]
+    # Verify execute() did not set stop_evt just because runner ended.
+    assert captured_evt[0] is not None and captured_evt[0].is_set() is False
+
+
+def test_execute_runner_crash_propagates_and_stops_algorithm(store: DummyLightningStore):
+    # Runner raises; algorithm should see stop_evt and exit.
+    started_a: List[str] = []
+    finished_a: List[str] = []
+    saw_stop: List[bool] = []
+
+    strat = SharedMemoryExecutionStrategy(n_runners=1, main_thread="algorithm", graceful_delay=0.02, join_timeout=0.3)
+
+    async def algo(store: LightningStore, event: Event) -> None:
+        started_a.append("algo")
+        # wait until stop evt set (by runner crash)
+        while not event.is_set():
+            await asyncio.sleep(0.005)
+        saw_stop.append(True)
+        finished_a.append("algo")
+
+    async def bad_runner(store: LightningStore, worker_id: int, event: Event) -> None:
+        await asyncio.sleep(0.02)
+        raise RuntimeError("boom")
+
+    with pytest.raises(RuntimeError):
+        strat.execute(algo, bad_runner, store)
+
+    assert started_a == ["algo"]
+    assert saw_stop == [True]
+    assert finished_a == ["algo"]
+
+
+def test_execute_ctrl_c_on_algorithm_stops_runners(store: DummyLightningStore, caplog: LogCaptureFixture):
+    caplog.set_level(logging.DEBUG)
+    # Simulate Ctrl+C by having the algorithm raise KeyboardInterrupt on the main thread.
+    runner_started: List[int] = []
+    runner_finished: List[int] = []
+
+    strat = SharedMemoryExecutionStrategy(n_runners=2, main_thread="algorithm", graceful_delay=0.02, join_timeout=0.3)
+
+    async def algo_kbi(store: LightningStore, event: Event) -> None:
+        await asyncio.sleep(0.01)
+        raise KeyboardInterrupt()
+
+    async def runner(store: LightningStore, worker_id: int, event: Event) -> None:
+        runner_started.append(worker_id)
+        while not event.is_set():
+            await asyncio.sleep(0.005)
+        runner_finished.append(worker_id)
+
+    strat.execute(algo_kbi, runner, store)
+
+    # Runners should have started and then exited due to stop signal
+    assert sorted(runner_started) == [0, 1]
+    assert sorted(runner_finished) == [0, 1]
+    assert any("KeyboardInterrupt received on main thread" in rec.message for rec in caplog.records)
+
+
+def test_execute_ctrl_c_on_runner_stops_algorithm(store: DummyLightningStore, caplog: LogCaptureFixture):
+    caplog.set_level(logging.DEBUG)
+    # Simulate Ctrl+C on the main thread when main_thread='runner'
+    algo_started: List[str] = []
+    algo_finished: List[str] = []
+
+    strat = SharedMemoryExecutionStrategy(n_runners=1, main_thread="runner", graceful_delay=0.02, join_timeout=0.3)
+
+    async def algo(store: LightningStore, event: Event) -> None:
+        algo_started.append("a")
+        # Await stop_evt after KBI from runner
+        while not event.is_set():
+            await asyncio.sleep(0.005)
+        algo_finished.append("a")
+
+    async def runner_kbi(store: LightningStore, worker_id: int, event: Event) -> None:
+        await asyncio.sleep(0.01)
+        raise KeyboardInterrupt()
+
+    strat.execute(algo, runner_kbi, store)
+
+    assert algo_started == ["a"]
+    assert algo_finished == ["a"]  # algorithm should finish after stop_evt set by KBI handler

tests/store/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Microsoft. All rights reserved.

tests/store/conftest.py

@@ -9,13 +9,13 @@ from opentelemetry.sdk.trace import ReadableSpan
 from agentlightning.store.memory import InMemoryLightningStore
 
 __all__ = [
-    "store",
+    "inmemory_store",
     "mock_readable_span",
 ]
 
 
 @pytest.fixture
-def store() -> InMemoryLightningStore:
+def inmemory_store() -> InMemoryLightningStore:
     """Create a fresh InMemoryLightningStore instance."""
     return InMemoryLightningStore()
 

tests/store/dummy_store.py

@@ -0,0 +1,176 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import Any, Dict, List, Literal, Optional, Sequence
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.store.base import UNSET, LightningStore
+from agentlightning.tracer import Span
+from agentlightning.types import (
+    Attempt,
+    AttemptedRollout,
+    AttemptStatus,
+    NamedResources,
+    ResourcesUpdate,
+    RolloutStatus,
+    RolloutV2,
+    TaskInput,
+)
+
+
+class DummyLightningStore(LightningStore):
+    def __init__(self, return_values: Dict[str, Any]) -> None:
+        super().__init__()
+        self.calls: List[tuple[str, tuple[Any, ...], Dict[str, Any]]] = []
+        self.return_values = return_values
+
+    async def start_rollout(
+        self,
+        input: TaskInput,
+        mode: Optional[str] = None,
+        resources_id: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> AttemptedRollout:
+        self.calls.append(("start_rollout", (input, mode, resources_id, metadata), {}))
+        return self.return_values["start_rollout"]
+
+    async def enqueue_rollout(
+        self,
+        input: TaskInput,
+        mode: Optional[str] = None,
+        resources_id: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> RolloutV2:
+        self.calls.append(("enqueue_rollout", (input, mode, resources_id, metadata), {}))
+        return self.return_values["enqueue_rollout"]
+
+    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+        self.calls.append(("dequeue_rollout", (), {}))
+        return self.return_values["dequeue_rollout"]
+
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+        self.calls.append(("start_attempt", (rollout_id,), {}))
+        return self.return_values["start_attempt"]
+
+    async def query_rollouts(
+        self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
+    ) -> List[RolloutV2]:
+        self.calls.append(("query_rollouts", (), {"status": status, "rollout_ids": rollout_ids}))
+        return self.return_values["query_rollouts"]
+
+    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
+        self.calls.append(("query_attempts", (rollout_id,), {}))
+        return self.return_values["query_attempts"]
+
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[RolloutV2]:
+        self.calls.append(("get_rollout_by_id", (rollout_id,), {}))
+        return self.return_values["get_rollout_by_id"]
+
+    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
+        self.calls.append(("get_latest_attempt", (rollout_id,), {}))
+        return self.return_values["get_latest_attempt"]
+
+    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
+        self.calls.append(("update_resources", (resources_id, resources), {}))
+        return self.return_values["update_resources"]
+
+    async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
+        self.calls.append(("get_resources_by_id", (resources_id,), {}))
+        return self.return_values["get_resources_by_id"]
+
+    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
+        self.calls.append(("get_latest_resources", (), {}))
+        return self.return_values["get_latest_resources"]
+
+    async def add_span(self, span: Span) -> Span:
+        self.calls.append(("add_span", (span,), {}))
+        return self.return_values["add_span"]
+
+    async def add_otel_span(
+        self,
+        rollout_id: str,
+        attempt_id: str,
+        readable_span: ReadableSpan,
+        sequence_id: Optional[int] = None,
+    ) -> Span:
+        self.calls.append(("add_otel_span", (rollout_id, attempt_id, readable_span, sequence_id), {}))
+        return self.return_values["add_otel_span"]
+
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[RolloutV2]:
+        self.calls.append(("wait_for_rollouts", (), {"rollout_ids": rollout_ids, "timeout": timeout}))
+        return self.return_values["wait_for_rollouts"]
+
+    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
+        self.calls.append(("get_next_span_sequence_id", (rollout_id, attempt_id), {}))
+        return self.return_values["get_next_span_sequence_id"]
+
+    async def query_spans(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"] | None = None,
+    ) -> List[Span]:
+        self.calls.append(("query_spans", (rollout_id, attempt_id), {}))
+        return self.return_values["query_spans"]
+
+    async def update_rollout(
+        self,
+        rollout_id: str,
+        input: TaskInput | Any = UNSET,
+        mode: Optional[str] | Any = UNSET,
+        resources_id: Optional[str] | Any = UNSET,
+        status: RolloutStatus | Any = UNSET,
+        config: Any = UNSET,
+        metadata: Optional[Dict[str, Any]] | Any = UNSET,
+    ) -> RolloutV2:
+        self.calls.append(
+            (
+                "update_rollout",
+                (rollout_id, input, mode, resources_id, status, config, metadata),
+                {},
+            )
+        )
+        return self.return_values["update_rollout"]
+
+    async def update_attempt(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"],
+        status: AttemptStatus | Any = UNSET,
+        worker_id: str | Any = UNSET,
+        last_heartbeat_time: float | Any = UNSET,
+        metadata: Optional[Dict[str, Any]] | Any = UNSET,
+    ) -> Attempt:
+        self.calls.append(
+            (
+                "update_attempt",
+                (rollout_id, attempt_id, status, worker_id, last_heartbeat_time, metadata),
+                {},
+            )
+        )
+        return self.return_values["update_attempt"]
+
+
+def minimal_dummy_store() -> DummyLightningStore:
+    # Provide minimal return values
+    return DummyLightningStore(
+        return_values={
+            "start_rollout": None,
+            "enqueue_rollout": None,
+            "dequeue_rollout": None,
+            "start_attempt": None,
+            "query_rollouts": [],
+            "query_attempts": [],
+            "get_rollout_by_id": None,
+            "get_latest_attempt": None,
+            "update_resources": None,
+            "get_resources_by_id": None,
+            "get_latest_resources": None,
+            "add_span": None,
+            "add_otel_span": None,
+            "wait_for_rollouts": [],
+            "get_next_span_sequence_id": 0,
+            "query_spans": [],
+            "update_rollout": None,
+            "update_attempt": None,
+        }
+    )

tests/store/test_client_server.py

@@ -2,7 +2,9 @@
 
 import asyncio
 import contextlib
+import multiprocessing
 import socket
+import sys
 from typing import AsyncGenerator, Tuple
 from unittest.mock import patch
 
@@ -241,3 +243,139 @@ async def test_concurrent_add_otel_span_sequence_ids_unique(
 
     stored_spans = await client.query_spans(rollout_id, attempt_id="latest")
     assert len(stored_spans) >= 2
+
+
+@pytest.mark.asyncio
+async def test_subprocess_operations_sync_via_http_automatically() -> None:
+    """
+    Test that LightningStoreServer automatically uses HTTP client in subprocesses.
+
+    When LightningStoreServer is passed to a subprocess, it detects it's in a different
+    process (via PID tracking) and automatically delegates to an HTTP client instead of
+    the local store. This ensures operations in the subprocess are reflected in the
+    main process via the HTTP server.
+    """
+    store = InMemoryLightningStore()
+    port = _get_free_port()
+    server = LightningStoreServer(store, "127.0.0.1", port)
+    await server.start()
+
+    try:
+        # Record initial state
+        initial_rollouts = await store.query_rollouts()
+        initial_count = len(initial_rollouts)
+
+        def subprocess_work(server_obj: LightningStoreServer) -> None:
+            """Subprocess that performs operations via the server object."""
+
+            async def do_work() -> None:
+                # The server detects we're in a different process and automatically
+                # uses HTTP client to communicate with the main process server
+                await server_obj.enqueue_rollout(input={"origin": "subprocess"})
+
+            asyncio.run(do_work())
+
+        # Spawn a subprocess to perform operations
+        ctx = multiprocessing.get_context()
+        process = ctx.Process(target=subprocess_work, args=(server,))
+        process.start()
+        await asyncio.to_thread(process.join, timeout=5.0)
+
+        assert process.exitcode == 0
+
+        # Allow time for HTTP request to complete
+        await asyncio.sleep(0.2)
+
+        # Subprocess operations ARE reflected in main process store
+        # because the server automatically used HTTP client in the subprocess
+        main_process_rollouts = await store.query_rollouts()
+        assert len(main_process_rollouts) == initial_count + 1, (
+            "Subprocess operations should be reflected in main process store " "via automatic HTTP client delegation"
+        )
+
+    finally:
+        await server.stop()
+
+
+@pytest.mark.asyncio
+async def test_subprocess_client_operations_work_but_direct_store_access_fails() -> None:
+    """
+    Demonstrate that:
+    1. Client operations via HTTP work correctly (data persists in main process)
+    2. Direct store access in subprocess does NOT work (data isolated to subprocess)
+    """
+    store = InMemoryLightningStore()
+    port = _get_free_port()
+    server = LightningStoreServer(store, "127.0.0.1", port)
+    await server.start()
+
+    try:
+        initial_rollouts = await store.query_rollouts()
+        initial_count = len(initial_rollouts)
+
+        def subprocess_client_work(endpoint: str) -> None:
+            """Subprocess using HTTP client - this WORKS."""
+
+            async def do_work() -> None:
+                client = LightningStoreClient(endpoint)
+                try:
+                    await client.enqueue_rollout(input={"origin": "subprocess-client"})
+                except Exception as e:
+                    print(f"Client subprocess error: {e}", file=sys.stderr)
+                    raise
+                finally:
+                    await client.close()
+
+            asyncio.run(do_work())
+
+        def subprocess_direct_store_work(server_obj: LightningStoreServer) -> None:
+            """Subprocess using direct store access - this does NOT work."""
+
+            async def do_work() -> None:
+                # This operates on the subprocess's copy of the store
+                await server_obj.enqueue_rollout(input={"origin": "subprocess-direct"})
+
+            asyncio.run(do_work())
+
+        # Test 1: Client operations via HTTP - should work
+        ctx = multiprocessing.get_context()
+        client_process = ctx.Process(target=subprocess_client_work, args=(server.endpoint,))
+        client_process.start()
+        await asyncio.to_thread(client_process.join, timeout=5.0)  # Add timeout
+
+        # Handle timeout case
+        if client_process.is_alive():
+            client_process.terminate()
+            client_process.join(timeout=1.0)
+            pytest.fail("Client subprocess hung and had to be terminated")
+
+        assert client_process.exitcode == 0, f"Client subprocess failed with exit code {client_process.exitcode}"
+
+        await asyncio.sleep(0.2)
+        after_client = await store.query_rollouts()
+        # Client operations WORK - the rollout is in the main process store
+        assert len(after_client) == initial_count + 1
+
+        # Test 2: Server object in subprocess - ALSO works now (auto-delegates to HTTP)
+        direct_process = ctx.Process(target=subprocess_direct_store_work, args=(server,))
+        direct_process.start()
+        await asyncio.to_thread(direct_process.join, timeout=5.0)
+
+        # Handle timeout case
+        if direct_process.is_alive():
+            direct_process.terminate()
+            direct_process.join(timeout=1.0)
+            pytest.fail("Server subprocess hung and had to be terminated")
+
+        assert direct_process.exitcode == 0, f"Server subprocess failed with exit code {direct_process.exitcode}"
+
+        await asyncio.sleep(0.2)
+        after_direct = await store.query_rollouts()
+        # With the fix: server object in subprocess ALSO works via auto HTTP delegation
+        # Both rollouts (client + server) should be in the store
+        assert (
+            len(after_direct) == initial_count + 2
+        ), "Both explicit client and server object operations should work via HTTP"
+
+    finally:
+        await server.stop()

tests/store/test_threading.py

@@ -2,7 +2,7 @@
 
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any, Dict, List, Literal, Optional, Sequence
+from typing import Any, Dict
 from unittest.mock import MagicMock
 
 import pytest
@@ -18,142 +18,11 @@ from agentlightning.types import (
     AttemptStatus,
     NamedResources,
     ResourcesUpdate,
-    RolloutStatus,
     RolloutV2,
     TaskInput,
 )
 
-
-class DummyLightningStore(LightningStore):
-    def __init__(self, return_values: Dict[str, Any]) -> None:
-        super().__init__()
-        self.calls: List[tuple[str, tuple[Any, ...], Dict[str, Any]]] = []
-        self.return_values = return_values
-
-    async def start_rollout(
-        self,
-        input: TaskInput,
-        mode: Optional[str] = None,
-        resources_id: Optional[str] = None,
-        metadata: Optional[Dict[str, Any]] = None,
-    ) -> AttemptedRollout:
-        self.calls.append(("start_rollout", (input, mode, resources_id, metadata), {}))
-        return self.return_values["start_rollout"]
-
-    async def enqueue_rollout(
-        self,
-        input: TaskInput,
-        mode: Optional[str] = None,
-        resources_id: Optional[str] = None,
-        metadata: Optional[Dict[str, Any]] = None,
-    ) -> RolloutV2:
-        self.calls.append(("enqueue_rollout", (input, mode, resources_id, metadata), {}))
-        return self.return_values["enqueue_rollout"]
-
-    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
-        self.calls.append(("dequeue_rollout", (), {}))
-        return self.return_values["dequeue_rollout"]
-
-    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
-        self.calls.append(("start_attempt", (rollout_id,), {}))
-        return self.return_values["start_attempt"]
-
-    async def query_rollouts(
-        self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
-    ) -> List[RolloutV2]:
-        self.calls.append(("query_rollouts", (), {"status": status, "rollout_ids": rollout_ids}))
-        return self.return_values["query_rollouts"]
-
-    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
-        self.calls.append(("query_attempts", (rollout_id,), {}))
-        return self.return_values["query_attempts"]
-
-    async def get_rollout_by_id(self, rollout_id: str) -> Optional[RolloutV2]:
-        self.calls.append(("get_rollout_by_id", (rollout_id,), {}))
-        return self.return_values["get_rollout_by_id"]
-
-    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
-        self.calls.append(("get_latest_attempt", (rollout_id,), {}))
-        return self.return_values["get_latest_attempt"]
-
-    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
-        self.calls.append(("update_resources", (resources_id, resources), {}))
-        return self.return_values["update_resources"]
-
-    async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
-        self.calls.append(("get_resources_by_id", (resources_id,), {}))
-        return self.return_values["get_resources_by_id"]
-
-    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
-        self.calls.append(("get_latest_resources", (), {}))
-        return self.return_values["get_latest_resources"]
-
-    async def add_span(self, span: Span) -> Span:
-        self.calls.append(("add_span", (span,), {}))
-        return self.return_values["add_span"]
-
-    async def add_otel_span(
-        self,
-        rollout_id: str,
-        attempt_id: str,
-        readable_span: ReadableSpan,
-        sequence_id: Optional[int] = None,
-    ) -> Span:
-        self.calls.append(("add_otel_span", (rollout_id, attempt_id, readable_span, sequence_id), {}))
-        return self.return_values["add_otel_span"]
-
-    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[RolloutV2]:
-        self.calls.append(("wait_for_rollouts", (), {"rollout_ids": rollout_ids, "timeout": timeout}))
-        return self.return_values["wait_for_rollouts"]
-
-    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
-        self.calls.append(("get_next_span_sequence_id", (rollout_id, attempt_id), {}))
-        return self.return_values["get_next_span_sequence_id"]
-
-    async def query_spans(
-        self,
-        rollout_id: str,
-        attempt_id: str | Literal["latest"] | None = None,
-    ) -> List[Span]:
-        self.calls.append(("query_spans", (rollout_id, attempt_id), {}))
-        return self.return_values["query_spans"]
-
-    async def update_rollout(
-        self,
-        rollout_id: str,
-        input: TaskInput | Any = UNSET,
-        mode: Optional[str] | Any = UNSET,
-        resources_id: Optional[str] | Any = UNSET,
-        status: RolloutStatus | Any = UNSET,
-        config: Any = UNSET,
-        metadata: Optional[Dict[str, Any]] | Any = UNSET,
-    ) -> RolloutV2:
-        self.calls.append(
-            (
-                "update_rollout",
-                (rollout_id, input, mode, resources_id, status, config, metadata),
-                {},
-            )
-        )
-        return self.return_values["update_rollout"]
-
-    async def update_attempt(
-        self,
-        rollout_id: str,
-        attempt_id: str | Literal["latest"],
-        status: AttemptStatus | Any = UNSET,
-        worker_id: str | Any = UNSET,
-        last_heartbeat_time: float | Any = UNSET,
-        metadata: Optional[Dict[str, Any]] | Any = UNSET,
-    ) -> Attempt:
-        self.calls.append(
-            (
-                "update_attempt",
-                (rollout_id, attempt_id, status, worker_id, last_heartbeat_time, metadata),
-                {},
-            )
-        )
-        return self.return_values["update_attempt"]
+from .dummy_store import DummyLightningStore
 
 
 class SlowAttemptStore(LightningStore):

tests/test_tracer.py

@@ -792,7 +792,9 @@ def test_run_with_agentops_tracer(agent_func):
     global _langchain_callback_handler
     _langchain_callback_handler = tracer.get_langchain_callback_handler()
 
+    loop = asyncio.new_event_loop()
     try:
+        asyncio.set_event_loop(loop)
         tracer.trace_run(
             run_one,
             agent_func,
@@ -831,6 +833,8 @@ def test_run_with_agentops_tracer(agent_func):
     finally:
         tracer.teardown_worker(0)
         tracer.teardown()
+        loop.close()
+        asyncio.set_event_loop(None)
 
 
 @pytest.mark.parametrize("agent_func", list(iterate_over_agents()), ids=lambda f: f.__name__)