From b5f0a68e41e9567e2606bd5114b98d99adc68dc5 Mon Sep 17 00:00:00 2001 From: harikrishna1410 Date: Wed, 15 Jul 2026 17:47:17 -0500 Subject: [PATCH 1/6] added ensembleexecutionbackend --- src/rhapsody/backends/__init__.py | 8 + src/rhapsody/backends/execution/__init__.py | 7 + src/rhapsody/backends/execution/el.py | 268 ++++++++++ tests/unit/test_backend_execution_el.py | 548 ++++++++++++++++++++ tutorials/el-backend-tutorial.ipynb | 213 ++++++++ 5 files changed, 1044 insertions(+) create mode 100644 src/rhapsody/backends/execution/el.py create mode 100644 tests/unit/test_backend_execution_el.py create mode 100644 tutorials/el-backend-tutorial.ipynb diff --git a/src/rhapsody/backends/__init__.py b/src/rhapsody/backends/__init__.py index b544c18..87ff379 100644 --- a/src/rhapsody/backends/__init__.py +++ b/src/rhapsody/backends/__init__.py @@ -59,6 +59,14 @@ except ImportError: pass +try: + from .execution import EnsembleBackend + + __all__.append(EnsembleBackend) + +except ImportError: + pass + # Try to import optional inference backends try: from .inference.vllm import DragonVllmInferenceBackend # noqa: F401 diff --git a/src/rhapsody/backends/execution/__init__.py b/src/rhapsody/backends/execution/__init__.py index 50ec4db..e311c95 100644 --- a/src/rhapsody/backends/execution/__init__.py +++ b/src/rhapsody/backends/execution/__init__.py @@ -25,6 +25,13 @@ except ImportError: pass +try: + from .el import EnsembleBackend # noqa: F401 + + __all__.append("EnsembleBackend") +except ImportError: + pass + try: from .dragon import DragonExecutionBackendV1 # noqa: F401 from .dragon import DragonExecutionBackendV2 # noqa: F401 diff --git a/src/rhapsody/backends/execution/el.py b/src/rhapsody/backends/execution/el.py new file mode 100644 index 0000000..77c882e --- /dev/null +++ b/src/rhapsody/backends/execution/el.py @@ -0,0 +1,268 @@ +"""Ensemble Launcher execution backend for distributed computing. + +This module provides a backend that executes tasks via the Ensemble Launcher framework, +supporting MPI-based distributed execution environments. +""" + +from __future__ import annotations + +import asyncio +import inspect +import logging +import os +import uuid +from typing import Any +from typing import Callable +from typing import Optional +from typing import Union + +from ensemble_launcher import EnsembleLauncher +from ensemble_launcher.config import LauncherConfig +from ensemble_launcher.config import MPIConfig +from ensemble_launcher.config import PolicyConfig +from ensemble_launcher.config import SystemConfig +from ensemble_launcher.ensemble import AsyncTask as AsyncELTask +from ensemble_launcher.ensemble import Task as ELTask +from ensemble_launcher.helper_functions import get_nodes +from ensemble_launcher.orchestrator import ClusterClient + +from ..base import BaseBackend +from ..constants import BackendMainStates +from ..constants import StateMapper + + +class EnsembleBackend(BaseBackend): + + def __init__(self, name: str | None = None, + child_executor_name: str = "async_mpi", + return_stdout: bool = True, + worker_logs: bool = True, + master_logs: bool = True, + gpu_selector: str = "ZE_AFFINITY_MASK", + children_scheduler_policy: str = "fixed_leafs_children_policy", + task_scheduler_policy: str = "large_resource_policy", + enable_workstealing: bool = False, + checkpoint_dir: str | None = None, + mpi_flavour: str = "mpich", + nlevels: int = 0, + nleafs: int | None = None, + cpus: list[int] | None = None, + gpus: list[int] | None = None, + client_only: bool = False, + node_id: str = "global", + ): + super().__init__(name=name) + + self.logger = logging.getLogger(__name__) + self._initialized = False + self._backend_state = BackendMainStates.INITIALIZED + self._callback_func: Callable = lambda t, s: None + self._client_only = client_only + self._node_id = node_id + + task_executor_name = ["async_loky", "async_mpi"] + + self._launcher_config = LauncherConfig( + child_executor_name=child_executor_name, + task_executor_name=task_executor_name, + return_stdout=return_stdout, + worker_logs=worker_logs, + master_logs=master_logs, + gpu_selector=gpu_selector, + children_scheduler_policy=children_scheduler_policy, + task_scheduler_policy=task_scheduler_policy, + policy_config=PolicyConfig(nlevels=nlevels, leaf_nodes=nleafs or len(get_nodes())), + enable_workstealing=enable_workstealing, + cluster=True, + checkpoint_dir=checkpoint_dir or os.path.join(os.getcwd(), f"ckpt_{uuid.uuid4()}"), + mpi_config=MPIConfig(flavor=mpi_flavour), + ) + cpus = cpus or list(range(os.cpu_count())) + ngpus = len(gpus) if gpus is not None else 0 + gpus = gpus or [] + self._sys_config = SystemConfig(name="cluster", ncpus=len(cpus), cpus=cpus, ngpus=ngpus, gpus=gpus) + self._el: EnsembleLauncher | None = None + self._client: ClusterClient | None = None + self.tasks: dict = {} + + def __await__(self): + """Make EnsembleBackend awaitable.""" + return self._async_init().__await__() + + async def _async_init(self): + if not self._initialized: + try: + self.logger.debug("Registering backend states...") + StateMapper.register_backend_states_with_defaults(backend=self) + + self.logger.debug("Registering task states...") + StateMapper.register_backend_tasks_states_with_defaults(backend=self) + + self._backend_state = BackendMainStates.INITIALIZED + self.logger.debug(f"Backend state set to: {self._backend_state.value}") + + await self._initialize() + self._initialized = True + + self.logger.info("Ensemble backend fully initialized and ready") + + except Exception as e: + self.logger.exception(f"Ensemble backend initialization failed: {e}") + self._initialized = False + raise + return self + + async def _initialize(self): + if not self._client_only: + self._el = EnsembleLauncher( + ensemble_file={}, + system_config=self._sys_config, + launcher_config=self._launcher_config, + ) + await asyncio.to_thread(self._el.start, wait_time=5) + + self._client = ClusterClient( + self._launcher_config.checkpoint_dir, + node_id=self._node_id, + ) + await asyncio.to_thread(self._client.start) + + def _ensure_initialized(self): + if not self._initialized: + raise RuntimeError( + "EnsembleBackend must be awaited before use. " + "Use: backend = await EnsembleBackend(...)" + ) + + async def _handle_task(self, task: dict) -> None: + try: + is_executable = task.get("executable", None) is not None + self._callback_func(task, "RUNNING") + el_task = self.build_task(task) + fut = self._client.submit(el_task) + task["future"] = fut + result = await asyncio.wrap_future(fut) + if is_executable: + task["return_value"] = "" + task["stdout"] = result.split(",")[0] ## EL returns "stdout,stderr" + task["stderr"] = result.split(",")[1] + else: + task["return_value"] = result + task["stdout"] = "" + task["stderr"] = "" + self._callback_func(task, "DONE") + except Exception as e: + task["exception"] = e + task["stdout"] = "" + task["stderr"] = str(e) + self._callback_func(task, "FAILED") + + async def submit_tasks(self, tasks: list[dict]) -> None: + self._ensure_initialized() + + if self._backend_state == BackendMainStates.SHUTDOWN: + raise RuntimeError("Cannot submit during shutdown") + + if self._backend_state != BackendMainStates.RUNNING: + self._backend_state = BackendMainStates.RUNNING + self.logger.debug(f"Backend state set to: {self._backend_state.value}") + + for task in tasks: + future = asyncio.create_task(self._handle_task(task)) + self.tasks[task["uid"]] = task + self.tasks[task["uid"]]["future"] = future + + async def shutdown(self) -> None: + self._backend_state = BackendMainStates.SHUTDOWN + self.logger.debug(f"Backend state set to: {self._backend_state.value}") + + try: + if self._client is not None: + self._client.teardown() + if self._el is not None: + self._el.stop() + except Exception as e: + self.logger.exception(f"Error during shutdown: {e}") + finally: + self._client = None + self._el = None + self.tasks.clear() + self._initialized = False + self.logger.info("Ensemble execution backend shutdown complete") + + async def state(self) -> str: + return self._backend_state.value + + def task_state_cb(self, task: dict, state: str) -> None: + self._callback_func(task, state) + + def register_callback(self, func: Callable[[dict[str, Any], str], None]) -> None: + self._callback_func = func + + def get_task_states_map(self) -> Any: + return StateMapper(backend=self) + + def build_task(self, task: dict) -> ELTask | AsyncELTask: + backend_kwargs = task.get("task_backend_specific_kwargs", {}) + nnodes = backend_kwargs.get("nnodes", 1) + ppn = backend_kwargs.get("ranks", 1) // nnodes + ngpus_per_process = backend_kwargs.get("gpus_per_rank", 0) + env = backend_kwargs.get("env", {}) + cpu_affinity = list(map(int, backend_kwargs.get("cpu_affinity").split(","))) if "cpu_affinity" in backend_kwargs else [] + gpu_affinity = list(map(int, backend_kwargs.get("gpu_affinity").split(","))) if "gpu_affinity" in backend_kwargs else [] + + print(cpu_affinity, gpu_affinity) + + func = task.get("function") + if func is not None and inspect.iscoroutinefunction(func): + task_class = AsyncELTask + else: + task_class = ELTask + + is_executable = task.get("executable", None) is not None + if is_executable: + ## When ELTask.executable is str, it ignore args and kwargs. + exec = task["executable"] + " " + " ".join([str(x) for x in task["arguments"]]) + else: + exec = func + + return task_class( + task_id=task["uid"], + nnodes=nnodes, + ppn=ppn, + ngpus_per_process=ngpus_per_process, + executable=exec, + args=task["args"], + kwargs=task["kwargs"], + executor_name="async_mpi" if is_executable else "async_loky", + env=env, + cpu_affinity=cpu_affinity, + gpu_affinity=gpu_affinity, + ) + + def link_implicit_data_deps(self, src_task: dict[str, Any], dst_task: dict[str, Any]) -> None: # noqa: B027 + pass + + def link_explicit_data_deps( # noqa: B027 + self, + src_task: dict[str, Any] | None = None, + dst_task: dict[str, Any] | None = None, + file_name: str | None = None, + file_path: str | None = None, + ) -> None: + pass + + async def cancel_task(self, uid: str) -> bool: + if uid in self.tasks: + future = self.tasks[uid].get("future") + if future: + return future.cancel() + return False + + 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() diff --git a/tests/unit/test_backend_execution_el.py b/tests/unit/test_backend_execution_el.py new file mode 100644 index 0000000..b3b3e04 --- /dev/null +++ b/tests/unit/test_backend_execution_el.py @@ -0,0 +1,548 @@ +"""Unit tests for Ensemble Launcher execution backend.""" + +import asyncio +from concurrent.futures import Future as ConcurrentFuture +from unittest.mock import MagicMock +from unittest.mock import patch + +import pytest + +from rhapsody import ComputeTask + +try: + from rhapsody.backends.execution.el import EnsembleBackend + + _el_available = True +except ImportError: + _el_available = False + +pytestmark = pytest.mark.skipif(not _el_available, reason="ensemble_launcher not installed") + + +# --------------------------------------------------------------------------- +# Import and class structure +# --------------------------------------------------------------------------- + + +def test_el_backend_import(): + from rhapsody.backends.execution import EnsembleBackend as EB + + assert EB is not None + + +def test_el_backend_inherits_base(): + from rhapsody.backends.base import BaseBackend + + assert issubclass(EnsembleBackend, BaseBackend) + + +# --------------------------------------------------------------------------- +# __init__ +# --------------------------------------------------------------------------- + + +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +def test_el_backend_init_defaults(mock_nodes): + backend = EnsembleBackend() + assert backend is not None + assert not backend._initialized + assert backend._client is None + assert backend._el is None + assert backend.tasks == {} + assert backend._client_only is False + assert backend._node_id == "global" + + +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +def test_el_backend_init_client_only(mock_nodes): + backend = EnsembleBackend(client_only=True, node_id="main.w0", checkpoint_dir="/tmp/ckpt") + assert backend._client_only is True + assert backend._node_id == "main.w0" + + +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +def test_el_backend_init_custom_name(mock_nodes): + backend = EnsembleBackend(name="my_el") + assert backend.name == "my_el" + + +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +def test_el_backend_init_default_name(mock_nodes): + backend = EnsembleBackend() + assert backend.name == "EnsembleBackend" + + +# --------------------------------------------------------------------------- +# Awaitable / async init +# --------------------------------------------------------------------------- + + +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +def test_el_backend_is_awaitable(mock_nodes): + backend = EnsembleBackend() + assert hasattr(backend, "__await__") + + +@pytest.mark.asyncio +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +async def test_el_backend_async_init(mock_nodes): + backend = EnsembleBackend() + + mock_el = MagicMock() + mock_client = MagicMock() + + with patch("rhapsody.backends.execution.el.EnsembleLauncher", return_value=mock_el), \ + patch("rhapsody.backends.execution.el.ClusterClient", return_value=mock_client), \ + patch("rhapsody.backends.execution.el.asyncio.to_thread", side_effect=_sync_to_thread): + result = await backend + assert result is backend + assert backend._initialized + mock_el.start.assert_called_once_with(wait_time=5) + mock_client.start.assert_called_once() + + +@pytest.mark.asyncio +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +async def test_el_backend_async_init_client_only(mock_nodes): + backend = EnsembleBackend(client_only=True, checkpoint_dir="/tmp/ckpt") + + mock_client = MagicMock() + + with patch("rhapsody.backends.execution.el.EnsembleLauncher") as mock_el_cls, \ + patch("rhapsody.backends.execution.el.ClusterClient", return_value=mock_client), \ + patch("rhapsody.backends.execution.el.asyncio.to_thread", side_effect=_sync_to_thread): + await backend + assert backend._initialized + assert backend._el is None + mock_el_cls.assert_not_called() + mock_client.start.assert_called_once() + + +# --------------------------------------------------------------------------- +# Context manager +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +async def test_el_backend_context_manager(mock_nodes): + backend = EnsembleBackend() + + mock_el = MagicMock() + mock_client = MagicMock() + + with patch("rhapsody.backends.execution.el.EnsembleLauncher", return_value=mock_el), \ + patch("rhapsody.backends.execution.el.ClusterClient", return_value=mock_client), \ + patch("rhapsody.backends.execution.el.asyncio.to_thread", side_effect=_sync_to_thread): + async with backend as b: + assert b._initialized + assert b is backend + + assert not backend._initialized + mock_client.teardown.assert_called_once() + mock_el.stop.assert_called_once() + + +# --------------------------------------------------------------------------- +# Callback registration +# --------------------------------------------------------------------------- + + +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +def test_el_backend_callback_registration(mock_nodes): + backend = EnsembleBackend() + + def my_cb(task, state): + pass + + backend.register_callback(my_cb) + assert backend._callback_func is my_cb + + +# --------------------------------------------------------------------------- +# State +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +async def test_el_backend_state(mock_nodes): + backend = EnsembleBackend() + state = await backend.state() + assert state == "INITIALIZED" + + +# --------------------------------------------------------------------------- +# State mapper +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +async def test_el_backend_state_mapper(mock_nodes): + backend = EnsembleBackend() + + with patch("rhapsody.backends.execution.el.EnsembleLauncher", return_value=MagicMock()), \ + patch("rhapsody.backends.execution.el.ClusterClient", return_value=MagicMock()), \ + patch("rhapsody.backends.execution.el.asyncio.to_thread", side_effect=_sync_to_thread): + await backend + + mapper = backend.get_task_states_map() + assert mapper is not None + + +# --------------------------------------------------------------------------- +# submit_tasks — not initialized guard +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +async def test_el_backend_submit_not_initialized(mock_nodes): + backend = EnsembleBackend() + + with pytest.raises(RuntimeError, match="EnsembleBackend must be awaited"): + await backend.submit_tasks([]) + + +# --------------------------------------------------------------------------- +# submit_tasks — happy path +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +async def test_el_backend_submit_tasks(mock_nodes): + backend = EnsembleBackend() + backend._initialized = True + from rhapsody.backends.constants import BackendMainStates + backend._backend_state = BackendMainStates.INITIALIZED + + captured = [] + backend.register_callback(lambda t, s: captured.append((t["uid"], s))) + + fut = ConcurrentFuture() + mock_client = MagicMock() + mock_client.submit.return_value = fut + backend._client = mock_client + + task = ComputeTask( + function=lambda x: x * 2, + args=(5,), + kwargs={}, + ) + + await backend.submit_tasks([task]) + + assert task["uid"] in backend.tasks + # future stored is the asyncio.Task wrapping _handle_task + assert isinstance(backend.tasks[task["uid"]]["future"], asyncio.Task) + state = await backend.state() + assert state == "RUNNING" + + # Resolve the underlying future so the task completes + fut.set_result(10) + await backend.tasks[task["uid"]]["future"] + + states = [s for _, s in captured] + assert "RUNNING" in states + assert "DONE" in states + + +@pytest.mark.asyncio +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +async def test_el_submit_multiple_tasks_complete(mock_nodes): + """All submitted tasks complete and fire callbacks.""" + backend = EnsembleBackend() + backend._initialized = True + from rhapsody.backends.constants import BackendMainStates + backend._backend_state = BackendMainStates.INITIALIZED + backend.register_callback(lambda t, s: None) + + mock_client = MagicMock() + futs = [ConcurrentFuture() for _ in range(3)] + mock_client.submit.side_effect = futs + backend._client = mock_client + + tasks = [ + ComputeTask(function=lambda: 1, args=(), kwargs={}) + for _ in range(3) + ] + + await backend.submit_tasks(tasks) + async_tasks = [backend.tasks[t["uid"]]["future"] for t in tasks] + + for f in futs: + f.set_result(42) + await asyncio.gather(*async_tasks) + + for t in tasks: + assert t["return_value"] == 42 + + +# --------------------------------------------------------------------------- +# submit_tasks — shutdown guard +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +async def test_el_backend_submit_after_shutdown(mock_nodes): + backend = EnsembleBackend() + backend._initialized = True + from rhapsody.backends.constants import BackendMainStates + backend._backend_state = BackendMainStates.SHUTDOWN + + with pytest.raises(RuntimeError, match="Cannot submit during shutdown"): + await backend.submit_tasks([]) + + +# --------------------------------------------------------------------------- +# build_task +# --------------------------------------------------------------------------- + + +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +def test_el_build_task_sync_function(mock_nodes): + from ensemble_launcher.ensemble import Task as ELTask + + backend = EnsembleBackend() + + def my_func(a, b): + return a + b + + task = ComputeTask( + function=my_func, + args=(1, 2), + kwargs={}, + task_backend_specific_kwargs={"nnodes": 2, "ranks": 4, "gpus_per_rank": 1}, + ) + + el_task = backend.build_task(task) + assert isinstance(el_task, ELTask) + assert el_task.nnodes == 2 + assert el_task.ppn == 2 + assert el_task.ngpus_per_process == 1 + + +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +def test_el_build_task_async_function(mock_nodes): + from ensemble_launcher.ensemble import AsyncTask as AsyncELTask + + backend = EnsembleBackend() + + async def my_async_func(a): + return a + + task = ComputeTask( + function=my_async_func, + args=(1,), + kwargs={}, + ) + + el_task = backend.build_task(task) + assert isinstance(el_task, AsyncELTask) + + +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +def test_el_build_task_executable(mock_nodes): + from ensemble_launcher.ensemble import Task as ELTask + + backend = EnsembleBackend() + + task = ComputeTask( + executable="/bin/echo", + args=("hello",), + kwargs={}, + ) + + el_task = backend.build_task(task) + assert isinstance(el_task, ELTask) + assert el_task.executable == "/bin/echo" + + +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +def test_el_build_task_default_resources(mock_nodes): + backend = EnsembleBackend() + + task = ComputeTask( + function=lambda: 1, + args=(), + kwargs={}, + ) + + el_task = backend.build_task(task) + assert el_task.nnodes == 1 + assert el_task.ppn == 1 + assert el_task.ngpus_per_process == 0 + + +# --------------------------------------------------------------------------- +# cancel_task +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +async def test_el_cancel_task_success(mock_nodes): + backend = EnsembleBackend() + + mock_future = MagicMock() + mock_future.cancel.return_value = True + backend.tasks["task-1"] = {"uid": "task-1", "future": mock_future} + + result = await backend.cancel_task("task-1") + assert result is True + mock_future.cancel.assert_called_once() + + +@pytest.mark.asyncio +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +async def test_el_cancel_task_not_found(mock_nodes): + backend = EnsembleBackend() + result = await backend.cancel_task("nonexistent") + assert result is False + + +# --------------------------------------------------------------------------- +# shutdown +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +async def test_el_shutdown(mock_nodes): + backend = EnsembleBackend() + backend._initialized = True + + mock_client = MagicMock() + mock_el = MagicMock() + backend._client = mock_client + backend._el = mock_el + backend.tasks = {"t1": {}} + + await backend.shutdown() + + mock_client.teardown.assert_called_once() + mock_el.stop.assert_called_once() + assert backend._client is None + assert backend._el is None + assert len(backend.tasks) == 0 + assert not backend._initialized + state = await backend.state() + assert state == "SHUTDOWN" + + +@pytest.mark.asyncio +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +async def test_el_shutdown_client_only(mock_nodes): + backend = EnsembleBackend(client_only=True, checkpoint_dir="/tmp/ckpt") + backend._initialized = True + + mock_client = MagicMock() + backend._client = mock_client + backend._el = None + + await backend.shutdown() + + mock_client.teardown.assert_called_once() + assert backend._client is None + assert not backend._initialized + + +@pytest.mark.asyncio +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +async def test_el_shutdown_without_init(mock_nodes): + backend = EnsembleBackend() + await backend.shutdown() + assert not backend._initialized + + +# --------------------------------------------------------------------------- +# Callback firing (done / failed) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +async def test_el_callback_done(mock_nodes): + backend = EnsembleBackend() + backend._initialized = True + from rhapsody.backends.constants import BackendMainStates + backend._backend_state = BackendMainStates.INITIALIZED + + captured = [] + backend.register_callback(lambda t, s: captured.append((t["uid"], s))) + + fut = ConcurrentFuture() + mock_client = MagicMock() + mock_client.submit.return_value = fut + backend._client = mock_client + + task = ComputeTask(function=lambda: 42, args=(), kwargs={}) + await backend.submit_tasks([task]) + + fut.set_result(42) + await backend.tasks[task["uid"]]["future"] + + states = [s for _, s in captured] + assert "RUNNING" in states + assert "DONE" in states + assert task["return_value"] == 42 + assert task["stdout"] == "" + assert task["stderr"] == "" + + +@pytest.mark.asyncio +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +async def test_el_callback_failed(mock_nodes): + backend = EnsembleBackend() + backend._initialized = True + from rhapsody.backends.constants import BackendMainStates + backend._backend_state = BackendMainStates.INITIALIZED + + captured = [] + backend.register_callback(lambda t, s: captured.append((t["uid"], s))) + + fut = ConcurrentFuture() + mock_client = MagicMock() + mock_client.submit.return_value = fut + backend._client = mock_client + + task = ComputeTask(function=lambda: 1, args=(), kwargs={}) + await backend.submit_tasks([task]) + + fut.set_exception(ValueError("boom")) + await backend.tasks[task["uid"]]["future"] + + states = [s for _, s in captured] + assert "RUNNING" in states + assert "FAILED" in states + assert isinstance(task["exception"], ValueError) + assert task["stdout"] == "" + assert "boom" in task["stderr"] + + +# --------------------------------------------------------------------------- +# task_state_cb +# --------------------------------------------------------------------------- + + +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +def test_el_task_state_cb(mock_nodes): + backend = EnsembleBackend() + captured = [] + backend.register_callback(lambda t, s: captured.append(s)) + + backend.task_state_cb({"uid": "x"}, "RUNNING") + assert captured == ["RUNNING"] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +async def _sync_to_thread(fn, /, *args, **kwargs): + """Drop-in replacement for asyncio.to_thread that runs synchronously.""" + return fn(*args, **kwargs) diff --git a/tutorials/el-backend-tutorial.ipynb b/tutorials/el-backend-tutorial.ipynb new file mode 100644 index 0000000..4833e18 --- /dev/null +++ b/tutorials/el-backend-tutorial.ipynb @@ -0,0 +1,213 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# RHAPSODY Ensemble Backend Tutorial\n", + "\n", + "This tutorial shows how to use the **`EnsembleExecutionBackend`** to run tasks on EnsembleLauncher cluster:" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import asyncio\n", + "\n", + "from rhapsody.api import ComputeTask, Session\n", + "from rhapsody.backends import EnsembleBackend" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "---\n", + "\n", + "## 1. Execute Functions" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "def compute_square(n):\n", + " \"\"\"Compute the square of a number.\"\"\"\n", + " return {\"input\": n, \"result\": n * n}" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Done: 20, Failed: 0\n", + " task.000001 -> {'input': 0, 'result': 0}\n", + " task.000002 -> {'input': 1, 'result': 1}\n", + " task.000003 -> {'input': 2, 'result': 4}\n" + ] + } + ], + "source": [ + "async def run_sync_functions():\n", + " async with EnsembleBackend() as backend:\n", + " session = Session(backends=[backend])\n", + " tasks = [ComputeTask(function=compute_square, args=(i,)) for i in range(20)]\n", + " async with session:\n", + " await session.submit_tasks(tasks)\n", + " await session.wait_tasks(tasks)\n", + "\n", + " done = sum(1 for t in tasks if t.state == \"DONE\")\n", + " failed = sum(1 for t in tasks if t.state == \"FAILED\")\n", + " print(f\"Done: {done}, Failed: {failed}\")\n", + " for t in tasks[:3]:\n", + " print(f\" {t.uid} -> {t.return_value}\")\n", + "\n", + "await run_sync_functions()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Expected output:**\n", + "```\n", + "Done: 20, Failed: 0\n", + " task.000001 -> {'input': 0, 'result': 0}\n", + " task.000002 -> {'input': 1, 'result': 1}\n", + " task.000003 -> {'input': 2, 'result': 4}\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "---\n", + "\n", + "## 2. Executable Tasks\n", + "\n", + "Executable task results are available via `task.stdout`, `task.stderr`." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "task.000001 | exit=None | stdout='task 0'\n", + "task.000002 | exit=None | stdout='task 1'\n", + "task.000003 | exit=None | stdout='task 2'\n", + "task.000004 | exit=None | stdout='task 3'\n", + "task.000005 | exit=None | stdout='task 4'\n" + ] + } + ], + "source": [ + "async def run_executables():\n", + " async with EnsembleBackend() as backend:\n", + " session = Session(backends=[backend])\n", + " tasks = [\n", + " ComputeTask(executable=\"/bin/echo\", arguments=[f\"task {i}\"])\n", + " for i in range(5)\n", + " ]\n", + " async with session:\n", + " await session.submit_tasks(tasks)\n", + " await session.wait_tasks(tasks)\n", + "\n", + " for t in tasks:\n", + " print(f\"{t.uid} | exit={t.exit_code} | stdout={t.stdout.strip()!r}\")\n", + "\n", + "await run_executables()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Expected output:**\n", + "```\n", + "task.000001 | exit=0 | stdout='task 0'\n", + "task.000002 | exit=0 | stdout='task 1'\n", + "...\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Resource awareness and pinning\n", + "EnsembleBackend supports resource-aware scheduling and pinning of tasks to specific resources. gpu pinning is controlled through `gpu_selector` argument of `EnsembleExecutionBackend`. The following example shows how to pin tasks to specific GPUs. The below example just shows that each task will have specific `gpu_selector` set correctly in its env. Note, `cpu_affinity` is only ensured when python `os` module allows it. Consequently, the `cpu_affinity` may not be set on some platforms, like macos." + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[0] [1]\n", + "[1] [0]\n", + "task.000007 | exit=None | stdout='1'\n", + "task.000008 | exit=None | stdout='0'\n" + ] + } + ], + "source": [ + "async def run_executables():\n", + " async with EnsembleBackend(gpu_selector=\"ZE_AFFINITY_MASK\", gpus=[0, 1]) as backend:\n", + " session = Session(backends=[backend])\n", + " tasks = [\n", + " ComputeTask(executable=\"printenv\", arguments=[\"ZE_AFFINITY_MASK\"], task_backend_specific_kwargs = {\"ranks\": 1, \"gpus_per_rank\": 1, \"cpu_affinity\":f\"{i}\", \"gpu_affinity\": f\"{gpu_id}\"})\n", + " for i, gpu_id in zip(list(range(2)),[1, 0]) #notice that the affinity is set opposite to the i\n", + " ]\n", + " async with session:\n", + " await session.submit_tasks(tasks)\n", + " await session.wait_tasks(tasks)\n", + "\n", + " for t in tasks:\n", + " print(f\"{t.uid} | exit={t.exit_code} | stdout={t.stdout.strip()!r}\")\n", + "\n", + "await run_executables()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.12" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} From 7f652ff6249976e2f7e3cfbd22d23c14368895af Mon Sep 17 00:00:00 2001 From: harikrishna1410 Date: Wed, 15 Jul 2026 18:13:10 -0500 Subject: [PATCH 2/6] Added MPI+serial example to demo --- src/rhapsody/backends/execution/el.py | 2 - tutorials/el-backend-tutorial.ipynb | 69 ++++++++++++++++++++++----- 2 files changed, 58 insertions(+), 13 deletions(-) diff --git a/src/rhapsody/backends/execution/el.py b/src/rhapsody/backends/execution/el.py index 77c882e..49883a2 100644 --- a/src/rhapsody/backends/execution/el.py +++ b/src/rhapsody/backends/execution/el.py @@ -211,8 +211,6 @@ def build_task(self, task: dict) -> ELTask | AsyncELTask: cpu_affinity = list(map(int, backend_kwargs.get("cpu_affinity").split(","))) if "cpu_affinity" in backend_kwargs else [] gpu_affinity = list(map(int, backend_kwargs.get("gpu_affinity").split(","))) if "gpu_affinity" in backend_kwargs else [] - print(cpu_affinity, gpu_affinity) - func = task.get("function") if func is not None and inspect.iscoroutinefunction(func): task_class = AsyncELTask diff --git a/tutorials/el-backend-tutorial.ipynb b/tutorials/el-backend-tutorial.ipynb index 4833e18..6aca089 100644 --- a/tutorials/el-backend-tutorial.ipynb +++ b/tutorials/el-backend-tutorial.ipynb @@ -32,7 +32,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 2, "metadata": {}, "outputs": [], "source": [ @@ -108,11 +108,11 @@ "name": "stdout", "output_type": "stream", "text": [ - "task.000001 | exit=None | stdout='task 0'\n", - "task.000002 | exit=None | stdout='task 1'\n", - "task.000003 | exit=None | stdout='task 2'\n", - "task.000004 | exit=None | stdout='task 3'\n", - "task.000005 | exit=None | stdout='task 4'\n" + "task.000021 | exit=None | stdout='task 0'\n", + "task.000022 | exit=None | stdout='task 1'\n", + "task.000023 | exit=None | stdout='task 2'\n", + "task.000024 | exit=None | stdout='task 3'\n", + "task.000025 | exit=None | stdout='task 4'\n" ] } ], @@ -156,17 +156,15 @@ }, { "cell_type": "code", - "execution_count": 12, + "execution_count": 5, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ - "[0] [1]\n", - "[1] [0]\n", - "task.000007 | exit=None | stdout='1'\n", - "task.000008 | exit=None | stdout='0'\n" + "task.000026 | exit=None | stdout='1'\n", + "task.000027 | exit=None | stdout='0'\n" ] } ], @@ -187,6 +185,55 @@ "\n", "await run_executables()" ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Heteregeneous tasks\n", + "The example shows how to execute MPI + CPU + GPU tasks." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "task.000028 | stdout='0\\n1' | return=\n", + "task.000029 | stdout='' | return=Hello CPU\n" + ] + } + ], + "source": [ + "def echo_hello():\n", + " return \"Hello CPU\"\n", + "\n", + "async def run_executables():\n", + " async with EnsembleBackend(mpi_flavour=\"test\") as backend:\n", + " session = Session(backends=[backend])\n", + " mpi_task = ComputeTask(executable=\"printenv\", arguments=[\"OMPI_COMM_WORLD_LOCAL_RANK\"], task_backend_specific_kwargs = {\"ranks\": 2})\n", + " serial_task = ComputeTask(function=echo_hello)\n", + " tasks = [mpi_task, serial_task]\n", + " async with session:\n", + " await session.submit_tasks(tasks)\n", + " await session.wait_tasks(tasks)\n", + "\n", + " for t in tasks:\n", + " print(f\"{t.uid} | stdout={t.stdout.strip()!r} | return={t.return_value}\")\n", + "\n", + "await run_executables()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { From e326f493c719d2d7e6687d6cca1e589534bc2b73 Mon Sep 17 00:00:00 2001 From: harikrishna1410 Date: Wed, 12 Aug 2026 09:33:01 -0500 Subject: [PATCH 3/6] Update format --- src/rhapsody/backends/execution/el.py | 61 ++++++++++++++----------- tests/unit/test_backend_execution_el.py | 42 ++++++++++------- 2 files changed, 61 insertions(+), 42 deletions(-) diff --git a/src/rhapsody/backends/execution/el.py b/src/rhapsody/backends/execution/el.py index 49883a2..2ddc4e4 100644 --- a/src/rhapsody/backends/execution/el.py +++ b/src/rhapsody/backends/execution/el.py @@ -1,7 +1,7 @@ """Ensemble Launcher execution backend for distributed computing. -This module provides a backend that executes tasks via the Ensemble Launcher framework, -supporting MPI-based distributed execution environments. +This module provides a backend that executes tasks via the Ensemble Launcher framework, supporting +MPI-based distributed execution environments. """ from __future__ import annotations @@ -13,8 +13,6 @@ import uuid from typing import Any from typing import Callable -from typing import Optional -from typing import Union from ensemble_launcher import EnsembleLauncher from ensemble_launcher.config import LauncherConfig @@ -32,24 +30,25 @@ class EnsembleBackend(BaseBackend): - - def __init__(self, name: str | None = None, - child_executor_name: str = "async_mpi", - return_stdout: bool = True, - worker_logs: bool = True, - master_logs: bool = True, - gpu_selector: str = "ZE_AFFINITY_MASK", - children_scheduler_policy: str = "fixed_leafs_children_policy", - task_scheduler_policy: str = "large_resource_policy", - enable_workstealing: bool = False, - checkpoint_dir: str | None = None, - mpi_flavour: str = "mpich", - nlevels: int = 0, - nleafs: int | None = None, - cpus: list[int] | None = None, - gpus: list[int] | None = None, - client_only: bool = False, - node_id: str = "global", + def __init__( + self, + name: str | None = None, + child_executor_name: str = "async_mpi", + return_stdout: bool = True, + worker_logs: bool = True, + master_logs: bool = True, + gpu_selector: str = "ZE_AFFINITY_MASK", + children_scheduler_policy: str = "fixed_leafs_children_policy", + task_scheduler_policy: str = "large_resource_policy", + enable_workstealing: bool = False, + checkpoint_dir: str | None = None, + mpi_flavour: str = "mpich", + nlevels: int = 0, + nleafs: int | None = None, + cpus: list[int] | None = None, + gpus: list[int] | None = None, + client_only: bool = False, + node_id: str = "global", ): super().__init__(name=name) @@ -80,7 +79,9 @@ def __init__(self, name: str | None = None, cpus = cpus or list(range(os.cpu_count())) ngpus = len(gpus) if gpus is not None else 0 gpus = gpus or [] - self._sys_config = SystemConfig(name="cluster", ncpus=len(cpus), cpus=cpus, ngpus=ngpus, gpus=gpus) + self._sys_config = SystemConfig( + name="cluster", ncpus=len(cpus), cpus=cpus, ngpus=ngpus, gpus=gpus + ) self._el: EnsembleLauncher | None = None self._client: ClusterClient | None = None self.tasks: dict = {} @@ -144,7 +145,7 @@ async def _handle_task(self, task: dict) -> None: result = await asyncio.wrap_future(fut) if is_executable: task["return_value"] = "" - task["stdout"] = result.split(",")[0] ## EL returns "stdout,stderr" + task["stdout"] = result.split(",")[0] ## EL returns "stdout,stderr" task["stderr"] = result.split(",")[1] else: task["return_value"] = result @@ -208,8 +209,16 @@ def build_task(self, task: dict) -> ELTask | AsyncELTask: ppn = backend_kwargs.get("ranks", 1) // nnodes ngpus_per_process = backend_kwargs.get("gpus_per_rank", 0) env = backend_kwargs.get("env", {}) - cpu_affinity = list(map(int, backend_kwargs.get("cpu_affinity").split(","))) if "cpu_affinity" in backend_kwargs else [] - gpu_affinity = list(map(int, backend_kwargs.get("gpu_affinity").split(","))) if "gpu_affinity" in backend_kwargs else [] + cpu_affinity = ( + list(map(int, backend_kwargs.get("cpu_affinity").split(","))) + if "cpu_affinity" in backend_kwargs + else [] + ) + gpu_affinity = ( + list(map(int, backend_kwargs.get("gpu_affinity").split(","))) + if "gpu_affinity" in backend_kwargs + else [] + ) func = task.get("function") if func is not None and inspect.iscoroutinefunction(func): diff --git a/tests/unit/test_backend_execution_el.py b/tests/unit/test_backend_execution_el.py index b3b3e04..245708f 100644 --- a/tests/unit/test_backend_execution_el.py +++ b/tests/unit/test_backend_execution_el.py @@ -91,9 +91,11 @@ async def test_el_backend_async_init(mock_nodes): mock_el = MagicMock() mock_client = MagicMock() - with patch("rhapsody.backends.execution.el.EnsembleLauncher", return_value=mock_el), \ - patch("rhapsody.backends.execution.el.ClusterClient", return_value=mock_client), \ - patch("rhapsody.backends.execution.el.asyncio.to_thread", side_effect=_sync_to_thread): + with ( + patch("rhapsody.backends.execution.el.EnsembleLauncher", return_value=mock_el), + patch("rhapsody.backends.execution.el.ClusterClient", return_value=mock_client), + patch("rhapsody.backends.execution.el.asyncio.to_thread", side_effect=_sync_to_thread), + ): result = await backend assert result is backend assert backend._initialized @@ -108,9 +110,11 @@ async def test_el_backend_async_init_client_only(mock_nodes): mock_client = MagicMock() - with patch("rhapsody.backends.execution.el.EnsembleLauncher") as mock_el_cls, \ - patch("rhapsody.backends.execution.el.ClusterClient", return_value=mock_client), \ - patch("rhapsody.backends.execution.el.asyncio.to_thread", side_effect=_sync_to_thread): + with ( + patch("rhapsody.backends.execution.el.EnsembleLauncher") as mock_el_cls, + patch("rhapsody.backends.execution.el.ClusterClient", return_value=mock_client), + patch("rhapsody.backends.execution.el.asyncio.to_thread", side_effect=_sync_to_thread), + ): await backend assert backend._initialized assert backend._el is None @@ -131,9 +135,11 @@ async def test_el_backend_context_manager(mock_nodes): mock_el = MagicMock() mock_client = MagicMock() - with patch("rhapsody.backends.execution.el.EnsembleLauncher", return_value=mock_el), \ - patch("rhapsody.backends.execution.el.ClusterClient", return_value=mock_client), \ - patch("rhapsody.backends.execution.el.asyncio.to_thread", side_effect=_sync_to_thread): + with ( + patch("rhapsody.backends.execution.el.EnsembleLauncher", return_value=mock_el), + patch("rhapsody.backends.execution.el.ClusterClient", return_value=mock_client), + patch("rhapsody.backends.execution.el.asyncio.to_thread", side_effect=_sync_to_thread), + ): async with backend as b: assert b._initialized assert b is backend @@ -182,9 +188,11 @@ async def test_el_backend_state(mock_nodes): async def test_el_backend_state_mapper(mock_nodes): backend = EnsembleBackend() - with patch("rhapsody.backends.execution.el.EnsembleLauncher", return_value=MagicMock()), \ - patch("rhapsody.backends.execution.el.ClusterClient", return_value=MagicMock()), \ - patch("rhapsody.backends.execution.el.asyncio.to_thread", side_effect=_sync_to_thread): + with ( + patch("rhapsody.backends.execution.el.EnsembleLauncher", return_value=MagicMock()), + patch("rhapsody.backends.execution.el.ClusterClient", return_value=MagicMock()), + patch("rhapsody.backends.execution.el.asyncio.to_thread", side_effect=_sync_to_thread), + ): await backend mapper = backend.get_task_states_map() @@ -216,6 +224,7 @@ async def test_el_backend_submit_tasks(mock_nodes): backend = EnsembleBackend() backend._initialized = True from rhapsody.backends.constants import BackendMainStates + backend._backend_state = BackendMainStates.INITIALIZED captured = [] @@ -256,6 +265,7 @@ async def test_el_submit_multiple_tasks_complete(mock_nodes): backend = EnsembleBackend() backend._initialized = True from rhapsody.backends.constants import BackendMainStates + backend._backend_state = BackendMainStates.INITIALIZED backend.register_callback(lambda t, s: None) @@ -264,10 +274,7 @@ async def test_el_submit_multiple_tasks_complete(mock_nodes): mock_client.submit.side_effect = futs backend._client = mock_client - tasks = [ - ComputeTask(function=lambda: 1, args=(), kwargs={}) - for _ in range(3) - ] + tasks = [ComputeTask(function=lambda: 1, args=(), kwargs={}) for _ in range(3)] await backend.submit_tasks(tasks) async_tasks = [backend.tasks[t["uid"]]["future"] for t in tasks] @@ -291,6 +298,7 @@ async def test_el_backend_submit_after_shutdown(mock_nodes): backend = EnsembleBackend() backend._initialized = True from rhapsody.backends.constants import BackendMainStates + backend._backend_state = BackendMainStates.SHUTDOWN with pytest.raises(RuntimeError, match="Cannot submit during shutdown"): @@ -469,6 +477,7 @@ async def test_el_callback_done(mock_nodes): backend = EnsembleBackend() backend._initialized = True from rhapsody.backends.constants import BackendMainStates + backend._backend_state = BackendMainStates.INITIALIZED captured = [] @@ -499,6 +508,7 @@ async def test_el_callback_failed(mock_nodes): backend = EnsembleBackend() backend._initialized = True from rhapsody.backends.constants import BackendMainStates + backend._backend_state = BackendMainStates.INITIALIZED captured = [] From 509628b0cb792f6a422d84c07167652770f79aee Mon Sep 17 00:00:00 2001 From: harikrishna1410 Date: Wed, 12 Aug 2026 09:41:43 -0500 Subject: [PATCH 4/6] changed EnsembleBackend to EnsembleExecutionBackend --- src/rhapsody/backends/__init__.py | 4 +- src/rhapsody/backends/execution/__init__.py | 4 +- src/rhapsody/backends/execution/el.py | 8 +-- tests/unit/test_backend_execution_el.py | 71 +++++++++++---------- 4 files changed, 45 insertions(+), 42 deletions(-) diff --git a/src/rhapsody/backends/__init__.py b/src/rhapsody/backends/__init__.py index 87ff379..deea9a0 100644 --- a/src/rhapsody/backends/__init__.py +++ b/src/rhapsody/backends/__init__.py @@ -60,9 +60,9 @@ pass try: - from .execution import EnsembleBackend + from .execution import EnsembleExecutionBackend - __all__.append(EnsembleBackend) + __all__.append(EnsembleExecutionBackend) except ImportError: pass diff --git a/src/rhapsody/backends/execution/__init__.py b/src/rhapsody/backends/execution/__init__.py index e311c95..17b053c 100644 --- a/src/rhapsody/backends/execution/__init__.py +++ b/src/rhapsody/backends/execution/__init__.py @@ -26,9 +26,9 @@ pass try: - from .el import EnsembleBackend # noqa: F401 + from .el import EnsembleExecutionBackend # noqa: F401 - __all__.append("EnsembleBackend") + __all__.append("EnsembleExecutionBackend") except ImportError: pass diff --git a/src/rhapsody/backends/execution/el.py b/src/rhapsody/backends/execution/el.py index 2ddc4e4..783195d 100644 --- a/src/rhapsody/backends/execution/el.py +++ b/src/rhapsody/backends/execution/el.py @@ -29,7 +29,7 @@ from ..constants import StateMapper -class EnsembleBackend(BaseBackend): +class EnsembleExecutionBackend(BaseBackend): def __init__( self, name: str | None = None, @@ -87,7 +87,7 @@ def __init__( self.tasks: dict = {} def __await__(self): - """Make EnsembleBackend awaitable.""" + """Make EnsembleExecutionBackend awaitable.""" return self._async_init().__await__() async def _async_init(self): @@ -131,8 +131,8 @@ async def _initialize(self): def _ensure_initialized(self): if not self._initialized: raise RuntimeError( - "EnsembleBackend must be awaited before use. " - "Use: backend = await EnsembleBackend(...)" + "EnsembleExecutionBackend must be awaited before use. " + "Use: backend = await EnsembleExecutionBackend(...)" ) async def _handle_task(self, task: dict) -> None: diff --git a/tests/unit/test_backend_execution_el.py b/tests/unit/test_backend_execution_el.py index 245708f..1c22dcf 100644 --- a/tests/unit/test_backend_execution_el.py +++ b/tests/unit/test_backend_execution_el.py @@ -10,7 +10,7 @@ from rhapsody import ComputeTask try: - from rhapsody.backends.execution.el import EnsembleBackend + from rhapsody.backends.execution.el import EnsembleExecutionBackend _el_available = True except ImportError: @@ -25,15 +25,16 @@ def test_el_backend_import(): - from rhapsody.backends.execution import EnsembleBackend as EB - - assert EB is not None + try: + from rhapsody.backends.execution import EnsembleExecutionBackend + except ModuleNotFoundError: + raise ModuleNotFoundError def test_el_backend_inherits_base(): from rhapsody.backends.base import BaseBackend - assert issubclass(EnsembleBackend, BaseBackend) + assert issubclass(EnsembleExecutionBackend, BaseBackend) # --------------------------------------------------------------------------- @@ -43,7 +44,7 @@ def test_el_backend_inherits_base(): @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) def test_el_backend_init_defaults(mock_nodes): - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() assert backend is not None assert not backend._initialized assert backend._client is None @@ -55,21 +56,23 @@ def test_el_backend_init_defaults(mock_nodes): @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) def test_el_backend_init_client_only(mock_nodes): - backend = EnsembleBackend(client_only=True, node_id="main.w0", checkpoint_dir="/tmp/ckpt") + backend = EnsembleExecutionBackend( + client_only=True, node_id="main.w0", checkpoint_dir="/tmp/ckpt" + ) assert backend._client_only is True assert backend._node_id == "main.w0" @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) def test_el_backend_init_custom_name(mock_nodes): - backend = EnsembleBackend(name="my_el") + backend = EnsembleExecutionBackend(name="my_el") assert backend.name == "my_el" @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) def test_el_backend_init_default_name(mock_nodes): - backend = EnsembleBackend() - assert backend.name == "EnsembleBackend" + backend = EnsembleExecutionBackend() + assert backend.name == "EnsembleExecutionBackend" # --------------------------------------------------------------------------- @@ -79,14 +82,14 @@ def test_el_backend_init_default_name(mock_nodes): @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) def test_el_backend_is_awaitable(mock_nodes): - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() assert hasattr(backend, "__await__") @pytest.mark.asyncio @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) async def test_el_backend_async_init(mock_nodes): - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() mock_el = MagicMock() mock_client = MagicMock() @@ -106,7 +109,7 @@ async def test_el_backend_async_init(mock_nodes): @pytest.mark.asyncio @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) async def test_el_backend_async_init_client_only(mock_nodes): - backend = EnsembleBackend(client_only=True, checkpoint_dir="/tmp/ckpt") + backend = EnsembleExecutionBackend(client_only=True, checkpoint_dir="/tmp/ckpt") mock_client = MagicMock() @@ -130,7 +133,7 @@ async def test_el_backend_async_init_client_only(mock_nodes): @pytest.mark.asyncio @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) async def test_el_backend_context_manager(mock_nodes): - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() mock_el = MagicMock() mock_client = MagicMock() @@ -156,7 +159,7 @@ async def test_el_backend_context_manager(mock_nodes): @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) def test_el_backend_callback_registration(mock_nodes): - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() def my_cb(task, state): pass @@ -173,7 +176,7 @@ def my_cb(task, state): @pytest.mark.asyncio @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) async def test_el_backend_state(mock_nodes): - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() state = await backend.state() assert state == "INITIALIZED" @@ -186,7 +189,7 @@ async def test_el_backend_state(mock_nodes): @pytest.mark.asyncio @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) async def test_el_backend_state_mapper(mock_nodes): - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() with ( patch("rhapsody.backends.execution.el.EnsembleLauncher", return_value=MagicMock()), @@ -207,9 +210,9 @@ async def test_el_backend_state_mapper(mock_nodes): @pytest.mark.asyncio @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) async def test_el_backend_submit_not_initialized(mock_nodes): - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() - with pytest.raises(RuntimeError, match="EnsembleBackend must be awaited"): + with pytest.raises(RuntimeError, match="EnsembleExecutionBackend must be awaited"): await backend.submit_tasks([]) @@ -221,7 +224,7 @@ async def test_el_backend_submit_not_initialized(mock_nodes): @pytest.mark.asyncio @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) async def test_el_backend_submit_tasks(mock_nodes): - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() backend._initialized = True from rhapsody.backends.constants import BackendMainStates @@ -262,7 +265,7 @@ async def test_el_backend_submit_tasks(mock_nodes): @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) async def test_el_submit_multiple_tasks_complete(mock_nodes): """All submitted tasks complete and fire callbacks.""" - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() backend._initialized = True from rhapsody.backends.constants import BackendMainStates @@ -295,7 +298,7 @@ async def test_el_submit_multiple_tasks_complete(mock_nodes): @pytest.mark.asyncio @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) async def test_el_backend_submit_after_shutdown(mock_nodes): - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() backend._initialized = True from rhapsody.backends.constants import BackendMainStates @@ -314,7 +317,7 @@ async def test_el_backend_submit_after_shutdown(mock_nodes): def test_el_build_task_sync_function(mock_nodes): from ensemble_launcher.ensemble import Task as ELTask - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() def my_func(a, b): return a + b @@ -337,7 +340,7 @@ def my_func(a, b): def test_el_build_task_async_function(mock_nodes): from ensemble_launcher.ensemble import AsyncTask as AsyncELTask - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() async def my_async_func(a): return a @@ -356,7 +359,7 @@ async def my_async_func(a): def test_el_build_task_executable(mock_nodes): from ensemble_launcher.ensemble import Task as ELTask - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() task = ComputeTask( executable="/bin/echo", @@ -371,7 +374,7 @@ def test_el_build_task_executable(mock_nodes): @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) def test_el_build_task_default_resources(mock_nodes): - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() task = ComputeTask( function=lambda: 1, @@ -393,7 +396,7 @@ def test_el_build_task_default_resources(mock_nodes): @pytest.mark.asyncio @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) async def test_el_cancel_task_success(mock_nodes): - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() mock_future = MagicMock() mock_future.cancel.return_value = True @@ -407,7 +410,7 @@ async def test_el_cancel_task_success(mock_nodes): @pytest.mark.asyncio @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) async def test_el_cancel_task_not_found(mock_nodes): - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() result = await backend.cancel_task("nonexistent") assert result is False @@ -420,7 +423,7 @@ async def test_el_cancel_task_not_found(mock_nodes): @pytest.mark.asyncio @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) async def test_el_shutdown(mock_nodes): - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() backend._initialized = True mock_client = MagicMock() @@ -444,7 +447,7 @@ async def test_el_shutdown(mock_nodes): @pytest.mark.asyncio @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) async def test_el_shutdown_client_only(mock_nodes): - backend = EnsembleBackend(client_only=True, checkpoint_dir="/tmp/ckpt") + backend = EnsembleExecutionBackend(client_only=True, checkpoint_dir="/tmp/ckpt") backend._initialized = True mock_client = MagicMock() @@ -461,7 +464,7 @@ async def test_el_shutdown_client_only(mock_nodes): @pytest.mark.asyncio @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) async def test_el_shutdown_without_init(mock_nodes): - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() await backend.shutdown() assert not backend._initialized @@ -474,7 +477,7 @@ async def test_el_shutdown_without_init(mock_nodes): @pytest.mark.asyncio @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) async def test_el_callback_done(mock_nodes): - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() backend._initialized = True from rhapsody.backends.constants import BackendMainStates @@ -505,7 +508,7 @@ async def test_el_callback_done(mock_nodes): @pytest.mark.asyncio @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) async def test_el_callback_failed(mock_nodes): - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() backend._initialized = True from rhapsody.backends.constants import BackendMainStates @@ -540,7 +543,7 @@ async def test_el_callback_failed(mock_nodes): @patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) def test_el_task_state_cb(mock_nodes): - backend = EnsembleBackend() + backend = EnsembleExecutionBackend() captured = [] backend.register_callback(lambda t, s: captured.append(s)) From 594112e3c63ecdafb0f1756a5ab496bd2f35c8d5 Mon Sep 17 00:00:00 2001 From: harikrishna1410 Date: Wed, 12 Aug 2026 09:50:16 -0500 Subject: [PATCH 5/6] extend words --- _typos.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/_typos.toml b/_typos.toml index fb4b095..dadcfc4 100644 --- a/_typos.toml +++ b/_typos.toml @@ -14,6 +14,7 @@ extend-ignore-identifiers-re = [ # Add specific words that should not be corrected # rhapsody = "rhapsody" HPE = "HPE" +leafs = "leafs" [files] extend-exclude = [ From b51a230c91a15f6a96f0908e8c2a87467a7ebc3c Mon Sep 17 00:00:00 2001 From: harikrishna1410 Date: Wed, 12 Aug 2026 10:51:27 -0500 Subject: [PATCH 6/6] fixed Gemini comments, Added doc strings --- src/rhapsody/backends/execution/el.py | 143 ++++++++++++++++++++++++-- 1 file changed, 137 insertions(+), 6 deletions(-) diff --git a/src/rhapsody/backends/execution/el.py b/src/rhapsody/backends/execution/el.py index 783195d..26cb17d 100644 --- a/src/rhapsody/backends/execution/el.py +++ b/src/rhapsody/backends/execution/el.py @@ -50,6 +50,28 @@ def __init__( client_only: bool = False, node_id: str = "global", ): + """Initialize the Ensemble Launcher execution backend. + + Args: + name: Optional name for the backend instance. + child_executor_name: Executor used for child processes. + return_stdout: Whether to capture and return stdout from tasks. + worker_logs: Enable logging on worker processes. + master_logs: Enable logging on the master process. + gpu_selector: Environment variable name used for GPU affinity selection. + children_scheduler_policy: Policy for scheduling child processes across nodes. + task_scheduler_policy: Policy for scheduling tasks onto resources. + enable_workstealing: Allow idle workers to steal tasks from busy ones. + checkpoint_dir: Directory for task checkpoints. Auto-generated if not provided. + mpi_flavour: MPI implementation to use (e.g., "mpich", "openmpi"). + nlevels: Number of hierarchy levels in the launcher tree. + nleafs: Number of leaf nodes. Defaults to the number of available nodes. + cpus: List of CPU IDs available for tasks. Defaults to all CPUs. + gpus: List of GPU IDs available for tasks. + client_only: If True, only start the client without launching the ensemble. + node_id: Scheduler node ID for the ClusterClient to connect to. + "global" (default) connects to the global master node. + """ super().__init__(name=name) self.logger = logging.getLogger(__name__) @@ -76,7 +98,7 @@ def __init__( checkpoint_dir=checkpoint_dir or os.path.join(os.getcwd(), f"ckpt_{uuid.uuid4()}"), mpi_config=MPIConfig(flavor=mpi_flavour), ) - cpus = cpus or list(range(os.cpu_count())) + cpus = cpus or list(range(os.cpu_count() or 1)) ngpus = len(gpus) if gpus is not None else 0 gpus = gpus or [] self._sys_config = SystemConfig( @@ -91,6 +113,17 @@ def __await__(self): return self._async_init().__await__() async def _async_init(self): + """Perform asynchronous initialization of the backend. + + Registers backend and task states, then starts the ensemble launcher + and cluster client. This method is idempotent. + + Returns: + The initialized EnsembleExecutionBackend instance. + + Raises: + Exception: If initialization fails, the backend remains uninitialized. + """ if not self._initialized: try: self.logger.debug("Registering backend states...") @@ -114,6 +147,11 @@ async def _async_init(self): return self async def _initialize(self): + """Start the EnsembleLauncher and ClusterClient. + + If ``client_only`` is False, starts the full ensemble launcher first. + Always starts a ClusterClient connected to the checkpoint directory. + """ if not self._client_only: self._el = EnsembleLauncher( ensemble_file={}, @@ -129,6 +167,7 @@ async def _initialize(self): await asyncio.to_thread(self._client.start) def _ensure_initialized(self): + """Raise RuntimeError if the backend has not been awaited yet.""" if not self._initialized: raise RuntimeError( "EnsembleExecutionBackend must be awaited before use. " @@ -136,12 +175,21 @@ def _ensure_initialized(self): ) async def _handle_task(self, task: dict) -> None: + """Submit a single task to the cluster and await its result. + + Builds an EL task, submits it via the cluster client, and populates + the task dict with return_value/stdout/stderr on success or + exception/stderr on failure. Invokes the registered callback with + the appropriate state ("RUNNING", "DONE", or "FAILED"). + + Args: + task: Mutable task dictionary. Updated in-place with results. + """ try: is_executable = task.get("executable", None) is not None self._callback_func(task, "RUNNING") el_task = self.build_task(task) fut = self._client.submit(el_task) - task["future"] = fut result = await asyncio.wrap_future(fut) if is_executable: task["return_value"] = "" @@ -159,6 +207,18 @@ async def _handle_task(self, task: dict) -> None: self._callback_func(task, "FAILED") async def submit_tasks(self, tasks: list[dict]) -> None: + """Submit a batch of tasks for asynchronous execution. + + Each task is scheduled as an independent asyncio task. The backend + transitions to RUNNING state on the first submission. + + Args: + tasks: List of task dictionaries, each containing at minimum a + "uid" key and either "function" or "executable". + + Raises: + RuntimeError: If the backend is not initialized or has been shut down. + """ self._ensure_initialized() if self._backend_state == BackendMainStates.SHUTDOWN: @@ -174,14 +234,19 @@ async def submit_tasks(self, tasks: list[dict]) -> None: self.tasks[task["uid"]]["future"] = future async def shutdown(self) -> None: + """Shut down the backend, tearing down the client and launcher. + + Clears all tracked tasks and resets initialization state. Safe to + call multiple times. + """ self._backend_state = BackendMainStates.SHUTDOWN self.logger.debug(f"Backend state set to: {self._backend_state.value}") try: if self._client is not None: - self._client.teardown() + await asyncio.to_thread(self._client.teardown) if self._el is not None: - self._el.stop() + await asyncio.to_thread(self._el.stop) except Exception as e: self.logger.exception(f"Error during shutdown: {e}") finally: @@ -192,20 +257,58 @@ async def shutdown(self) -> None: self.logger.info("Ensemble execution backend shutdown complete") async def state(self) -> str: + """Return the current backend state as a string.""" return self._backend_state.value def task_state_cb(self, task: dict, state: str) -> None: + """Invoke the registered callback with a task and its new state. + + Args: + task: The task dictionary whose state changed. + state: The new state string (e.g., "RUNNING", "DONE", "FAILED"). + """ self._callback_func(task, state) def register_callback(self, func: Callable[[dict[str, Any], str], None]) -> None: + """Register a callback to be invoked on task state transitions. + + Args: + func: A callable accepting (task_dict, state_string). + """ self._callback_func = func def get_task_states_map(self) -> Any: + """Return a StateMapper instance for this backend's task states.""" return StateMapper(backend=self) def build_task(self, task: dict) -> ELTask | AsyncELTask: + """Convert a task dictionary into an Ensemble Launcher Task object. + + Selects ``AsyncELTask`` if the task's function is a coroutine, + otherwise uses ``ELTask``. For executable-based tasks, the executable + and arguments are combined into a single command string. + + Args: + task: Task dictionary containing "uid" and either "function" with + "args"/"kwargs" or "executable" with "arguments", plus optional + "task_backend_specific_kwargs" for resource configuration: + + - **nnodes** (int): Number of nodes to run on. Defaults to 1. + - **ranks** (int): Total number of MPI ranks. Divided by + ``nnodes`` to get processes-per-node (ppn). Defaults to 1. + - **gpus_per_rank** (int): GPUs allocated per MPI rank. + Defaults to 0. + - **env** (dict): Extra environment variables passed to the task. + - **cpu_affinity** (str): Comma-separated CPU IDs for pinning + (e.g., ``"0,1,2,3"``). + - **gpu_affinity** (str): Comma-separated GPU IDs for pinning + (e.g., ``"0,1"``). + + Returns: + An ``ELTask`` or ``AsyncELTask`` configured for submission. + """ backend_kwargs = task.get("task_backend_specific_kwargs", {}) - nnodes = backend_kwargs.get("nnodes", 1) + nnodes = max(backend_kwargs.get("nnodes", 1),1) ppn = backend_kwargs.get("ranks", 1) // nnodes ngpus_per_process = backend_kwargs.get("gpus_per_rank", 0) env = backend_kwargs.get("env", {}) @@ -229,7 +332,7 @@ def build_task(self, task: dict) -> ELTask | AsyncELTask: is_executable = task.get("executable", None) is not None if is_executable: ## When ELTask.executable is str, it ignore args and kwargs. - exec = task["executable"] + " " + " ".join([str(x) for x in task["arguments"]]) + exec = task["executable"] + " " + " ".join([str(x) for x in task.get("arguments",[])]) else: exec = func @@ -248,6 +351,14 @@ def build_task(self, task: dict) -> ELTask | AsyncELTask: ) def link_implicit_data_deps(self, src_task: dict[str, Any], dst_task: dict[str, Any]) -> None: # noqa: B027 + """Register an implicit data dependency between two tasks. + + Not implemented for this backend; this is a no-op. + + Args: + src_task: The upstream task that produces data. + dst_task: The downstream task that consumes data. + """ pass def link_explicit_data_deps( # noqa: B027 @@ -257,9 +368,27 @@ def link_explicit_data_deps( # noqa: B027 file_name: str | None = None, file_path: str | None = None, ) -> None: + """Register an explicit file-based data dependency between two tasks. + + Not implemented for this backend; this is a no-op. + + Args: + src_task: The upstream task that produces the file. + dst_task: The downstream task that consumes the file. + file_name: Name of the shared file. + file_path: Path to the shared file. + """ pass async def cancel_task(self, uid: str) -> bool: + """Cancel a previously submitted task by its UID. + + Args: + uid: Unique identifier of the task to cancel. + + Returns: + True if the task was found and successfully cancelled, False otherwise. + """ if uid in self.tasks: future = self.tasks[uid].get("future") if future: @@ -267,9 +396,11 @@ async def cancel_task(self, uid: str) -> bool: return False async def __aenter__(self): + """Enter the async context manager, initializing the backend if needed.""" if not self._initialized: await self._async_init() return self async def __aexit__(self, exc_type, exc_val, exc_tb): + """Exit the async context manager, shutting down the backend.""" await self.shutdown()