mantisdk 0.1.0__py3-none-any.whl

mantisdk/execution/shared_memory.py

@@ -0,0 +1,282 @@
+# 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 mantisdk.env_var import LightningEnvVar, resolve_bool_env_var
+from mantisdk.store.base import LightningStore
+from mantisdk.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`][mantisdk.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

mantisdk/instrumentation/__init__.py

@@ -0,0 +1,119 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import warnings
+
+AGENTOPS_INSTALLED: bool = False
+AGENTOPS_LANGCHAIN_INSTALLED: bool = False
+LITELLM_INSTALLED: bool = False
+VLLM_INSTALLED: bool = False
+WEAVE_INSTALLED: bool = False
+OPENAI_OTEL_INSTALLED: bool = False
+
+try:
+    from . import agentops  # type: ignore
+
+    AGENTOPS_INSTALLED = True  # type: ignore
+except ImportError:
+    pass
+
+try:
+    from . import litellm  # type: ignore
+
+    LITELLM_INSTALLED = True  # type: ignore
+except ImportError:
+    pass
+
+try:
+    from opentelemetry.instrumentation.openai import OpenAIInstrumentor  # type: ignore
+
+    OPENAI_OTEL_INSTALLED = True
+except ImportError:
+    pass
+
+
+def instrument_all():
+    """Instrument all the instrumentation libraries."""
+    if AGENTOPS_INSTALLED:
+        from .agentops import instrument_agentops
+
+        instrument_agentops()
+    else:
+        warnings.warn("agentops is not installed. It's therefore not instrumented.")
+
+    if LITELLM_INSTALLED:
+        from .litellm import instrument_litellm
+
+        instrument_litellm()
+    else:
+        warnings.warn("litellm is not installed. It's therefore not instrumented.")
+
+    if VLLM_INSTALLED:
+        from .vllm import instrument_vllm
+
+        instrument_vllm()
+    else:
+        warnings.warn("vllm is not installed. It's therefore not instrumented.")
+
+    if AGENTOPS_LANGCHAIN_INSTALLED:
+        from .agentops_langchain import instrument_agentops_langchain
+
+        instrument_agentops_langchain()
+    else:
+        warnings.warn("Agentops-langchain integration is not installed. It's therefore not instrumented.")
+    
+    # Enable standard OpenAI instrumentation if available
+    # This allows client-side spans to be generated for OpenAI SDK calls
+    if OPENAI_OTEL_INSTALLED:
+        try:
+            OpenAIInstrumentor().instrument()  # type: ignore
+        except Exception as e:
+            warnings.warn(f"Failed to instrument OpenAI SDK: {e}")
+
+
+def uninstrument_all():
+    """Uninstrument all the instrumentation libraries."""
+    if AGENTOPS_INSTALLED:
+        try:
+            from .agentops import uninstrument_agentops
+
+            uninstrument_agentops()
+        except ImportError:
+            warnings.warn("agentops is installed but uninstrument_agentops could not be imported.")
+    else:
+        warnings.warn("agentops is not installed. It's therefore not uninstrumented.")
+
+    if LITELLM_INSTALLED:
+        try:
+            from .litellm import uninstrument_litellm
+
+            uninstrument_litellm()
+        except ImportError:
+            warnings.warn("litellm is installed but uninstrument_litellm could not be imported.")
+    else:
+        warnings.warn("litellm is not installed. It's therefore not uninstrumented.")
+
+    if VLLM_INSTALLED:
+        try:
+            from .vllm import uninstrument_vllm
+
+            uninstrument_vllm()
+        except ImportError:
+            warnings.warn("vllm is installed but uninstrument_vllm could not be imported.")
+    else:
+        warnings.warn("vllm is not installed. It's therefore not uninstrumented.")
+
+    if AGENTOPS_LANGCHAIN_INSTALLED:
+        try:
+            from .agentops_langchain import uninstrument_agentops_langchain
+
+            uninstrument_agentops_langchain()
+        except ImportError:
+            warnings.warn("agentops_langchain is installed but uninstrument_agentops_langchain could not be imported.")
+    else:
+        warnings.warn("Agentops-langchain integration is not installed. It's therefore not uninstrumented.")
+    
+    if OPENAI_OTEL_INSTALLED:
+        try:
+            OpenAIInstrumentor().uninstrument()  # type: ignore
+        except Exception:
+            pass