mantisdk 0.1.0__py3-none-any.whl

mantisdk/trainer/trainer.py

@@ -0,0 +1,618 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import functools
+import logging
+import warnings
+from typing import Any, Callable, Dict, Optional, Sequence, TypeVar, Union
+
+from mantisdk.adapter import TraceAdapter, TracerTraceToTriplet
+from mantisdk.algorithm import Algorithm, Baseline, FastAlgorithm
+from mantisdk.client import MantisdkClient
+from mantisdk.execution.base import ExecutionStrategy
+from mantisdk.execution.client_server import ClientServerExecutionStrategy
+from mantisdk.execution.events import ExecutionEvent
+from mantisdk.litagent import LitAgent
+from mantisdk.llm_proxy import LLMProxy
+from mantisdk.runner import LitAgentRunner, Runner
+from mantisdk.store.base import LightningStore
+from mantisdk.store.memory import InMemoryLightningStore
+from mantisdk.tracer.agentops import AgentOpsTracer
+from mantisdk.tracer.base import Tracer
+from mantisdk.tracer.otel import OtelTracer
+from mantisdk.types import Dataset, Hook, NamedResources
+
+from .init_utils import build_component, instantiate_component
+from .legacy import TrainerLegacy
+from .registry import ExecutionStrategyRegistry
+
+logger = logging.getLogger(__name__)
+
+T_co = TypeVar("T_co", covariant=True)
+T = TypeVar("T")
+
+ComponentSpec = Union[T, type[T], Callable[[], T], str, Dict[str, Any], None]
+
+
+class Trainer(TrainerLegacy):
+    """High-level orchestration layer that wires Algorithm <-> Runner <-> Store.
+
+    A [`Trainer`][mantisdk.Trainer] packages the moving parts of Agent-Lightning's
+    training loop into a single entry point:
+
+    * **Algorithm lifecycle:** Instantiates or accepts an [`Algorithm`][mantisdk.Algorithm],
+      attaches the current [`LightningStore`][mantisdk.LightningStore], adapter, and
+      initial resources, then executes the algorithm role inside the configured execution strategy.
+    * **Runner fleet:** Spawns one or more [`Runner`][mantisdk.Runner] instances (defaulting
+      to [`LitAgentRunner`][mantisdk.LitAgentRunner]) that hydrate a [`LitAgent`][mantisdk.LitAgent],
+      claim rollouts, stream spans, and respect graceful termination signals from the execution strategy.
+    * **Execution strategy:** Delegates process management to an
+      [`ExecutionStrategy`][mantisdk.ExecutionStrategy] (shared memory, client/server, etc.),
+      so advanced users can swap orchestration backends without changing trainer code.
+    * **Telemetry plumbing:** Ensures tracers, adapters, and optional [`LLMProxy`][mantisdk.LLMProxy]
+      are wired into both algorithm and runners so telemetry flows back into the store.
+
+    The trainer exposes two convenience entry points:
+    [`fit()`][mantisdk.Trainer.fit] for full training and
+    [`dev()`][mantisdk.Trainer.dev] for fast, reproducible dry-runs. See the
+    [Train the First Agent](../how-to/train-first-agent.md) and
+    [Write the First Algorithm](../how-to/write-first-algorithm.md) tutorials for the broader context.
+    """
+
+    algorithm: Optional[Algorithm]
+    """An instance of [`Algorithm`][mantisdk.Algorithm] to use for training."""
+
+    store: LightningStore
+    """An instance of [`LightningStore`][mantisdk.LightningStore] to use for storing tasks and traces."""
+
+    runner: Runner[Any]
+    """An instance of [`Runner`][mantisdk.Runner] to use for running the agent."""
+
+    initial_resources: Optional[NamedResources]
+    """An instance of [`NamedResources`][mantisdk.NamedResources] to use for bootstrapping the fit/dev process.
+
+    The resources will be handed over to the algorithm. Note that not all algorithms support seeding resources.
+    """
+
+    n_runners: int
+    """Number of agent runners to run in parallel."""
+
+    max_rollouts: Optional[int]
+    """Maximum number of rollouts to process per runner. If None, workers run until no more rollouts are available."""
+
+    strategy: ExecutionStrategy
+    """An instance of [`ExecutionStrategy`][mantisdk.ExecutionStrategy] to use for spawning the algorithm and runners."""
+
+    tracer: Tracer
+    """A tracer instance, or a string pointing to the class full name or a dictionary with a 'type' key
+    that specifies the class full name and other initialization parameters.
+    If None, a default [`AgentOpsTracer`][mantisdk.AgentOpsTracer] will be created with the current settings."""
+
+    hooks: Sequence[Hook]
+    """A sequence of [`Hook`][mantisdk.Hook] instances to be called at various lifecycle stages (e.g., `on_trace_start`,
+    `on_trace_end`, `on_rollout_start`, `on_rollout_end`)."""
+
+    adapter: TraceAdapter[Any]
+    """An instance of [`TraceAdapter`][mantisdk.TraceAdapter] to export data consumble by algorithms from traces."""
+
+    llm_proxy: Optional[LLMProxy]
+    """An instance of [`LLMProxy`][mantisdk.LLMProxy] to use for intercepting the LLM calls.
+    If not provided, algorithm may create one on its own."""
+
+    n_workers: int
+    """Number of agent workers to run in parallel. Deprecated in favor of `n_runners`."""
+
+    max_tasks: Optional[int]
+    """Maximum number of tasks to process per runner. Deprecated in favor of `max_rollouts`."""
+
+    daemon: bool
+    """Whether worker processes should be daemons. Daemon processes
+    are terminated automatically when the main process exits. Deprecated.
+    Only have effect with `fit_v0`."""
+
+    triplet_exporter: TraceAdapter[Any]
+    """An instance of [`TracerTraceToTriplet`][mantisdk.TracerTraceToTriplet] to export triplets from traces,
+    or a dictionary with the initialization parameters for the exporter.
+    Deprecated. Use [`adapter`][mantisdk.Trainer.adapter] instead."""
+
+    port: Optional[int]
+    """Port forwarded to [`ClientServerExecutionStrategy`][mantisdk.ClientServerExecutionStrategy]."""
+
+    def __init__(
+        self,
+        *,
+        dev: bool = False,
+        n_runners: Optional[int] = None,
+        max_rollouts: Optional[int] = None,
+        initial_resources: Optional[NamedResources] = None,
+        tracer: ComponentSpec[Tracer] = None,
+        adapter: ComponentSpec[TraceAdapter[Any]] = None,
+        store: ComponentSpec[LightningStore] = None,
+        runner: ComponentSpec[Runner[Any]] = None,
+        strategy: ComponentSpec[ExecutionStrategy] = None,
+        port: Optional[int] = None,
+        algorithm: ComponentSpec[Algorithm] = None,
+        llm_proxy: ComponentSpec[LLMProxy] = None,
+        n_workers: Optional[int] = None,
+        max_tasks: Optional[int] = None,
+        daemon: bool = True,
+        triplet_exporter: ComponentSpec[TracerTraceToTriplet] = None,
+        hooks: Optional[Union[Hook, Sequence[Hook]]] = None,
+    ):
+        """Configure the trainer and resolve user-provided component specifications.
+
+        Each keyword accepts either a concrete instance, a class, a callable factory, a
+        registry string, or a lightweight configuration dictionary (see
+        [`build_component()`][mantisdk.trainer.init_utils.build_component]).
+
+        When ``port`` is provided it is forwarded to
+        [`ClientServerExecutionStrategy`][mantisdk.ClientServerExecutionStrategy]
+        instances constructed (or supplied) for the trainer.
+        """
+        # Do not call super().__init__() here.
+        # super().__init__() will call TrainerLegacy's initialization, which is not intended.
+        self.worker_id: Optional[int] = None
+
+        if dev:
+            warnings.warn(
+                "Trainer(dev=True) is deprecated and will be removed in future versions. "
+                "Please use Trainer.dev(...) instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        self._dev = dev
+        self.daemon = daemon
+        self._client: MantisdkClient | None = None  # Will be initialized in fit or fit_v0
+
+        if n_workers is not None:
+            warnings.warn(
+                "`n_workers` is deprecated. Please use `n_runners`.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+        if n_runners is None:
+            n_runners = n_workers if n_workers is not None else 1
+        else:
+            if n_workers is not None and n_workers != n_runners:
+                warnings.warn(
+                    "`n_workers` is ignored when `n_runners` is provided.",
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
+
+        self.n_runners = n_runners
+        self.n_workers = n_runners  # Backwards compatibility for fit_v0
+
+        if max_tasks is not None:
+            warnings.warn(
+                "`max_tasks` is deprecated. Please use `max_rollouts`.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+        if max_rollouts is None:
+            max_rollouts = max_tasks
+        elif max_tasks is not None and max_tasks != max_rollouts:
+            warnings.warn(
+                "`max_tasks` is ignored when `max_rollouts` is provided.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+        self.max_rollouts = max_rollouts
+        self.max_tasks = max_tasks if max_tasks is not None else max_rollouts
+
+        # Note: tracer is created after store to enable auto-detection of OTLP-capable stores
+        self._tracer_spec = tracer  # Store spec for later resolution
+
+        if adapter is not None and triplet_exporter is not None:
+            warnings.warn(
+                "`triplet_exporter` is deprecated and ignored because `adapter` is provided.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+        adapter_spec = adapter if adapter is not None else triplet_exporter
+        self.adapter = self._make_adapter(adapter_spec)
+        self.triplet_exporter = self.adapter  # Backwards compatibility
+
+        self.algorithm = self._make_algorithm(algorithm)
+
+        # We might be able to support a list of resources in future.
+        self.initial_resources = initial_resources
+
+        self.port = port
+
+        self.strategy = self._make_strategy(
+            strategy,
+            n_runners=self.n_runners,
+            port=port,
+        )
+
+        # The active store for the current execution context
+        self.store = self._make_store(store, self.strategy)
+
+        # Create tracer after store so we can auto-detect OTLP-capable stores
+        self.tracer = self._make_tracer(self._tracer_spec, self.store)
+
+        self.runner = self._make_runner(runner)
+
+        if hasattr(self.strategy, "n_runners"):
+            strategy_runners = getattr(self.strategy, "n_runners")
+            if isinstance(strategy_runners, int) and strategy_runners > 0:
+                self.n_runners = strategy_runners
+                self.n_workers = strategy_runners
+
+        self.llm_proxy = self._make_llm_proxy(llm_proxy, store=self.store)
+
+        self.hooks = self._normalize_hooks(hooks)
+
+        if not self.daemon:
+            logger.warning(
+                "daemon=False. Worker processes are non-daemonic. "
+                "The worker processes will NOT be terminated when the main process exits. "
+                "The cleanup must be handled manually."
+            )
+
+    def _make_tracer(
+        self, tracer: ComponentSpec[Tracer], store: Optional[LightningStore] = None
+    ) -> Tracer:
+        """Resolve the tracer component from user input.
+
+        Auto-detects OTLP-capable stores (like InsightLightningStore) and uses OtelTracer
+        to enable native OTLP trace export. Falls back to AgentOpsTracer otherwise.
+        """
+        # Check if store supports OTLP traces (e.g., InsightLightningStore)
+        store_supports_otlp = False
+        if store is not None:
+            # Check store's capabilities
+            if hasattr(store, "capabilities"):
+                store_supports_otlp = store.capabilities.get("otlp_traces", False)
+            # Also check listeners on the store for OTLP capability
+            elif hasattr(store, "_listeners"):
+                for listener in getattr(store, "_listeners", []):
+                    if hasattr(listener, "capabilities"):
+                        if listener.capabilities.get("otlp_traces", False):
+                            store_supports_otlp = True
+                            break
+
+        if store_supports_otlp:
+            logger.info("OTLP-capable store detected. Using OtelTracer for native trace export.")
+            default_factory = lambda: OtelTracer()
+        else:
+            default_factory = lambda: AgentOpsTracer(
+                agentops_managed=True,
+                instrument_managed=True,
+                daemon=self.daemon,
+            )
+
+        return build_component(
+            tracer,
+            expected_type=Tracer,
+            spec_name="tracer",
+            default_factory=default_factory,
+            dict_requires_type=True,
+            invalid_spec_error_fmt="Invalid tracer type: {actual_type}. Expected Tracer, str, dict, or None.",
+            type_error_fmt="Tracer factory returned {type_name}, which is not a Tracer subclass.",
+        )
+
+    def _make_algorithm(self, algorithm: ComponentSpec[Algorithm]) -> Optional[Algorithm]:
+        """Resolve the algorithm component, allowing `None` for dev-mode dry runs."""
+        return build_component(
+            algorithm,
+            expected_type=Algorithm,
+            spec_name="algorithm",
+            allow_none=True,
+            invalid_spec_error_fmt="Invalid algorithm type: {actual_type}. Expected Algorithm, str, dict, or None.",
+            type_error_fmt="Algorithm factory returned {type_name}, which is not a Algorithm subclass.",
+        )
+
+    def _make_adapter(self, adapter: ComponentSpec[TraceAdapter[Any]]) -> TraceAdapter[Any]:
+        """Resolve the adapter used to transform spans into algorithm-ready payloads."""
+        return build_component(
+            adapter,
+            expected_type=TraceAdapter,
+            spec_name="adapter",
+            default_factory=TracerTraceToTriplet,
+            dict_requires_type=False,
+            dict_default_cls=TracerTraceToTriplet,
+            invalid_spec_error_fmt="Invalid adapter type: {actual_type}. Expected TraceAdapter, dict, or None.",
+            type_error_fmt="Adapter factory returned {type_name}, which is not a TraceAdapter subclass.",
+        )
+
+    def _make_store(self, store: ComponentSpec[LightningStore], strategy: ExecutionStrategy) -> LightningStore:
+        """Resolve the store implementation backing rollouts, attempts, spans, and resources.
+
+        By default, it's always a in-memory store. If using a client/server execution strategy,
+        the in-memory store will be initialized in a thread-safe manner.
+        """
+        is_client_server = isinstance(strategy, ClientServerExecutionStrategy)
+        default_store_factory = lambda: InMemoryLightningStore(thread_safe=is_client_server)
+        return build_component(
+            store,
+            expected_type=LightningStore,
+            spec_name="store",
+            default_factory=default_store_factory,
+            invalid_spec_error_fmt="Invalid store type: {actual_type}. Expected LightningStore, str, dict, or None.",
+            type_error_fmt="Store factory returned {type_name}, which is not a LightningStore subclass.",
+        )
+
+    def _make_strategy(
+        self,
+        strategy: ComponentSpec[ExecutionStrategy],
+        *,
+        n_runners: int,
+        port: Optional[int] = None,
+    ) -> ExecutionStrategy:
+        """Resolve the execution strategy and seed defaults such as `n_runners`."""
+        if isinstance(strategy, ExecutionStrategy):
+            if port is not None and isinstance(strategy, ClientServerExecutionStrategy):
+                strategy.server_port = port
+            return strategy
+        optional_defaults: Dict[str, Callable[[], Any]] = {"n_runners": lambda: n_runners}
+        if port is not None:
+            optional_defaults["server_port"] = lambda: port
+
+        def default_factory() -> ExecutionStrategy:
+            if port is not None:
+                return ClientServerExecutionStrategy(n_runners=n_runners, server_port=port)
+            return ClientServerExecutionStrategy(n_runners=n_runners)
+
+        return build_component(
+            strategy,
+            expected_type=ExecutionStrategy,
+            spec_name="strategy",
+            default_factory=default_factory,
+            optional_defaults=optional_defaults,
+            invalid_spec_error_fmt="Invalid strategy type: {actual_type}. Expected ExecutionStrategy, str, dict, or None.",
+            type_error_fmt="Strategy factory returned {type_name}, which is not an ExecutionStrategy subclass.",
+            registry=ExecutionStrategyRegistry,
+        )
+
+    def _make_llm_proxy(
+        self,
+        llm_proxy: ComponentSpec[LLMProxy],
+        *,
+        store: LightningStore,
+    ) -> Optional[LLMProxy]:
+        """Resolve an optional LLM proxy and ensure it shares the trainer's store instance."""
+        if isinstance(llm_proxy, LLMProxy):
+            return llm_proxy
+
+        optional_defaults: Dict[str, Callable[[], Any]] = {"store": lambda: store}
+        
+        # Extract OTLP endpoint and headers from store if it supports OTLP
+        if hasattr(store, "capabilities") and store.capabilities.get("otlp_traces", False):
+            try:
+                otlp_endpoint = store.otlp_traces_endpoint()
+                otlp_headers = store.get_otlp_headers() if hasattr(store, "get_otlp_headers") else {}
+                optional_defaults["otlp_endpoint"] = lambda: otlp_endpoint
+                optional_defaults["otlp_headers"] = lambda: otlp_headers
+                logger.info(f"Trainer: Store provides OTLP endpoint for proxy: {otlp_endpoint}")
+            except Exception as e:
+                logger.warning(f"Trainer: Failed to get OTLP endpoint from store: {e}")
+        
+        if isinstance(llm_proxy, dict):
+            llm_proxy = {**llm_proxy}
+            llm_proxy.setdefault("store", store)
+            # Pass OTLP config to proxy if store provides it and it's not already set
+            if "otlp_endpoint" in optional_defaults and "otlp_endpoint" not in llm_proxy:
+                llm_proxy.setdefault("otlp_endpoint", optional_defaults["otlp_endpoint"]())
+            if "otlp_headers" in optional_defaults and "otlp_headers" not in llm_proxy:
+                llm_proxy.setdefault("otlp_headers", optional_defaults["otlp_headers"]())
+
+        return build_component(
+            llm_proxy,
+            expected_type=LLMProxy,
+            spec_name="llm_proxy",
+            allow_none=True,
+            optional_defaults=optional_defaults,
+            invalid_spec_error_fmt="Invalid llm_proxy type: {actual_type}. Expected LLMProxy, dict, str, or None.",
+            type_error_fmt="llm_proxy factory returned {type_name}, which is not an LLMProxy subclass.",
+        )
+
+    def _make_runner(self, runner: ComponentSpec[Runner[Any]]) -> Runner[Any]:
+        """Resolve the runner responsible for executing the agent inside each worker."""
+        optional_defaults: Dict[str, Callable[[], Any]] = {"tracer": lambda: self.tracer}
+        if self.max_rollouts is not None:
+            optional_defaults["max_rollouts"] = lambda: self.max_rollouts
+
+        def default_runner_factory() -> Runner[Any]:
+            return instantiate_component(LitAgentRunner, optional_defaults=optional_defaults)
+
+        return build_component(
+            runner,
+            expected_type=Runner,
+            spec_name="runner",
+            default_factory=default_runner_factory,
+            optional_defaults=optional_defaults,
+            invalid_spec_error_fmt="Invalid runner type: {actual_type}. Expected Runner, callable, str, dict, or None.",
+            type_error_fmt="Runner factory returned {type_name}, which is not a Runner subclass.",
+        )
+
+    def _normalize_hooks(self, hooks: Optional[Union[Hook, Sequence[Hook]]]) -> Sequence[Hook]:
+        """Coerce hook inputs into an immutable sequence for runner initialization."""
+        if hooks is None:
+            return ()
+        if isinstance(hooks, Hook):
+            return (hooks,)
+        return tuple(hooks)
+
+    def fit(
+        self,
+        agent: LitAgent[T_co],
+        train_dataset: Optional[Dataset[T_co]] = None,
+        *,
+        val_dataset: Optional[Dataset[T_co]] = None,
+    ) -> None:
+        """Execute the full algorithm/runner training loop.
+
+        [`Trainer.fit`][mantisdk.Trainer.fit] packages the algorithm and runner bundles,
+        then hands them to the active [`ExecutionStrategy`][mantisdk.ExecutionStrategy].
+        The strategy rarely returns until:
+
+        * The algorithm exhausts the dataset(s) and stops enqueuing rollouts.
+        * `max_rollouts` causes individual runners to exit.
+        * An exception or interrupt cancels the shared [`ExecutionEvent`][mantisdk.ExecutionEvent].
+
+        Args:
+            agent: [`LitAgent`][mantisdk.LitAgent] implementation executed by runners.
+            train_dataset: Optional iterable of rollout inputs consumed by the algorithm.
+            val_dataset: Optional iterable consumed by validation passes.
+        """
+        if isinstance(train_dataset, str):
+            logger.warning(
+                "Trainer.fit will no longer accepts a string URL in future version. "
+                "To continue using a string URL, please use Trainer.fit_v0 instead. "
+                "See documentation for how to migrate to latest version: https://microsoft.github.io/mantisdk/stable/"
+            )
+            return self.fit_v0(  # type: ignore
+                agent,
+                train_dataset,
+                val_dataset,  # type: ignore
+            )
+
+        agent.set_trainer(self)
+
+        algorithm_bundle = functools.partial(
+            self._algorithm_bundle,
+            train_dataset=train_dataset,
+            val_dataset=val_dataset,
+            algorithm=self.algorithm,
+        )
+        runner_bundle = functools.partial(self._runner_bundle, agent=agent)
+
+        self.strategy.execute(algorithm_bundle, runner_bundle, self.store)
+
+    def dev(
+        self,
+        agent: LitAgent[T_co],
+        train_dataset: Optional[Dataset[T_co]] = None,
+        *,
+        val_dataset: Optional[Dataset[T_co]] = None,
+    ) -> None:
+        """Exercise the infrastructure using a fast, synchronous algorithm.
+
+        [`Trainer.dev`][mantisdk.Trainer.dev] mirrors [`fit()`][mantisdk.Trainer.fit] but
+        insists on an [`Algorithm`][mantisdk.Algorithm] subtype that also derives from
+        [`FastAlgorithm`][mantisdk.FastAlgorithm]. This keeps the loop responsive for
+        debugging while still touching the same store, runners, hooks, and tracer plumbing.
+
+        If no algorithm is provided, a default [`Baseline`][mantisdk.Baseline] algorithm will be used.
+
+        Args:
+            agent: [`LitAgent`][mantisdk.LitAgent] implementation to execute.
+            train_dataset: Optional iterable passed to the algorithm.
+            val_dataset: Optional iterable passed to the algorithm.
+
+        Raises:
+            TypeError: If the configured algorithm does not inherit from `FastAlgorithm`.
+        """
+        agent.set_trainer(self)
+
+        # Sanity check
+        if self.algorithm is None:
+            algorithm = Baseline()
+        else:
+            algorithm = self.algorithm
+
+        if not isinstance(algorithm, FastAlgorithm):
+            raise TypeError(
+                "Trainer.dev() requires an algorithm that inherits from FastAlgorithm. "
+                f"Received {type(algorithm).__name__}."
+            )
+
+        algorithm_bundle = functools.partial(
+            self._algorithm_bundle,
+            train_dataset=train_dataset,
+            val_dataset=val_dataset,
+            algorithm=algorithm,
+        )
+        runner_bundle = functools.partial(self._runner_bundle, agent=agent)
+        self.strategy.execute(algorithm_bundle, runner_bundle, self.store)
+
+    async def _algorithm_bundle(
+        self,
+        store: LightningStore,
+        event: ExecutionEvent,
+        train_dataset: Optional[Dataset[T_co]],
+        val_dataset: Optional[Dataset[T_co]],
+        algorithm: Optional[Algorithm],
+    ) -> None:
+        """Internal entry point executed by the strategy for the algorithm role.
+
+        This coroutine is scheduled inside the strategy's process/thread and is responsible
+        for binding algorithm dependencies (store, adapter, initial resources, proxy) before
+        invoking [`Algorithm.run`][mantisdk.Algorithm.run].
+        When `algorithm` is `None` the bundle simply waits for the
+        shared `event` to signal shutdown so runners can still execute (useful for manual queue
+        seeding or external algorithms).
+        """
+        if algorithm is not None:
+            algorithm.set_trainer(self)
+            algorithm.set_store(store)
+            algorithm.set_adapter(self.adapter)
+            if self.initial_resources is not None:
+                algorithm.set_initial_resources(self.initial_resources)
+            if self.llm_proxy is not None:
+                self.llm_proxy.set_store(store)
+                algorithm.set_llm_proxy(self.llm_proxy)
+
+        if algorithm is None:
+            while not event.is_set():
+                await asyncio.sleep(0.1)
+            return
+        try:
+            if algorithm.is_async():
+                await algorithm.run(  # type: ignore
+                    train_dataset=train_dataset,
+                    val_dataset=val_dataset,
+                )
+            else:
+                # This will block the event loop to maximize the debugging experience
+                # It's the responsibility of the execution strategy to enable async execution
+                algorithm.run(
+                    train_dataset=train_dataset,
+                    val_dataset=val_dataset,
+                )
+        except Exception:
+            logger.exception("Algorithm bundle encountered an error.")
+            raise
+
+    async def _runner_bundle(
+        self, store: LightningStore, worker_id: int, event: ExecutionEvent, agent: LitAgent[T_co]
+    ) -> None:
+        """Internal entry point executed by the strategy for each runner role.
+
+        The bundle materializes the configured runner, binds the agent and hooks, associates
+        the worker with the shared store, and then drives the runner's [`iter`][mantisdk.Runner.iter]
+        loop until the execution event is set or an exception occurs. Cleanup mirrors the initialization
+        sequence to keep tracer state, hooks, and agent resources consistent across restarts.
+        """
+        runner_instance: Runner[Any] | None = None
+        runner_initialized = False
+        worker_initialized = False
+        try:
+            # If not using shm execution strategy, we are already in the forked process
+            runner_instance = self.runner
+            runner_instance.init(agent=agent, hooks=self.hooks)
+            runner_initialized = True
+            runner_instance.init_worker(worker_id, store)
+            worker_initialized = True
+            await runner_instance.iter(event=event)
+        except Exception:
+            logger.exception("Runner bundle encountered an error (worker_id=%s).", worker_id)
+            raise
+        finally:
+            if runner_instance is not None:
+                if worker_initialized:
+                    try:
+                        runner_instance.teardown_worker(worker_id)
+                    except Exception:
+                        logger.exception("Error during runner worker teardown (worker_id=%s).", worker_id)
+                if runner_initialized:
+                    try:
+                        runner_instance.teardown()
+                    except Exception:
+                        logger.exception("Error during runner teardown (worker_id=%s).", worker_id)

mantisdk/types/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .core import *
+from .resources import *
+from .tracer import *
+from .tracing import TracingConfig