Skip to content
Closed
Show file tree
Hide file tree
Changes from 46 commits
Commits
Show all changes
51 commits
Select commit Hold shift + click to select a range
100eeee
basics for RM move from RP
andre-merzky Nov 10, 2025
30494a6
snap
andre-merzky Nov 10, 2025
3bae4e8
fix slurm rm
andre-merzky Nov 20, 2025
33c3df9
Merge branch 'main' into feature/rm
andre-merzky Nov 26, 2025
bc1055b
snap
andre-merzky Nov 26, 2025
fd8b3b2
Merge branch 'improve/execution_backends' into feature/rm
andre-merzky Jan 13, 2026
dc63bd9
linting, dataclass
andre-merzky Jan 13, 2026
9492971
linting
andre-merzky Jan 13, 2026
c41983e
snap
andre-merzky Jan 13, 2026
0c10406
linting
andre-merzky Jan 13, 2026
169e675
linting
andre-merzky Jan 16, 2026
dbea907
add partitioning to RM
andre-merzky Jan 17, 2026
33752d5
remove rc dependency
andre-merzky Jan 19, 2026
20222e6
fix RM bugs: align _initialize interface, fix logger calls, fix _get_…
andre-merzky Feb 3, 2026
72186a5
add hostlist compactify/expand methods, refactor hostlist parsing
andre-merzky Feb 5, 2026
2c7739b
add partition environment variable support for all RMs
andre-merzky Feb 5, 2026
5b81177
add partition support to execution backends, fix RM bugs
andre-merzky Feb 5, 2026
f50496f
clean up resource_manager: remove redundant comments, fix docstrings
andre-merzky Feb 5, 2026
53ac4c6
Refactor Resource Manager code
andre-merzky Feb 5, 2026
86bb452
snap
andre-merzky Feb 5, 2026
4297d9a
Merge branch 'main' into feature/rm
andre-merzky Feb 6, 2026
f7ac6f9
Merge branch 'main' into feature/rm
andre-merzky Feb 20, 2026
bd0df7a
first shot at edge backend
andre-merzky Apr 8, 2026
41228d7
add missing files
andre-merzky Apr 8, 2026
c814d23
backend refactoring
andre-merzky Apr 8, 2026
0e34f2f
snap
andre-merzky Apr 9, 2026
dc81e05
fix task routing and improve error reporting
andre-merzky Apr 10, 2026
c002ef5
Add noop execution backend for performance benchmarking
andre-merzky Apr 10, 2026
b976560
snap
andre-merzky Apr 10, 2026
60da2ad
snap
andre-merzky Apr 10, 2026
ca7fbd7
snap
andre-merzky Apr 10, 2026
0aef47d
fix silent import error
andre-merzky Apr 26, 2026
6a656bc
edge backend: fail-fast on Python version skew for cloudpickled tasks
andre-merzky Apr 27, 2026
a9522eb
edge backend: auto-select edge, simplify
andre-merzky Apr 30, 2026
1bcf587
Merge branch 'main' into feature/edge
andre-merzky May 1, 2026
d4c3fe6
collect stderr in case of failure for task diagnosis
andre-merzky May 1, 2026
606e232
revert v1 fix, v3 fix for missing stderr
andre-merzky May 1, 2026
735e6d5
another v3 fix for missing stderr
andre-merzky May 1, 2026
4d3eaf1
another v3 fix for missing stderr
andre-merzky May 1, 2026
d0f60ca
another v3 fix for missing stderr
andre-merzky May 2, 2026
4547c73
merge snapshot - untested
andre-merzky May 12, 2026
05ae574
merge from main
andre-merzky May 18, 2026
9eff60c
migrate rm to new repo
andre-merzky May 26, 2026
379ce85
Merge remote-tracking branch 'origin/dev' into feature/rm-reconcile
andre-merzky Jul 1, 2026
c5f2689
backends: drop edge backend, superseded by orbit
andre-merzky Jul 1, 2026
fa1da76
address gemini review — Makefile typo, RM partition safety
andre-merzky Jul 2, 2026
b290390
style: satisfy pre-commit formatters (ruff-format/docformatter)
andre-merzky Jul 2, 2026
ef1ce1f
address gemini review — RM resources retry-safety and noop task lifec…
andre-merzky Jul 2, 2026
9fbfccd
address gemini review round 2 — noop/task (shared) + radical_pilot
andre-merzky Jul 2, 2026
4af2cd1
concurrent: add `from __future__ import annotations` for Python 3.9
andre-merzky Jul 2, 2026
45c4cc7
noop: address gemini review — task_state_cb signature + create() factory
andre-merzky Jul 2, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/rhapsody/api/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
7 changes: 4 additions & 3 deletions src/rhapsody/api/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Comment thread
andre-merzky marked this conversation as resolved.
Outdated
Comment thread
andre-merzky marked this conversation as resolved.
Outdated
else:
raise TaskValidationError(
Expand Down
1 change: 1 addition & 0 deletions src/rhapsody/backends/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ class BackendMainStates(Enum):
SHUTDOWN = "SHUTDOWN"



class TasksMainStates(Enum):
"""Enumeration of standard task states used across all backends.

Expand Down
3 changes: 2 additions & 1 deletion src/rhapsody/backends/execution/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
7 changes: 6 additions & 1 deletion src/rhapsody/backends/execution/concurrent.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,10 +34,15 @@ 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()
Expand Down
4 changes: 4 additions & 0 deletions src/rhapsody/backends/execution/dask_parallel.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__()
Expand Down
4 changes: 4 additions & 0 deletions src/rhapsody/backends/execution/dragon.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
105 changes: 105 additions & 0 deletions src/rhapsody/backends/execution/noop.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,105 @@
"""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._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]]) -> list[asyncio.Task]:
if self._backend_state != BackendMainStates.RUNNING:
self._backend_state = BackendMainStates.RUNNING

submitted = []
for task in tasks:
task.update(
{
"return_value": True,
"stdout": "",
"stderr": "",
"exit_code": 0,
}
)
self.tasks[task["uid"]] = task
future = asyncio.create_task(self._complete(task))
submitted.append(future)
Comment thread
andre-merzky marked this conversation as resolved.
Outdated
return submitted

async def _complete(self, task: dict) -> None:
self._callback_func(task, "DONE")

async def cancel_task(self, uid: str) -> bool:
return uid in self.tasks

async def cancel_all_tasks(self) -> int:
n = len(self.tasks)
self.tasks.clear()
return n

async def shutdown(self) -> None:
self._backend_state = BackendMainStates.SHUTDOWN
self.tasks.clear()
self.logger.info("Noop execution backend shutdown")
Comment thread
andre-merzky marked this conversation as resolved.

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
Comment thread
andre-merzky marked this conversation as resolved.
Outdated

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()
Comment thread
andre-merzky marked this conversation as resolved.
43 changes: 39 additions & 4 deletions src/rhapsody/backends/execution/radical_pilot.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import copy
import logging
import os
import shlex
import threading
from collections.abc import Generator
from typing import Any
Expand Down Expand Up @@ -234,16 +235,50 @@ 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

# Extract and remove partition info from resources. Copy first so
# we don't mutate the caller-provided dict in place.
self.resources = dict(self.resources)
partition = self.resources.pop("partition", {})
partition_nodelist = partition.get("nodelist", [])
partition_env = partition.get("env", {})
Comment thread
andre-merzky marked this conversation as resolved.
Outdated

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(self.resources)
Comment thread
andre-merzky marked this conversation as resolved.
Outdated

# 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)
Expand Down
2 changes: 1 addition & 1 deletion tests/performance/test_api_performance.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
9 changes: 9 additions & 0 deletions tox.ini
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading