From d1c1673a44bbb64d851b23d925b45b557d28b431 Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Tue, 18 Aug 2026 15:27:40 +0000 Subject: [PATCH 1/6] Add backend-independent Store abstraction with RedisStore and DragonStore Introduces radex.store (Store/Endpoint base + RedisStore/DragonStore) so RADEX can provision and tear down its own Redis or Dragon DDict backend instead of requiring callers to do it manually, with idempotent async lifecycle (start/ready/endpoints/shutdown), plus tests and examples. --- dev-resources/radex-config.toml | 4 +- dev-resources/requirements-dev.txt | 4 +- example/py-store-exchange/dragon/driver.py | 88 +++++++ example/py-store-exchange/redis/driver.py | 79 ++++++ src/python/pyproject.toml | 4 + src/python/setup.py | 3 +- src/python/src/radex/__init__.py | 29 +++ src/python/src/radex/store/__init__.py | 27 ++ src/python/src/radex/store/base.py | 148 +++++++++++ src/python/src/radex/store/dragon_store.py | 150 +++++++++++ src/python/src/radex/store/redis_store.py | 242 ++++++++++++++++++ tests/python/store/conftest.py | 29 +++ tests/python/store/test_base_state_machine.py | 143 +++++++++++ tests/python/store/test_dragon_store.py | 60 +++++ tests/python/store/test_redis_store.py | 140 ++++++++++ 15 files changed, 1147 insertions(+), 3 deletions(-) create mode 100644 example/py-store-exchange/dragon/driver.py create mode 100644 example/py-store-exchange/redis/driver.py create mode 100644 src/python/src/radex/store/__init__.py create mode 100644 src/python/src/radex/store/base.py create mode 100644 src/python/src/radex/store/dragon_store.py create mode 100644 src/python/src/radex/store/redis_store.py create mode 100644 tests/python/store/conftest.py create mode 100644 tests/python/store/test_base_state_machine.py create mode 100644 tests/python/store/test_dragon_store.py create mode 100644 tests/python/store/test_redis_store.py diff --git a/dev-resources/radex-config.toml b/dev-resources/radex-config.toml index 9239999..38acc02 100644 --- a/dev-resources/radex-config.toml +++ b/dev-resources/radex-config.toml @@ -4,9 +4,11 @@ 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" + "compiled: Test has a component compiled at run time", + "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..215c1bb --- /dev/null +++ b/example/py-store-exchange/dragon/driver.py @@ -0,0 +1,88 @@ +""" +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.client() -> a real radex.clients.core.DragonClient bound + to this store's DDict + DragonStore.shutdown() -> destroys the DDict + + 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 + +from dragon.native.process import Process, ProcessTemplate + +from radex import DragonStore +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. A ready-to-use client, from the Store itself ────────────────────── + client = store.client(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) + + print("[Driver] Looking for keys") + print_scalar(client, "cpp-double") + print_scalar(client, "cpp-int") + 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}") + + +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..03216f2 --- /dev/null +++ b/example/py-store-exchange/redis/driver.py @@ -0,0 +1,79 @@ +""" +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, and each endpoint gets its own + `.client()` (a plain redis-py `redis.Redis`). + + 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}", + ) + + Store never mutates os.environ or the C++ SmartRedis env-var convention + (RADEX_STORE/RADEX_STORE_OPTS) itself -- if you need the compiled + `radex::redis::smartredis::Client` to attach to a node this Store + launched, wire that through explicitly yourself, e.g.: + + os.environ["RADEX_STORE"] = store.endpoints[0].serialize() + +Run with: + python driver.py +""" + +import asyncio + +from radex import RedisStore + + +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 = store.client() + client.set("greeting", "hello from RADEX") + print(f"[Driver] Got back: {client.get('greeting').decode()}") + 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 = endpoint.client() + client.set("node-id", i) + print(f"[Driver] node {i}: {endpoint.serialize()} -> " + f"node-id={client.get('node-id').decode()}") + 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/setup.py b/src/python/setup.py index dc7beba..d9ad341 100644 --- a/src/python/setup.py +++ b/src/python/setup.py @@ -1,4 +1,4 @@ -from setuptools import setup, Extension +from setuptools import setup, Extension, find_packages from Cython.Build import cythonize, build_ext import pathlib import os @@ -63,6 +63,7 @@ ] ), package_dir={"": "src"}, + packages=find_packages(where="src"), ) finally: os.chdir(ORIGINAL_DIR) 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..24c1a29 --- /dev/null +++ b/src/python/src/radex/store/base.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +import abc +import asyncio +import enum +from typing import Any + + +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`/`.client()` 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.""" + + @abc.abstractmethod + def client(self, **kwargs: Any) -> Any: + """Construct a ready-to-use client bound to this endpoint. The + concrete return type is backend-specific.""" + + +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) -> None: + 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._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 + + async def shutdown(self) -> None: + async with self._lock: + if self._state is StoreState.SHUTDOWN: + return + try: + await self._do_shutdown() + finally: + self._state = StoreState.SHUTDOWN + self._endpoints = [] + + def client(self, index: int = 0, **kwargs: Any) -> Any: + return self.endpoints[index].client(**kwargs) + + @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..4cafce3 --- /dev/null +++ b/src/python/src/radex/store/dragon_store.py @@ -0,0 +1,150 @@ +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. + """ + + descriptor: str + + def serialize(self) -> str: + return self.descriptor + + def client(self, *, timeout: int = 5) -> Any: + # Deferred import: radex.store must stay importable even when the + # compiled radex.clients.core extension isn't available/built. + from radex.clients.core import DragonClient + + return DragonClient(descriptor=self.descriptor, timeout=timeout) + + +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. + """ + + 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..29b09c1 --- /dev/null +++ b/src/python/src/radex/store/redis_store.py @@ -0,0 +1,242 @@ +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 that compiled client to attach to a node this + Store launched, e.g.:: + + os.environ["RADEX_STORE"] = endpoint.serialize() + os.environ["RADEX_STORE_OPTS"] = "" + """ + + host: str + port: int + + def serialize(self) -> str: + return f"{self.host}:{self.port}" + + def client(self, **kwargs: Any) -> "redis.Redis": + if redis is None: + raise ImportError( + "redis-py is required for RedisEndpoint.client(); " + "install with `pip install radex[redis]`." + ) + return redis.Redis(host=self.host, port=self.port, **kwargs) + + +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. + """ + + 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 _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..d9491c4 --- /dev/null +++ b/tests/python/store/test_base_state_machine.py @@ -0,0 +1,143 @@ +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 + + def client(self, **kwargs): + return {"tag": self.tag, **kwargs} + + +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_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_store_client_delegates_and_gates_on_readiness(): + store = _FakeStore() + with pytest.raises(StoreNotReadyError): + store.client() + await store.start() + assert store.client(extra=1) == {"tag": "fake", "extra": 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..7a51b2a --- /dev/null +++ b/tests/python/store/test_dragon_store.py @@ -0,0 +1,60 @@ +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.handles.handles import IncomingHandle, OutgoingHandle + + store = DragonStore(managers_per_node=1, n_nodes=1) + await store.start() + try: + client = store.client(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..597dd61 --- /dev/null +++ b/tests/python/store/test_redis_store.py @@ -0,0 +1,140 @@ +import asyncio +import time + +import pytest + +from radex.store.base import StoreStartupError, StoreState +from radex.store.redis_store import RedisStore +import radex.store.redis_store as redis_store_module + + +@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() From 2b6dcec2ec807c044a08be32618b0b49eacd6236 Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Tue, 18 Aug 2026 19:15:20 +0000 Subject: [PATCH 2/6] fix corner case --- src/python/src/radex/store/base.py | 15 +++++++++++---- src/python/src/radex/store/dragon_store.py | 5 +++++ tests/python/store/test_base_state_machine.py | 12 ++++++++++++ 3 files changed, 28 insertions(+), 4 deletions(-) diff --git a/src/python/src/radex/store/base.py b/src/python/src/radex/store/base.py index 24c1a29..5f2a03d 100644 --- a/src/python/src/radex/store/base.py +++ b/src/python/src/radex/store/base.py @@ -93,7 +93,10 @@ async def ready(self) -> bool: return False return await self._do_ready() - async def start(self, wait: bool = True) -> None: + 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( @@ -102,7 +105,7 @@ async def start(self, wait: bool = True) -> None: "construct a new Store instance instead." ) if self._state is StoreState.READY: - return + return self self._state = StoreState.STARTING try: endpoints = await self._do_start(wait=wait) @@ -119,16 +122,20 @@ async def start(self, wait: bool = True) -> None: else: self._endpoints = list(endpoints) self._state = StoreState.READY + return self - async def shutdown(self) -> None: + 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 + return self try: await self._do_shutdown() finally: self._state = StoreState.SHUTDOWN self._endpoints = [] + return self def client(self, index: int = 0, **kwargs: Any) -> Any: return self.endpoints[index].client(**kwargs) diff --git a/src/python/src/radex/store/dragon_store.py b/src/python/src/radex/store/dragon_store.py index 4cafce3..3df81fe 100644 --- a/src/python/src/radex/store/dragon_store.py +++ b/src/python/src/radex/store/dragon_store.py @@ -53,6 +53,11 @@ class DragonStore(Store): 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 client from `.client()` 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__( diff --git a/tests/python/store/test_base_state_machine.py b/tests/python/store/test_base_state_machine.py index d9491c4..30e7b3e 100644 --- a/tests/python/store/test_base_state_machine.py +++ b/tests/python/store/test_base_state_machine.py @@ -52,6 +52,18 @@ async def test_endpoints_before_start_raises(): _ = 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() From 7a830b32391ae11ead1b0a873ea90b7f936c42e1 Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Fri, 21 Aug 2026 05:05:05 +0000 Subject: [PATCH 3/6] Decouple Store/Endpoint lifecycle from client construction Store and Endpoint now only launch backends and expose connection info (.endpoints/.serialize()) -- callers build the typed RADEX client (DragonClient/RedisClient) themselves from an endpoint, giving clean N-clients-per-endpoint / clients-across-M-endpoints support instead of Store guessing cardinality. RedisStore keeps a .client() convenience for a raw redis-py client only; DragonStore has no client() at all. Updated both examples and the affected tests to match. --- example/py-store-exchange/dragon/driver.py | 36 ++++++++-- example/py-store-exchange/redis/driver.py | 63 +++++++++++++----- example/py-store-exchange/redis/dump.rdb | Bin 0 -> 121 bytes src/python/src/radex/store/base.py | 11 +-- src/python/src/radex/store/dragon_store.py | 23 +++---- src/python/src/radex/store/redis_store.py | 34 +++++++--- tests/python/store/test_base_state_machine.py | 11 --- tests/python/store/test_dragon_store.py | 3 +- 8 files changed, 118 insertions(+), 63 deletions(-) create mode 100644 example/py-store-exchange/redis/dump.rdb diff --git a/example/py-store-exchange/dragon/driver.py b/example/py-store-exchange/dragon/driver.py index 215c1bb..315ed67 100644 --- a/example/py-store-exchange/dragon/driver.py +++ b/example/py-store-exchange/dragon/driver.py @@ -9,10 +9,12 @@ DragonStore.start() -> constructs the DDict, blocks until ready DragonStore.endpoints -> [DragonEndpoint(descriptor=...)] - DragonStore.client() -> a real radex.clients.core.DragonClient bound - to this store's DDict 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. @@ -26,9 +28,11 @@ 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() @@ -42,8 +46,9 @@ async def main() -> int: await store.start() print(f"[Driver] DragonStore ready: {store.endpoints[0].serialize()[:32]}...") - # ── 2. A ready-to-use client, from the Store itself ────────────────────── - client = store.client(timeout=5) + # ── 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 @@ -65,9 +70,26 @@ async def main() -> int: 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: @@ -84,5 +106,11 @@ def print_scalar(client, key): 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 index 03216f2..d5094ea 100644 --- a/example/py-store-exchange/redis/driver.py +++ b/example/py-store-exchange/redis/driver.py @@ -5,8 +5,22 @@ ──────────── 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, and each endpoint gets its own - `.client()` (a plain redis-py `redis.Redis`). + 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 @@ -20,20 +34,30 @@ cmd="srun --nodelist={host} redis-server --port {port}", ) - Store never mutates os.environ or the C++ SmartRedis env-var convention - (RADEX_STORE/RADEX_STORE_OPTS) itself -- if you need the compiled - `radex::redis::smartredis::Client` to attach to a node this Store - launched, wire that through explicitly yourself, e.g.: - - os.environ["RADEX_STORE"] = store.endpoints[0].serialize() - 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: @@ -44,9 +68,16 @@ async def single_node_demo() -> None: endpoint = store.endpoints[0] print(f"[Driver] RedisStore ready at {endpoint.serialize()}") - client = store.client() - client.set("greeting", "hello from RADEX") - print(f"[Driver] Got back: {client.get('greeting').decode()}") + 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() @@ -61,10 +92,10 @@ async def multi_node_demo() -> None: 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 = endpoint.client() - client.set("node-id", i) - print(f"[Driver] node {i}: {endpoint.serialize()} -> " - f"node-id={client.get('node-id').decode()}") + 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() diff --git a/example/py-store-exchange/redis/dump.rdb b/example/py-store-exchange/redis/dump.rdb new file mode 100644 index 0000000000000000000000000000000000000000..f8f41fafbe5503c3a2abe9bce320617910aa6ec9 GIT binary patch literal 121 zcmWG?b@2=~FfcUz#aWb^l3A=5xUrsf{Hki^UIiz6{VO*bjAIQ0O-KZf6o3=ACUMX9MJnR)2~8L2rr`3h-8 W`MC-~jxMed|IH4ytuiSO_5c9q`!2Hp literal 0 HcmV?d00001 diff --git a/src/python/src/radex/store/base.py b/src/python/src/radex/store/base.py index 5f2a03d..cf0281d 100644 --- a/src/python/src/radex/store/base.py +++ b/src/python/src/radex/store/base.py @@ -3,7 +3,6 @@ import abc import asyncio import enum -from typing import Any class StoreState(enum.Enum): @@ -30,7 +29,7 @@ class StoreStateError(StoreError): class StoreNotReadyError(StoreStateError): - """Raised by `.endpoints`/`.client()` before a successful start().""" + """Raised by `.endpoints` before a successful start().""" class StoreTerminatedError(StoreStateError): @@ -51,11 +50,6 @@ def serialize(self) -> str: kwarg, task arg) to reconnect a client elsewhere. Store never sets this into os.environ itself -- callers own that decision.""" - @abc.abstractmethod - def client(self, **kwargs: Any) -> Any: - """Construct a ready-to-use client bound to this endpoint. The - concrete return type is backend-specific.""" - class Store(abc.ABC): """Backend-independent lifecycle for a RADEX data-exchange backend. @@ -137,9 +131,6 @@ async def shutdown(self) -> "Store": self._endpoints = [] return self - def client(self, index: int = 0, **kwargs: Any) -> Any: - return self.endpoints[index].client(**kwargs) - @abc.abstractmethod async def _do_start(self, wait: bool) -> list[Endpoint]: """Launch the backend and return its endpoint(s). Raise on failure diff --git a/src/python/src/radex/store/dragon_store.py b/src/python/src/radex/store/dragon_store.py index 3df81fe..16d517c 100644 --- a/src/python/src/radex/store/dragon_store.py +++ b/src/python/src/radex/store/dragon_store.py @@ -23,7 +23,12 @@ class DragonEndpoint(Endpoint): 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. + 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 @@ -31,13 +36,6 @@ class DragonEndpoint(Endpoint): def serialize(self) -> str: return self.descriptor - def client(self, *, timeout: int = 5) -> Any: - # Deferred import: radex.store must stay importable even when the - # compiled radex.clients.core extension isn't available/built. - from radex.clients.core import DragonClient - - return DragonClient(descriptor=self.descriptor, timeout=timeout) - class DragonStore(Store): """A Store backed by a single Dragon DDict. @@ -54,10 +52,11 @@ class DragonStore(Store): `wait_for_keys=False`, so any other value would guarantee every downstream C++/Cython client fails to construct. - Client lifetime note: drop/`del` any client from `.client()` 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. + 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__( diff --git a/src/python/src/radex/store/redis_store.py b/src/python/src/radex/store/redis_store.py index 29b09c1..734fc37 100644 --- a/src/python/src/radex/store/redis_store.py +++ b/src/python/src/radex/store/redis_store.py @@ -28,13 +28,18 @@ class RedisEndpoint(Endpoint): 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 that compiled client to attach to a node this + 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 @@ -43,14 +48,6 @@ class RedisEndpoint(Endpoint): def serialize(self) -> str: return f"{self.host}:{self.port}" - def client(self, **kwargs: Any) -> "redis.Redis": - if redis is None: - raise ImportError( - "redis-py is required for RedisEndpoint.client(); " - "install with `pip install radex[redis]`." - ) - return redis.Redis(host=self.host, port=self.port, **kwargs) - class RedisStore(Store): """A Store backed by N independent per-node `redis-server` instances. @@ -65,6 +62,12 @@ class RedisStore(Store): 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__( @@ -103,6 +106,19 @@ def __init__( 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 diff --git a/tests/python/store/test_base_state_machine.py b/tests/python/store/test_base_state_machine.py index 30e7b3e..50ef285 100644 --- a/tests/python/store/test_base_state_machine.py +++ b/tests/python/store/test_base_state_machine.py @@ -19,9 +19,6 @@ def __init__(self, tag: str = "fake"): def serialize(self) -> str: return self.tag - def client(self, **kwargs): - return {"tag": self.tag, **kwargs} - class _FakeStore(Store): def __init__(self, fail_start: bool = False, fail_ready: bool = False): @@ -131,14 +128,6 @@ async def test_shutdown_after_failure_is_noop_and_invokes_hook_once(): assert store.shutdown_calls == 1 -async def test_store_client_delegates_and_gates_on_readiness(): - store = _FakeStore() - with pytest.raises(StoreNotReadyError): - store.client() - await store.start() - assert store.client(extra=1) == {"tag": "fake", "extra": 1} - - async def test_ready_reflects_lifecycle(): store = _FakeStore() assert await store.ready() is False diff --git a/tests/python/store/test_dragon_store.py b/tests/python/store/test_dragon_store.py index 7a51b2a..c5c6f07 100644 --- a/tests/python/store/test_dragon_store.py +++ b/tests/python/store/test_dragon_store.py @@ -30,12 +30,13 @@ async def test_start_ready_shutdown_roundtrip(): 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 = store.client(timeout=5) + 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: From 17127d7bfcebdb27d22d3eb15a8b91471f0b80ea Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Fri, 21 Aug 2026 05:08:27 +0000 Subject: [PATCH 4/6] deleting db file --- example/py-store-exchange/redis/dump.rdb | Bin 121 -> 0 bytes 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 example/py-store-exchange/redis/dump.rdb diff --git a/example/py-store-exchange/redis/dump.rdb b/example/py-store-exchange/redis/dump.rdb deleted file mode 100644 index f8f41fafbe5503c3a2abe9bce320617910aa6ec9..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 121 zcmWG?b@2=~FfcUz#aWb^l3A=5xUrsf{Hki^UIiz6{VO*bjAIQ0O-KZf6o3=ACUMX9MJnR)2~8L2rr`3h-8 W`MC-~jxMed|IH4ytuiSO_5c9q`!2Hp From 6ad708de7fc203dc63774321b9850a655dffe32e Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Sat, 22 Aug 2026 20:12:13 +0000 Subject: [PATCH 5/6] fix style --- example/py-store-exchange/dragon/driver.py | 4 +++- tests/python/store/test_redis_store.py | 8 ++++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/example/py-store-exchange/dragon/driver.py b/example/py-store-exchange/dragon/driver.py index 315ed67..862c7fd 100644 --- a/example/py-store-exchange/dragon/driver.py +++ b/example/py-store-exchange/dragon/driver.py @@ -76,7 +76,9 @@ async def main() -> int: time.sleep(3) print("[Driver] Setting Int Tensor") - client.put_tensor(OutgoingHandle("py-int-tensor"), np.arange(4, dtype=np.int32)) + client.put_tensor( + OutgoingHandle("py-int-tensor"), np.arange(4, dtype=np.int32) + ) time.sleep(3) print("[Driver] Setting Float Tensor") diff --git a/tests/python/store/test_redis_store.py b/tests/python/store/test_redis_store.py index 597dd61..d9f3b4c 100644 --- a/tests/python/store/test_redis_store.py +++ b/tests/python/store/test_redis_store.py @@ -48,7 +48,9 @@ async def test_multi_node_concurrent_launch(_requires_redis_server): @pytest.mark.redis -async def test_partial_failure_rolls_back_all_nodes(_requires_redis_server, monkeypatch): +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 @@ -116,7 +118,9 @@ async def wait(self): @pytest.mark.redis -async def test_client_raises_import_error_without_redis_py(_requires_redis_server, monkeypatch): +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() From a01a87ad97df48ed097d1e7009b5660e84a5404b Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Sat, 22 Aug 2026 20:19:47 +0000 Subject: [PATCH 6/6] fix isort --- tests/python/store/test_redis_store.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/python/store/test_redis_store.py b/tests/python/store/test_redis_store.py index d9f3b4c..6b9d29b 100644 --- a/tests/python/store/test_redis_store.py +++ b/tests/python/store/test_redis_store.py @@ -3,9 +3,9 @@ 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 -import radex.store.redis_store as redis_store_module @pytest.mark.redis