data backends for radex

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; \

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())

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())

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]

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.<uid> 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:

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