Resource Manager

src/rhapsody/api/session.py

@@ -234,7 +234,7 @@ async def submit_tasks(self, tasks: list[dict | BaseTask]) -> list[asyncio.Futur
         if submission_tasks:
             await asyncio.gather(*submission_tasks)
 
-        logger.info(f"Successfully submitted {len(tasks)} tasks")
+        logger.debug(f"Successfully submitted {len(tasks)} tasks")
 
         return futures
 

src/rhapsody/api/task.py

@@ -244,10 +244,11 @@ def from_dict(cls, data: dict[str, Any]) -> BaseTask:
         Raises:
             TaskValidationError: If required fields are missing or invalid
         """
-        # Determine task type based on fields present
-        if "prompt" in data:
+        # Determine task type based on fields present (use .get() so that
+        # None-valued keys don't misroute the task class selection).
+        if data.get("prompt"):
             return AITask(**data)
-        elif "executable" in data or "function" in data:
+        elif data.get("executable") or data.get("function"):
             return ComputeTask(**data)
         else:
             raise TaskValidationError(

src/rhapsody/backends/execution/__init__.py

@@ -7,8 +7,9 @@
 from __future__ import annotations
 
 from .concurrent import ConcurrentExecutionBackend  # noqa: F401
+from .noop import NoopExecutionBackend  # noqa: F401
 
-__all__ = ["ConcurrentExecutionBackend"]
+__all__ = ["ConcurrentExecutionBackend", "NoopExecutionBackend"]
 
 # Try to import optional backends
 try:

src/rhapsody/backends/execution/concurrent.py

@@ -34,10 +34,17 @@ def _get_logger() -> logging.Logger:
 class ConcurrentExecutionBackend(BaseBackend):
     """Simple async-only concurrent execution backend."""
 
-    def __init__(self, executor: Executor = None, name: str = "concurrent"):
+    def __init__(
+        self, executor: Executor = None, name: str = "concurrent", resources: dict | None = None
+    ):
         super().__init__(name=name)
 
         self.logger = _get_logger()
+        self._resources = resources or {}
+
+        # Concurrent backend does not support partitions
+        if self._resources.get("partition"):
+            raise ValueError("ConcurrentExecutionBackend does not support partitions")
 
         if not executor:
             executor = ThreadPoolExecutor()

src/rhapsody/backends/execution/dask_parallel.py

@@ -126,6 +126,10 @@ def __init__(
         self._initialized = False
         self._backend_state = BackendMainStates.INITIALIZED
 
+        # Dask backend does not support partitions
+        if self._resources.get("partition"):
+            raise ValueError("DaskExecutionBackend does not support partitions")
+
     def __await__(self):
         """Make DaskExecutionBackend awaitable like Dask Client."""
         return self._async_init().__await__()

src/rhapsody/backends/execution/dragon.py

@@ -3118,13 +3118,17 @@ def __init__(
         self,
         batch_kwargs: Optional[dict] = None,
         name: Optional[str] = "dragon",
+        resources: Optional[dict] = None,
     ):
         if not Batch:
             raise RuntimeError("Dragon Batch not available")
 
         super().__init__(name=name)
 
         self.logger = _get_logger()
+        self._resources = resources or {}
+        if self._resources:
+            raise NotImplementedError("DragonExecutionBackendV3 does not yet support resources")
         self.batch = Batch(**(batch_kwargs or {}))
 
         self._backend_state = BackendMainStates.INITIALIZED

src/rhapsody/backends/execution/noop.py

@@ -0,0 +1,124 @@
+"""No-op execution backend for performance benchmarking.
+
+Tasks are immediately marked as DONE without executing anything.
+"""
+
+import asyncio
+import logging
+from typing import Any
+from typing import Callable
+
+from ..base import BaseBackend
+from ..constants import BackendMainStates
+from ..constants import StateMapper
+
+
+def _get_logger() -> logging.Logger:
+    return logging.getLogger(__name__)
+
+
+class NoopExecutionBackend(BaseBackend):
+    """Backend that completes every task instantly.
+
+    Useful for measuring Orbit/bridge/client overhead without any actual task execution cost.
+    """
+
+    def __init__(self, name: str = "noop"):
+        super().__init__(name=name)
+        self.logger = _get_logger()
+        self.tasks: dict[str, dict] = {}
+        self._futures: dict[str, asyncio.Task] = {}
+        self._callback_func: Callable = lambda t, s: None
+        self._initialized = False
+        self._backend_state = BackendMainStates.INITIALIZED
+
+    def __await__(self):
+        return self._async_init().__await__()
+
+    async def _async_init(self):
+        if not self._initialized:
+            StateMapper.register_backend_states_with_defaults(backend=self)
+            StateMapper.register_backend_tasks_states_with_defaults(backend=self)
+            self._backend_state = BackendMainStates.INITIALIZED
+            self._initialized = True
+            self.logger.info("Noop execution backend started")
+        return self
+
+    def get_task_states_map(self):
+        return StateMapper(backend=self)
+
+    async def submit_tasks(self, tasks: list[dict[str, Any]]) -> None:
+        if self._backend_state != BackendMainStates.RUNNING:
+            self._backend_state = BackendMainStates.RUNNING
+
+        for task in tasks:
+            task.update(
+                {
+                    "return_value": True,
+                    "stdout": "",
+                    "stderr": "",
+                    "exit_code": 0,
+                }
+            )
+            uid = task["uid"]
+            self.tasks[uid] = task
+            # Track the completion future so it can be cancelled via
+            # cancel_task / cancel_all_tasks / shutdown.
+            self._futures[uid] = asyncio.create_task(self._complete(task))
+
+    async def _complete(self, task: dict) -> None:
+        try:
+            task["state"] = "DONE"
+            self._callback_func(task, "DONE")
+        finally:
+            self._futures.pop(task["uid"], None)
+
+    async def cancel_task(self, uid: str) -> bool:
+        if uid not in self.tasks:
+            return False
+
+        future = self._futures.pop(uid, None)
+        if future is not None and not future.done():
+            future.cancel()
+
+        task = self.tasks[uid]
+        if task.get("state") not in ("DONE", "FAILED", "CANCELED"):
+            task["state"] = "CANCELED"
+            self._callback_func(task, "CANCELED")
+        return True
+
+    async def cancel_all_tasks(self) -> int:
+        uids = list(self.tasks)
+        for uid in uids:
+            await self.cancel_task(uid)
+        return len(uids)
+
+    async def shutdown(self) -> None:
+        await self.cancel_all_tasks()
+        self._backend_state = BackendMainStates.SHUTDOWN
+        self.tasks.clear()
+        self._futures.clear()
+        self.logger.info("Noop execution backend shutdown")
+
+    def build_task(self, uid, task_desc, task_specific_kwargs):
+        pass
+
+    def link_explicit_data_deps(self, src_task=None, dst_task=None, file_name=None, file_path=None):
+        pass
+
+    def link_implicit_data_deps(self, src_task, dst_task):
+        pass
+
+    async def state(self) -> str:
+        return self._backend_state.value
+
+    def task_state_cb(self):
+        pass
+
+    async def __aenter__(self):
+        if not self._initialized:
+            await self._async_init()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.shutdown()

src/rhapsody/backends/execution/radical_pilot.py

@@ -10,6 +10,7 @@
 import copy
 import logging
 import os
+import shlex
 import threading
 from collections.abc import Generator
 from typing import Any
@@ -234,16 +235,52 @@ async def _async_init(self):
         return self
 
     async def _initialize(self) -> None:
-        """Initialize Radical Pilot components."""
+        """Initialize Radical Pilot components.
+
+        If partition info is provided in resources, configures the pilot to:
+        - Use only the specified number of nodes
+        - Set partition environment variables via prepare_env
+        """
         try:
             self.tasks = {}
             self.raptor_mode = False
+
+            # Work on a local copy and strip "partition" from it, so the pilot
+            # description doesn't receive partition config while self.resources
+            # stays intact (caller's dict is not mutated, and a retry of
+            # _initialize still sees the partition entry).
+            resources = dict(self.resources)
+            partition = resources.pop("partition", {})
+            partition_nodelist = partition.get("nodelist", [])
+            partition_env = partition.get("env", {})
+
             self.session = rp.Session(uid=ru.generate_id("rhapsody.session", mode=ru.ID_PRIVATE))
             self.task_manager = rp.TaskManager(self.session)
             self.pilot_manager = rp.PilotManager(self.session)
-            self.resource_pilot = self.pilot_manager.submit_pilots(
-                rp.PilotDescription(self.resources)
-            )
+
+            # Create pilot description
+            pd = rp.PilotDescription(resources)
+
+            # Configure partition if specified
+            if partition_nodelist:
+                pd.nodes = len(partition_nodelist)
+
+            if partition_env:
+                # Create shell environment with export directives.  Quote keys
+                # and values with shlex to avoid shell injection / breakage on
+                # spaces or special characters.
+                export_cmds = [
+                    f"export {shlex.quote(str(k))}={shlex.quote(str(v))}"
+                    for k, v in partition_env.items()
+                ]
+                pd.prepare_env = {
+                    "partition": {
+                        "type": "shell",
+                        "pre_exec": export_cmds,
+                    }
+                }
+
+            self.resource_pilot = self.pilot_manager.submit_pilots(pd)
             self.pilot_manager.register_callback(self.handle_pilot_state_callback)
 
             self.task_manager.add_pilots(self.resource_pilot)

tests/performance/test_api_performance.py

@@ -69,7 +69,7 @@ def test_task_creation_performance(self):
         duration = time.time() - start
 
         print(f"\nTask Creation (100K): {duration:.4f}s ({duration / n * 1e6:.2f} μs/task)")
-        assert duration < 1.0  # Should be well under 1s
+        assert duration < 2.0  # Should be well under 2s
 
     def test_serialization_performance(self):
         """Benchmark pickling 100,000 task objects."""

tox.ini

@@ -88,6 +88,15 @@ deps =
 commands =
     pytest tests/integration/ {posargs}
 
+[testenv:unit]
+deps =
+    pytest
+    pytest-asyncio
+    psycopg2-binary
+    .[dev,tracing,metrics,serialization]
+commands =
+    pytest tests/unit/ {posargs}
+
 ; [testenv:notebooks]
 ; deps =
 ;     nbconvert