diff --git a/dev-resources/radex-config.toml b/dev-resources/radex-config.toml index d41daa8..d3d8817 100644 --- a/dev-resources/radex-config.toml +++ b/dev-resources/radex-config.toml @@ -4,10 +4,12 @@ addopts = [ ] strict_config = true strict_markers = true +asyncio_mode = "auto" markers = [ "slow: Test may be slow to run", "compiled: Test has a component compiled at run time", "example: Test is running one of the examples", + "redis: Test requires a real redis-server binary on PATH" ] [tool.black] diff --git a/dev-resources/requirements-dev.txt b/dev-resources/requirements-dev.txt index 19e68a3..3b273c7 100644 --- a/dev-resources/requirements-dev.txt +++ b/dev-resources/requirements-dev.txt @@ -1,3 +1,5 @@ pytest +pytest-asyncio black -isort \ No newline at end of file +isort +redis diff --git a/example/py-store-exchange/dragon/driver.py b/example/py-store-exchange/dragon/driver.py new file mode 100644 index 0000000..862c7fd --- /dev/null +++ b/example/py-store-exchange/dragon/driver.py @@ -0,0 +1,118 @@ +""" +DragonStore-managed data exchange -- single-file example. + +Architecture +──────────── + This mirrors example/py-cpp-exchange/dragon/driver.py, but instead of the + driver manually constructing/serializing/destroying a `dragon.data.ddict.DDict` + itself, RADEX's `DragonStore` owns that lifecycle: + + DragonStore.start() -> constructs the DDict, blocks until ready + DragonStore.endpoints -> [DragonEndpoint(descriptor=...)] + DragonStore.shutdown() -> destroys the DDict + + DragonStore never constructs a client itself -- a real + radex.clients.core.DragonClient is built directly from the endpoint's + serialized descriptor, below. + + The serialized descriptor (`store.endpoints[0].serialize()`) is what you'd + hand to a separately-launched process (env var, task kwarg, etc.) -- the + Store itself never touches os.environ, so that handoff is always explicit. + +Run with: + dragon -s -- python driver.py +""" + +import asyncio +import os +import pathlib +import time + +import numpy as np +from dragon.native.process import Process, ProcessTemplate + +from radex import DragonStore +from radex.clients.core import DragonClient +from radex.handles.handles import IncomingHandle, OutgoingHandle + +HERE = pathlib.Path(__file__).parent.absolute() +ROOT = HERE.parent.parent.parent +EXAMPLES_BIN_DIR = ROOT / "install" / "bin" / "examples" + + +async def main() -> int: + # ── 1. RADEX starts and owns the DDict-backed store ───────────────────── + store = DragonStore(managers_per_node=1, n_nodes=1) + await store.start() + print(f"[Driver] DragonStore ready: {store.endpoints[0].serialize()[:32]}...") + + # ── 2. RADEX client, built directly from the endpoint's descriptor -- + # Store never constructs clients itself ───────────────────────── + client = DragonClient(descriptor=store.endpoints[0].serialize(), timeout=5) + + try: + # ── 3. Hand the serialized descriptor to a separately-launched + # process -- the store never sets env vars for you. + app_tmpl = ProcessTemplate( + target=os.fspath(EXAMPLES_BIN_DIR / "dragon-cpp-with-py"), + env={"SERIALIZED_DDICT": store.endpoints[0].serialize()}, + ) + app = Process.from_template(app_tmpl) + + print("[Driver] Starting C++ app") + app.start() + try: + time.sleep(3) + print("[Driver] Setting Int") + client.put_scalar(OutgoingHandle("py-int"), 123) + + time.sleep(3) + print("[Driver] Setting Double") + client.put_scalar(OutgoingHandle("py-double"), 9.87) + + time.sleep(3) + print("[Driver] Setting Numpy Float") + client.put_scalar(OutgoingHandle("py-np-float"), np.float32(45.6)) + + time.sleep(3) + print("[Driver] Setting Int Tensor") + client.put_tensor( + OutgoingHandle("py-int-tensor"), np.arange(4, dtype=np.int32) + ) + + time.sleep(3) + print("[Driver] Setting Float Tensor") + client.put_tensor( + OutgoingHandle("py-float-tensor"), + np.arange(12, dtype=np.float64).reshape((6, 2)), + ) + + print("[Driver] Looking for keys") + print_scalar(client, "cpp-double") + print_scalar(client, "cpp-int") + print_tensor(client, "cpp-double-tensor") + print_tensor(client, "cpp-long-tensor") + finally: + app.join() + finally: + # ── 4. RADEX owns teardown too -- idempotent, safe to call again. + await store.shutdown() + print(f"[Driver] Store state: {store.state.name}") + + return 0 + + +def print_scalar(client, key): + print(f"[Driver] Waiting for scalar key `{key}`") + scalar = client.wait_for_scalar(IncomingHandle(key), 10) + print(f"[Driver] Got scalar: {scalar}") + + +def print_tensor(client, key): + print(f"[Driver] Waiting for tensor key `{key}`") + tensor = client.wait_for_tensor(IncomingHandle(key), 10) + print(f"[Driver] Got tensor: {tensor.ravel()}") + + +if __name__ == "__main__": + raise SystemExit(asyncio.run(main())) diff --git a/example/py-store-exchange/redis/driver.py b/example/py-store-exchange/redis/driver.py new file mode 100644 index 0000000..d5094ea --- /dev/null +++ b/example/py-store-exchange/redis/driver.py @@ -0,0 +1,110 @@ +""" +RedisStore-managed data exchange -- single-file example. + +Architecture +──────────── + RedisStore models N *independent* per-node Redis instances -- not a Redis + Cluster. Each node is its own isolated keyspace; `store.endpoints` is a + list with one RedisEndpoint per node. + + RedisStore never constructs a client itself. The typed RADEX client + (`radex.clients.core.RedisClient` -- `put_scalar`/`get_scalar`/ + `put_tensor`/`get_tensor`, the same API `DragonClient` exposes) only + supports env-based construction (no host/port constructor args), so a + client for one specific endpoint is built by pointing + `RADEX_STORE`/`RADEX_STORE_OPTS` at that endpoint first, then + constructing `RedisClient()`: + + os.environ["RADEX_STORE"] = endpoint.serialize() + os.environ["RADEX_STORE_OPTS"] = "Standalone" + client = RedisClient() + + (For a raw redis-py client instead -- direct SET/GET, not RADEX's typed + API -- use `RedisStore.client(index=...)`.) + + Locally, RedisStore() with no arguments spawns a single `redis-server` on + an auto-picked free port -- this is what the first half of this example + uses. The second half shows the same API scaled out to several + independent local nodes, and how you'd point it at an HPC launcher + instead (commented out, since it needs a real Slurm allocation to run): + + RedisStore( + hosts=["nid00001", "nid00002", "nid00003"], + port=6380, + cmd="srun --nodelist={host} redis-server --port {port}", + ) + +Run with: + python driver.py +""" + +import asyncio +import os + +import numpy as np + +from radex import RedisStore +from radex.clients.core import RedisClient +from radex.handles.handles import IncomingHandle, OutgoingHandle +from radex.store.redis_store import RedisEndpoint + + +def client_for(endpoint: RedisEndpoint) -> RedisClient: + """Build a typed RADEX client bound to one specific endpoint. + + `RedisClient()` only constructs from the environment, so this points + `RADEX_STORE`/`RADEX_STORE_OPTS` at `endpoint` first. + """ + os.environ["RADEX_STORE"] = endpoint.serialize() + os.environ["RADEX_STORE_OPTS"] = "Standalone" + return RedisClient() + + +async def single_node_demo() -> None: + print("── Single local node ──────────────────────────────────────────") + store = RedisStore() + await store.start() + try: + endpoint = store.endpoints[0] + print(f"[Driver] RedisStore ready at {endpoint.serialize()}") + + client = client_for(endpoint) + + client.put_scalar(OutgoingHandle("greeting-count"), 1) + count = client.get_scalar(IncomingHandle("greeting-count")) + print(f"[Driver] Got scalar back: {count}") + + client.put_tensor(OutgoingHandle("samples"), np.arange(6, dtype=np.float64)) + samples = client.get_tensor(IncomingHandle("samples")) + print(f"[Driver] Got tensor back: {samples}") + + print(f"[Driver] ready(): {await store.ready()}") + finally: + await store.shutdown() + print(f"[Driver] Store state: {store.state.name}") + + +async def multi_node_demo() -> None: + print("\n── Multiple independent local nodes ───────────────────────────") + store = RedisStore(hosts=["localhost", "localhost", "localhost"]) + await store.start() + try: + print(f"[Driver] {len(store.endpoints)} independent nodes:") + for i, endpoint in enumerate(store.endpoints): + # Each node is its own keyspace -- write a distinct value to each. + client = client_for(endpoint) + client.put_scalar(OutgoingHandle("node-id"), i) + node_id = client.get_scalar(IncomingHandle("node-id")) + print(f"[Driver] node {i}: {endpoint.serialize()} -> node-id={node_id}") + finally: + await store.shutdown() + + +async def main() -> int: + await single_node_demo() + await multi_node_demo() + return 0 + + +if __name__ == "__main__": + raise SystemExit(asyncio.run(main())) diff --git a/src/python/pyproject.toml b/src/python/pyproject.toml index 91ef8d5..015cd2d 100644 --- a/src/python/pyproject.toml +++ b/src/python/pyproject.toml @@ -13,3 +13,7 @@ dependencies = [ "cloudpickle", "numpy", ] + +[project.optional-dependencies] +redis = ["redis>=5.0"] +dragon = [] diff --git a/src/python/src/radex/__init__.py b/src/python/src/radex/__init__.py index e69de29..b617e50 100644 --- a/src/python/src/radex/__init__.py +++ b/src/python/src/radex/__init__.py @@ -0,0 +1,29 @@ +from radex.store import ( + DragonEndpoint, + DragonStore, + Endpoint, + RedisEndpoint, + RedisStore, + Store, + StoreError, + StoreNotReadyError, + StoreStartupError, + StoreState, + StoreStateError, + StoreTerminatedError, +) + +__all__ = [ + "DragonEndpoint", + "DragonStore", + "Endpoint", + "RedisEndpoint", + "RedisStore", + "Store", + "StoreError", + "StoreNotReadyError", + "StoreStartupError", + "StoreState", + "StoreStateError", + "StoreTerminatedError", +] diff --git a/src/python/src/radex/store/__init__.py b/src/python/src/radex/store/__init__.py new file mode 100644 index 0000000..f0fcd9c --- /dev/null +++ b/src/python/src/radex/store/__init__.py @@ -0,0 +1,27 @@ +from radex.store.base import ( + Endpoint, + Store, + StoreError, + StoreNotReadyError, + StoreStartupError, + StoreState, + StoreStateError, + StoreTerminatedError, +) +from radex.store.dragon_store import DragonEndpoint, DragonStore +from radex.store.redis_store import RedisEndpoint, RedisStore + +__all__ = [ + "DragonEndpoint", + "DragonStore", + "Endpoint", + "RedisEndpoint", + "RedisStore", + "Store", + "StoreError", + "StoreNotReadyError", + "StoreStartupError", + "StoreState", + "StoreStateError", + "StoreTerminatedError", +] diff --git a/src/python/src/radex/store/base.py b/src/python/src/radex/store/base.py new file mode 100644 index 0000000..cf0281d --- /dev/null +++ b/src/python/src/radex/store/base.py @@ -0,0 +1,146 @@ +from __future__ import annotations + +import abc +import asyncio +import enum + + +class StoreState(enum.Enum): + CREATED = "CREATED" + STARTING = "STARTING" + READY = "READY" + FAILED = "FAILED" + SHUTDOWN = "SHUTDOWN" + + +_TERMINAL_STATES = frozenset({StoreState.FAILED, StoreState.SHUTDOWN}) + + +class StoreError(Exception): + """Base class for all radex.store errors.""" + + +class StoreStartupError(StoreError): + """Raised when a Store fails to reach READY during start().""" + + +class StoreStateError(StoreError): + """Raised when a Store method is invoked in an invalid lifecycle state.""" + + +class StoreNotReadyError(StoreStateError): + """Raised by `.endpoints` before a successful start().""" + + +class StoreTerminatedError(StoreStateError): + """Raised by start() on a Store that is already FAILED or SHUTDOWN.""" + + +class Endpoint(abc.ABC): + """Backend-specific connection information for one Store instance/node. + + Deliberately generic: a Redis node has a host/port, a Dragon DDict has an + opaque serialized descriptor, and neither concept is assumed by callers + that only depend on this interface. + """ + + @abc.abstractmethod + def serialize(self) -> str: + """Return a string a caller can pass through explicitly (env var, + kwarg, task arg) to reconnect a client elsewhere. Store never sets + this into os.environ itself -- callers own that decision.""" + + +class Store(abc.ABC): + """Backend-independent lifecycle for a RADEX data-exchange backend. + + State machine: CREATED -> STARTING -> {READY, FAILED}; + {CREATED, READY, FAILED} -> SHUTDOWN. FAILED and SHUTDOWN are terminal -- + a Store cannot be restarted once it lands in either; construct a new + instance instead. + + start()/shutdown() each run their body under a single per-instance lock, + which is what makes repeated/concurrent calls idempotent and the + terminal-state guard correct without any state duplicated in subclasses. + """ + + def __init__(self) -> None: + self._state: StoreState = StoreState.CREATED + self._endpoints: list[Endpoint] = [] + self._lock = asyncio.Lock() + + @property + def state(self) -> StoreState: + return self._state + + @property + def endpoints(self) -> list[Endpoint]: + if self._state is not StoreState.READY: + raise StoreNotReadyError( + f"{type(self).__name__} is not ready (state={self._state.name}); " + "call `await store.start()` first." + ) + return list(self._endpoints) + + async def ready(self) -> bool: + if self._state is not StoreState.READY: + return False + return await self._do_ready() + + async def start(self, wait: bool = True) -> "Store": + """Start the backend and return self, so both + `await store.start()` and `store = await RedisStore(...).start()` + work.""" + async with self._lock: + if self._state in _TERMINAL_STATES: + raise StoreTerminatedError( + f"{type(self).__name__} is in terminal state " + f"{self._state.name} and cannot be started again; " + "construct a new Store instance instead." + ) + if self._state is StoreState.READY: + return self + self._state = StoreState.STARTING + try: + endpoints = await self._do_start(wait=wait) + except BaseException as exc: + self._state = StoreState.FAILED + self._endpoints = [] + if isinstance(exc, StoreError) or not isinstance(exc, Exception): + # Already a StoreError, or a BaseException we must not + # mask (CancelledError, KeyboardInterrupt, SystemExit). + raise + raise StoreStartupError( + f"{type(self).__name__} failed to start" + ) from exc + else: + self._endpoints = list(endpoints) + self._state = StoreState.READY + return self + + async def shutdown(self) -> "Store": + """Tear down the backend and return self, for the same fluent + usage as start().""" + async with self._lock: + if self._state is StoreState.SHUTDOWN: + return self + try: + await self._do_shutdown() + finally: + self._state = StoreState.SHUTDOWN + self._endpoints = [] + return self + + @abc.abstractmethod + async def _do_start(self, wait: bool) -> list[Endpoint]: + """Launch the backend and return its endpoint(s). Raise on failure + after cleaning up any partially-started state.""" + + @abc.abstractmethod + async def _do_shutdown(self) -> None: + """Tear down the backend. Must tolerate being called with empty or + partial internal state (never-started, or a failed start).""" + + @abc.abstractmethod + async def _do_ready(self) -> bool: + """Live readiness check, never cached.""" diff --git a/src/python/src/radex/store/dragon_store.py b/src/python/src/radex/store/dragon_store.py new file mode 100644 index 0000000..16d517c --- /dev/null +++ b/src/python/src/radex/store/dragon_store.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +import asyncio +import dataclasses +from typing import Any, Mapping + +from radex.store.base import Endpoint, Store, StoreStartupError + +try: + import dragon # noqa: F401 + from dragon.data.ddict import DDict as _DDict +except ImportError: # pragma: no cover - environment without Dragon + dragon = None + _DDict = None + + +_PROBE_KEY = "__radex_store_liveness_probe__" + + +@dataclasses.dataclass(frozen=True) +class DragonEndpoint(Endpoint): + """Connection information for a Dragon DDict-backed store. + + Unlike Redis, there is no host/port here -- `serialize()` returns the + opaque base64 descriptor produced by `DDict.serialize()`, which is what + both the compiled `radex::drg::ddict::Client` and `DDict.attach()` + expect. `DragonEndpoint` never constructs a client itself -- build one + directly from the descriptor, e.g.:: + + from radex.clients.core import DragonClient + client = DragonClient(descriptor=endpoint.serialize(), timeout=5) + """ + + descriptor: str + + def serialize(self) -> str: + return self.descriptor + + +class DragonStore(Store): + """A Store backed by a single Dragon DDict. + + Constructing a `dragon.data.ddict.DDict` is itself the blocking startup + call -- it spins up the orchestrator and manager processes and blocks + until ready, with no separate "start" step. This wraps that blocking + call in `asyncio.to_thread`, matching Dragon's own established + convention for calling blocking Dragon primitives from asyncio. + + `wait_for_keys` is hard-enforced to True: the compiled + `radex::drg::ddict::Client` (`include/radex/dragon.hpp`, `_validate_ddict` + in `src/cpp/dragon.cpp`) refuses to attach to a DDict created with + `wait_for_keys=False`, so any other value would guarantee every + downstream C++/Cython client fails to construct. + + Client lifetime note: drop/`del` any `DragonClient` constructed from + this store's endpoint before calling `shutdown()`. Destroying the DDict + first and letting a client outlive it is harmless but prints a + DRAGON_OBJECT_DESTROYED message from the compiled client's destructor + when it later tries to detach. + """ + + def __init__( + self, + managers_per_node: int = 1, + n_nodes: int = 1, + total_mem: int | None = None, + *, + working_set_size: int = 2, + wait_for_keys: bool = True, + wait_for_writers: bool = False, + policy: Any = None, + managers_per_policy: int = 1, + orc_policy: Any = None, + persist_freq: int = 0, + name: str = "", + timeout: float | None = None, + trace: bool = False, + streams_per_manager: int = 5, + manager_pool_full_thresh: float = 0.9, + extra_ddict_kwargs: Mapping[str, Any] | None = None, + ) -> None: + super().__init__() + if _DDict is None: + raise ImportError( + "The 'dragon' package is required to use DragonStore. It is " + "not pip-installable from PyPI; install it per your Dragon " + "distribution/environment first." + ) + if wait_for_keys is not True: + raise ValueError( + "DragonStore requires wait_for_keys=True: the compiled " + "radex::drg::ddict::Client (include/radex/dragon.hpp, " + "_validate_ddict) refuses to attach to a DDict created with " + "wait_for_keys=False, so every downstream C++/Cython client " + "would immediately fail to construct." + ) + if working_set_size < 2: + raise ValueError( + "DragonStore requires working_set_size >= 2 when combined " + "with the (always-forced) wait_for_keys=True: DDict itself " + "rejects wait_for_keys=True with working_set_size < 2." + ) + + self._ddict_kwargs: dict[str, Any] = dict( + managers_per_node=managers_per_node, + n_nodes=n_nodes, + working_set_size=working_set_size, + wait_for_keys=True, + wait_for_writers=wait_for_writers, + policy=policy, + managers_per_policy=managers_per_policy, + orc_policy=orc_policy, + persist_freq=persist_freq, + name=name, + timeout=timeout, + trace=trace, + streams_per_manager=streams_per_manager, + manager_pool_full_thresh=manager_pool_full_thresh, + ) + if total_mem is not None: + self._ddict_kwargs["total_mem"] = total_mem + self._ddict_kwargs.update(extra_ddict_kwargs or {}) + self._ddict: Any = None + + async def _do_start(self, wait: bool) -> list[Endpoint]: + # `wait` is accepted for interface parity with Store.start() but has + # no effect: DDict.__init__ is already atomically blocking-until- + # ready -- there is no separate "start" step to skip waiting on. + try: + self._ddict = await asyncio.to_thread(_DDict, **self._ddict_kwargs) + except Exception as exc: + self._ddict = None + raise StoreStartupError( + f"DragonStore failed to construct DDict: {exc}" + ) from exc + return [DragonEndpoint(descriptor=self._ddict.serialize())] + + async def _do_shutdown(self) -> None: + if self._ddict is not None: + await asyncio.to_thread(self._ddict.destroy) + self._ddict = None + + async def _do_ready(self) -> bool: + if self._ddict is None: + return False + + def _probe() -> bool: + try: + _ = _PROBE_KEY in self._ddict + return True + except Exception: + return False + + return await asyncio.to_thread(_probe) diff --git a/src/python/src/radex/store/redis_store.py b/src/python/src/radex/store/redis_store.py new file mode 100644 index 0000000..734fc37 --- /dev/null +++ b/src/python/src/radex/store/redis_store.py @@ -0,0 +1,258 @@ +from __future__ import annotations + +import asyncio +import dataclasses +import shlex +import shutil +import socket +import time +from typing import Any, Mapping, Sequence + +from radex.store.base import Endpoint, Store, StoreStartupError + +try: + import redis +except ImportError: # pragma: no cover - environment without redis-py + redis = None + + +_PING = b"*1\r\n$4\r\nPING\r\n" +_PONG_PREFIXES = (b"+PONG", b"+PONG\r\n") + + +@dataclasses.dataclass(frozen=True) +class RedisEndpoint(Endpoint): + """Connection information for one independent Redis node. + + `serialize()` returns a `host:port` string compatible with the SmartRedis + SSDB env-var format read by the existing compiled + `radex::redis::smartredis::Client` (`include/radex/smartredis.hpp`) -- + RadexStore never sets this into os.environ itself; wire it through + yourself if you need the typed RADEX client to attach to a node this + Store launched, e.g.:: + + os.environ["RADEX_STORE"] = endpoint.serialize() + os.environ["RADEX_STORE_OPTS"] = "" + from radex.clients.core import RedisClient + client = RedisClient() + + `RedisEndpoint` never constructs a client itself -- for a raw + redis-py client instead, use `RedisStore.client()`. + """ + + host: str + port: int + + def serialize(self) -> str: + return f"{self.host}:{self.port}" + + +class RedisStore(Store): + """A Store backed by N independent per-node `redis-server` instances. + + This models independent, unrelated keyspaces -- not a Redis Cluster. + `store.endpoints` is a list with one `RedisEndpoint` per host in `hosts`. + + `RedisStore()` with no arguments spawns a single local `redis-server` on + an auto-picked free port -- it works out of the box locally. Multi-node + (HPC) usage requires an explicit `cmd` launch-command template (e.g. + `"srun --nodelist={host} redis-server --port {port}"`, formatted per + host/port and executed directly, never via a shell) plus an explicit + `port`, since a free port picked on the launching host says nothing + about availability on a remote target host. + + `RedisStore.client(index=0)` is a convenience for a raw redis-py client + on one node, for direct redis-py interaction. It is not RADEX's typed + client -- for `put_scalar`/`get_scalar`/`put_tensor`/`get_tensor`, + construct `radex.clients.core.RedisClient` yourself from an endpoint's + `serialize()` (see `RedisEndpoint`'s docstring). + """ + + def __init__( + self, + *, + hosts: Sequence[str] | None = None, + port: int | None = None, + cmd: str | None = None, + redis_server_path: str = "redis-server", + extra_args: Sequence[str] = (), + env: Mapping[str, str] | None = None, + connect_timeout: float = 5.0, + startup_timeout: float = 30.0, + poll_interval: float = 0.2, + shutdown_grace_period: float = 5.0, + ) -> None: + super().__init__() + self._hosts = list(hosts) if hosts is not None else ["localhost"] + if not self._hosts: + raise ValueError("`hosts` must be non-empty if provided") + if cmd is not None and port is None: + raise ValueError( + "`port` must be given explicitly when `cmd` is provided: " + "RedisStore cannot safely auto-pick a free port on a remote " + "host from the launching process." + ) + self._cmd_template = cmd + self._explicit_port = port + self._redis_server_path = redis_server_path + self._extra_args = list(extra_args) + self._env_overrides = dict(env) if env is not None else None + self._connect_timeout = connect_timeout + self._startup_timeout = startup_timeout + self._poll_interval = poll_interval + self._shutdown_grace_period = shutdown_grace_period + self._processes: dict[int, asyncio.subprocess.Process] = {} + self._planned: list[tuple[str, int]] = [] + + def client(self, index: int = 0, **kwargs: Any) -> "redis.Redis": + """Raw redis-py client for one node (default: the first), for + direct redis-py interaction. This is not RADEX's typed client -- + for that, construct `radex.clients.core.RedisClient` yourself (see + `RedisEndpoint`'s docstring).""" + if redis is None: + raise ImportError( + "redis-py is required for RedisStore.client(); " + "install with `pip install radex[redis]`." + ) + endpoint = self.endpoints[index] + return redis.Redis(host=endpoint.host, port=endpoint.port, **kwargs) + + def _resolve_port(self) -> int: + if self._explicit_port is not None: + return self._explicit_port + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("", 0)) + return sock.getsockname()[1] + + def _build_argv(self, host: str, port: int) -> list[str]: + if self._cmd_template is not None: + try: + formatted = self._cmd_template.format(host=host, port=port) + except (KeyError, IndexError) as exc: + raise ValueError( + f"Invalid `cmd` template {self._cmd_template!r}: {exc}" + ) from exc + return shlex.split(formatted) + return [self._redis_server_path, "--port", str(port), *self._extra_args] + + async def _do_start(self, wait: bool) -> list[Endpoint]: + if self._cmd_template is None and shutil.which(self._redis_server_path) is None: + raise StoreStartupError( + f"'{self._redis_server_path}' not found on PATH; install " + "redis-server, or pass redis_server_path=/cmd= explicitly." + ) + + self._processes = {} + self._planned = [(host, self._resolve_port()) for host in self._hosts] + + launch_results = await asyncio.gather( + *(self._launch_one(i, h, p) for i, (h, p) in enumerate(self._planned)), + return_exceptions=True, + ) + launch_errors = [r for r in launch_results if isinstance(r, BaseException)] + if launch_errors: + await self._terminate_all() + raise StoreStartupError( + f"Failed to launch {len(launch_errors)}/{len(self._planned)} " + f"redis node(s): {launch_errors}" + ) + + if not wait: + return [RedisEndpoint(host=h, port=p) for h, p in self._planned] + + ready_results = await asyncio.gather( + *( + self._wait_ready_one(i, h, p) + for i, (h, p) in enumerate(self._planned) + ), + return_exceptions=True, + ) + ready_errors = [r for r in ready_results if isinstance(r, BaseException)] + if ready_errors: + await self._terminate_all() + raise StoreStartupError( + f"{len(ready_errors)}/{len(self._planned)} redis node(s) " + f"failed to become ready: {ready_errors}" + ) + return list(ready_results) + + async def _launch_one(self, index: int, host: str, port: int) -> None: + argv = self._build_argv(host, port) + proc = await asyncio.create_subprocess_exec( + *argv, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + env=self._env_overrides, + ) + self._processes[index] = proc + + async def _wait_ready_one(self, index: int, host: str, port: int) -> RedisEndpoint: + deadline = time.monotonic() + self._startup_timeout + proc = self._processes[index] + while True: + if proc.returncode is not None: + tail = await self._read_stderr_tail(proc) + raise StoreStartupError( + f"redis-server for node {index} ({host}:{port}) exited " + f"early with code {proc.returncode}: {tail}" + ) + if await self._ping(host, port): + return RedisEndpoint(host=host, port=port) + if time.monotonic() >= deadline: + raise StoreStartupError( + f"Timed out after {self._startup_timeout}s waiting for " + f"redis node {index} ({host}:{port}) to become ready" + ) + await asyncio.sleep(self._poll_interval) + + async def _read_stderr_tail(self, proc: asyncio.subprocess.Process, n: int = 2000) -> str: + try: + assert proc.stderr is not None + data = await asyncio.wait_for(proc.stderr.read(n), timeout=1.0) + return data.decode(errors="replace") + except Exception: + return "" + + async def _ping(self, host: str, port: int) -> bool: + def _do_ping() -> bool: + try: + with socket.create_connection( + (host, port), timeout=self._connect_timeout + ) as sock: + sock.sendall(_PING) + reply = sock.recv(64) + return reply.startswith(b"+PONG") + except OSError: + return False + + return await asyncio.to_thread(_do_ping) + + async def _do_shutdown(self) -> None: + await self._terminate_all() + + async def _terminate_all(self) -> None: + await asyncio.gather( + *(self._terminate_one(p) for p in self._processes.values()), + return_exceptions=True, + ) + self._processes = {} + + async def _terminate_one(self, proc: asyncio.subprocess.Process) -> None: + if proc.returncode is not None: + return + proc.terminate() + try: + await asyncio.wait_for(proc.wait(), timeout=self._shutdown_grace_period) + except asyncio.TimeoutError: + proc.kill() + await proc.wait() + + async def _do_ready(self) -> bool: + if not self._endpoints: + return False + results = await asyncio.gather( + *(self._ping(ep.host, ep.port) for ep in self._endpoints) + ) + return all(results) diff --git a/tests/python/store/conftest.py b/tests/python/store/conftest.py new file mode 100644 index 0000000..fa4defe --- /dev/null +++ b/tests/python/store/conftest.py @@ -0,0 +1,29 @@ +import shutil + +import pytest + + +@pytest.fixture(scope="session") +def _requires_redis_server(): + if shutil.which("redis-server") is None: + pytest.skip("redis-server binary not found on PATH") + yield + + +@pytest.fixture(scope="session") +def _requires_dragon_runtime(): + try: + import dragon # noqa: F401 + from dragon.globalservices.api_setup import get_gs_ret_cuid + + get_gs_ret_cuid() + except Exception: + in_dragon_session = False + else: + in_dragon_session = True + + if not in_dragon_session: + pytest.skip( + "Dragon session not detected! Try running with `dragon -s -- -m pytest`?" + ) + yield diff --git a/tests/python/store/test_base_state_machine.py b/tests/python/store/test_base_state_machine.py new file mode 100644 index 0000000..50ef285 --- /dev/null +++ b/tests/python/store/test_base_state_machine.py @@ -0,0 +1,144 @@ +import asyncio + +import pytest + +from radex.store.base import ( + Endpoint, + Store, + StoreNotReadyError, + StoreStartupError, + StoreState, + StoreTerminatedError, +) + + +class _FakeEndpoint(Endpoint): + def __init__(self, tag: str = "fake"): + self.tag = tag + + def serialize(self) -> str: + return self.tag + + +class _FakeStore(Store): + def __init__(self, fail_start: bool = False, fail_ready: bool = False): + super().__init__() + self.fail_start = fail_start + self.fail_ready = fail_ready + self.start_calls = 0 + self.shutdown_calls = 0 + self.shutdown_seen_states = [] + + async def _do_start(self, wait: bool): + self.start_calls += 1 + if self.fail_start: + raise RuntimeError("boom") + return [_FakeEndpoint()] + + async def _do_shutdown(self) -> None: + self.shutdown_calls += 1 + self.shutdown_seen_states.append(self.state) + + async def _do_ready(self) -> bool: + return not self.fail_ready + + +async def test_endpoints_before_start_raises(): + store = _FakeStore() + with pytest.raises(StoreNotReadyError): + _ = store.endpoints + + +async def test_start_and_shutdown_return_self_for_fluent_usage(): + # Guards against `store = await RedisStore(...).start()` silently + # assigning None -- start()/shutdown() must return the Store instance. + store = await _FakeStore().start() + assert isinstance(store, _FakeStore) + assert store.state is StoreState.READY + + returned = await store.shutdown() + assert returned is store + assert store.state is StoreState.SHUTDOWN + + +async def test_start_then_endpoints(): + store = _FakeStore() + await store.start() + assert store.state is StoreState.READY + eps = store.endpoints + assert len(eps) == 1 + assert eps[0].serialize() == "fake" + + +async def test_repeated_start_is_idempotent(): + store = _FakeStore() + await store.start() + await store.start() + await store.start() + assert store.start_calls == 1 + + +async def test_repeated_shutdown_is_idempotent(): + store = _FakeStore() + await store.start() + await store.shutdown() + await store.shutdown() + await store.shutdown() + assert store.shutdown_calls == 1 + + +async def test_shutdown_never_started_is_noop_but_calls_hook_once(): + store = _FakeStore() + await store.shutdown() + assert store.state is StoreState.SHUTDOWN + assert store.shutdown_calls == 1 + assert store.shutdown_seen_states == [StoreState.CREATED] + + +async def test_start_after_shutdown_raises(): + store = _FakeStore() + await store.start() + await store.shutdown() + with pytest.raises(StoreTerminatedError): + await store.start() + + +async def test_failing_start_sets_failed_and_raises_with_cause(): + store = _FakeStore(fail_start=True) + with pytest.raises(StoreStartupError) as excinfo: + await store.start() + assert store.state is StoreState.FAILED + assert isinstance(excinfo.value.__cause__, RuntimeError) + + +async def test_start_after_failure_raises_terminated(): + store = _FakeStore(fail_start=True) + with pytest.raises(StoreStartupError): + await store.start() + with pytest.raises(StoreTerminatedError): + await store.start() + + +async def test_shutdown_after_failure_is_noop_and_invokes_hook_once(): + store = _FakeStore(fail_start=True) + with pytest.raises(StoreStartupError): + await store.start() + await store.shutdown() + assert store.state is StoreState.SHUTDOWN + assert store.shutdown_calls == 1 + + +async def test_ready_reflects_lifecycle(): + store = _FakeStore() + assert await store.ready() is False + await store.start() + assert await store.ready() is True + await store.shutdown() + assert await store.ready() is False + + +async def test_concurrent_start_calls_do_start_once(): + store = _FakeStore() + await asyncio.gather(store.start(), store.start(), store.start()) + assert store.start_calls == 1 + assert store.state is StoreState.READY diff --git a/tests/python/store/test_dragon_store.py b/tests/python/store/test_dragon_store.py new file mode 100644 index 0000000..c5c6f07 --- /dev/null +++ b/tests/python/store/test_dragon_store.py @@ -0,0 +1,61 @@ +import pytest + +from radex.store.base import StoreState, StoreTerminatedError +from radex.store.dragon_store import DragonStore + +pytestmark = pytest.mark.usefixtures("_requires_dragon_runtime") + + +def test_wait_for_keys_false_raises(): + with pytest.raises(ValueError, match=r"wait_for_keys"): + DragonStore(wait_for_keys=False) + + +def test_default_construction_does_not_raise(): + DragonStore() + + +async def test_start_ready_shutdown_roundtrip(): + store = DragonStore(managers_per_node=1, n_nodes=1) + await store.start() + try: + assert store.state is StoreState.READY + eps = store.endpoints + assert len(eps) == 1 + assert eps[0].descriptor + assert await store.ready() is True + finally: + await store.shutdown() + assert store.state is StoreState.SHUTDOWN + + +async def test_client_roundtrip_put_get_scalar(): + from radex.clients.core import DragonClient + from radex.handles.handles import IncomingHandle, OutgoingHandle + + store = DragonStore(managers_per_node=1, n_nodes=1) + await store.start() + try: + client = DragonClient(descriptor=store.endpoints[0].serialize(), timeout=5) + client.put_scalar(OutgoingHandle("store-test-key"), 0.5) + assert client.get_scalar(IncomingHandle("store-test-key")) == 0.5 + finally: + await store.shutdown() + + +async def test_start_after_shutdown_raises(): + store = DragonStore(managers_per_node=1, n_nodes=1) + await store.start() + await store.shutdown() + with pytest.raises(StoreTerminatedError): + await store.start() + + +async def test_wait_false_is_noop_for_dragon(): + store = DragonStore(managers_per_node=1, n_nodes=1) + await store.start(wait=False) + try: + assert store.state is StoreState.READY + assert store.endpoints[0].descriptor + finally: + await store.shutdown() diff --git a/tests/python/store/test_redis_store.py b/tests/python/store/test_redis_store.py new file mode 100644 index 0000000..6b9d29b --- /dev/null +++ b/tests/python/store/test_redis_store.py @@ -0,0 +1,144 @@ +import asyncio +import time + +import pytest + +import radex.store.redis_store as redis_store_module +from radex.store.base import StoreStartupError, StoreState +from radex.store.redis_store import RedisStore + + +@pytest.mark.redis +async def test_zero_arg_local_roundtrip(_requires_redis_server): + store = RedisStore() + await store.start() + try: + assert store.state is StoreState.READY + eps = store.endpoints + assert len(eps) == 1 + assert eps[0].host == "localhost" + assert await store.ready() is True + finally: + await store.shutdown() + assert store.state is StoreState.SHUTDOWN + + +@pytest.mark.redis +async def test_multi_node_concurrent_launch(_requires_redis_server): + single = RedisStore() + t0 = time.monotonic() + await single.start() + single_elapsed = time.monotonic() - t0 + await single.shutdown() + + store = RedisStore(hosts=["localhost"] * 4) + t0 = time.monotonic() + await store.start() + multi_elapsed = time.monotonic() - t0 + try: + eps = store.endpoints + assert len(eps) == 4 + assert len({ep.port for ep in eps}) == 4 + assert await store.ready() is True + # Loose bound: concurrent launch of 4 nodes should be nowhere near + # 4x a single node's startup time. + assert multi_elapsed < 3 * max(single_elapsed, 0.5) + finally: + await store.shutdown() + + +@pytest.mark.redis +async def test_partial_failure_rolls_back_all_nodes( + _requires_redis_server, monkeypatch +): + store = RedisStore(hosts=["localhost"] * 3, startup_timeout=1.0, poll_interval=0.05) + seen: dict[tuple[str, int], int] = {} + orig_ping = RedisStore._ping + + async def fake_ping(self, host, port): + key = (host, port) + if key not in seen: + seen[key] = len(seen) + if seen[key] == 1: + return False + return await orig_ping(self, host, port) + + monkeypatch.setattr(RedisStore, "_ping", fake_ping) + + with pytest.raises(StoreStartupError): + await store.start() + + assert store.state is StoreState.FAILED + assert store._processes == {} + + +async def test_terminate_then_kill_after_grace_period(): + store = RedisStore(shutdown_grace_period=0.05) + + class _StubProc: + def __init__(self): + self.returncode = None + self.terminate_called = False + self.kill_called = False + + def terminate(self): + self.terminate_called = True + + def kill(self): + self.kill_called = True + self.returncode = -9 + + async def wait(self): + if not self.kill_called: + await asyncio.sleep(9999) + return self.returncode + + proc = _StubProc() + await store._terminate_one(proc) + assert proc.terminate_called is True + assert proc.kill_called is True + + +async def test_terminate_one_skips_already_exited_process(): + store = RedisStore() + + class _ExitedProc: + returncode = 0 + + def terminate(self): + raise AssertionError("should not be called") + + def kill(self): + raise AssertionError("should not be called") + + async def wait(self): + raise AssertionError("should not be called") + + await store._terminate_one(_ExitedProc()) + + +@pytest.mark.redis +async def test_client_raises_import_error_without_redis_py( + _requires_redis_server, monkeypatch +): + monkeypatch.setattr(redis_store_module, "redis", None) + store = RedisStore() + await store.start() + try: + with pytest.raises(ImportError): + store.client() + finally: + await store.shutdown() + + +@pytest.mark.redis +async def test_client_pings_real_node(_requires_redis_server): + if redis_store_module.redis is None: + pytest.skip("redis-py not installed") + store = RedisStore() + await store.start() + try: + client = store.client() + assert client.ping() is True + finally: + await store.shutdown()