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 = [ diff --git a/src/rhapsody/backends/__init__.py b/src/rhapsody/backends/__init__.py index b544c18..deea9a0 100644 --- a/src/rhapsody/backends/__init__.py +++ b/src/rhapsody/backends/__init__.py @@ -59,6 +59,14 @@ except ImportError: pass +try: + from .execution import EnsembleExecutionBackend + + __all__.append(EnsembleExecutionBackend) + +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..17b053c 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 EnsembleExecutionBackend # noqa: F401 + + __all__.append("EnsembleExecutionBackend") +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..26cb17d --- /dev/null +++ b/src/rhapsody/backends/execution/el.py @@ -0,0 +1,406 @@ +"""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 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 EnsembleExecutionBackend(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", + ): + """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__) + 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() or 1)) + 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 EnsembleExecutionBackend awaitable.""" + 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...") + 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): + """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={}, + 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): + """Raise RuntimeError if the backend has not been awaited yet.""" + if not self._initialized: + raise RuntimeError( + "EnsembleExecutionBackend must be awaited before use. " + "Use: backend = await EnsembleExecutionBackend(...)" + ) + + 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) + 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: + """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: + 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: + """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: + await asyncio.to_thread(self._client.teardown) + if self._el is not None: + await asyncio.to_thread(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 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 = 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", {}) + 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): + 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.get("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 + """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 + 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: + """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: + return future.cancel() + 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() diff --git a/tests/unit/test_backend_execution_el.py b/tests/unit/test_backend_execution_el.py new file mode 100644 index 0000000..1c22dcf --- /dev/null +++ b/tests/unit/test_backend_execution_el.py @@ -0,0 +1,561 @@ +"""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 EnsembleExecutionBackend + + _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(): + 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(EnsembleExecutionBackend, BaseBackend) + + +# --------------------------------------------------------------------------- +# __init__ +# --------------------------------------------------------------------------- + + +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +def test_el_backend_init_defaults(mock_nodes): + backend = EnsembleExecutionBackend() + 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 = 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 = 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 = EnsembleExecutionBackend() + assert backend.name == "EnsembleExecutionBackend" + + +# --------------------------------------------------------------------------- +# Awaitable / async init +# --------------------------------------------------------------------------- + + +@patch("rhapsody.backends.execution.el.get_nodes", return_value=["node0"]) +def test_el_backend_is_awaitable(mock_nodes): + 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 = EnsembleExecutionBackend() + + 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 = EnsembleExecutionBackend(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 = EnsembleExecutionBackend() + + 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 = EnsembleExecutionBackend() + + 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 = EnsembleExecutionBackend() + 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 = EnsembleExecutionBackend() + + 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 = EnsembleExecutionBackend() + + with pytest.raises(RuntimeError, match="EnsembleExecutionBackend 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 = EnsembleExecutionBackend() + 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 = EnsembleExecutionBackend() + 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 = EnsembleExecutionBackend() + 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 = EnsembleExecutionBackend() + + 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 = EnsembleExecutionBackend() + + 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 = EnsembleExecutionBackend() + + 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 = EnsembleExecutionBackend() + + 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 = EnsembleExecutionBackend() + + 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 = EnsembleExecutionBackend() + 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 = EnsembleExecutionBackend() + 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 = EnsembleExecutionBackend(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 = EnsembleExecutionBackend() + 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 = EnsembleExecutionBackend() + 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 = EnsembleExecutionBackend() + 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 = EnsembleExecutionBackend() + 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..6aca089 --- /dev/null +++ b/tutorials/el-backend-tutorial.ipynb @@ -0,0 +1,260 @@ +{ + "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": 2, + "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.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" + ] + } + ], + "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": 5, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "task.000026 | exit=None | stdout='1'\n", + "task.000027 | 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()" + ] + }, + { + "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": { + "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 +}