gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/store/redis_stream.py

@@ -0,0 +1,517 @@
+# Copyright (c) Gateway. All rights reserved.
+
+"""RedisStreamLightningStore — distributed rollout coordination via Redis Streams.
+
+This store wraps InMemoryLightningStore for all metadata (rollouts, attempts,
+spans, resources) and replaces only the rollout queue with Redis Streams.
+
+Redis Stream protocol:
+  - Stream key: ``gateway:run:{run_id}:rollouts``
+  - Consumer group: ``gateway-runners``
+  - Enqueue: ``XADD`` with rollout_id, input_json, resources_id, mode
+  - Dequeue: ``XREADGROUP GROUP gateway-runners {worker_id} COUNT 1 BLOCK {timeout_ms}``
+  - Ack: ``XACK`` after rollout reaches terminal status
+
+Usage::
+
+    store = RedisStreamLightningStore(
+        redis_url="redis://:myredissecret@localhost:6379",
+        run_id="run-abc123",
+    )
+    await store.initialize()
+
+    # Enqueue rollouts (algorithm side)
+    rollout = await store.enqueue_rollout(input={"task": "solve"})
+
+    # Dequeue rollouts (runner side)
+    attempted = await store.dequeue_rollout(worker_id="runner-1")
+
+    await store.close()
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+from typing import Any, Dict, List, Optional, Sequence
+
+from gatewaysdk.store.base import (
+    UNSET,
+    LightningStore,
+    LightningStoreCapabilities,
+    LightningStoreStatistics,
+    Unset,
+)
+from gatewaysdk.store.memory import InMemoryLightningStore
+from gatewaysdk.types import (
+    Attempt,
+    AttemptedRollout,
+    AttemptStatus,
+    EnqueueRolloutRequest,
+    NamedResources,
+    PaginatedResult,
+    ResourcesUpdate,
+    Rollout,
+    RolloutConfig,
+    RolloutMode,
+    RolloutStatus,
+    Span,
+    TaskInput,
+    Worker,
+    WorkerStatus,
+)
+from gatewaysdk.utils.redact import redact_url
+
+try:
+    import redis.asyncio as aioredis
+except ImportError:
+    aioredis = None  # type: ignore[assignment]
+
+logger = logging.getLogger(__name__)
+
+
+class RedisStreamLightningStore(LightningStore):
+    """DEPRECATED: Use LightningStoreServer + LightningStoreClient instead.
+
+    This store wraps InMemoryLightningStore with Redis Streams for distributed
+    rollout coordination. It has been superseded by the StoreServer/Client
+    architecture which avoids the rollout desync bug (separate memory stores
+    on algorithm and agent sides can't see each other's state).
+
+    For platform runs, the worker now creates an InMemoryLightningStore
+    wrapped in a LightningStoreServer. Agent containers connect via
+    LightningStoreClient (HTTP). See platform/worker.py AlgorithmContext.
+    """
+
+    @classmethod
+    def from_env(cls) -> "RedisStreamLightningStore":
+        """Create a store from environment variables set by DockerOrchestrator.
+
+        Expected env vars:
+            ``GATEWAY_RUN_ID``: The GatewayRun ID.
+            ``REDIS_HOST``: Redis hostname (default ``localhost``).
+            ``REDIS_PORT``: Redis port (default ``6379``).
+            ``REDIS_AUTH``: Redis password (optional).
+
+        Returns:
+            A configured but not-yet-initialized store.
+            Caller must ``await store.initialize()`` before use.
+
+        Raises:
+            EnvironmentError: If ``GATEWAY_RUN_ID`` is not set.
+        """
+        import os
+
+        run_id = os.environ.get("GATEWAY_RUN_ID")
+        if not run_id:
+            raise EnvironmentError(
+                "GATEWAY_RUN_ID environment variable is required. "
+                "This is set automatically by the Gateway platform orchestrator."
+            )
+        host = os.environ.get("REDIS_HOST", "localhost")
+        port = os.environ.get("REDIS_PORT", "6379")
+        auth = os.environ.get("REDIS_AUTH", "")
+        redis_url = f"redis://:{auth}@{host}:{port}" if auth else f"redis://{host}:{port}"
+        return cls(redis_url=redis_url, run_id=run_id)
+
+    def __init__(
+        self,
+        *,
+        redis_url: str,
+        run_id: str,
+        consumer_group: str = "gateway-runners",
+        thread_safe: bool = False,
+        listeners: Optional[Sequence] = None,
+        block_timeout_ms: int = 5000,
+    ):
+        """
+        Args:
+            redis_url: Redis connection string (e.g., redis://:password@host:6379/0).
+            run_id: The GatewayRun ID. Used to namespace the Redis stream key.
+            consumer_group: Consumer group name for XREADGROUP.
+            thread_safe: Whether the underlying InMemoryLightningStore is thread-safe.
+            listeners: StorageListener instances to attach.
+            block_timeout_ms: Timeout for XREADGROUP BLOCK (milliseconds).
+        """
+        self._redis_url = redis_url
+        self._run_id = run_id
+        self._consumer_group = consumer_group
+        self._block_timeout_ms = block_timeout_ms
+        self._redis = None
+        self._stream_key = f"gateway:run:{run_id}:rollouts"
+
+        # Create the inner in-memory store for all metadata
+        self._inner = InMemoryLightningStore(
+            thread_safe=thread_safe,
+            listeners=list(listeners) if listeners else None,
+        )
+
+    @property
+    def run_id(self) -> str:
+        return self._run_id
+
+    @property
+    def capabilities(self) -> LightningStoreCapabilities:
+        caps = self._inner.capabilities
+        caps["async_safe"] = True
+        caps["thread_safe"] = True  # Redis handles concurrency at server level
+        return caps
+
+    async def initialize(self) -> None:
+        """Create Redis connection and consumer group."""
+        if aioredis is None:
+            raise ImportError(
+                "redis[asyncio] is required for RedisStreamLightningStore. Install with: pip install redis>=5.0.0"
+            )
+
+        try:
+            self._redis = aioredis.from_url(self._redis_url, decode_responses=True)
+            await self._redis.ping()
+        except Exception as e:
+            self._redis = None
+            raise ConnectionError(f"Failed to connect to Redis at {redact_url(self._redis_url)}: {e}") from e
+
+        # Create consumer group (MKSTREAM creates the stream if it doesn't exist)
+        try:
+            await self._redis.xgroup_create(self._stream_key, self._consumer_group, id="0", mkstream=True)
+            logger.info(f"Created consumer group '{self._consumer_group}' on stream '{self._stream_key}'")
+        except Exception as e:
+            if "BUSYGROUP" in str(e):
+                logger.debug(f"Consumer group '{self._consumer_group}' already exists")
+            else:
+                raise
+
+    async def close(self, *, cleanup: bool = True) -> None:
+        """Close Redis connection and optionally clean up stream resources.
+
+        Args:
+            cleanup: If True (default), delete the stream key and consumer group
+                to prevent Redis memory growth from orphaned streams.
+        """
+        if self._redis:
+            if cleanup:
+                try:
+                    await self._redis.xgroup_destroy(self._stream_key, self._consumer_group)
+                    logger.debug(f"Destroyed consumer group '{self._consumer_group}' " f"on '{self._stream_key}'")
+                except Exception:
+                    pass  # Group may not exist or Redis may be closing
+
+                try:
+                    await self._redis.delete(self._stream_key)
+                    logger.debug(f"Deleted stream key '{self._stream_key}'")
+                except Exception:
+                    pass
+
+                # Clean up cancellation signal key if it exists
+                try:
+                    await self._redis.delete(f"gateway:run:{self._run_id}:cancel")
+                except Exception:
+                    pass
+
+            await self._redis.aclose()
+            self._redis = None
+
+    # =========================================================================
+    # Rollout Queue — Redis Stream backed
+    # =========================================================================
+
+    async def enqueue_rollout(
+        self,
+        input: TaskInput,
+        *,
+        mode: Optional[RolloutMode] = None,
+        resources_id: Optional[str] = None,
+        config: Optional[RolloutConfig] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> Rollout:
+        """Create rollout in memory, then XADD to Redis stream."""
+        # Create rollout in the inner store (status="queuing")
+        rollout = await self._inner.enqueue_rollout(
+            input=input,
+            mode=mode,
+            resources_id=resources_id,
+            config=config,
+            metadata=metadata,
+        )
+
+        # Add to Redis stream for distributed consumption
+        if self._redis:
+            try:
+                await self._redis.xadd(
+                    self._stream_key,
+                    {
+                        "rollout_id": rollout.rollout_id,
+                        "input_json": json.dumps(input) if not isinstance(input, str) else input,
+                        "resources_id": resources_id or "",
+                        "mode": mode or "",
+                    },
+                )
+                logger.debug(f"XADD rollout {rollout.rollout_id} to {self._stream_key}")
+            except Exception as e:
+                logger.warning(
+                    f"Failed to XADD rollout {rollout.rollout_id} to Redis: {e}. "
+                    f"Rollout exists in memory but not on Redis stream."
+                )
+                raise RuntimeError(f"Failed to enqueue rollout {rollout.rollout_id} to Redis stream: {e}") from e
+
+        return rollout
+
+    async def enqueue_rollouts(
+        self,
+        requests: Sequence[EnqueueRolloutRequest],
+    ) -> List[Rollout]:
+        """Batch enqueue rollouts with Redis pipelining (F12)."""
+        import json as _json
+
+        # 1. Create all rollouts in inner (memory) store
+        rollouts = []
+        for req in requests:
+            rollout = await self._inner.enqueue_rollout(
+                input=req.input,
+                mode=req.mode,
+                resources_id=req.resources_id,
+                config=req.config,
+                metadata=req.metadata,
+            )
+            rollouts.append(rollout)
+
+        # 2. Pipeline all XADD commands in one round-trip
+        if self._redis and rollouts:
+            pipe = self._redis.pipeline()
+            for rollout, req in zip(rollouts, requests):
+                pipe.xadd(
+                    self._stream_key,
+                    {
+                        "rollout_id": rollout.rollout_id,
+                        "input_json": _json.dumps(req.input) if not isinstance(req.input, str) else req.input,
+                        "resources_id": req.resources_id or "",
+                        "mode": req.mode or "",
+                    },
+                )
+            try:
+                await pipe.execute()
+            except Exception as e:
+                logger.error(f"Failed to pipeline enqueue {len(rollouts)} rollouts: {e}")
+                raise RuntimeError(f"Failed to pipeline enqueue {len(rollouts)} rollouts: {e}") from e
+
+        return rollouts
+
+    async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
+        """XREADGROUP from Redis stream, then create attempt via inner store."""
+        if not self._redis:
+            return await self._inner.dequeue_rollout(worker_id=worker_id)
+
+        consumer_name = worker_id or f"consumer-{id(self)}"
+
+        # Read one message from the stream
+        try:
+            result = await self._redis.xreadgroup(
+                self._consumer_group,
+                consumer_name,
+                {self._stream_key: ">"},
+                count=1,
+                block=self._block_timeout_ms,
+            )
+        except Exception as e:
+            logger.error(f"Redis XREADGROUP failed: {e}. No fallback in distributed mode.")
+            return None
+
+        if not result:
+            return None
+
+        # result format: [[stream_key, [(message_id, {field: value})]]]
+        stream_messages = result[0][1]
+        if not stream_messages:
+            return None
+
+        message_id, fields = stream_messages[0]
+        rollout_id = fields.get("rollout_id")
+
+        if not rollout_id:
+            logger.warning(f"Dequeued message without rollout_id: {fields}")
+            await self._redis.xack(self._stream_key, self._consumer_group, message_id)
+            return None
+
+        # Use the inner store's dequeue to create the attempt
+        # The rollout is already in memory from enqueue_rollout()
+        attempted = await self._inner.dequeue_rollout(worker_id=worker_id)
+
+        if attempted and attempted.rollout.rollout_id != rollout_id:
+            logger.warning(
+                f"Redis delivered rollout {rollout_id} but inner store "
+                f"returned {attempted.rollout.rollout_id}. Possible desync."
+            )
+
+        if attempted:
+            # Store message_id for later ACK
+            attempted.rollout._redis_message_id = message_id  # type: ignore[attr-defined]
+            logger.debug(f"Dequeued rollout {rollout_id} for worker {consumer_name}")
+        else:
+            # Rollout was in Redis but not in memory — ACK it to prevent redelivery
+            await self._redis.xack(self._stream_key, self._consumer_group, message_id)
+            logger.warning(f"Rollout {rollout_id} in Redis stream but not in memory store")
+
+        return attempted
+
+    async def start_rollout(
+        self,
+        input: TaskInput,
+        *,
+        mode: Optional[RolloutMode] = None,
+        resources_id: Optional[str] = None,
+        config: Optional[RolloutConfig] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> AttemptedRollout:
+        """Create rollout + attempt immediately (bypasses queue)."""
+        return await self._inner.start_rollout(
+            input=input,
+            mode=mode,
+            resources_id=resources_id,
+            config=config,
+            metadata=metadata,
+        )
+
+    async def wait_for_rollouts(
+        self,
+        *,
+        rollout_ids: Optional[List[str]] = None,
+        timeout: Optional[float] = None,
+    ) -> List[Rollout]:
+        """Wait for rollouts to reach terminal status."""
+        return await self._inner.wait_for_rollouts(rollout_ids=rollout_ids, timeout=timeout)
+
+    # =========================================================================
+    # Rollout/Attempt updates — ACK on terminal status
+    # =========================================================================
+
+    async def update_rollout(
+        self,
+        rollout_id: str,
+        *,
+        status: Optional[RolloutStatus] = None,
+        input: TaskInput | Unset = UNSET,
+        resources_id: Optional[str] | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Rollout:
+        """Update rollout, ACK Redis message on terminal status."""
+        rollout = await self._inner.update_rollout(
+            rollout_id,
+            status=status,
+            input=input,
+            resources_id=resources_id,
+            metadata=metadata,
+        )
+
+        # ACK the Redis message when rollout reaches terminal status
+        if self._redis and status and status in ("succeeded", "failed", "cancelled"):
+            msg_id = getattr(rollout, "_redis_message_id", None)
+            if msg_id:
+                try:
+                    await self._redis.xack(self._stream_key, self._consumer_group, msg_id)
+                    logger.debug(f"XACK rollout {rollout_id} (status={status})")
+                except Exception as e:
+                    logger.warning(f"Failed to XACK rollout {rollout_id}: {e}")
+
+        return rollout
+
+    async def update_attempt(
+        self,
+        rollout_id: str,
+        attempt_id: str,
+        *,
+        status: Optional[AttemptStatus] = None,
+        worker_id: Optional[str] | Unset = UNSET,
+        heartbeat: bool = False,
+    ) -> Attempt:
+        """Update attempt via inner store."""
+        return await self._inner.update_attempt(
+            rollout_id,
+            attempt_id,
+            status=status,
+            worker_id=worker_id,
+            heartbeat=heartbeat,
+        )
+
+    # =========================================================================
+    # All other methods — delegate to InMemoryLightningStore
+    # =========================================================================
+
+    async def add_span(self, span: Span) -> Optional[Span]:
+        return await self._inner.add_span(span)
+
+    async def add_otel_span(
+        self,
+        rollout_id: str,
+        attempt_id: str,
+        readable_span: Any,
+    ) -> Optional[Span]:
+        return await self._inner.add_otel_span(rollout_id, attempt_id, readable_span)
+
+    async def query_rollouts(self, **kwargs: Any) -> PaginatedResult[Rollout]:
+        return await self._inner.query_rollouts(**kwargs)
+
+    async def query_attempts(self, rollout_id: str, **kwargs: Any) -> PaginatedResult[Attempt]:
+        return await self._inner.query_attempts(rollout_id, **kwargs)
+
+    async def query_spans(self, rollout_id: str, attempt_id: str, **kwargs: Any) -> PaginatedResult[Span]:
+        return await self._inner.query_spans(rollout_id, attempt_id, **kwargs)
+
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
+        return await self._inner.get_rollout_by_id(rollout_id)
+
+    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
+        return await self._inner.get_latest_attempt(rollout_id)
+
+    async def add_resources(
+        self,
+        resources: ResourcesUpdate,
+        *,
+        default: bool = False,
+    ) -> ResourcesUpdate:
+        return await self._inner.add_resources(resources)
+
+    async def update_resources(
+        self,
+        resources_id: str,
+        resources: ResourcesUpdate,
+        *,
+        default: bool = False,
+    ) -> ResourcesUpdate:
+        return await self._inner.update_resources(resources_id, resources)
+
+    async def get_latest_resources(
+        self,
+        resource_type: Optional[str] = None,
+    ) -> Optional[NamedResources]:
+        return await self._inner.get_latest_resources(resource_type=resource_type)
+
+    async def get_resources_by_id(self, resources_id: str) -> Optional[NamedResources]:
+        return await self._inner.get_resources_by_id(resources_id)
+
+    def get_statistics(self) -> LightningStoreStatistics:
+        return self._inner.get_statistics()
+
+    async def register_worker(
+        self,
+        worker_id: str,
+        *,
+        status: WorkerStatus = "idle",
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> Worker:
+        return await self._inner.register_worker(worker_id, status=status, metadata=metadata)
+
+    async def update_worker(
+        self,
+        worker_id: str,
+        *,
+        status: Optional[WorkerStatus] = None,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Worker:
+        return await self._inner.update_worker(worker_id, status=status, metadata=metadata)
+
+    async def get_worker(self, worker_id: str) -> Optional[Worker]:
+        return await self._inner.get_worker(worker_id)
+
+    async def query_workers(self, **kwargs: Any) -> PaginatedResult[Worker]:
+        return await self._inner.query_workers(**kwargs)

gatewaysdk/store/sqlite.py

@@ -0,0 +1,3 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# TODO: Implement this