gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/platform/orchestrator.py

@@ -0,0 +1,901 @@
+# Copyright (c) Gateway. All rights reserved.
+
+"""Orchestrator ABC and DockerOrchestrator implementation.
+
+The ``Orchestrator`` base class defines the contract for compute lifecycle
+management.  Each backend (Docker, Nomad, KubeRay) implements this contract.
+``ComputeHandle`` is an opaque token that the orchestrator returns from
+``start()`` and accepts back in every subsequent call.
+
+Usage (via AlgorithmContext)::
+
+    orch = DockerOrchestrator()
+    handle = orch.start(
+        run_id="run-abc",
+        image="my-agent:v1",
+        store_url="http://host.docker.internal:9876",
+        count=3,
+        lifecycle=ContainerLifecycle.PERSISTENT,
+    )
+    orch.wait_ready(handle, timeout=120.0)
+    # ... algorithm enqueues rollouts, N agents process them ...
+    orch.stop(handle)
+"""
+
+from __future__ import annotations
+
+import logging
+import time as _time
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# Enums & data classes
+# ═══════════════════════════════════════════════════════════════════════
+
+
+class ComputeStatus(str, Enum):
+    """Lifecycle state of orchestrated compute."""
+
+    PENDING = "pending"
+    RUNNING = "running"
+    SUCCEEDED = "succeeded"
+    FAILED = "failed"
+    UNKNOWN = "unknown"
+
+
+class ContainerLifecycle(str, Enum):
+    """How agent containers are recycled during a run.
+
+    PERSISTENT: containers stay alive for the entire run duration.
+               Best for stateless agents -- avoids startup overhead.
+    EPHEMERAL:  containers are replaced after each rollout batch.
+               Guarantees clean state for stateful agents (e.g. CRM sim).
+    """
+
+    PERSISTENT = "persistent"
+    EPHEMERAL = "ephemeral"
+
+
+@dataclass
+class ComputeHandle:
+    """Opaque handle returned by start(), passed back to stop()/status().
+
+    Each orchestrator subclass stores whatever it needs here.
+    Docker stores container IDs; Nomad stores a job ID; KubeRay stores a
+    RayJob name; etc.
+    """
+
+    run_id: str
+    backend: str  # "docker", "nomad", "kuberay", ...
+    identifiers: Dict[str, str] = field(default_factory=dict)
+    # e.g. {"agent-0": "abc123", "agent-1": "def456"}  (Docker)
+    # e.g. {"job_id": "gateway-run-42"}                  (Nomad)
+    # e.g. {"rayjob": "gateway-verl-run-42"}             (KubeRay)
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# Orchestrator ABC
+# ═══════════════════════════════════════════════════════════════════════
+
+
+class Orchestrator(ABC):
+    """Base class for compute orchestrators.
+
+    An orchestrator manages the lifecycle of compute resources for a
+    platform run.  Subclasses handle the specifics of Docker, Nomad,
+    KubeRay, etc.
+
+    Lifecycle (called by AlgorithmContext)::
+
+        handle = orch.start(...)        # __aenter__
+        orch.wait_ready(handle)         # block until healthy
+        # ... algorithm runs ...
+        orch.status(handle)             # health checks during run
+        orch.stop(handle)               # __aexit__
+    """
+
+    @property
+    @abstractmethod
+    def available(self) -> bool:
+        """Can this orchestrator function in the current environment?"""
+        ...
+
+    @property
+    @abstractmethod
+    def backend_name(self) -> str:
+        """Short identifier: 'docker', 'nomad', 'kuberay', etc."""
+        ...
+
+    @abstractmethod
+    def start(
+        self,
+        *,
+        run_id: str,
+        image: str,
+        store_url: str,
+        count: int = 1,
+        lifecycle: ContainerLifecycle = ContainerLifecycle.PERSISTENT,
+        env: Optional[Dict[str, str]] = None,
+        labels: Optional[Dict[str, str]] = None,
+        runner_config: Optional[Dict[str, Any]] = None,
+        resource_limits: Optional[Dict[str, Any]] = None,
+        # --- reactive scaling (optional) ---
+        store: Optional[Any] = None,
+        n_runners_per_unit: int = 1,
+        max_units: int = 20,
+        # --- extended fields (used by KubeRay, ignored by Docker/Nomad) ---
+        trainer_resources: Optional[Dict[str, Any]] = None,
+        services: Optional[Dict[str, Any]] = None,
+        algorithm_config: Optional[Dict[str, Any]] = None,
+    ) -> ComputeHandle:
+        """Provision compute for a run.
+
+        For Docker:  starts ``count`` agent containers
+        For Nomad:   submits a job spec with task group ``count``
+        For KubeRay: creates a RayJob CR (head + workers)
+
+        If ``store`` is provided, the orchestrator registers a listener
+        for reactive scaling (start containers when rollouts are enqueued)
+        and ephemeral recycling (replace containers after rollout completion
+        when lifecycle is EPHEMERAL).
+
+        Args:
+            run_id: GatewayRun ID.
+            image: Container image to run.
+            store_url: HTTP URL of the worker's LightningStoreServer.
+            count: Number of agent instances to start immediately.
+                Use 0 for fully reactive (containers start on demand).
+            lifecycle: Whether containers persist or are replaced per batch.
+            env: Additional environment variables.
+            labels: Metadata labels for the containers.
+            runner_config: Runner lifecycle config from agent manifest.
+            resource_limits: Container resource limits from agent manifest.
+            store: LightningStore instance. If provided, enables reactive
+                scaling via store listener hooks.
+            n_runners_per_unit: Rollouts a single compute unit handles
+                concurrently (used for capacity calculation).
+            max_units: Hard cap on number of compute units.
+            trainer_resources: GPU/memory for trainer (KubeRay only).
+            services: Sidecar services (KubeRay only).
+            algorithm_config: Algorithm-specific config (KubeRay only).
+
+        Returns:
+            ComputeHandle with identifiers for the provisioned compute.
+
+        Raises:
+            RuntimeError: If compute cannot be provisioned.
+        """
+        ...
+
+    @abstractmethod
+    def stop(self, handle: ComputeHandle, timeout: int = 30) -> None:
+        """Tear down all compute for a run.
+
+        Must be idempotent -- safe to call multiple times or after
+        compute has already exited.
+        """
+        ...
+
+    @abstractmethod
+    def status(self, handle: ComputeHandle) -> ComputeStatus:
+        """Check aggregate health of the compute for a run.
+
+        Returns RUNNING if all members are healthy, FAILED if any have
+        crashed, PENDING if still starting up.
+        """
+        ...
+
+    @abstractmethod
+    def wait_ready(self, handle: ComputeHandle, timeout: float = 120.0) -> None:
+        """Block until all provisioned compute is healthy.
+
+        Docker: poll container.status until all report 'running'.
+        Nomad:  poll deployment health.
+        KubeRay: poll RayJob status.
+
+        Raises:
+            RuntimeError: If timeout exceeded before all members are ready.
+        """
+        ...
+
+    @abstractmethod
+    def scale(self, handle: ComputeHandle, count: int) -> None:
+        """Adjust the number of running instances.
+
+        Docker: start/stop delta containers, update handle.identifiers.
+        Nomad:  update job count, scheduler reconciles.
+
+        Args:
+            handle: ComputeHandle to scale.
+            count: New desired count.
+        """
+        ...
+
+    @abstractmethod
+    def replace_member(
+        self,
+        handle: ComputeHandle,
+        member_key: str,
+    ) -> Optional[str]:
+        """Replace one member of the compute group (ephemeral mode).
+
+        Docker: stop container, start fresh one, update handle.
+        Nomad:  ``nomad alloc restart``, or no-op (scheduler handles).
+
+        Args:
+            handle: ComputeHandle containing the member.
+            member_key: Key in handle.identifiers (e.g. "agent-0").
+
+        Returns:
+            New container/alloc ID, or None if replacement failed.
+        """
+        ...
+
+    def members(self, handle: ComputeHandle) -> List[Dict[str, Any]]:
+        """List individual members of the compute group.
+
+        Optional -- returns empty list by default.
+        """
+        return []
+
+    def get_logs(self, handle: ComputeHandle, tail: int = 100) -> str:
+        """Fetch recent logs. Optional -- returns empty string by default."""
+        return ""
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# DockerOrchestrator
+# ═══════════════════════════════════════════════════════════════════════
+
+
+class DockerOrchestrator(Orchestrator):
+    """Manages agent containers via Docker SDK.
+
+    For local development, uses the Docker socket to start/stop containers.
+    Gracefully degrades if Docker is not available (e.g., in CI without Docker).
+    """
+
+    def __init__(self, docker_host: Optional[str] = None):
+        """
+        Args:
+            docker_host: Docker daemon URL. Defaults to unix:///var/run/docker.sock.
+        """
+        self._client = None
+        self._stop_timeouts: Dict[str, int] = {}
+        try:
+            import docker
+
+            self._client = docker.from_env() if not docker_host else docker.DockerClient(base_url=docker_host)
+            # Quick connectivity check
+            self._client.ping()
+            logger.info("DockerOrchestrator connected to Docker daemon")
+        except ImportError:
+            logger.warning("Docker SDK not installed. Install with: pip install docker")
+        except Exception as e:
+            logger.warning(f"Could not connect to Docker daemon: {e}")
+            self._client = None
+
+    # ── Orchestrator properties ─────────────────────────────────────
+
+    @property
+    def available(self) -> bool:
+        return self._client is not None
+
+    @property
+    def backend_name(self) -> str:
+        return "docker"
+
+    # ── Private helpers ─────────────────────────────────────────────
+
+    def _build_container_env(
+        self,
+        *,
+        run_id: str,
+        store_url: str,
+        env: Optional[Dict[str, str]],
+        runner_config: Optional[Dict[str, Any]],
+    ) -> Dict[str, str]:
+        """Build the environment dict for an agent container."""
+        import os as _os
+
+        container_env: Dict[str, str] = {
+            "GATEWAY_RUN_ID": run_id,
+            "GATEWAY_STORE_URL": store_url,
+        }
+
+        # Forward Gateway tracing config from worker env unconditionally
+        for _key in (
+            "GATEWAY_HOST",
+            "GATEWAY_PUBLIC_KEY",
+            "GATEWAY_SECRET_KEY",
+        ):
+            _val = _os.environ.get(_key)
+            if _val:
+                if _key == "GATEWAY_HOST" and ("localhost" in _val or "127.0.0.1" in _val):
+                    import platform as _plat
+
+                    _access = "host.docker.internal" if _plat.system() == "Darwin" else _val
+                    _val = _val.replace("localhost", _access).replace("127.0.0.1", _access)
+                container_env[_key] = _val
+
+        # Fallback: forward LLM API keys from worker env ONLY if not already
+        # provided via ProjectSecret (caller-provided env takes precedence).
+        for _key in (
+            "ANTHROPIC_API_KEY",
+            "OPENAI_API_KEY",
+            "AZURE_API_KEY",
+            "GOOGLE_API_KEY",
+            "COHERE_API_KEY",
+        ):
+            if _key not in container_env:
+                _val = _os.environ.get(_key)
+                if _val:
+                    container_env[_key] = _val
+
+        # Merge caller-provided env (takes precedence)
+        if env:
+            container_env.update(env)
+
+        # Inject runner config as env vars (read by LitAgentRunner)
+        if runner_config:
+            if runner_config.get("max_rollouts") is not None:
+                container_env["GATEWAY_MAX_ROLLOUTS"] = str(runner_config["max_rollouts"])
+            if runner_config.get("poll_interval") is not None:
+                container_env["GATEWAY_POLL_INTERVAL"] = str(runner_config["poll_interval"])
+            if runner_config.get("n_runners") is not None:
+                container_env["GATEWAY_N_RUNNERS"] = str(runner_config["n_runners"])
+
+        return container_env
+
+    def _detect_network(self) -> str:
+        """Auto-detect the Docker network to attach agent containers to."""
+        if not self._client:
+            return "bridge"
+        try:
+            networks = [n.name for n in self._client.networks.list()]
+            for candidate in (
+                "langfuse-network",
+                "gateway_default",
+                "gateway_default",
+            ):
+                if candidate in networks:
+                    logger.info(f"Auto-detected Docker network: {candidate}")
+                    return candidate
+        except Exception:
+            pass
+        return "bridge"
+
+    def _build_docker_kwargs(
+        self,
+        *,
+        container_env: Dict[str, str],
+        container_labels: Dict[str, str],
+        network: str,
+        resource_limits: Optional[Dict[str, Any]],
+        name: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Build the kwargs dict for ``docker.containers.run()``."""
+        docker_kwargs: Dict[str, Any] = {
+            "detach": True,
+            "network": network,
+            "environment": container_env,
+            "labels": container_labels,
+            "auto_remove": False,
+        }
+
+        if name:
+            docker_kwargs["name"] = name
+
+        if resource_limits:
+            if resource_limits.get("memory"):
+                docker_kwargs["mem_limit"] = resource_limits["memory"]
+            if resource_limits.get("cpu"):
+                docker_kwargs["nano_cpus"] = int(float(resource_limits["cpu"]) * 1e9)
+            gpu = resource_limits.get("gpu")
+            if gpu and gpu is not False:
+                try:
+                    import docker as _docker
+
+                    gpu_count = -1 if gpu is True else int(gpu)
+                    docker_kwargs["device_requests"] = [
+                        _docker.types.DeviceRequest(count=gpu_count, capabilities=[["gpu"]])
+                    ]
+                except (ImportError, Exception) as e:
+                    logger.warning(f"GPU device request failed: {e}")
+
+        return docker_kwargs
+
+    def _start_one_container(
+        self,
+        *,
+        image: str,
+        run_id: str,
+        index: int,
+        container_env: Dict[str, str],
+        network: str,
+        labels: Optional[Dict[str, str]],
+        resource_limits: Optional[Dict[str, Any]],
+    ) -> Optional[str]:
+        """Start a single agent container. Returns container ID or None."""
+        if not self._client:
+            return None
+
+        name = f"gateway-{run_id[:8]}-agent-{index}"
+        container_labels = {
+            "gateway.run_id": run_id,
+            "gateway.component": "agent",
+            "gateway.agent_index": str(index),
+            **(labels or {}),
+        }
+
+        docker_kwargs = self._build_docker_kwargs(
+            container_env=container_env,
+            container_labels=container_labels,
+            network=network,
+            resource_limits=resource_limits,
+            name=name,
+        )
+
+        try:
+            container = self._client.containers.run(image, **docker_kwargs)
+            logger.info(f"Started container {name} ({container.short_id}) " f"from {image} for run {run_id}")
+            return container.id
+        except Exception as e:
+            logger.error(f"Failed to start container {name}: {e}")
+            return None
+
+    # ── Orchestrator interface ──────────────────────────────────────
+
+    def start(
+        self,
+        *,
+        run_id: str,
+        image: str,
+        store_url: str,
+        count: int = 0,
+        lifecycle: ContainerLifecycle = ContainerLifecycle.PERSISTENT,
+        env: Optional[Dict[str, str]] = None,
+        labels: Optional[Dict[str, str]] = None,
+        runner_config: Optional[Dict[str, Any]] = None,
+        resource_limits: Optional[Dict[str, Any]] = None,
+        store: Optional[Any] = None,
+        n_runners_per_unit: int = 1,
+        max_units: int = 20,
+        trainer_resources: Optional[Dict[str, Any]] = None,
+        services: Optional[Dict[str, Any]] = None,
+        algorithm_config: Optional[Dict[str, Any]] = None,
+    ) -> ComputeHandle:
+        if not self._client:
+            raise RuntimeError("Docker not available, cannot start containers")
+
+        container_env = self._build_container_env(
+            run_id=run_id,
+            store_url=store_url,
+            env=env,
+            runner_config=runner_config,
+        )
+        network = self._detect_network()
+
+        # Store stop timeout for use in stop()
+        stop_timeout = (resource_limits or {}).get("stop_timeout")
+        if stop_timeout:
+            self._stop_timeouts[run_id] = int(stop_timeout)
+
+        identifiers: Dict[str, str] = {}
+        for i in range(count):
+            cid = self._start_one_container(
+                image=image,
+                run_id=run_id,
+                index=i,
+                container_env=container_env,
+                network=network,
+                labels=labels,
+                resource_limits=resource_limits,
+            )
+            if cid is None:
+                # Clean up any containers that were already started
+                for key, started_id in identifiers.items():
+                    self._stop_container(started_id)
+                raise RuntimeError(f"Failed to start container agent-{i} " f"({i}/{count} started successfully)")
+            identifiers[f"agent-{i}"] = cid
+
+        if count > 0:
+            logger.info(f"Started {count} agent container(s) for run {run_id} " f"(lifecycle={lifecycle.value})")
+
+        handle = ComputeHandle(
+            run_id=run_id,
+            backend=self.backend_name,
+            identifiers=identifiers,
+            metadata={
+                "lifecycle": lifecycle.value,
+                "desired_count": count,
+                "image": image,
+                "store_url": store_url,
+                "network": network,
+                "env": env or {},
+                "runner_config": runner_config or {},
+                "resource_limits": resource_limits or {},
+                "labels": labels or {},
+            },
+        )
+
+        # If a store is provided, set up reactive scaling automatically.
+        if store is not None:
+            self._setup_reactive_listener(
+                handle,
+                store,
+                n_runners_per_unit,
+                max_units,
+            )
+
+        return handle
+
+    def stop(self, handle: ComputeHandle, timeout: int = 30) -> None:
+        if not self._client:
+            return
+
+        # Remove reactive listener first so no more events fire.
+        self._teardown_reactive_listener(handle)
+
+        stopped = 0
+        for key, cid in list(handle.identifiers.items()):
+            self._stop_container(cid)
+            stopped += 1
+
+        if stopped:
+            logger.info(f"Killed and removed {stopped} container(s) for " f"run {handle.run_id}")
+
+    def status(self, handle: ComputeHandle) -> ComputeStatus:
+        if not self._client:
+            return ComputeStatus.UNKNOWN
+
+        statuses = []
+        for key, cid in handle.identifiers.items():
+            s = self._get_container_status(cid)
+            statuses.append(s)
+
+        if not statuses:
+            return ComputeStatus.UNKNOWN
+        if all(s == "running" for s in statuses):
+            return ComputeStatus.RUNNING
+        if any(s in ("exited", "dead") for s in statuses):
+            return ComputeStatus.FAILED
+        if all(s in ("created", "restarting") for s in statuses):
+            return ComputeStatus.PENDING
+        # Mixed state (some running, some starting) -- still consider running
+        if any(s == "running" for s in statuses):
+            return ComputeStatus.RUNNING
+        return ComputeStatus.UNKNOWN
+
+    def wait_ready(
+        self,
+        handle: ComputeHandle,
+        timeout: float = 120.0,
+    ) -> None:
+        if not self._client:
+            raise RuntimeError("Docker not available")
+
+        deadline = _time.monotonic() + timeout
+        poll_interval = 2.0
+        target_count = len(handle.identifiers)
+
+        while _time.monotonic() < deadline:
+            ready = 0
+            for key, cid in handle.identifiers.items():
+                s = self._get_container_status(cid)
+                if s == "running":
+                    ready += 1
+                elif s in ("exited", "dead"):
+                    raise RuntimeError(f"Container {key} ({cid[:12]}) exited before " f"becoming ready")
+
+            if ready >= target_count:
+                logger.info(f"All {target_count} container(s) ready for " f"run {handle.run_id}")
+                return
+
+            logger.debug(
+                f"Waiting for containers: {ready}/{target_count} ready "
+                f"({timeout - (_time.monotonic() - (deadline - timeout)):.0f}s remaining)"
+            )
+            _time.sleep(poll_interval)
+
+        raise RuntimeError(
+            f"Timeout ({timeout}s) waiting for {target_count} container(s) " f"to become ready for run {handle.run_id}"
+        )
+
+    def scale(self, handle: ComputeHandle, count: int) -> None:
+        if not self._client:
+            raise RuntimeError("Docker not available")
+
+        current = len(handle.identifiers)
+        if count == current:
+            return
+
+        meta = handle.metadata
+
+        if count > current:
+            # Scale up: start additional containers
+            container_env = self._build_container_env(
+                run_id=handle.run_id,
+                store_url=meta["store_url"],
+                env=meta.get("env"),
+                runner_config=meta.get("runner_config"),
+            )
+            network = meta.get("network", self._detect_network())
+
+            for i in range(current, count):
+                cid = self._start_one_container(
+                    image=meta["image"],
+                    run_id=handle.run_id,
+                    index=i,
+                    container_env=container_env,
+                    network=network,
+                    labels=meta.get("labels"),
+                    resource_limits=meta.get("resource_limits"),
+                )
+                if cid is None:
+                    logger.error(f"Failed to scale up: could not start agent-{i}")
+                    break
+                handle.identifiers[f"agent-{i}"] = cid
+
+            logger.info(f"Scaled up from {current} to {len(handle.identifiers)} " f"for run {handle.run_id}")
+        else:
+            # Scale down: stop highest-indexed containers
+            keys_sorted = sorted(
+                handle.identifiers.keys(),
+                key=lambda k: int(k.split("-")[-1]),
+                reverse=True,
+            )
+            to_remove = keys_sorted[: current - count]
+            for key in to_remove:
+                cid = handle.identifiers.pop(key)
+                self._stop_container(cid)
+
+            logger.info(f"Scaled down from {current} to {len(handle.identifiers)} " f"for run {handle.run_id}")
+
+        handle.metadata["desired_count"] = count
+
+    def replace_member(
+        self,
+        handle: ComputeHandle,
+        member_key: str,
+    ) -> Optional[str]:
+        if not self._client:
+            return None
+
+        old_cid = handle.identifiers.get(member_key)
+        if old_cid:
+            self._stop_container(old_cid)
+
+        meta = handle.metadata
+        index = int(member_key.split("-")[-1])
+
+        container_env = self._build_container_env(
+            run_id=handle.run_id,
+            store_url=meta["store_url"],
+            env=meta.get("env"),
+            runner_config=meta.get("runner_config"),
+        )
+        network = meta.get("network", self._detect_network())
+
+        new_cid = self._start_one_container(
+            image=meta["image"],
+            run_id=handle.run_id,
+            index=index,
+            container_env=container_env,
+            network=network,
+            labels=meta.get("labels"),
+            resource_limits=meta.get("resource_limits"),
+        )
+
+        if new_cid:
+            handle.identifiers[member_key] = new_cid
+            logger.info(
+                f"Replaced {member_key} ({old_cid[:12] if old_cid else '?'}) "
+                f"with {new_cid[:12]} for run {handle.run_id}"
+            )
+        else:
+            logger.error(f"Failed to replace {member_key} for run {handle.run_id}")
+
+        return new_cid
+
+    def members(self, handle: ComputeHandle) -> List[Dict[str, Any]]:
+        if not self._client:
+            return []
+
+        result = []
+        for key, cid in sorted(handle.identifiers.items()):
+            index = int(key.split("-")[-1])
+            status = self._get_container_status(cid) or "unknown"
+            result.append(
+                {
+                    "key": key,
+                    "id": cid,
+                    "short_id": cid[:12],
+                    "index": index,
+                    "status": status,
+                    "name": f"gateway-{handle.run_id[:8]}-agent-{index}",
+                }
+            )
+        return result
+
+    def get_logs(self, handle: ComputeHandle, tail: int = 100) -> str:
+        if not self._client:
+            return ""
+
+        parts: List[str] = []
+        for key, cid in sorted(handle.identifiers.items()):
+            try:
+                container = self._client.containers.get(cid)
+                logs = container.logs(tail=tail, timestamps=False)
+                if isinstance(logs, bytes):
+                    logs = logs.decode("utf-8", errors="replace")
+                parts.append(f"=== {key} ({cid[:12]}) ===\n{logs}")
+            except Exception as e:
+                parts.append(f"=== {key} ({cid[:12]}) === [error: {e}]")
+        return "\n".join(parts)
+
+    # ── Low-level container helpers ─────────────────────────────────
+
+    def _stop_container(self, container_id: str) -> None:
+        """Kill and remove a container immediately."""
+        if not self._client:
+            return
+        try:
+            container = self._client.containers.get(container_id)
+            try:
+                container.kill()
+            except Exception:
+                pass  # Already stopped / exited
+            container.remove(force=True)
+            logger.debug(f"Killed and removed container {container_id[:12]}")
+        except Exception as e:
+            logger.warning(f"Could not kill/remove container {container_id[:12]}: {e}")
+
+    def _get_container_status(self, container_id: str) -> Optional[str]:
+        """Get container status (running, exited, etc.)."""
+        if not self._client:
+            return None
+        try:
+            container = self._client.containers.get(container_id)
+            return container.status
+        except Exception as e:
+            logger.debug(f"Could not get container status: {e}")
+            return None
+
+    # ── Legacy helpers (used by stop_all_for_run fallback) ──────────
+
+    def list_run_containers(self, run_id: str) -> List[Dict[str, Any]]:
+        """List all containers associated with a run (by label)."""
+        if not self._client:
+            return []
+        try:
+            containers = self._client.containers.list(
+                all=True,
+                filters={"label": f"gateway.run_id={run_id}"},
+            )
+            return [
+                {
+                    "id": c.id,
+                    "short_id": c.short_id,
+                    "name": c.name,
+                    "status": c.status,
+                    "image": str(c.image),
+                }
+                for c in containers
+            ]
+        except Exception as e:
+            logger.debug(f"Could not list containers for run {run_id}: {e}")
+            return []
+
+    def stop_all_for_run(self, run_id: str) -> int:
+        """Kill all containers associated with a run (label-based fallback)."""
+        containers = self.list_run_containers(run_id)
+        stopped = 0
+        for c in containers:
+            self._stop_container(c["id"])
+            stopped += 1
+        return stopped
+
+    # ── Reactive scaling / ephemeral recycling ───────────────────
+
+    def _setup_reactive_listener(
+        self,
+        handle: ComputeHandle,
+        store: Any,
+        n_runners_per_unit: int,
+        max_units: int,
+    ) -> None:
+        """Register a store listener for on-demand scaling and ephemeral replacement.
+
+        Called automatically from ``start()`` when a ``store`` is provided.
+        """
+        import math
+
+        lifecycle_str = handle.metadata.get("lifecycle", "persistent")
+        is_ephemeral = lifecycle_str == ContainerLifecycle.EPHEMERAL.value
+        orch = self
+
+        class _ReactiveListener:
+            """Store listener that starts/replaces containers on demand."""
+
+            @property
+            def capabilities(self) -> Dict[str, bool]:
+                return {}
+
+            def otlp_traces_endpoint(self) -> Optional[str]:
+                return None
+
+            def get_otlp_headers(self) -> Dict[str, str]:
+                return {}
+
+            async def on_rollout_created(self, rollout: Any) -> None:
+                current_units = len(handle.identifiers)
+                # Scale to one more container than we currently have,
+                # but only if all existing slots are occupied.
+                # Don't call wait_ready() here -- it blocks the event
+                # loop and causes cascading scale-ups.  The container
+                # will become ready asynchronously; the runner inside
+                # polls the store and picks up work once healthy.
+                if current_units < max_units:
+                    needed = current_units + 1
+                    try:
+                        orch.scale(handle, needed)
+                        logger.info(
+                            f"Reactive: started container agent-{current_units} "
+                            f"for run {handle.run_id} "
+                            f"({needed} total)"
+                        )
+                    except Exception as e:
+                        logger.warning(f"Reactive scale-up failed for " f"run {handle.run_id}: {e}")
+
+            async def on_rollout_updated(self, rollout: Any) -> None:
+                if not is_ephemeral:
+                    return
+                status = getattr(rollout, "status", None)
+                if status not in ("succeeded", "failed"):
+                    return
+
+                # Replace the first container that is no longer running
+                # (exited after processing its rollout).
+                for key, cid in list(handle.identifiers.items()):
+                    cstatus = orch._get_container_status(cid)
+                    if cstatus in ("exited", "dead", None):
+                        try:
+                            new_cid = orch.replace_member(handle, key)
+                            if new_cid:
+                                logger.info(
+                                    f"Ephemeral recycle: replaced {key} "
+                                    f"({cid[:12]}) with {new_cid[:12]} "
+                                    f"for run {handle.run_id}"
+                                )
+                        except Exception as e:
+                            logger.warning(f"Ephemeral recycle failed for {key}: {e}")
+                        break
+
+        listener = _ReactiveListener()
+        if not hasattr(self, "_reactive_listeners"):
+            self._reactive_listeners: Dict[str, Any] = {}
+        self._reactive_listeners[handle.run_id] = (listener, store)
+
+        store.add_listener(listener)
+        logger.info(
+            f"Reactive listener registered for run {handle.run_id} "
+            f"(ephemeral={is_ephemeral}, n_runners={n_runners_per_unit}, "
+            f"max_units={max_units})"
+        )
+
+    def _teardown_reactive_listener(self, handle: ComputeHandle) -> None:
+        """Remove the reactive listener from the store."""
+        listeners = getattr(self, "_reactive_listeners", {})
+        entry = listeners.pop(handle.run_id, None)
+        if entry:
+            listener, store = entry
+            store.remove_listener(listener)
+            logger.debug(f"Reactive listener removed for run {handle.run_id}")