diff --git a/Makefile b/Makefile index 97da283..3246b85 100644 --- a/Makefile +++ b/Makefile @@ -52,11 +52,12 @@ test-unit-no-dragon: ## Run unit tests excluding Dragon (uses regular pytest, f @echo "Running non-Dragon unit tests with pytest..." @$(PYTHON) -m pytest tests/unit/ \ --ignore=tests/unit/test_backend_execution_dragon.py \ + --ignore=tests/unit/test_backend_data_dragon.py \ -xvs test-unit-dragon: ## Run Dragon unit tests (requires 'dragon pytest', may hang other tests) @echo "Running Dragon unit tests with dragon pytest..." - @dragon $(PYTHON) -m pytest tests/unit/test_backend_execution_dragon.py tests/unit/telemetry/test_adapters_dragon.py -xvs + @dragon $(PYTHON) -m pytest tests/unit/test_backend_execution_dragon.py tests/unit/test_backend_data_dragon.py tests/unit/telemetry/test_adapters_dragon.py -xvs test-integration: ## Run all integration tests (non-Dragon + Dragon) test-integration: test-integration-no-dragon test-integration-dragon @@ -76,12 +77,13 @@ test-regular: ## Run all regular tests (excludes Dragon tests, uses pytest) @echo "Running all regular tests (excluding Dragon)..." @$(PYTHON) -m pytest tests/unit/ tests/integration/ \ --ignore=tests/unit/test_backend_execution_dragon.py \ + --ignore=tests/unit/test_backend_data_dragon.py \ -xvs test-dragon-only: ## Run Dragon tests only (requires 'dragon' launcher) @echo "Running Dragon tests..." @if command -v dragon >/dev/null 2>&1; then \ - dragon $(PYTHON) -m pytest tests/unit/test_backend_execution_dragon.py -xvs; \ + dragon $(PYTHON) -m pytest tests/unit/test_backend_execution_dragon.py tests/unit/test_backend_data_dragon.py -xvs; \ else \ echo "Dragon launcher not available. Install Dragon and run with 'dragon' command."; \ exit 1; \ diff --git a/examples/data/00-producer-consumer-redis.py b/examples/data/00-producer-consumer-redis.py new file mode 100644 index 0000000..450d1eb --- /dev/null +++ b/examples/data/00-producer-consumer-redis.py @@ -0,0 +1,101 @@ +import asyncio +import logging +import os + +import rhapsody + +from concurrent.futures import ProcessPoolExecutor + +from rhapsody.api import ComputeTask +from rhapsody.api import Session +from rhapsody.backends import ConcurrentExecutionBackend +from rhapsody.backends.data import RedisDataBackend + +rhapsody.enable_logging(level=logging.INFO) + + +# NOTE: task functions below are each fully self-contained. A task may be +# invoked in a completely separate process/node with no knowledge of this +# module or anything else defined here -- every import and every bit of +# setup a task needs must live inside that task's own function body, never +# factored into a shared helper or relying on driver-scope state. The only +# input a task gets is whatever is explicitly passed as an argument +# (`descriptor`, here) -- RedisDataBackend hands that back from `.start()`, +# it never constructs a client itself. + + +# func1 (producer) and func2 (consumer) are submitted together, with no +# ordering guarantee between them -- func2 uses wait_for_*, not get_*, so +# it correctly blocks until func1's data actually lands instead of racing +# it. + + +def func1(descriptor): + import os + + import numpy as np + + from radex.clients.core import RedisClient + from radex.handles.handles import OutgoingHandle + + os.environ["RADEX_STORE"] = descriptor + os.environ["RADEX_STORE_OPTS"] = "Standalone" + client = RedisClient() + + samples = np.arange(10, dtype=np.float64) ** 2 # [0, 1, 4, 9, ..., 81] + client.put_tensor(OutgoingHandle("samples"), samples) + client.put_scalar(OutgoingHandle("sample-count"), len(samples)) + return len(samples) + + +def func2(descriptor): + import os + + from radex.clients.core import RedisClient + from radex.handles.handles import IncomingHandle + + os.environ["RADEX_STORE"] = descriptor + os.environ["RADEX_STORE_OPTS"] = "Standalone" + client = RedisClient() + + samples = client.wait_for_tensor(IncomingHandle("samples"), 10) + count = client.wait_for_scalar(IncomingHandle("sample-count"), 10) + return {"count": int(count), "sum": float(samples.sum()), "mean": float(samples.mean())} + + +async def main(): + # RHAPSODY owns launching the Redis infrastructure; RADEX only ever + # sees the resulting endpoint, never the launch mechanism. + data_backend = await RedisDataBackend( + redis_server_path="redis-stable/src/redis-server" # or export redis-server on $PATH + ) + exec_backend = await ConcurrentExecutionBackend(ProcessPoolExecutor()) + + session = Session([exec_backend, data_backend]) + + descriptor = data_backend.endpoints[0].serialize() + + # Define tasks (UIDs auto-generated!) + tasks = [ + ComputeTask(function=func1, args=(descriptor,)), + ComputeTask(function=func2, args=(descriptor,)), + ] + + # Submit tasks + futures = await session.submit_tasks(tasks) + + # Wait for all tasks to complete (no manual callback needed!) + results = await asyncio.gather(*futures) + + # Access task results - tasks are updated in-place + for task in tasks: + print(f"Task {task.uid} in {task.state} state.") + print(f"Output: {task.return_value}") + + # Cleanup + await data_backend.shutdown() + await exec_backend.shutdown() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/data/01-producer-consumer-dragon.py b/examples/data/01-producer-consumer-dragon.py new file mode 100644 index 0000000..2c3682f --- /dev/null +++ b/examples/data/01-producer-consumer-dragon.py @@ -0,0 +1,96 @@ +"""Dragon counterpart of 00-workload-native-api.py. + +Run with: + dragon -s -- python3 00-workload-native-api-dragon.py +""" + +import asyncio +import logging + +import rhapsody + +from rhapsody.api import ComputeTask +from rhapsody.api import Session +from rhapsody.backends import DragonExecutionBackend +from rhapsody.backends.data import DragonDataBackend + +rhapsody.enable_logging(level=logging.INFO) + + +# NOTE: task functions below are each fully self-contained. A task may be +# invoked in a completely separate process/node with no knowledge of this +# module or anything else defined here -- every import and every bit of +# setup a task needs must live inside that task's own function body, never +# factored into a shared helper or relying on driver-scope state. The only +# input a task gets is whatever is explicitly passed as an argument +# (`descriptor`, here) -- DragonDataBackend hands that back from +# `.start()`, it never constructs a client itself. Unlike Redis, the +# Dragon client takes the descriptor directly as a constructor argument -- +# no environment variables involved. + + +# func1 (producer) and func2 (consumer) are submitted together, with no +# ordering guarantee between them -- func2 uses wait_for_*, not get_*, so +# it correctly blocks until func1's data actually lands instead of racing +# it. + + +def func1(descriptor): + import numpy as np + + from radex.clients.core import DragonClient + from radex.handles.handles import OutgoingHandle + + client = DragonClient(descriptor=descriptor, timeout=5) + + samples = np.arange(10, dtype=np.float64) ** 2 # [0, 1, 4, 9, ..., 81] + client.put_tensor(OutgoingHandle("samples"), samples) + client.put_scalar(OutgoingHandle("sample-count"), len(samples)) + return len(samples) + + +def func2(descriptor): + from radex.clients.core import DragonClient + from radex.handles.handles import IncomingHandle + + client = DragonClient(descriptor=descriptor, timeout=5) + + samples = client.wait_for_tensor(IncomingHandle("samples"), 10) + count = client.wait_for_scalar(IncomingHandle("sample-count"), 10) + return {"count": int(count), "sum": float(samples.sum()), "mean": float(samples.mean())} + + +async def main(): + # RHAPSODY owns launching the Dragon DDict; RADEX only ever sees the + # resulting endpoint, never the launch mechanism. + data_backend = await DragonDataBackend(managers_per_node=1, n_nodes=1) + exec_backend = await DragonExecutionBackend() + + session = Session([exec_backend, data_backend]) + + descriptor = data_backend.endpoints[0].serialize() + + # Define tasks (UIDs auto-generated!) + tasks = [ + ComputeTask(function=func1, args=(descriptor,)), + ComputeTask(function=func2, args=(descriptor,)), + ] + + # Submit tasks + futures = await session.submit_tasks(tasks) + + # Wait for all tasks to complete (no manual callback needed!) + results = await asyncio.gather(*futures) + + # Access task results - tasks are updated in-place + for task in tasks: + print(f"Task {task.uid} in {task.state} state.") + print(f"Output: {task.return_value}") + + # Cleanup + await data_backend.shutdown() + await exec_backend.shutdown() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/pyproject.toml b/pyproject.toml index 35b75c5..110f6d1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -131,7 +131,8 @@ markers = [ "slow: marks tests as slow-running", "radical_pilot: marks tests that use RADICAL-Pilot backend", "performance: marks tests as performance benchmarks", - "result_contract: marks tests that verify the task callback contract — required fields, types, and state transitions — across all backends" + "result_contract: marks tests that verify the task callback contract — required fields, types, and state transitions — across all backends", + "redis: marks tests that require a real redis-server binary on PATH" ] [tool.ruff] diff --git a/src/rhapsody/api/session.py b/src/rhapsody/api/session.py index e9f9213..7dc85e3 100644 --- a/src/rhapsody/api/session.py +++ b/src/rhapsody/api/session.py @@ -12,6 +12,7 @@ if TYPE_CHECKING: from rhapsody.api.task import BaseTask from rhapsody.backends.base import BaseBackend + from rhapsody.backends.data.base import DataBackend from rhapsody.telemetry.manager import TelemetryManager @@ -110,14 +111,17 @@ class Session: def __init__( self, - backends: list[BaseBackend] | None = None, + backends: list[BaseBackend | DataBackend] | None = None, uid: str | None = None, work_dir: str | None = None, ): """Initialize a new session. Args: - backends: List of execution backends to use. If None, no backends are configured. + backends: List of backends to use -- task-executing backends + (ConcurrentExecutionBackend, DragonExecutionBackend, ...) and/or + DataBackend instances (RedisDataBackend, DragonDataBackend, ...). + If None, no backends are configured. uid: Optional unique identifier for the session. work_dir: working directory (default: cwd). """ @@ -130,23 +134,40 @@ def __init__( # Register callbacks with all provided backends backends_list = backends or [] - self.backends: dict[str, BaseBackend] = {} + self.backends: dict[str, BaseBackend | DataBackend] = {} + # Task-executing subset of self.backends, used for routing in + # submit_tasks() -- kept separate so routing stays O(1) per task + # regardless of how many DataBackend instances are also registered. + self._exec_backends: dict[str, BaseBackend] = {} for backend in backends_list: self.add_backend(backend) - def add_backend(self, backend: BaseBackend) -> None: + def add_backend(self, backend: BaseBackend | DataBackend) -> None: """Add a backend to the session and register callbacks. Args: - backend: The execution or inference backend to add. + backend: The execution/inference backend, or DataBackend, to add. """ + self.backends[backend.name] = backend + + if not hasattr(backend, "submit_tasks"): + # Infrastructure backend (e.g. DataBackend) -- no tasks, no + # callbacks, no task-state map, and no Session-assigned + # _work_dir: it resolves (and may already be using) its own + # work_dir before ever reaching a Session, so stamping a fresh + # rhapsody.session. directory here would just create an + # empty, unused one. Registered only for inclusion in + # Session.close()/telemetry. + logger.debug(f"Registered data backend '{backend.name}' with Session '{self.uid}'") + return + backend._work_dir = os.path.join(self.work_dir, self.uid) os.makedirs(backend._work_dir, exist_ok=True) + + self._exec_backends[backend.name] = backend backend.is_attached = True backend.attached_to.append(self.uid) - self.backends[backend.name] = backend - # Register state manager callback backend.register_callback(self._state_manager.update_task) @@ -178,8 +199,8 @@ async def submit_tasks(self, tasks: list[dict | BaseTask]) -> list[asyncio.Futur self._state_manager.bind_loop(asyncio.get_running_loop()) except RuntimeError: pass - if not self.backends: - raise RuntimeError("No backends configured in Session") + if not self._exec_backends: + raise RuntimeError("No task-executing backend configured in Session") # Group tasks by their explicit backend target tasks_by_backend: dict[str, list] = {} @@ -204,16 +225,18 @@ async def submit_tasks(self, tasks: list[dict | BaseTask]) -> list[asyncio.Futur # Routing decision target_name = task.get("backend") if not target_name: - # If no backend specified, use the first one as default - target_name = next(iter(self.backends)) + # If no backend specified, use the first task-executing one as + # default (DataBackend instances registered in the same + # Session are never eligible here). + target_name = next(iter(self._exec_backends)) task["backend"] = target_name # Ensure it's recorded # Emit TaskSubmitted AFTER routing so task["backend"] is always set. if self._telemetry is not None: self._telemetry._on_task_submitted(task) - if target_name not in self.backends: - available = list(self.backends.keys()) + if target_name not in self._exec_backends: + available = list(self._exec_backends.keys()) raise ValueError( f"Backend '{target_name}' requested by task {uid} not found in Session. " f"Available backends: {available}" @@ -224,7 +247,7 @@ async def submit_tasks(self, tasks: list[dict | BaseTask]) -> list[asyncio.Futur # Submit each group to its respective backend concurrently submission_tasks = [] for name, backend_tasks in tasks_by_backend.items(): - backend = self.backends[name] + backend = self._exec_backends[name] # Emit TaskQueued at the backend boundary (after routing, before execution) if self._telemetry is not None: for task in backend_tasks: diff --git a/src/rhapsody/backends/data/__init__.py b/src/rhapsody/backends/data/__init__.py new file mode 100644 index 0000000..6368db9 --- /dev/null +++ b/src/rhapsody/backends/data/__init__.py @@ -0,0 +1,42 @@ +"""Data infrastructure subsystem for Rhapsody. + +This module provides backends that launch and own the lifecycle of data infrastructure (a Redis +server, a Dragon DDict) and hand back connection endpoints, mirroring how execution/inference +backends launch and own compute/inference infrastructure. +""" + +from __future__ import annotations + +from .base import DataBackend +from .base import DataBackendError +from .base import DataBackendNotReadyError +from .base import DataBackendStartupError +from .base import DataBackendState +from .base import DataBackendStateError +from .base import DataBackendTerminatedError +from .base import Endpoint +from .redis import RedisDataBackend +from .redis import RedisEndpoint + +__all__ = [ + "DataBackend", + "DataBackendError", + "DataBackendNotReadyError", + "DataBackendStartupError", + "DataBackendState", + "DataBackendStateError", + "DataBackendTerminatedError", + "Endpoint", + "RedisDataBackend", + "RedisEndpoint", +] + +# Try to import the optional Dragon-backed data backend +try: + from .dragon import DragonDataBackend # noqa: F401 + from .dragon import DragonEndpoint # noqa: F401 + + __all__.append("DragonDataBackend") + __all__.append("DragonEndpoint") +except ImportError: + pass diff --git a/src/rhapsody/backends/data/base.py b/src/rhapsody/backends/data/base.py new file mode 100644 index 0000000..32dbec1 --- /dev/null +++ b/src/rhapsody/backends/data/base.py @@ -0,0 +1,162 @@ +"""Backend-independent lifecycle abstraction for RHAPSODY data infrastructure.""" + +from __future__ import annotations + +import abc +import asyncio +import enum + + +class DataBackendState(enum.Enum): + """Lifecycle state of a `DataBackend`.""" + + CREATED = "CREATED" + STARTING = "STARTING" + READY = "READY" + FAILED = "FAILED" + SHUTDOWN = "SHUTDOWN" + + +_TERMINAL_STATES = frozenset({DataBackendState.FAILED, DataBackendState.SHUTDOWN}) + + +class DataBackendError(Exception): + """Base class for all rhapsody.data errors.""" + + +class DataBackendStartupError(DataBackendError): + """Raised when a DataBackend fails to reach READY during start().""" + + +class DataBackendStateError(DataBackendError): + """Raised when a DataBackend method is invoked in an invalid lifecycle state.""" + + +class DataBackendNotReadyError(DataBackendStateError): + """Raised by `.endpoints` before a successful start().""" + + +class DataBackendTerminatedError(DataBackendStateError): + """Raised by start() on a DataBackend that is already FAILED or SHUTDOWN.""" + + +class Endpoint(abc.ABC): + """Connection information for one DataBackend 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. An `Endpoint` never constructs a + client itself -- callers build whatever client they need from + `serialize()`'s output. + """ + + @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.""" + + +class DataBackend(abc.ABC): + """Backend-independent lifecycle for a RHAPSODY data-infrastructure backend. + + State machine: CREATED -> STARTING -> {READY, FAILED}; + {CREATED, READY, FAILED} -> SHUTDOWN. FAILED and SHUTDOWN are terminal -- + a DataBackend 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, *, name: str | None = None) -> None: + self._name = name + self._state: DataBackendState = DataBackendState.CREATED + self._endpoints: list[Endpoint] = [] + self._lock = asyncio.Lock() + + @property + def name(self) -> str: + return self._name or type(self).__name__ + + @property + def state(self) -> DataBackendState: + return self._state + + @property + def endpoints(self) -> list[Endpoint]: + if self._state is not DataBackendState.READY: + raise DataBackendNotReadyError( + f"{type(self).__name__} is not ready (state={self._state.name}); " + "call `await backend.start()` first." + ) + return list(self._endpoints) + + async def ready(self) -> bool: + if self._state is not DataBackendState.READY: + return False + return await self._do_ready() + + def __await__(self): + """Allow `await RedisDataBackend(...)` to start it and return self, matching BaseBackend's + `await ConcurrentExecutionBackend(...)` convention.""" + return self.start().__await__() + + async def start(self, wait: bool = True) -> DataBackend: + """Start the backend and return self, so both `await backend.start()` and `backend = await + RedisDataBackend(...).start()` work.""" + async with self._lock: + if self._state in _TERMINAL_STATES: + raise DataBackendTerminatedError( + f"{type(self).__name__} is in terminal state " + f"{self._state.name} and cannot be started again; " + "construct a new DataBackend instance instead." + ) + if self._state is DataBackendState.READY: + return self + self._state = DataBackendState.STARTING + try: + endpoints = await self._do_start(wait=wait) + except BaseException as exc: + self._state = DataBackendState.FAILED + self._endpoints = [] + if isinstance(exc, DataBackendError) or not isinstance(exc, Exception): + # Already a DataBackendError, or a BaseException we must + # not mask (CancelledError, KeyboardInterrupt, SystemExit). + raise + raise DataBackendStartupError(f"{type(self).__name__} failed to start") from exc + else: + self._endpoints = list(endpoints) + self._state = DataBackendState.READY + return self + + async def shutdown(self) -> DataBackend: + """Tear down the backend and return self, for the same fluent usage as start().""" + async with self._lock: + if self._state is DataBackendState.SHUTDOWN: + return self + try: + await self._do_shutdown() + finally: + self._state = DataBackendState.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/rhapsody/backends/data/dragon.py b/src/rhapsody/backends/data/dragon.py new file mode 100644 index 0000000..20115a2 --- /dev/null +++ b/src/rhapsody/backends/data/dragon.py @@ -0,0 +1,162 @@ +"""DragonDataBackend: launches and owns the lifecycle of a Dragon DDict.""" + +from __future__ import annotations + +import asyncio +import dataclasses +import logging +from typing import Any + +from rhapsody.backends.data.base import DataBackend +from rhapsody.backends.data.base import DataBackendStartupError +from rhapsody.backends.data.base import Endpoint + +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 = "__rhapsody_data_backend_liveness_probe__" + + +def _get_logger() -> logging.Logger: + """Get logger for the dragon data backend module. + + This function provides lazy logger evaluation, ensuring the logger is created after the user has + configured logging, not at module import time. + """ + return logging.getLogger(__name__) + + +@dataclasses.dataclass(frozen=True) +class DragonEndpoint(Endpoint): + """Connection information for a Dragon DDict-backed backend. + + Unlike Redis, there is no host/port here -- `serialize()` returns the + opaque base64 descriptor produced by `DDict.serialize()`. + `DragonEndpoint` never constructs a client itself -- build one directly + from the descriptor, e.g. `radex.clients.core.DragonClient(descriptor=endpoint.serialize(), timeout=5)` + or `DDict.attach(endpoint.serialize())`. + """ + + descriptor: str + + def serialize(self) -> str: + return self.descriptor + + +class DragonDataBackend(DataBackend): + """A DataBackend 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 constructor parameter is + named `_wait_for_keys` (leading underscore) to make clear it isn't + meant to be set by normal callers -- it exists so tests can exercise + the rejection path. + + All other `DDict` construction arguments are accepted as arbitrary + keyword arguments and forwarded to `DDict(...)` unchanged -- this + class does not mirror or re-validate `DDict.__init__`'s own parameter + list. If you pass something `DDict` doesn't like, `DDict` raises its + own error; see `dragon.data.ddict.DDict` for the full set of accepted + keyword arguments. The one adjustment made on your behalf: if you + don't pass `working_set_size`, it defaults to `2` (`DDict`'s own + default of `1` is incompatible with the `wait_for_keys=True` this + class always forces) -- pass it explicitly to override. + """ + + def __init__( + self, *, name: str = "dragon", _wait_for_keys: bool = True, **ddict_kwargs: Any + ) -> None: + """Initialize a DragonDataBackend. + + Args: + name: Name this backend is registered under when attached to a + Session. + _wait_for_keys: Must be True (the default) -- exists only so + tests can exercise the rejection path. Do not set this. + **ddict_kwargs: Forwarded directly to `dragon.data.ddict.DDict`. + See its docstring for the full set of accepted arguments. + + Raises: + ImportError: If the `dragon` package is not importable. + ValueError: If `_wait_for_keys` is not True, or `ddict_kwargs` + tries to set `wait_for_keys`. + """ + super().__init__(name=name) + self.logger = _get_logger() + if _DDict is None: + raise ImportError( + "The 'dragon' package is required to use DragonDataBackend. " + "It is not pip-installable from PyPI; install it per your " + "Dragon distribution/environment first." + ) + if _wait_for_keys is not True: + raise ValueError( + "DragonDataBackend requires wait_for_keys=True: downstream " + "clients (e.g. the compiled radex::drg::ddict::Client) " + "refuse to attach to a DDict created with " + "wait_for_keys=False." + ) + if "wait_for_keys" in ddict_kwargs: + raise ValueError( + "'wait_for_keys' must not be passed as a keyword argument: " + "it would silently override the enforced " + "wait_for_keys=True above. Use the _wait_for_keys " + "constructor parameter instead (and only for tests " + "exercising the rejection path)." + ) + + self._ddict_kwargs: dict[str, Any] = dict(ddict_kwargs) + self._ddict_kwargs["wait_for_keys"] = True + self._ddict_kwargs.setdefault("working_set_size", 2) + self._ddict: Any = None + + async def _do_start(self, wait: bool) -> list[Endpoint]: + # `wait` is accepted for interface parity with DataBackend.start() + # but has no effect: DDict.__init__ is already atomically + # blocking-until-ready -- there is no separate "start" step to skip + # waiting on. + self.logger.info("Starting DragonDataBackend (constructing DDict)...") + try: + self._ddict = await asyncio.to_thread(_DDict, **self._ddict_kwargs) + except Exception as exc: + self._ddict = None + self.logger.error("DragonDataBackend failed to construct DDict: %s", exc) + raise DataBackendStartupError( + f"DragonDataBackend failed to construct DDict: {exc}" + ) from exc + descriptor = self._ddict.serialize() + self.logger.info("DragonDataBackend ready at %s...", descriptor[:32]) + return [DragonEndpoint(descriptor=descriptor)] + + async def _do_shutdown(self) -> None: + self.logger.info("Shutting down DragonDataBackend...") + if self._ddict is not None: + await asyncio.to_thread(self._ddict.destroy) + self._ddict = None + self.logger.info("DragonDataBackend shutdown complete") + + async def _do_ready(self) -> bool: + if self._ddict is None: + return False + + # TODO: this probe-key check is a stand-in for a real health check. + # Wire in a more direct Dragon-native liveness API here once one is + # available, instead of a synthetic containment lookup. + 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/rhapsody/backends/data/redis.py b/src/rhapsody/backends/data/redis.py new file mode 100644 index 0000000..d8762bf --- /dev/null +++ b/src/rhapsody/backends/data/redis.py @@ -0,0 +1,327 @@ +"""RedisDataBackend: launches and owns the lifecycle of one or more +independent `redis-server` instances.""" + +from __future__ import annotations + +import asyncio +import dataclasses +import logging +import os +import shlex +import shutil +import socket +import time +import uuid +from collections.abc import Mapping +from collections.abc import Sequence + +from rhapsody.backends.data.base import DataBackend +from rhapsody.backends.data.base import DataBackendStartupError +from rhapsody.backends.data.base import Endpoint + +_PING = b"*1\r\n$4\r\nPING\r\n" + + +def _get_logger() -> logging.Logger: + """Get logger for the redis data backend module. + + This function provides lazy logger evaluation, ensuring the logger is created after the user has + configured logging, not at module import time. + """ + return logging.getLogger(__name__) + + +@dataclasses.dataclass(frozen=True) +class RedisEndpoint(Endpoint): + """Connection information for one independent Redis node. + + `serialize()` returns a `host:port` string. `RedisEndpoint` never + constructs a client itself -- build whichever client you need directly, + e.g. `redis.Redis(host=endpoint.host, port=endpoint.port)`. + """ + + host: str + port: int + + def serialize(self) -> str: + return f"{self.host}:{self.port}" + + +class RedisDataBackend(DataBackend): + """A DataBackend backed by N independent per-node `redis-server` instances. + + This models independent, unrelated keyspaces -- not a Redis Cluster. + `backend.endpoints` is a list with one `RedisEndpoint` per host in + `hosts`. + + `RedisDataBackend()` 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. + + Launching is done via a plain `asyncio` subprocess for now. + """ + + def __init__( + self, + *, + name: str = "redis", + 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, + work_dir: 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: + """Initialize a RedisDataBackend. + + Args: + name: Name this backend is registered under when attached to a + Session. + hosts: Hostnames to launch a `redis-server` on, one per entry. + Defaults to a single `"localhost"` node. + port: Port to bind on every host. Required when `cmd` is given + (a free port picked on the launching host says nothing + about a remote target host); auto-picked per host + otherwise. + cmd: Launch-command template, formatted with `host`/`port` and + executed directly (never via a shell), e.g. + `"srun --nodelist={host} redis-server --port {port}"`. + redis_server_path: Path to the `redis-server` executable, used + when `cmd` is not given. + extra_args: Extra CLI arguments appended when `cmd` is not + given. + env: Environment variables to add/override for the launched + process(es); the parent's own environment is preserved + underneath (PATH included), not replaced. + work_dir: Directory for per-node `redis.node{index}.log` files + (redis-server's stdout/stderr, redirected there instead of + left as an undrained pipe). Defaults to a fresh + `rhapsody.data.` directory under the cwd. + connect_timeout: Per-attempt timeout for the readiness PING. + startup_timeout: Overall timeout to wait for each node to + become ready. + poll_interval: Delay between readiness poll attempts. + shutdown_grace_period: Time to wait after SIGTERM before + escalating to SIGKILL. + + Raises: + ValueError: If `hosts` is empty, or `cmd` is given without an + explicit `port`. + """ + super().__init__(name=name) + self.logger = _get_logger() + 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: " + "RedisDataBackend 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 + # Merge (not replace) the parent environment: passing `env=` should + # add/override variables, not strip PATH and everything else out + # from under the child (which would break resolving a bare + # `redis-server` via PATH even though shutil.which() just found it). + self._env = ( + {**os.environ, **self._env_overrides} if self._env_overrides is not None else None + ) + self._work_dir = work_dir or os.path.join( + os.getcwd(), f"rhapsody.data.{uuid.uuid4().hex[:8]}" + ) + 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]] = [] + self._log_paths: dict[int, str] = {} + + 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 DataBackendStartupError( + f"'{self._redis_server_path}' not found on PATH; install " + "redis-server, or pass redis_server_path=/cmd= explicitly." + ) + + os.makedirs(self._work_dir, exist_ok=True) + self.logger.info( + "Starting %d redis-server node(s) (logs under %s)...", + len(self._hosts), + self._work_dir, + ) + 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() + self.logger.error( + "Failed to launch %d/%d redis node(s): %s", + len(launch_errors), + len(self._planned), + launch_errors, + ) + raise DataBackendStartupError( + f"Failed to launch {len(launch_errors)}/{len(self._planned)} " + f"redis node(s): {launch_errors}" + ) + + if not wait: + self.logger.info("RedisDataBackend launched (not waiting for ready)") + 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() + self.logger.error( + "%d/%d redis node(s) failed to become ready: %s", + len(ready_errors), + len(self._planned), + ready_errors, + ) + raise DataBackendStartupError( + f"{len(ready_errors)}/{len(self._planned)} redis node(s) " + f"failed to become ready: {ready_errors}" + ) + self.logger.info("RedisDataBackend ready: %d node(s)", len(ready_results)) + return list(ready_results) + + async def _launch_one(self, index: int, host: str, port: int) -> None: + argv = self._build_argv(host, port) + log_path = os.path.join(self._work_dir, f"redis.node{index}.log") + self._log_paths[index] = log_path + log_file = open(log_path, "wb") + try: + proc = await asyncio.create_subprocess_exec( + *argv, + stdout=log_file, + stderr=log_file, + env=self._env, + ) + finally: + # The child has its own fd copy after fork/exec; safe to close + # the parent's handle immediately. + log_file.close() + self._processes[index] = proc + + async def _tail_log(self, index: int, n: int = 4000) -> str: + path = self._log_paths.get(index) + if not path: + return "" + + def _read() -> str: + try: + with open(path, "rb") as f: + f.seek(0, os.SEEK_END) + size = f.tell() + f.seek(max(0, size - n)) + return f.read().decode(errors="replace") + except OSError: + return "" + + return await asyncio.to_thread(_read) + + 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._tail_log(index) + raise DataBackendStartupError( + f"redis-server for node {index} ({host}:{port}) exited " + f"early with code {proc.returncode}; see " + f"{self._log_paths.get(index)}:\n{tail}" + ) + if await self._ping(host, port): + self.logger.debug("redis node %d (%s:%d) ready", index, host, port) + return RedisEndpoint(host=host, port=port) + if time.monotonic() >= deadline: + tail = await self._tail_log(index) + raise DataBackendStartupError( + f"Timed out after {self._startup_timeout}s waiting for " + f"redis node {index} ({host}:{port}) to become ready; " + f"see {self._log_paths.get(index)}:\n{tail}" + ) + await asyncio.sleep(self._poll_interval) + + 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: + self.logger.info("Shutting down RedisDataBackend (%d node(s))...", len(self._processes)) + await self._terminate_all() + self.logger.info("RedisDataBackend shutdown complete") + + 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) + self.logger.debug("redis process (pid=%s) terminated", proc.pid) + except asyncio.TimeoutError: + proc.kill() + await proc.wait() + self.logger.debug("redis process (pid=%s) killed after grace period", proc.pid) + + 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/unit/test_backend_data_base.py b/tests/unit/test_backend_data_base.py new file mode 100644 index 0000000..fd6804b --- /dev/null +++ b/tests/unit/test_backend_data_base.py @@ -0,0 +1,151 @@ +"""Unit tests for the DataBackend/Endpoint lifecycle state machine. + +No real infrastructure is needed -- these exercise rhapsody.data.base's state machine via an in-file +fake DataBackend subclass. +""" + +import asyncio + +import pytest + +from rhapsody.backends.data.base import DataBackend +from rhapsody.backends.data.base import DataBackendNotReadyError +from rhapsody.backends.data.base import DataBackendStartupError +from rhapsody.backends.data.base import DataBackendState +from rhapsody.backends.data.base import DataBackendTerminatedError +from rhapsody.backends.data.base import Endpoint + + +class _FakeEndpoint(Endpoint): + def __init__(self, tag: str = "fake"): + self.tag = tag + + def serialize(self) -> str: + return self.tag + + +class _FakeDataBackend(DataBackend): + 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 + + +def test_data_backend_is_abstract(): + with pytest.raises(TypeError): + DataBackend() + + +async def test_endpoints_before_start_raises(): + backend = _FakeDataBackend() + with pytest.raises(DataBackendNotReadyError): + _ = backend.endpoints + + +async def test_start_then_endpoints(): + backend = _FakeDataBackend() + await backend.start() + assert backend.state is DataBackendState.READY + eps = backend.endpoints + assert len(eps) == 1 + assert eps[0].serialize() == "fake" + + +async def test_start_and_shutdown_return_self_for_fluent_usage(): + backend = await _FakeDataBackend().start() + assert isinstance(backend, _FakeDataBackend) + assert backend.state is DataBackendState.READY + + returned = await backend.shutdown() + assert returned is backend + assert backend.state is DataBackendState.SHUTDOWN + + +async def test_repeated_start_is_idempotent(): + backend = _FakeDataBackend() + await backend.start() + await backend.start() + await backend.start() + assert backend.start_calls == 1 + + +async def test_repeated_shutdown_is_idempotent(): + backend = _FakeDataBackend() + await backend.start() + await backend.shutdown() + await backend.shutdown() + await backend.shutdown() + assert backend.shutdown_calls == 1 + + +async def test_shutdown_never_started_is_noop_but_calls_hook_once(): + backend = _FakeDataBackend() + await backend.shutdown() + assert backend.state is DataBackendState.SHUTDOWN + assert backend.shutdown_calls == 1 + assert backend.shutdown_seen_states == [DataBackendState.CREATED] + + +async def test_start_after_shutdown_raises(): + backend = _FakeDataBackend() + await backend.start() + await backend.shutdown() + with pytest.raises(DataBackendTerminatedError): + await backend.start() + + +async def test_failing_start_sets_failed_and_raises_with_cause(): + backend = _FakeDataBackend(fail_start=True) + with pytest.raises(DataBackendStartupError) as excinfo: + await backend.start() + assert backend.state is DataBackendState.FAILED + assert isinstance(excinfo.value.__cause__, RuntimeError) + + +async def test_start_after_failure_raises_terminated(): + backend = _FakeDataBackend(fail_start=True) + with pytest.raises(DataBackendStartupError): + await backend.start() + with pytest.raises(DataBackendTerminatedError): + await backend.start() + + +async def test_shutdown_after_failure_is_noop_and_invokes_hook_once(): + backend = _FakeDataBackend(fail_start=True) + with pytest.raises(DataBackendStartupError): + await backend.start() + await backend.shutdown() + assert backend.state is DataBackendState.SHUTDOWN + assert backend.shutdown_calls == 1 + + +async def test_ready_reflects_lifecycle(): + backend = _FakeDataBackend() + assert await backend.ready() is False + await backend.start() + assert await backend.ready() is True + await backend.shutdown() + assert await backend.ready() is False + + +async def test_concurrent_start_calls_do_start_once(): + backend = _FakeDataBackend() + await asyncio.gather(backend.start(), backend.start(), backend.start()) + assert backend.start_calls == 1 + assert backend.state is DataBackendState.READY diff --git a/tests/unit/test_backend_data_dragon.py b/tests/unit/test_backend_data_dragon.py new file mode 100644 index 0000000..adb0748 --- /dev/null +++ b/tests/unit/test_backend_data_dragon.py @@ -0,0 +1,75 @@ +"""Tests for DragonDataBackend. + +Run with: + dragon python -m pytest tests/unit/test_backend_data_dragon.py -v +""" + +import pytest + +# Skip the entire module when the Dragon runtime is not installed. +pytest.importorskip("dragon", reason="Dragon is required for Dragon data backend tests") + +from rhapsody.backends.data.base import DataBackendState # noqa: E402 +from rhapsody.backends.data.base import DataBackendTerminatedError # noqa: E402 +from rhapsody.backends.data.dragon import DragonDataBackend # noqa: E402 + + +def test_wait_for_keys_false_raises(): + with pytest.raises(ValueError, match=r"wait_for_keys"): + DragonDataBackend(_wait_for_keys=False) + + +def test_wait_for_keys_via_kwargs_raises(): + with pytest.raises(ValueError, match=r"wait_for_keys"): + DragonDataBackend(wait_for_keys=False) + + +def test_default_construction_does_not_raise(): + DragonDataBackend() + + +async def test_zero_arg_start_ready_shutdown_roundtrip(): + # Covers the working_set_size default-injection: DDict's own default + # (1) is incompatible with the wait_for_keys=True this class always + # forces, so a bare DragonDataBackend() must still construct cleanly. + backend = DragonDataBackend() + await backend.start() + try: + assert backend.state is DataBackendState.READY + assert backend.endpoints[0].descriptor + assert await backend.ready() is True + finally: + await backend.shutdown() + assert backend.state is DataBackendState.SHUTDOWN + + +async def test_start_ready_shutdown_roundtrip(): + backend = DragonDataBackend(managers_per_node=1, n_nodes=1) + await backend.start() + try: + assert backend.state is DataBackendState.READY + eps = backend.endpoints + assert len(eps) == 1 + assert eps[0].descriptor + assert await backend.ready() is True + finally: + await backend.shutdown() + assert backend.state is DataBackendState.SHUTDOWN + + +async def test_start_after_shutdown_raises(): + backend = DragonDataBackend(managers_per_node=1, n_nodes=1) + await backend.start() + await backend.shutdown() + with pytest.raises(DataBackendTerminatedError): + await backend.start() + + +async def test_wait_false_is_noop_for_dragon(): + backend = DragonDataBackend(managers_per_node=1, n_nodes=1) + await backend.start(wait=False) + try: + assert backend.state is DataBackendState.READY + assert backend.endpoints[0].descriptor + finally: + await backend.shutdown() diff --git a/tests/unit/test_backend_data_redis.py b/tests/unit/test_backend_data_redis.py new file mode 100644 index 0000000..4b620a1 --- /dev/null +++ b/tests/unit/test_backend_data_redis.py @@ -0,0 +1,130 @@ +"""Unit tests for RedisDataBackend. + +Tests requiring a real `redis-server` binary are marked `redis` and skip +cleanly when it isn't on PATH. +""" + +import asyncio +import shutil +import time + +import pytest + +from rhapsody.backends.data.base import DataBackendStartupError +from rhapsody.backends.data.base import DataBackendState +from rhapsody.backends.data.redis import RedisDataBackend + + +@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.mark.redis +async def test_zero_arg_local_roundtrip(_requires_redis_server): + backend = RedisDataBackend() + await backend.start() + try: + assert backend.state is DataBackendState.READY + eps = backend.endpoints + assert len(eps) == 1 + assert eps[0].host == "localhost" + assert await backend.ready() is True + finally: + await backend.shutdown() + assert backend.state is DataBackendState.SHUTDOWN + + +@pytest.mark.redis +async def test_multi_node_concurrent_launch(_requires_redis_server): + single = RedisDataBackend() + t0 = time.monotonic() + await single.start() + single_elapsed = time.monotonic() - t0 + await single.shutdown() + + backend = RedisDataBackend(hosts=["localhost"] * 4) + t0 = time.monotonic() + await backend.start() + multi_elapsed = time.monotonic() - t0 + try: + eps = backend.endpoints + assert len(eps) == 4 + assert len({ep.port for ep in eps}) == 4 + assert await backend.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 backend.shutdown() + + +@pytest.mark.redis +async def test_partial_failure_rolls_back_all_nodes(_requires_redis_server, monkeypatch): + backend = RedisDataBackend(hosts=["localhost"] * 3, startup_timeout=1.0, poll_interval=0.05) + seen: dict[tuple[str, int], int] = {} + orig_ping = RedisDataBackend._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(RedisDataBackend, "_ping", fake_ping) + + with pytest.raises(DataBackendStartupError): + await backend.start() + + assert backend.state is DataBackendState.FAILED + assert backend._processes == {} + + +async def test_terminate_then_kill_after_grace_period(): + backend = RedisDataBackend(shutdown_grace_period=0.05) + + class _StubProc: + def __init__(self): + self.pid = 12345 + 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 backend._terminate_one(proc) + assert proc.terminate_called is True + assert proc.kill_called is True + + +async def test_terminate_one_skips_already_exited_process(): + backend = RedisDataBackend() + + 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 backend._terminate_one(_ExitedProc())