microsoft--agent-lightning/agentlightning/execution/shared_memory.py

# 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.env_var import LightningEnvVar, resolve_bool_env_var
from agentlightning.store.base import LightningStore
from agentlightning.store.threading import LightningStoreThreaded

from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle
from .events import ExecutionEvent, ThreadingEvent

logger = logging.getLogger(__name__)


class SharedMemoryExecutionStrategy(ExecutionStrategy):
    """Execute bundles in a single process with cooperative worker threads.

    Stop Model:

    - All bundles share one [`ThreadingEvent`][agentlightning.ThreadingEvent]
      named `stop_evt`.
    - Only the main thread receives `KeyboardInterrupt`. When Ctrl+C occurs we
      set `stop_evt`.
    - Any exception raised inside a bundle sets `stop_evt` so other threads can
      unwind cooperatively.
    - Once the bundle running on the main thread exits successfully the
      treatment depends on `main_thread`:
        - `"algorithm"`: the runners are asked to stop by setting `stop_evt`.
        - `"runner"`: the algorithm keeps running until it exits naturally.
    - Background threads are marked as daemons. We join them briefly and log any
      stragglers before shutting down.

    !!! note
        Signals other than `SIGINT` (such as `SIGTERM`) are not intercepted;
        Python's default behavior for those signals is preserved.
    """

    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,
        managed_store: bool | None = None,
    ) -> 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. "
                "Either use 'algorithm' on the main thread or set n_runners to 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
        self.managed_store = resolve_bool_env_var(
            LightningEnvVar.AGL_MANAGED_STORE, override=managed_store, fallback=True
        )

    async def _run_until_completed_or_canceled(self, coro: Awaitable[Any], stop_evt: ExecutionEvent) -> Any:
        """Run `coro` until it finishes or a cooperative stop is requested.

        Control flow:

        1. Start the bundle coroutine as `task`.
        2. Launch a watcher that polls `stop_evt` without blocking the loop.
        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. Cancel the bundle task if it is still running after the grace
               period.
        4. Await both tasks and swallow `CancelledError` where appropriate.

        This is a *backup* mechanism for bundles that might not poll the event
        frequently; cooperative shutdown (checking `stop_evt` inside the
        bundle) remains the preferred approach.
        """
        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: ExecutionEvent,
        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: ExecutionEvent,
        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()
        if self.managed_store:
            thread_safe_store = LightningStoreThreaded(store)
        else:
            thread_safe_store = 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