From cee014c028bcd6be2a854c5ab2f53fdf6114527c Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Wed, 26 Aug 2026 18:33:20 +0000 Subject: [PATCH 1/5] follow Dask best practice --- .../backends/execution/dask_parallel.py | 425 +++++++++++------- .../test_backend_execution_dask_parallel.py | 328 +++++++++++++- tests/unit/test_dask_custom_cluster.py | 94 ++++ 3 files changed, 678 insertions(+), 169 deletions(-) diff --git a/src/rhapsody/backends/execution/dask_parallel.py b/src/rhapsody/backends/execution/dask_parallel.py index 87552d6..4a811bb 100644 --- a/src/rhapsody/backends/execution/dask_parallel.py +++ b/src/rhapsody/backends/execution/dask_parallel.py @@ -2,17 +2,36 @@ This module provides a backend that executes tasks on Dask clusters, supporting both local and distributed execution environments. + +Execution model +---------------- +Dask workers already natively distinguish sync and async callables: a submitted callable that +is a coroutine function runs directly on the worker's event loop, while a plain callable runs in +the worker's thread pool (see ``dask._task_spec.Task.is_coro`` / +``distributed.worker.Worker._maybe_deserialize_task``). Because of this, RHAPSODY submits +``task["function"]`` to Dask exactly as the caller defined it (optionally wrapped in +``functools.partial`` only to pre-bind keyword arguments) for *both* sync and async callables. +No RHAPSODY-side closure/adapter is used to "convert" an async callable into something Dask can +run — that used to be done via a `functools.wraps`-decorated local closure, which is exactly what +caused pickling failures (the closure's `__qualname__`/`__module__` were copied from the original +function, so pickling-by-reference resolved to a *different* object at that name and raised +``PicklingError: ... it's not the same object as ...``). Submitting the real callable directly +removes the synthetic function whose identity could ever mismatch. """ from __future__ import annotations import asyncio +import inspect import logging +import time +from dataclasses import dataclass from functools import partial -from functools import wraps from typing import Any from typing import Callable +from rhapsody.api.errors import BackendError + from ..base import BaseBackend from ..constants import BackendMainStates from ..constants import StateMapper @@ -80,6 +99,21 @@ def _run_executable( return result.stdout.decode(), result.stderr.decode(), result.returncode +@dataclass +class _TaskRuntime: + """Backend-private bookkeeping for a submitted task. + + Deliberately kept separate from the RHAPSODY task dict (``DaskExecutionBackend.tasks``), which + is shared with and owned by the caller: nothing Dask/runtime-specific (the Dask ``Future``, + submission bookkeeping) should be written onto that shared object. + """ + + uid: str + kind: str # "function" or "executable" — diagnostics only + future: Any | None = None # dask.distributed.Future, set once client.submit() succeeds + submitted_at: float = 0.0 + + class DaskExecutionBackend(BaseBackend): """A Dask execution backend for distributed task execution. @@ -87,6 +121,11 @@ class DaskExecutionBackend(BaseBackend): for distributed task execution using Dask. Supports async functions, sync functions, and executable tasks. + Client/cluster ownership: if the caller supplies ``client``, this backend never closes it + (``shutdown()`` leaves it open). If the caller supplies only ``cluster``, this backend creates + and owns a ``Client`` for it (closed on ``shutdown()``) but never touches the caller's cluster. + If neither is supplied, this backend creates and owns both. + Usage: backend = await DaskExecutionBackend(resources) # or @@ -107,8 +146,11 @@ def __init__( resources: Dictionary of resource requirements for tasks. Contains configuration parameters for the Dask client initialization. name: Name of the backend. - cluster: Optional preconfigured Dask Cluster object. - client: Optional preconfigured Dask Client object. + cluster: Optional preconfigured Dask Cluster object. Not closed by + shutdown() — the caller retains ownership. + client: Optional preconfigured Dask Client object, which must have been + created with asynchronous=True. Not closed by shutdown() — the + caller retains ownership. """ if dask is None: @@ -117,12 +159,15 @@ def __init__( super().__init__(name=name) self.logger = _get_logger() - self.tasks = {} + self.tasks: dict[str, dict[str, Any]] = {} + self._runtime: dict[str, _TaskRuntime] = {} self._client = None self._callback_func: Callable = lambda t, s: None self._resources = resources or {} self._cluster_provided = cluster self._client_provided = client + # Ownership: only close resources this backend created itself. + self._owns_client = client is None self._initialized = False self._backend_state = BackendMainStates.INITIALIZED @@ -169,12 +214,20 @@ async def _initialize(self) -> None: """Initialize the Dask client and set up worker environments. Raises: + ValueError: If an externally-provided client is not asynchronous. Exception: If Dask client initialization fails. """ try: - if self._client_provided: + if self._client_provided is not None: + if not getattr(self._client_provided, "asynchronous", False): + raise ValueError( + "DaskExecutionBackend requires an externally-provided Client to be " + "created with asynchronous=True (e.g. Client(..., asynchronous=True)). " + f"Got a client with asynchronous=" + f"{getattr(self._client_provided, 'asynchronous', None)!r}." + ) self._client = self._client_provided - elif self._cluster_provided: + elif self._cluster_provided is not None: self._client = await dask.Client( self._cluster_provided, asynchronous=True, **self._resources ) @@ -192,7 +245,7 @@ def register_callback(self, func: Callable) -> None: Args: func: Function to be called when task states change. Should accept - task and state parameters. + task and state parameters. May be sync or async. """ self._callback_func = func @@ -211,24 +264,28 @@ async def cancel_task(self, uid: str) -> bool: uid (str): The UID of the task to cancel. Returns: - bool: True if the task was found and cancellation was attempted, - False otherwise. + bool: True if the task was found (still tracked, with a submitted Dask + future) and cancellation was attempted, False for an unknown uid, a task + that hasn't been submitted to Dask yet, or one that already reached a + terminal state (and was purged from tracking). """ self._ensure_initialized() - if uid in self.tasks: - task = self.tasks[uid] - future = task.get("future") - if future: - await future.cancel() - return True - return False + runtime = self._runtime.get(uid) + if runtime is None or runtime.future is None: + return False + try: + await runtime.future.cancel() + except Exception: + self.logger.exception(f"Error cancelling task '{uid}'") + return False + return True async def submit_tasks(self, tasks: list[dict[str, Any]]) -> None: """Submit tasks to the Dask cluster. Dispatches each task to the appropriate submission method based on its type: - executable tasks run via subprocess, async functions via an async wrapper, - and sync functions are submitted directly to Dask workers. + executable tasks run via subprocess, function tasks (sync or async) are + submitted directly to Dask workers. Args: tasks: List of task dictionaries containing: @@ -240,11 +297,19 @@ async def submit_tasks(self, tasks: list[dict[str, Any]]) -> None: - arguments: CLI arguments for executable tasks - task_backend_specific_kwargs: Passed directly to client.submit() - Note: - Future objects are filtered out from args as they are not picklable. + Raises: + BackendError: If the Dask client is not in a usable state (the whole + backend is unavailable — this is not attributable to any one task). + ValueError: If a task specifies neither 'function' nor 'executable'. """ self._ensure_initialized() + if self._client is None or self._client.status != "running": + raise BackendError( + self.name, + f"Dask client is not usable (status={getattr(self._client, 'status', None)!r})", + ) + if self._backend_state != BackendMainStates.RUNNING: self._backend_state = BackendMainStates.RUNNING self.logger.debug(f"Backend state set to: {self._backend_state.value}") @@ -252,66 +317,56 @@ async def submit_tasks(self, tasks: list[dict[str, Any]]) -> None: for task in tasks: is_func_task = bool(task.get("function")) is_exec_task = bool(task.get("executable")) + if not is_func_task and not is_exec_task: + raise ValueError("Task must specify either 'function' or 'executable'") + + uid = task["uid"] + self.tasks[uid] = task + self._runtime[uid] = _TaskRuntime( + uid=uid, + kind="executable" if is_exec_task else "function", + submitted_at=time.monotonic(), + ) - self.tasks[task["uid"]] = task - - # Filter out future objects as they are not picklable for Dask workers - filtered_args = [ - arg for arg in task.get("args", ()) if not isinstance(arg, asyncio.Future) - ] - task["args"] = tuple(filtered_args) - - try: - if is_exec_task: - await self._submit_executable(task) - elif is_func_task and asyncio.iscoroutinefunction(task["function"]): - await self._submit_async_function(task) - elif is_func_task: - await self._submit_sync_function(task) - else: - raise ValueError("Task must specify either 'function' or 'executable'") - except Exception as e: - task["exception"] = e - task["stdout"] = "" - task["stderr"] = str(e) - self._callback_func(task, "FAILED") + if is_exec_task: + await self._submit_executable(task) + else: + await self._submit_function(task) - async def _submit_to_dask(self, task: dict[str, Any], fn: Callable, *args) -> None: - """Submit function to Dask and register completion callback. + async def _submit_and_track( + self, + task: dict[str, Any], + fn: Callable, + args: tuple, + *, + kind: str, + fn_kwargs: dict[str, Any] | None = None, + ) -> None: + """Submit ``fn(*args, **fn_kwargs)`` to Dask and schedule its completion handler. - Submits the wrapped function to Dask client and registers a callback - to handle task completion or failure. + This is the single low-level submission primitive shared by function and + executable tasks. ``fn``/``args``/``fn_kwargs`` are the user-visible callable + and its arguments; Dask submission options (resources, key, retries, ...) come + only from ``task["task_backend_specific_kwargs"]``, kept strictly separate so a + user's own function kwarg can never collide with a Dask-reserved submit() kwarg. Args: - task: Task dictionary containing task metadata and configuration. - fn: The async function to submit to Dask. - *args: Arguments to pass to the function. + task: The RHAPSODY task dictionary (mutated in place with result/state fields). + fn: The callable to submit (never a synthetic wrapper — see module docstring). + args: Positional arguments to submit alongside ``fn``. + kind: "function" or "executable" — controls result-shape handling in `_on_done`. + fn_kwargs: Keyword arguments to submit alongside ``fn`` (executable path only; + function-task kwargs are pre-bound via `functools.partial` by the caller). """ - - async def on_done(f): - task_uid = task["uid"] - try: - self._callback_func(task, "RUNNING") - result = await f - task["return_value"] = result - task["stdout"] = "" - task["stderr"] = "" - self._callback_func(task, "DONE") - except dask.client.FutureCancelledError: - self._callback_func(task, "CANCELED") - except Exception as e: - task["exception"] = e - task["stdout"] = "" - task["stderr"] = str(e) - self._callback_func(task, "FAILED") - finally: - # Clean up the future reference once task is complete - if task_uid in self.tasks: - del self.tasks[task_uid] - - backend_kwargs = dict(task.get("task_backend_specific_kwargs", {})) - dask_resources = backend_kwargs.get("resources", {}) - if dask_resources and not self._check_resources_satisfiable(dask_resources): + uid = task["uid"] + submit_kwargs = dict(task.get("task_backend_specific_kwargs", {})) + if kind == "executable": + submit_kwargs = { + k: v for k, v in submit_kwargs.items() if k not in ("cwd", "shell", "env") + } + + dask_resources = submit_kwargs.get("resources", {}) + if dask_resources and not await self._check_resources_satisfiable(dask_resources): task["exception"] = RuntimeError( f"No worker can satisfy resources {dask_resources}. " f"Workers must be started with matching --resources flags " @@ -320,46 +375,108 @@ async def on_done(f): task["stdout"] = "" task["stderr"] = str(task["exception"]) task["exit_code"] = 1 - self._callback_func(task, "FAILED") + await self._invoke_callback(task, "FAILED") + self._purge(uid) return - dask_future = self._client.submit(fn, *args, **backend_kwargs) + # Use the RHAPSODY uid as the Dask key unless the caller explicitly set one. + # Without this, client.submit()'s default (tokenize(func, kwargs, *args) under + # pure=True) can make two distinct tasks that call the same function with the + # same arguments collide onto the same Dask key, silently sharing one Future. + submit_kwargs.setdefault("key", uid) - # Store the future for potential cancellation - self.tasks[task["uid"]]["future"] = dask_future + try: + dask_future = self._client.submit(fn, *args, **(fn_kwargs or {}), **submit_kwargs) + except Exception as e: + task["exception"] = e + task["stdout"] = "" + task["stderr"] = str(e) + await self._invoke_callback(task, "FAILED") + self._purge(uid) + return - # Schedule the callback to run when future completes - asyncio.create_task(on_done(dask_future)) + self._runtime[uid].future = dask_future + asyncio.create_task(self._on_done(task, dask_future, kind)) - async def _submit_async_function(self, task: dict[str, Any]) -> None: - """Submit async function to Dask. + async def _on_done(self, task: dict[str, Any], f: Any, kind: str) -> None: + """Shared completion handler for both function and executable Dask futures.""" + uid = task["uid"] + try: + await self._invoke_callback(task, "RUNNING") + if kind == "executable": + stdout, stderr, returncode = await f + task["stdout"] = stdout + task["stderr"] = stderr + task["exit_code"] = returncode + state = "DONE" if returncode == 0 else "FAILED" + else: + result = await f + task["return_value"] = result + task["stdout"] = "" + task["stderr"] = "" + state = "DONE" + await self._invoke_callback(task, state) + except dask.client.FutureCancelledError: + await self._invoke_callback(task, "CANCELED") + except Exception as e: + task["exception"] = e + task["stdout"] = "" + task["stderr"] = str(e) + await self._invoke_callback(task, "FAILED") + finally: + self._purge(uid) - Creates an async wrapper that preserves the original function name - for better visibility in the Dask dashboard. + async def _invoke_callback(self, task: dict[str, Any], state: str) -> None: + """Invoke the registered callback, tolerating both sync and async callables. - Args: - task: Task dictionary containing the async function and its parameters. + A callback failure is caught and logged here so it can never corrupt task state or interrupt + the completion handler that called it. """ + try: + result = self._callback_func(task, state) + if inspect.isawaitable(result): + await result + except Exception: + self.logger.exception( + f"Callback raised while handling state '{state}' for task '{task.get('uid')}'" + ) - # Preserve the real task name in dask dashboard - @wraps(task["function"]) - async def async_wrapper(): - return await task["function"](*task["args"], **task["kwargs"]) + def _purge(self, uid: str) -> None: + """Drop the backend-internal record for a task that has reached a terminal state. - await self._submit_to_dask(task, async_wrapper) + Safe to call once the terminal-state callback has already been invoked for this + uid: results live on the original task dict, which the caller retains + independently of this backend's own bookkeeping, and the callback has already + resolved any waiter before this runs. After this, `cancel_task(uid)` returns + False and `uid not in self.tasks`. + """ + self.tasks.pop(uid, None) + self._runtime.pop(uid, None) - async def _submit_sync_function(self, task: dict[str, Any]) -> None: - """Submit a sync (non-coroutine) function to Dask. + async def _submit_function(self, task: dict[str, Any]) -> None: + """Submit a Python callable (sync or async) directly to Dask. - Dask workers run sync functions natively — no async wrapper needed. + Dask's own worker executor detects whether the submitted callable is a + coroutine function and runs it on the worker's event loop if so, or in the + worker's thread pool otherwise — no RHAPSODY-side adapter is needed for + either case (see module docstring). The callable is submitted exactly as the + caller defined it, optionally pre-bound via `functools.partial` for kwargs; + it is never renamed or re-wrapped, so it can never claim an identity that + doesn't match what pickling finds at that name. Args: - task: Task dictionary containing the sync function and its parameters. + task: Task dictionary containing the function and its parameters. """ fn = task["function"] - if task.get("kwargs"): - fn = partial(fn, **task["kwargs"]) - await self._submit_to_dask(task, fn, *task["args"]) + kwargs = task.get("kwargs") or {} + if kwargs: + # Bound via partial (not passed as client.submit(**kwargs)) because + # client.submit reserves kwarg names like "resources"/"retries"/"key" for + # its own submission options; a user function kwarg with the same name + # must never collide with those. + fn = partial(fn, **kwargs) + args = tuple(task.get("args") or ()) + await self._submit_and_track(task, fn, args, kind="function") async def _submit_executable(self, task: dict[str, Any]) -> None: """Submit an executable task to run via subprocess inside a Dask worker. @@ -368,67 +485,36 @@ async def _submit_executable(self, task: dict[str, Any]) -> None: task: Task dictionary containing executable path, arguments, and metadata. """ bksp = task.get("task_backend_specific_kwargs", {}) - backend_kwargs = {k: v for k, v in bksp.items() if k not in ("cwd", "shell", "env")} - dask_resources = backend_kwargs.get("resources", {}) - if dask_resources and not self._check_resources_satisfiable(dask_resources): - msg = ( - f"No worker can satisfy resources {dask_resources}. " - f"Workers must be started with matching --resources flags " - f'(e.g. dask worker --resources "GPU=1").' - ) - task["stderr"] = msg - task["stdout"] = "" - task["exit_code"] = 1 - self._callback_func(task, "FAILED") - return - - dask_future = self._client.submit( - _run_executable, - task["executable"], - task.get("arguments", []), - cwd=bksp.get("cwd"), - env=bksp.get("env"), - shell=bksp.get("shell", False), - capture_stdio=task.get("capture_stdio", False), - output_dir=self._work_dir, - uid=task["uid"], - **backend_kwargs, + fn_kwargs = { + "cwd": bksp.get("cwd"), + "env": bksp.get("env"), + "shell": bksp.get("shell", False), + "capture_stdio": task.get("capture_stdio", False), + "output_dir": self._work_dir, + "uid": task["uid"], + } + args = (task["executable"], task.get("arguments", [])) + await self._submit_and_track( + task, _run_executable, args, kind="executable", fn_kwargs=fn_kwargs ) - self.tasks[task["uid"]]["future"] = dask_future - async def on_done(f): - task_uid = task["uid"] - try: - self._callback_func(task, "RUNNING") - stdout, stderr, returncode = await f - task["stdout"] = stdout - task["stderr"] = stderr - task["exit_code"] = returncode - state = "DONE" if returncode == 0 else "FAILED" - self._callback_func(task, state) - except dask.client.FutureCancelledError: - self._callback_func(task, "CANCELED") - except Exception as e: - task["exception"] = e - task["stdout"] = "" - task["stderr"] = str(e) - self._callback_func(task, "FAILED") - finally: - if task_uid in self.tasks: - del self.tasks[task_uid] - - asyncio.create_task(on_done(dask_future)) - - def _check_resources_satisfiable(self, resources: dict) -> bool: + async def _check_resources_satisfiable(self, resources: dict) -> bool: """Return True if at least one connected worker can satisfy all resource constraints. + Uses `client.scheduler.identity()` rather than `client.scheduler_info()`: for an + asynchronous Client, `scheduler_info()` returns a cached snapshot captured once + at connect time with an always-empty "workers" mapping (see its own docstring, + which recommends this exact alternative), so it can never see workers that + joined afterward or their resource advertisements. + Args: resources: Dict of resource requirements (e.g. {"GPU": 1}). Returns: True if a qualifying worker exists, False otherwise. """ - workers = self._client.scheduler_info().get("workers", {}) + info = await self._client.scheduler.identity(n_workers=-1) + workers = info.get("workers", {}) return any( all(w.get("resources", {}).get(k, 0) >= v for k, v in resources.items()) for w in workers.values() @@ -457,7 +543,16 @@ def link_explicit_data_deps( file_name: str | None = None, file_path: str | None = None, ) -> None: - """Handle explicit data dependencies between tasks. + """Intentional no-op: Dask has no file-staging directive analogous to this hook. + + RHAPSODY's other backends that implement this (e.g. RADICAL-Pilot) use + backend-native staging/transfer directives executed before a task runs. Dask + has no equivalent at the flat `submit_tasks(list)` batch-submission level used + here; Dask's native dependency mechanism is instead passing one task's Future + as an argument to another `client.submit()` call, which requires a task-graph + API this module doesn't currently expose. Should RHAPSODY grow an explicit + dependency-graph API, this backend should wire dependencies that way rather + than via file staging. Args: src_task: The source task that produces the dependency. @@ -468,7 +563,7 @@ def link_explicit_data_deps( pass def link_implicit_data_deps(self, src_task: dict[str, Any], dst_task: dict[str, Any]) -> None: - """Handle implicit data dependencies for a task. + """Intentional no-op — see `link_explicit_data_deps` for rationale. Args: src_task: The source task that produces data. @@ -487,6 +582,10 @@ async def state(self) -> str: async def task_state_cb(self, task: dict, state: str) -> None: """Callback function invoked when a task's state changes. + Intentional no-op: state notification is delivered through the callback + registered via `register_callback` (invoked from `_on_done`/`_invoke_callback`), + matching the convention used by every sibling backend. + Args: task: Dictionary containing task information and metadata. state: The new state of the task. @@ -494,7 +593,11 @@ async def task_state_cb(self, task: dict, state: str) -> None: pass async def build_task(self, task: dict) -> None: - """Build or prepare a task for execution. + """Intentional no-op: Dask has no separate "build description, then submit" phase. + + Task construction happens inline during `submit_tasks`/`_submit_function`/ + `_submit_executable` — there is no intermediate native task-description object + to build ahead of time the way e.g. RADICAL-Pilot's `TaskDescription` requires. Args: task: Dictionary containing task definition, parameters, and metadata @@ -505,8 +608,9 @@ async def build_task(self, task: dict) -> None: async def shutdown(self) -> None: """Shutdown the Dask client and clean up resources. - Closes the Dask client connection, clears task storage, and handles any cleanup exceptions - gracefully. + Cancels all outstanding tasks unconditionally. Closes the Dask client only if this backend + created it (see class docstring on ownership); an externally provided client/cluster is left + open for the caller to manage. """ # Set backend state to SHUTDOWN self._backend_state = BackendMainStates.SHUTDOWN @@ -514,12 +618,16 @@ async def shutdown(self) -> None: if self._client is not None: try: - # Cancel all running tasks first + # Cancel all running tasks first, regardless of ownership. await self.cancel_all_tasks() - # Close the client - await self._client.close() - self.logger.info("Dask client shutdown complete") + if self._owns_client: + await self._client.close() + self.logger.info("Dask client shutdown complete") + else: + self.logger.info( + "Externally-provided Dask client left open (not owned by this backend)" + ) except Exception as e: self.logger.exception(f"Error during shutdown: {str(e)}") finally: @@ -528,6 +636,7 @@ async def shutdown(self) -> None: # Always clean up state regardless of client presence self.tasks.clear() + self._runtime.clear() self._initialized = False def _ensure_initialized(self): diff --git a/tests/unit/test_backend_execution_dask_parallel.py b/tests/unit/test_backend_execution_dask_parallel.py index 4274036..8fb55f9 100644 --- a/tests/unit/test_backend_execution_dask_parallel.py +++ b/tests/unit/test_backend_execution_dask_parallel.py @@ -11,6 +11,19 @@ from rhapsody import ComputeTask +async def _pickling_regression_target(n): + """Module-level async target used by test_wraps_closure_pickling_regression.""" + return n + + +def _create_task_that_closes_coro(coro, **kwargs): + """Stand-in for `asyncio.create_task` in tests that never let the scheduled completion coroutine + run — closes it immediately instead of leaking it, which otherwise trips a "coroutine was never + awaited" RuntimeWarning at GC time.""" + coro.close() + return None + + def test_dask_backend_import(): """Test that DaskExecutionBackend can be imported.""" try: @@ -44,6 +57,7 @@ def test_dask_backend_init(): assert not backend._initialized assert backend._client is None assert backend.tasks == {} + assert backend._runtime == {} # Test initialization with resources resources = {"n_workers": 2, "threads_per_worker": 1} @@ -184,12 +198,15 @@ async def test_dask_backend_task_submission_routing(): """Test that tasks are routed to the correct submission methods.""" try: from unittest.mock import AsyncMock + from unittest.mock import MagicMock from unittest.mock import patch from rhapsody.backends import DaskExecutionBackend backend = DaskExecutionBackend() backend._initialized = True + backend._client = MagicMock() + backend._client.status = "running" # Executable tasks route to _submit_executable (not FAILED) with patch.object(backend, "_submit_executable", new_callable=AsyncMock) as mock_exec: @@ -197,25 +214,26 @@ async def test_dask_backend_task_submission_routing(): await backend.submit_tasks([executable_task]) mock_exec.assert_called_once() - # Sync function tasks route to _submit_sync_function (not FAILED) - with patch.object(backend, "_submit_sync_function", new_callable=AsyncMock) as mock_sync: + # Both sync and async function tasks route to _submit_function — Dask itself + # (not a RHAPSODY-side wrapper) distinguishes sync/async execution once the + # callable reaches a worker, so there is only one dispatch path here. + with patch.object(backend, "_submit_function", new_callable=AsyncMock) as mock_fn: def sync_fn(): return "sync" sync_task = ComputeTask(function=sync_fn, args=[], kwargs={}) await backend.submit_tasks([sync_task]) - mock_sync.assert_called_once() + mock_fn.assert_called_once() - # Async function tasks route to _submit_async_function - with patch.object(backend, "_submit_async_function", new_callable=AsyncMock) as mock_async: + with patch.object(backend, "_submit_function", new_callable=AsyncMock) as mock_fn: async def async_fn(): return "async" async_task = ComputeTask(function=async_fn, args=[], kwargs={}) await backend.submit_tasks([async_task]) - mock_async.assert_called_once() + mock_fn.assert_called_once() except ImportError: pytest.skip("Dask dependencies not available") @@ -272,11 +290,13 @@ async def test_dask_backend_shutdown(): backend = DaskExecutionBackend() backend._initialized = True backend.tasks = {"test": "task"} + backend._runtime = {"test": "runtime"} await backend.shutdown() assert backend._client is None assert not backend._initialized assert len(backend.tasks) == 0 + assert len(backend._runtime) == 0 except ImportError: pytest.skip("Dask dependencies not available") @@ -330,6 +350,7 @@ async def test_dask_submit_executable_passes_capture_stdio(tmp_path): from unittest.mock import patch from rhapsody.backends import DaskExecutionBackend + from rhapsody.backends.execution.dask_parallel import _TaskRuntime backend = DaskExecutionBackend() backend._initialized = True @@ -347,8 +368,9 @@ def fake_submit(fn, *args, **kwargs): task = ComputeTask(executable="/bin/echo", arguments=["hi"], capture_stdio=True) backend.tasks[task["uid"]] = task + backend._runtime[task["uid"]] = _TaskRuntime(uid=task["uid"], kind="executable") - with patch("asyncio.create_task"): + with patch("asyncio.create_task", side_effect=_create_task_that_closes_coro): await backend._submit_executable(task) assert captured.get("capture_stdio") is True @@ -415,6 +437,7 @@ async def test_dask_submit_executable_cwd_from_bksp(): from unittest.mock import patch from rhapsody.backends import DaskExecutionBackend + from rhapsody.backends.execution.dask_parallel import _TaskRuntime backend = DaskExecutionBackend() backend._initialized = True @@ -436,8 +459,9 @@ def fake_submit(fn, *args, **kwargs): task_backend_specific_kwargs={"cwd": "/tmp"}, ) backend.tasks[task["uid"]] = task + backend._runtime[task["uid"]] = _TaskRuntime(uid=task["uid"], kind="executable") - with patch("asyncio.create_task"): + with patch("asyncio.create_task", side_effect=_create_task_that_closes_coro): await backend._submit_executable(task) assert captured.get("cwd") == "/tmp" @@ -454,6 +478,7 @@ async def test_dask_submit_executable_no_cwd(): from unittest.mock import patch from rhapsody.backends import DaskExecutionBackend + from rhapsody.backends.execution.dask_parallel import _TaskRuntime backend = DaskExecutionBackend() backend._initialized = True @@ -470,8 +495,9 @@ def fake_submit(fn, *args, **kwargs): task = ComputeTask(executable="/bin/pwd") backend.tasks[task["uid"]] = task + backend._runtime[task["uid"]] = _TaskRuntime(uid=task["uid"], kind="executable") - with patch("asyncio.create_task"): + with patch("asyncio.create_task", side_effect=_create_task_that_closes_coro): await backend._submit_executable(task) assert captured.get("cwd") is None @@ -493,6 +519,7 @@ async def test_dask_function_done_stdout_is_string(): from unittest.mock import patch from rhapsody.backends import DaskExecutionBackend + from rhapsody.backends.execution.dask_parallel import _TaskRuntime except ImportError: pytest.skip("Dask not available") @@ -503,11 +530,12 @@ async def test_dask_function_done_stdout_is_string(): task = ComputeTask(function=lambda: 99, args=[]) backend.tasks[task["uid"]] = {} + backend._runtime[task["uid"]] = _TaskRuntime(uid=task["uid"], kind="function") fut = asyncio.get_event_loop().create_future() with patch.object(backend, "_client") as mc: mc.submit.return_value = fut - asyncio.create_task(backend._submit_async_function(task)) + asyncio.create_task(backend._submit_function(task)) await asyncio.sleep(0) fut.set_result(99) await asyncio.sleep(0) @@ -526,6 +554,7 @@ async def test_dask_function_failed_stdout_is_string(): from unittest.mock import patch from rhapsody.backends import DaskExecutionBackend + from rhapsody.backends.execution.dask_parallel import _TaskRuntime except ImportError: pytest.skip("Dask not available") @@ -536,11 +565,12 @@ async def test_dask_function_failed_stdout_is_string(): task = ComputeTask(function=lambda: 1 / 0, args=[]) backend.tasks[task["uid"]] = {} + backend._runtime[task["uid"]] = _TaskRuntime(uid=task["uid"], kind="function") fut = asyncio.get_event_loop().create_future() with patch.object(backend, "_client") as mc: mc.submit.return_value = fut - asyncio.create_task(backend._submit_async_function(task)) + asyncio.create_task(backend._submit_function(task)) await asyncio.sleep(0) fut.set_exception(ZeroDivisionError("div by zero")) await asyncio.sleep(0) @@ -548,3 +578,279 @@ async def test_dask_function_failed_stdout_is_string(): failed = [(t, s) for t, s in captured if s == "FAILED"] assert failed, "FAILED callback never fired" assert isinstance(failed[0][0].get("stdout"), str) + + +# --------------------------------------------------------------------------- +# Pickling regression — proves the old @wraps-closure pattern was the bug, +# and that submitting the real callable directly is not. +# --------------------------------------------------------------------------- + + +def test_wraps_closure_pickling_regression(): + """Regression test for the original `PicklingError`. + + The deleted `_submit_async_function` wrapped every async task callable in a + local closure decorated with `@wraps(task["function"])`. `@wraps` copies the + original function's `__module__`/`__qualname__` onto the closure, so pickling + it by reference resolves to a *different* object living at that name and + raises `PicklingError: ... it's not the same object as ...` — this is the + exact shape of the original bug (`Can't pickle : it's not the same object as __main__.infer_batch_task`). + + Proves (a) the old wrapping pattern really does break pickling, and (b) the + new pattern — submitting the real callable directly, optionally bound via + `functools.partial` for kwargs — does not. + """ + import pickle + from functools import partial + from functools import wraps + + # Old (deleted) pattern: a closure decorated with @wraps(original). + @wraps(_pickling_regression_target) + async def async_wrapper(): + return await _pickling_regression_target(1) + + with pytest.raises(pickle.PicklingError, match="not the same object as"): + pickle.dumps(async_wrapper) + + # New pattern: submit the real module-level function directly, or via partial + # to pre-bind kwargs — never through a renamed/re-wrapped closure. + restored_fn = pickle.loads(pickle.dumps(_pickling_regression_target)) + assert restored_fn is _pickling_regression_target + + bound = partial(_pickling_regression_target, 1) + restored_partial = pickle.loads(pickle.dumps(bound)) + assert restored_partial.func is _pickling_regression_target + assert restored_partial.args == (1,) + + +# --------------------------------------------------------------------------- +# Client/cluster ownership and validation +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dask_external_sync_client_rejected(): + """A caller-supplied Client that isn't asynchronous=True must be rejected at init.""" + try: + from unittest.mock import MagicMock + + from rhapsody.backends import DaskExecutionBackend + except ImportError: + pytest.skip("Dask dependencies not available") + + sync_client = MagicMock() + sync_client.asynchronous = False + + with pytest.raises(ValueError, match="asynchronous=True"): + await DaskExecutionBackend(client=sync_client) + + +# --------------------------------------------------------------------------- +# Callback handling +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dask_async_callback_registered_and_invoked(): + """_invoke_callback awaits an async callback rather than firing-and-forgetting it.""" + try: + from rhapsody.backends import DaskExecutionBackend + except ImportError: + pytest.skip("Dask dependencies not available") + + backend = DaskExecutionBackend() + backend._initialized = True + + calls = [] + + async def async_callback(task, state): + await asyncio.sleep(0) + calls.append((task["uid"], state)) + + backend.register_callback(async_callback) + + task = ComputeTask(function=lambda: 1, args=[]) + await backend._invoke_callback(task, "RUNNING") + + assert calls == [(task["uid"], "RUNNING")] + + +@pytest.mark.asyncio +async def test_dask_callback_exception_does_not_break_task_completion(): + """A raising callback must not crash `_on_done` or corrupt task state.""" + try: + from rhapsody.backends import DaskExecutionBackend + from rhapsody.backends.execution.dask_parallel import _TaskRuntime + except ImportError: + pytest.skip("Dask dependencies not available") + + backend = DaskExecutionBackend() + backend._initialized = True + + def bad_callback(task, state): + raise RuntimeError("callback boom") + + backend.register_callback(bad_callback) + + task = ComputeTask(function=lambda: 42, args=[]) + backend.tasks[task["uid"]] = task + backend._runtime[task["uid"]] = _TaskRuntime(uid=task["uid"], kind="function") + + fut = asyncio.get_event_loop().create_future() + fut.set_result(42) + + # Must not raise despite the callback raising on every invocation. + await backend._on_done(task, fut, "function") + + assert task["return_value"] == 42 + assert task["uid"] not in backend.tasks + + +# --------------------------------------------------------------------------- +# Data dependency hooks — documented no-ops +# --------------------------------------------------------------------------- + + +def test_dask_link_data_deps_are_safe_noops(): + """link_explicit_data_deps / link_implicit_data_deps are documented no-ops.""" + try: + from rhapsody.backends import DaskExecutionBackend + except ImportError: + pytest.skip("Dask dependencies not available") + + backend = DaskExecutionBackend() + + src = ComputeTask(function=lambda: 1, args=[]) + dst = ComputeTask(function=lambda: 2, args=[]) + src_before, dst_before = dict(src), dict(dst) + + assert ( + backend.link_explicit_data_deps( + src_task=src, dst_task=dst, file_name="x", file_path="/tmp/x" + ) + is None + ) + assert backend.link_implicit_data_deps(src, dst) is None + + assert dict(src) == src_before + assert dict(dst) == dst_before + assert backend.tasks == {} + + +# --------------------------------------------------------------------------- +# Argument handling — no more silent mutation/filtering of the caller's task +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dask_args_kwargs_not_mutated_on_caller_task(): + """submit_tasks must not rewrite the caller's task['args']/task['kwargs'].""" + try: + from unittest.mock import MagicMock + from unittest.mock import patch + + from rhapsody.backends import DaskExecutionBackend + except ImportError: + pytest.skip("Dask dependencies not available") + + backend = DaskExecutionBackend() + backend._initialized = True + backend._client = MagicMock() + backend._client.status = "running" + backend._client.submit = MagicMock(return_value=MagicMock()) + + task = ComputeTask(function=lambda a, x=None: (a, x), args=(1, 2), kwargs={"x": 1}) + args_before = task["args"] + kwargs_before = task["kwargs"] + + with patch("asyncio.create_task", side_effect=_create_task_that_closes_coro): + await backend.submit_tasks([task]) + + assert task["args"] == args_before + assert task["kwargs"] == kwargs_before + + +# --------------------------------------------------------------------------- +# Resource pre-check — must use a live scheduler snapshot, not the stale cache +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_check_resources_satisfiable_uses_scheduler_identity(): + """_check_resources_satisfiable must use scheduler.identity(), not the always- empty + scheduler_info() cache on an asynchronous client.""" + try: + from unittest.mock import AsyncMock + from unittest.mock import MagicMock + + from rhapsody.backends import DaskExecutionBackend + except ImportError: + pytest.skip("Dask dependencies not available") + + backend = DaskExecutionBackend() + backend._initialized = True + backend._client = MagicMock() + backend._client.scheduler = MagicMock() + backend._client.scheduler.identity = AsyncMock( + return_value={"workers": {"w1": {"resources": {"GPU": 2}}}} + ) + + assert await backend._check_resources_satisfiable({"GPU": 1}) is True + assert await backend._check_resources_satisfiable({"GPU": 4}) is False + backend._client.scheduler.identity.assert_called_with(n_workers=-1) + + +# --------------------------------------------------------------------------- +# Dask key collision — proves distinct tasks never silently share a Future +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dask_duplicate_function_args_get_distinct_keys(): + """Two tasks with the same function/args must get distinct Dask submission keys. + + Regression test: `client.submit()` defaults to `pure=True` with no explicit `key`, + deriving the Dask key from `tokenize(func, kwargs, *args)`. Two RHAPSODY tasks + calling the same function with the same args used to tokenize to the identical + key, so the second `submit()` call would silently return the first task's + `Future` (`distributed.client.Client.submit`: `if key in self.futures: return + Future(key, self)`) instead of doing independent work. Passing `key=task["uid"]` + explicitly bypasses the tokenize path entirely. + """ + try: + from unittest.mock import MagicMock + from unittest.mock import patch + + from rhapsody.backends import DaskExecutionBackend + except ImportError: + pytest.skip("Dask dependencies not available") + + backend = DaskExecutionBackend() + backend._initialized = True + backend._client = MagicMock() + backend._client.status = "running" + + submitted_keys = [] + + def fake_submit(fn, *args, **kwargs): + submitted_keys.append(kwargs.get("key")) + return MagicMock() + + backend._client.submit = fake_submit + + def shared_fn(n): + return n + + tasks = [ + ComputeTask(function=shared_fn, args=(1,)), + ComputeTask(function=shared_fn, args=(1,)), + ] + + with patch("asyncio.create_task", side_effect=_create_task_that_closes_coro): + await backend.submit_tasks(tasks) + + assert len(submitted_keys) == 2 + assert submitted_keys[0] != submitted_keys[1] + assert submitted_keys[0] == tasks[0]["uid"] + assert submitted_keys[1] == tasks[1]["uid"] diff --git a/tests/unit/test_dask_custom_cluster.py b/tests/unit/test_dask_custom_cluster.py index 40faf65..f973b64 100644 --- a/tests/unit/test_dask_custom_cluster.py +++ b/tests/unit/test_dask_custom_cluster.py @@ -29,6 +29,45 @@ async def test_dask_preconfigured_client(): await backend.shutdown() +# --------------------------------------------------------------------------- +# Client/cluster ownership — shutdown() must not close resources it didn't create +# --------------------------------------------------------------------------- + + +async def test_dask_shutdown_does_not_close_external_client(): + """A caller-provided Client is left open after backend.shutdown().""" + async with LocalCluster(n_workers=1, threads_per_worker=1, asynchronous=True) as cluster: + async with Client(cluster, asynchronous=True) as client: + backend = await DaskExecutionBackend(client=client) + await backend.shutdown() + + assert client.status == "running" + fut = client.submit(lambda: 1) + assert await fut == 1 + + +async def test_dask_shutdown_does_not_close_external_cluster(): + """A caller-provided Cluster stays usable after backend.shutdown() closes its own client.""" + async with LocalCluster(n_workers=1, threads_per_worker=1, asynchronous=True) as cluster: + backend = await DaskExecutionBackend(cluster=cluster) + await backend.shutdown() + + async with Client(cluster, asynchronous=True) as client: + fut = client.submit(lambda: 2) + assert await fut == 2 + + +async def test_dask_shutdown_closes_owned_client(): + """A backend-created Client is closed by backend.shutdown().""" + backend = await DaskExecutionBackend(resources={"n_workers": 1, "threads_per_worker": 1}) + owned_client = backend._client + + await backend.shutdown() + + assert backend._client is None + assert owned_client.status != "running" + + # --------------------------------------------------------------------------- # End-to-end task execution # --------------------------------------------------------------------------- @@ -128,6 +167,61 @@ async def test_dask_unmet_resources_executable_sets_stderr(): assert tasks[0].exit_code == 1 +async def test_dask_satisfiable_resources_succeed(): + """Tasks with satisfiable resource constraints must succeed, not always fail. + + Regression test: the resource pre-check used to call `Client.scheduler_info()`, + which for an asynchronous client always returns an empty `workers` mapping — so + the check could never detect a *satisfiable* request, only ever "fail fast" on + unsatisfiable ones. This exercises the case that used to be impossible to pass. + """ + async with LocalCluster( + n_workers=1, threads_per_worker=1, resources={"GPU": 1}, asynchronous=True + ) as cluster: + async with Client(cluster, asynchronous=True) as client: + backend = await DaskExecutionBackend(client=client) + session = Session(backends=[backend]) + tasks = [ + ComputeTask( + function=lambda: "ok", + task_backend_specific_kwargs={"resources": {"GPU": 1}}, + ) + ] + async with session: + await session.submit_tasks(tasks) + await session.wait_tasks(tasks) + + assert tasks[0].state == "DONE" + assert tasks[0].return_value == "ok" + + +async def test_dask_duplicate_function_args_do_not_share_a_future(): + """Two tasks calling the same function with the same args both complete correctly end-to-end. + + See `test_dask_duplicate_function_args_get_distinct_keys` in + test_backend_execution_dask_parallel.py for the precise regression proof (that each + task is submitted with its own distinct Dask key, rather than silently sharing one + via `client.submit()`'s default `pure=True` tokenization). + """ + + def make_marker(n): + return n + + async with DaskExecutionBackend(resources={"n_workers": 1, "threads_per_worker": 1}) as backend: + session = Session(backends=[backend]) + tasks = [ + ComputeTask(function=make_marker, args=(1,)), + ComputeTask(function=make_marker, args=(1,)), + ] + async with session: + await session.submit_tasks(tasks) + await session.wait_tasks(tasks) + + assert all(t.state == "DONE" for t in tasks) + assert tasks[0].uid != tasks[1].uid + assert [t.return_value for t in tasks] == [1, 1] + + if __name__ == "__main__": asyncio.run(test_dask_preconfigured_cluster()) asyncio.run(test_dask_preconfigured_client()) From 526b9658ffbf2ca22211da7e57cf374844cbe3d5 Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Wed, 26 Aug 2026 19:53:02 +0000 Subject: [PATCH 2/5] addign Dask Slurm cluster example --- examples/07-dask-backend-slrum-cluster.py | 61 +++++++++++++++++++++++ examples/README.md | 27 ++++++++++ 2 files changed, 88 insertions(+) create mode 100644 examples/07-dask-backend-slrum-cluster.py diff --git a/examples/07-dask-backend-slrum-cluster.py b/examples/07-dask-backend-slrum-cluster.py new file mode 100644 index 0000000..7769af3 --- /dev/null +++ b/examples/07-dask-backend-slrum-cluster.py @@ -0,0 +1,61 @@ +"""Example: DaskExecutionBackend against a Slurm cluster via dask_jobqueue. + +Requires `dask_jobqueue` installed and a Slurm scheduler reachable from this host +(salloc/sbatch on $PATH). Update queue/account/cores/memory/walltime for your allocation. + +SLURMCluster must be constructed with asynchronous=True (and entered via `async with`) +so it shares this script's event loop instead of spinning up its own background-thread +loop. Skipping that makes the Client we build on top of it (cluster=...) inherit a +mismatched loop: every await on a Future or on shutdown then silently falls back to +blocking-sync mode and breaks with confusing TypeError/AttributeError failures. +""" + +import asyncio +import logging + +import rhapsody +from dask_jobqueue import SLURMCluster +from rhapsody.api import ComputeTask +from rhapsody.api import Session +from rhapsody.backends.execution.dask_parallel import DaskExecutionBackend + +rhapsody.enable_logging(level=logging.DEBUG) + +logger = logging.getLogger(__name__) + + +def compute_task(n: int) -> int: + return n * n + + +async def main(): + # tested on purdue anvil + async with SLURMCluster( + queue="wholenode", + account="dmrxxx", # user must provide this + cores=16, # cores per Slurm job (worker) + memory="16GB", + walltime="00:30:00", + # job_extra_directives=["--gres=gpu:1"], # e.g. for GPU-constrained jobs + asynchronous=True, + ) as cluster: + await cluster.scale(jobs=1) # submit 1 Slurm job hosting one worker + + # cluster= (not client=) -> backend creates+owns the Client, not the cluster. + backend = await DaskExecutionBackend(cluster=cluster) + session = Session(backends=[backend]) + + tasks = [ComputeTask(function=compute_task, args=(i,)) for i in range(10)] + + async with session: + await session.submit_tasks(tasks) + await session.wait_tasks(tasks) + + for t in tasks: + print(t.uid, t.state, t.return_value) + # cluster is closed automatically on exit from `async with` — you created it, + # so it's yours to close, but the async form does it via await for you. + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/README.md b/examples/README.md index 34569f9..bee3be5 100644 --- a/examples/README.md +++ b/examples/README.md @@ -9,6 +9,7 @@ Different examples require different launchers depending on the backend they use | Backend | Launcher | Command | |---------|----------|---------| | `ConcurrentExecutionBackend` | Standard Python | `python example.py` | +| `DaskExecutionBackend` | Standard Python (+ Dask cluster) | `python example.py` | | `DragonExecutionBackend` | Dragon runtime | `dragon example.py` | | `DragonVllmInferenceBackend` | Dragon runtime + GPU | `dragon example.py` | @@ -91,6 +92,18 @@ Integrates RHAPSODY with [RADICAL AsyncFlow](https://github.com/radical-cybertoo --- +### 05 — Dask Backend (Local Cluster) + +```bash +python 05-dask-backend.py +``` + +Runs sync functions, async functions, and executables on `DaskExecutionBackend`, backed by an automatically-started local Dask cluster. Pass `cluster=` or `client=` to target Slurm, Kubernetes, or any other Dask-compatible deployment instead — see example 07 for the Slurm variant and a critical gotcha it demonstrates. + +**What you'll learn:** `DaskExecutionBackend`, sync/async function dispatch on the same code path, executable tasks, mixed task types in one submission. + +--- + ### 06 — Multi-Service AI-HPC Workflow (Dragon + vLLM) ```bash @@ -102,3 +115,17 @@ Runs **two independent `DragonVllmInferenceBackend` services** at once (one per **Requires:** GPU access (2 GPUs), vLLM installed (`dragonhpc[ai]`), and a locally downloaded model directory (`HF_HUB_OFFLINE=1` recommended — see the model download notes for this backend). **What you'll learn:** Running multiple inference services in one allocation, `use_service=True` HTTP endpoints, `get_endpoint()`, mixing direct (`AITask`) and service-style (`ComputeTask` + `aiohttp`) access to the same backend type. + +--- + +### 07 — Dask Backend on Slurm (dask_jobqueue) + +```bash +python 07-dask-backend-slrum-cluster.py +``` + +Runs `DaskExecutionBackend` against a real Slurm allocation via `dask_jobqueue.SLURMCluster`. Demonstrates the `cluster=`/`client=` ownership model and a critical gotcha: the cluster must be constructed with `asynchronous=True` and entered via `async with`, or the Dask `Client` built on top of it silently inherits a mismatched event loop — every `await` on a task result or on shutdown then breaks with confusing `TypeError`/`AttributeError` failures. + +**Requires:** `dask_jobqueue` installed and a Slurm scheduler reachable from this host (`sbatch` on `$PATH`). Update `queue`/`account`/`cores`/`memory`/`walltime` for your allocation. + +**What you'll learn:** `SLURMCluster(asynchronous=True)`, `cluster=` ownership semantics, running RHAPSODY against a real HPC batch scheduler. From 6ed70a40fc55353f95f93d3cb4044060fdea14a7 Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Tue, 8 Sep 2026 17:07:38 +0000 Subject: [PATCH 3/5] docs: fix stale Dask examples that reproduce the async-loop-mismatch bug The old @wraps-closure wording ("wrapped transparently") and the plain sync SLURMCluster()/KubeCluster() construction shown in advanced-usage.md, resource-specification.md, configuration.md, both dask examples, and the tutorial notebook no longer match the refactored backend and, worse, are exactly the pattern that silently breaks task results/shutdown via a mismatched event loop. Update all of them to asynchronous=True + async with, add a CHANGELOG entry, and cross-link the new 07 Slurm example. --- CHANGELOG.md | 69 ++++++++++++++ docs/getting-started/advanced-usage.md | 30 +++++-- docs/getting-started/configuration.md | 2 +- .../getting-started/resource-specification.md | 17 +++- examples/05-dask-backend.py | 3 +- examples/07-dask-backend-slrum-cluster.py | 2 +- tutorials/dask-backend-tutorial.ipynb | 90 +------------------ 7 files changed, 114 insertions(+), 99 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8a0e2fb..efcce6f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,74 @@ # Changelog +## [Unreleased] + +### Fixed + +- **`DaskExecutionBackend` PicklingError on async function tasks** — `submit_tasks` + wrapped every async callable in a local closure decorated with + `@functools.wraps(task["function"])` before submitting it. `@wraps` copies the + original function's `__module__`/`__qualname__` onto the closure, so pickling it + by reference resolved to a *different* object living at that name and raised + `PicklingError: Can't pickle : it's not the same object as + module.name`. Dask workers already natively detect `iscoroutinefunction()` on a + submitted callable and run it on the worker's own event loop (no thread pool) — + sync and async callables are now both submitted directly (optionally through + `functools.partial` to pre-bind kwargs), eliminating the wrapper entirely rather + than patching around it. +- **`DaskExecutionBackend._check_resources_satisfiable()` could never detect a + satisfiable resource request** — it called `Client.scheduler_info()`, which for + an asynchronous client always returns a cached snapshot with an empty `workers` + mapping (see that method's own docstring). Every resource-constrained task + failed regardless of whether a matching worker actually existed. Fixed to use + `await client.scheduler.identity(n_workers=-1)`, Dask's own documented + alternative for a live per-worker view. +- **`DaskExecutionBackend` could silently share one Dask `Future` across two + distinct tasks** — `client.submit()` defaults to `pure=True` with no explicit + `key`, deriving the Dask key from `tokenize(func, kwargs, *args)`. Two RHAPSODY + tasks calling the same function with the same arguments tokenized to the + identical key, so the second `submit()` silently returned the first task's + `Future` instead of doing independent work. Fixed by always passing + `key=task["uid"]` (unless the caller already set one via + `task_backend_specific_kwargs`). +- **`DaskExecutionBackend.shutdown()` always closed the Dask `Client`**, even one + the caller supplied via `client=`/`cluster=` — the constructor already tracked + `_client_provided`/`_cluster_provided` but never consulted them at shutdown. + Ownership is now tracked explicitly (`_owns_client`) and `shutdown()` only + closes a client this backend created itself. + +### Changed + +- **`DaskExecutionBackend`** no longer mutates the caller's task dict for + submission bookkeeping: `task["args"]` is no longer rewritten, and + `asyncio.Future` arguments are no longer silently filtered out of `args` + (nothing in RHAPSODY's task contract puts one there; a genuinely unpicklable + argument now surfaces as a real, attributable submission failure on that task + instead of silently shifting positional args). The Dask `Future` handle moved + off the shared task dict into a private per-task runtime record. +- **`DaskExecutionBackend.submit_tasks()`** now raises `ValueError` immediately for + a task specifying neither `function` nor `executable`, instead of recording it + as a per-task `FAILED` callback — this is a caller programming error, not a + runtime submission failure. Also raises `BackendError` if the Dask client + itself is unusable (e.g. scheduler connection lost), rather than attributing a + whole-backend outage to whichever task happened to be submitting at the time. +- **`DaskExecutionBackend`** now requires an externally-supplied `client=` to have + been constructed with `asynchronous=True`; raises `ValueError` at init time + otherwise instead of failing confusingly later. +- Registered callbacks on `DaskExecutionBackend` may now be sync or async + callables; a raising callback is caught and logged instead of corrupting task + completion state. + +### Added + +- `examples/07-dask-backend-slrum-cluster.py` — `DaskExecutionBackend` against a + real Slurm allocation via `dask_jobqueue.SLURMCluster`, including the + `asynchronous=True`/`async with` construction required for non-`LocalCluster` + cluster managers driven from async code (see `docs/getting-started/advanced-usage.md#dask-distributed-backend`). + Documented in `examples/README.md`. +- Regression tests for all four fixes above, plus client/cluster + ownership-on-shutdown tests, async-callback and callback-isolation tests, and a + test proving the deleted `@wraps`-closure pattern really did break pickling. + ## [0.5.0] - 2026-08-20 ### Added diff --git a/docs/getting-started/advanced-usage.md b/docs/getting-started/advanced-usage.md index 5592694..a5cf247 100644 --- a/docs/getting-started/advanced-usage.md +++ b/docs/getting-started/advanced-usage.md @@ -147,7 +147,7 @@ so they have their own event loop and do not share the session's loop. |---|---|---|---| | `ConcurrentExecutionBackend` (default) | `ThreadPoolExecutor` | called directly | run via `asyncio.run` | | `ConcurrentExecutionBackend` | `ProcessPoolExecutor` | called directly | run via `asyncio.run` | -| `DaskExecutionBackend` | Dask workers | submitted natively | wrapped transparently | +| `DaskExecutionBackend` | Dask workers | submitted natively | submitted natively — Dask runs coroutine functions on the worker's own event loop | | `OrbitExecutionBackend` | remote ORBIT endpoint | shipped via `cloudpickle` | shipped via `cloudpickle` | !!! note "ProcessPoolExecutor requires cloudpickle" @@ -658,7 +658,8 @@ from rhapsody.backends import DaskExecutionBackend def compute_square(n): return n * n -# Async function — wrapped transparently, name visible in Dask dashboard +# Async function — submitted directly; Dask runs it natively on the worker's +# own event loop, no RHAPSODY-side wrapper involved async def fetch_data(n): await asyncio.sleep(0.01) return n * 2 @@ -691,15 +692,19 @@ The task submission code is unchanged: # SLURM (requires dask-jobqueue) from dask_jobqueue import SLURMCluster -cluster = SLURMCluster(cores=4, memory="8GB", walltime="01:00:00") -cluster.scale(jobs=4) -backend = await DaskExecutionBackend(cluster=cluster) +async with SLURMCluster( + cores=4, memory="8GB", walltime="01:00:00", asynchronous=True +) as cluster: + await cluster.scale(jobs=4) + backend = await DaskExecutionBackend(cluster=cluster) + ... # Kubernetes (requires dask-kubernetes) from dask_kubernetes.operator import KubeCluster -cluster = KubeCluster(name="rhapsody-workers", n_workers=8) -backend = await DaskExecutionBackend(cluster=cluster) +async with KubeCluster(name="rhapsody-workers", n_workers=8, asynchronous=True) as cluster: + backend = await DaskExecutionBackend(cluster=cluster) + ... # Pre-existing Client from dask.distributed import Client @@ -710,6 +715,17 @@ backend = await DaskExecutionBackend(client=client) !!! tip "Default cluster" If `cluster` and `client` are both omitted, Rhapsody creates a `LocalCluster` using the `resources` dict (e.g. `{"n_workers": 4}`). +!!! warning "Cluster objects must be constructed with asynchronous=True" + `SLURMCluster`, `KubeCluster`, and other non-`LocalCluster` cluster managers + spin up their own background-thread event loop unless built with + `asynchronous=True` (and entered via `async with`, as above). The `Client` + `DaskExecutionBackend` creates around a `cluster=` you pass in inherits that + cluster's loop — if it's the wrong one, every `await` on a task result or on + `shutdown()` silently falls back to blocking-sync mode and fails with + confusing `TypeError`/`AttributeError` errors. See + [`examples/07-dask-backend-slrum-cluster.py`](https://github.com/radical-cybertools/rhapsody/blob/main/examples/07-dask-backend-slrum-cluster.py) + for a complete, working example. + ### GPU and CPU resource scheduling Pass Dask resource constraints via `task_backend_specific_kwargs={"resources": {...}}`. diff --git a/docs/getting-started/configuration.md b/docs/getting-started/configuration.md index 1bde367..7ecac56 100644 --- a/docs/getting-started/configuration.md +++ b/docs/getting-started/configuration.md @@ -56,7 +56,7 @@ ComputeTask( ``` !!! tip "Preconfigured Clusters" - If both `cluster` and `client` are omitted, Rhapsody creates a new `LocalCluster` using the provided `resources`. Pass `cluster=` to use SLURM, Kubernetes, or any other Dask cluster type. + If both `cluster` and `client` are omitted, Rhapsody creates a new `LocalCluster` using the provided `resources`. Pass `cluster=` to use SLURM, Kubernetes, or any other Dask cluster type — construct it with `asynchronous=True` (entered via `async with`); see [Cluster injection](advanced-usage.md#cluster-injection) for why this matters and a full example. ### Dragon Backend High-performance execution using the Dragon runtime. diff --git a/docs/getting-started/resource-specification.md b/docs/getting-started/resource-specification.md index cf9e37b..1f40db0 100644 --- a/docs/getting-started/resource-specification.md +++ b/docs/getting-started/resource-specification.md @@ -131,11 +131,22 @@ pass it via `cluster=` or `client=` — the task code is unchanged: from dask_jobqueue import SLURMCluster from rhapsody.backends import DaskExecutionBackend -cluster = SLURMCluster(cores=4, memory="8GB", walltime="01:00:00") -cluster.scale(jobs=4) -backend = await DaskExecutionBackend(cluster=cluster) +async with SLURMCluster( + cores=4, memory="8GB", walltime="01:00:00", asynchronous=True +) as cluster: + await cluster.scale(jobs=4) + backend = await DaskExecutionBackend(cluster=cluster) ``` +!!! warning "asynchronous=True is required" + Construct `SLURMCluster` (and other non-`LocalCluster` cluster managers) with + `asynchronous=True` and enter it via `async with`, as above — otherwise it runs + its own background-thread event loop, the `Client` built around it inherits the + mismatched loop, and task results / shutdown fail with confusing + `TypeError`/`AttributeError` errors. See + [Cluster injection](advanced-usage.md#cluster-injection) for details and + `examples/07-dask-backend-slrum-cluster.py` for a full working example. + --- ## Dragon Backend (V3) diff --git a/examples/05-dask-backend.py b/examples/05-dask-backend.py index fdb5783..13ca306 100644 --- a/examples/05-dask-backend.py +++ b/examples/05-dask-backend.py @@ -26,7 +26,8 @@ def compute_square_sync(n): async def compute_square_async(n): - """Async function — wrapped transparently, name visible in Dask dashboard.""" + """Async function — submitted directly; Dask runs it natively on the worker's own event loop, no + RHAPSODY-side wrapper involved.""" import asyncio await asyncio.sleep(0.1) diff --git a/examples/07-dask-backend-slrum-cluster.py b/examples/07-dask-backend-slrum-cluster.py index 7769af3..ec3a323 100644 --- a/examples/07-dask-backend-slrum-cluster.py +++ b/examples/07-dask-backend-slrum-cluster.py @@ -32,7 +32,7 @@ async def main(): # tested on purdue anvil async with SLURMCluster( queue="wholenode", - account="dmrxxx", # user must provide this + account="dmrxxx", # user must provide this cores=16, # cores per Slurm job (worker) memory="16GB", walltime="00:30:00", diff --git a/tutorials/dask-backend-tutorial.ipynb b/tutorials/dask-backend-tutorial.ipynb index 5e96e13..ce901eb 100644 --- a/tutorials/dask-backend-tutorial.ipynb +++ b/tutorials/dask-backend-tutorial.ipynb @@ -93,14 +93,7 @@ { "cell_type": "markdown", "metadata": {}, - "source": [ - "---\n", - "\n", - "## 2. Async Functions\n", - "\n", - "Async functions are wrapped transparently and run inside the Dask worker event loop.\n", - "The function name is preserved in the Dask dashboard for easy monitoring." - ] + "source": "---\n\n## 2. Async Functions\n\nAsync functions are submitted directly, exactly like sync ones — Dask detects that the\ncallable is a coroutine function and runs it natively on the worker's own event loop\n(no RHAPSODY-side wrapper involved). The function's real name is what shows up in the\nDask dashboard, since nothing renames or re-wraps it." }, { "cell_type": "code", @@ -198,29 +191,7 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": [ - "# from dask_jobqueue import SLURMCluster\n", - "#\n", - "# cluster = SLURMCluster(\n", - "# cores=4,\n", - "# memory=\"8GB\",\n", - "# walltime=\"01:00:00\",\n", - "# job_extra_directives=[\"--partition=compute\"],\n", - "# )\n", - "# cluster.scale(jobs=4) # Submit 4 SLURM jobs as Dask workers\n", - "#\n", - "# async def run_on_slurm():\n", - "# async with DaskExecutionBackend(cluster=cluster) as backend:\n", - "# session = Session(backends=[backend])\n", - "# tasks = [ComputeTask(function=compute_square, args=(i,)) for i in range(100)]\n", - "# async with session:\n", - "# await session.submit_tasks(tasks)\n", - "# await session.wait_tasks(tasks)\n", - "# done = sum(1 for t in tasks if t.state == \"DONE\")\n", - "# print(f\"{done}/100 tasks completed on SLURM\")\n", - "#\n", - "# asyncio.run(run_on_slurm())" - ] + "source": "# from dask_jobqueue import SLURMCluster\n#\n# # asynchronous=True + async with: the cluster runs on *our* event loop instead of\n# # spinning up its own background-thread loop. Without this, the Client we build\n# # on top of it (cluster=...) inherits that mismatched loop, and every await on a\n# # Future/close() silently falls back to blocking-sync mode and breaks.\n# async def run_on_slurm():\n# async with SLURMCluster(\n# cores=4,\n# memory=\"8GB\",\n# walltime=\"01:00:00\",\n# job_extra_directives=[\"--partition=compute\"],\n# asynchronous=True,\n# ) as cluster:\n# await cluster.scale(jobs=4) # Submit 4 SLURM jobs as Dask workers\n#\n# async with DaskExecutionBackend(cluster=cluster) as backend:\n# session = Session(backends=[backend])\n# tasks = [ComputeTask(function=compute_square, args=(i,)) for i in range(100)]\n# async with session:\n# await session.submit_tasks(tasks)\n# await session.wait_tasks(tasks)\n# done = sum(1 for t in tasks if t.state == \"DONE\")\n# print(f\"{done}/100 tasks completed on SLURM\")\n#\n# await run_on_slurm()" }, { "cell_type": "markdown", @@ -240,23 +211,7 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": [ - "# from dask_kubernetes.operator import KubeCluster\n", - "#\n", - "# cluster = KubeCluster(name=\"rhapsody-workers\", n_workers=4)\n", - "#\n", - "# async def run_on_k8s():\n", - "# async with DaskExecutionBackend(cluster=cluster) as backend:\n", - "# session = Session(backends=[backend])\n", - "# tasks = [ComputeTask(function=compute_square, args=(i,)) for i in range(50)]\n", - "# async with session:\n", - "# await session.submit_tasks(tasks)\n", - "# await session.wait_tasks(tasks)\n", - "# done = sum(1 for t in tasks if t.state == \"DONE\")\n", - "# print(f\"{done}/50 tasks completed on Kubernetes\")\n", - "#\n", - "# asyncio.run(run_on_k8s())" - ] + "source": "# from dask_kubernetes.operator import KubeCluster\n#\n# # asynchronous=True + async with, same reason as the SLURM cluster above: it\n# # keeps the cluster on our event loop instead of a mismatched background one.\n# async def run_on_k8s():\n# async with KubeCluster(name=\"rhapsody-workers\", n_workers=4, asynchronous=True) as cluster:\n# async with DaskExecutionBackend(cluster=cluster) as backend:\n# session = Session(backends=[backend])\n# tasks = [ComputeTask(function=compute_square, args=(i,)) for i in range(50)]\n# async with session:\n# await session.submit_tasks(tasks)\n# await session.wait_tasks(tasks)\n# done = sum(1 for t in tasks if t.state == \"DONE\")\n# print(f\"{done}/50 tasks completed on Kubernetes\")\n#\n# await run_on_k8s()" }, { "cell_type": "markdown", @@ -279,44 +234,7 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": [ - "# async def gpu_task(x):\n", - "# # In real usage: import cupy, torch, etc.\n", - "# return x ** 2\n", - "#\n", - "# async def run_gpu_tasks():\n", - "# # Workers must advertise GPU resources. Options:\n", - "# # a) LocalCUDACluster — auto-detects and sets GPU resources per worker\n", - "# # from dask_cuda import LocalCUDACluster\n", - "# # cluster = LocalCUDACluster()\n", - "# #\n", - "# # b) SLURMCluster — must pass --resources explicitly:\n", - "# # cluster = SLURMCluster(\n", - "# # gres=\"gpu:1\",\n", - "# # worker_extra_args=[\"--resources\", \"GPU=1\"],\n", - "# # )\n", - "# #\n", - "# # c) Pre-existing workers launched with:\n", - "# # dask worker --resources \"GPU=1\"\n", - "#\n", - "# async with DaskExecutionBackend() as backend:\n", - "# session = Session(backends=[backend])\n", - "# tasks = [\n", - "# ComputeTask(\n", - "# function=gpu_task,\n", - "# args=(i,),\n", - "# task_backend_specific_kwargs={\"resources\": {\"GPU\": 1}},\n", - "# )\n", - "# for i in range(4)\n", - "# ]\n", - "# async with session:\n", - "# await session.submit_tasks(tasks)\n", - "# await session.wait_tasks(tasks)\n", - "# for t in tasks:\n", - "# print(f\"{t.uid} -> {t.return_value if t.state == 'DONE' else t.exception}\")\n", - "#\n", - "# await run_gpu_tasks()" - ] + "source": "# async def gpu_task(x):\n# # In real usage: import cupy, torch, etc.\n# return x ** 2\n#\n# async def run_gpu_tasks():\n# # Workers must advertise GPU resources. Options:\n# # a) LocalCUDACluster — auto-detects and sets GPU resources per worker\n# # from dask_cuda import LocalCUDACluster\n# # cluster = LocalCUDACluster()\n# #\n# # b) SLURMCluster — must pass --resources explicitly (and asynchronous=True,\n# # entered via `async with`, as in section 4 — otherwise the Client built\n# # on top of it inherits a mismatched event loop):\n# # cluster = SLURMCluster(\n# # gres=\"gpu:1\",\n# # worker_extra_args=[\"--resources\", \"GPU=1\"],\n# # asynchronous=True,\n# # )\n# #\n# # c) Pre-existing workers launched with:\n# # dask worker --resources \"GPU=1\"\n#\n# async with DaskExecutionBackend() as backend:\n# session = Session(backends=[backend])\n# tasks = [\n# ComputeTask(\n# function=gpu_task,\n# args=(i,),\n# task_backend_specific_kwargs={\"resources\": {\"GPU\": 1}},\n# )\n# for i in range(4)\n# ]\n# async with session:\n# await session.submit_tasks(tasks)\n# await session.wait_tasks(tasks)\n# for t in tasks:\n# print(f\"{t.uid} -> {t.return_value if t.state == 'DONE' else t.exception}\")\n#\n# await run_gpu_tasks()" }, { "cell_type": "markdown", From da42e94e44a935c397641f29470d1feaca44c4cf Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Tue, 8 Sep 2026 17:35:59 +0000 Subject: [PATCH 4/5] fix tests: make tests backward compatible --- CHANGELOG.md | 6 ++++-- src/rhapsody/backends/execution/dask_parallel.py | 8 +++++++- tests/unit/test_backend_execution_dask_parallel.py | 5 +++-- 3 files changed, 14 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index efcce6f..c793b2b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,8 +20,10 @@ an asynchronous client always returns a cached snapshot with an empty `workers` mapping (see that method's own docstring). Every resource-constrained task failed regardless of whether a matching worker actually existed. Fixed to use - `await client.scheduler.identity(n_workers=-1)`, Dask's own documented - alternative for a live per-worker view. + `await client.scheduler.identity()`, Dask's own documented alternative for a + live per-worker view (called with no `n_workers=` kwarg — that parameter is + absent on older `distributed` releases still within this project's + `dask[distributed]>=2023.0.0` support range and raises `TypeError` there). - **`DaskExecutionBackend` could silently share one Dask `Future` across two distinct tasks** — `client.submit()` defaults to `pure=True` with no explicit `key`, deriving the Dask key from `tokenize(func, kwargs, *args)`. Two RHAPSODY diff --git a/src/rhapsody/backends/execution/dask_parallel.py b/src/rhapsody/backends/execution/dask_parallel.py index 4a811bb..be016ba 100644 --- a/src/rhapsody/backends/execution/dask_parallel.py +++ b/src/rhapsody/backends/execution/dask_parallel.py @@ -507,13 +507,19 @@ async def _check_resources_satisfiable(self, resources: dict) -> bool: which recommends this exact alternative), so it can never see workers that joined afterward or their resource advertisements. + No `n_workers=` kwarg is passed: newer `distributed` defaults it to -1 (all + workers) on the scheduler side, but older `distributed` releases (still within + this project's `dask[distributed]>=2023.0.0` support range) don't accept that + parameter at all and raise `TypeError` if it's passed. Omitting it returns all + workers on every supported version. + Args: resources: Dict of resource requirements (e.g. {"GPU": 1}). Returns: True if a qualifying worker exists, False otherwise. """ - info = await self._client.scheduler.identity(n_workers=-1) + info = await self._client.scheduler.identity() workers = info.get("workers", {}) return any( all(w.get("resources", {}).get(k, 0) >= v for k, v in resources.items()) diff --git a/tests/unit/test_backend_execution_dask_parallel.py b/tests/unit/test_backend_execution_dask_parallel.py index 8fb55f9..f088c77 100644 --- a/tests/unit/test_backend_execution_dask_parallel.py +++ b/tests/unit/test_backend_execution_dask_parallel.py @@ -779,7 +779,8 @@ async def test_dask_args_kwargs_not_mutated_on_caller_task(): @pytest.mark.asyncio async def test_check_resources_satisfiable_uses_scheduler_identity(): """_check_resources_satisfiable must use scheduler.identity(), not the always- empty - scheduler_info() cache on an asynchronous client.""" + scheduler_info() cache on an asynchronous client. Called with no n_workers= kwarg + since older `distributed` releases don't accept that parameter at all.""" try: from unittest.mock import AsyncMock from unittest.mock import MagicMock @@ -798,7 +799,7 @@ async def test_check_resources_satisfiable_uses_scheduler_identity(): assert await backend._check_resources_satisfiable({"GPU": 1}) is True assert await backend._check_resources_satisfiable({"GPU": 4}) is False - backend._client.scheduler.identity.assert_called_with(n_workers=-1) + backend._client.scheduler.identity.assert_called_with() # --------------------------------------------------------------------------- From a67a65c2bf0e8bfdb250cedd6dc9fdc00d5a4b69 Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Tue, 8 Sep 2026 17:36:35 +0000 Subject: [PATCH 5/5] ruff --- tests/unit/test_backend_execution_dask_parallel.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_backend_execution_dask_parallel.py b/tests/unit/test_backend_execution_dask_parallel.py index f088c77..2ec9a42 100644 --- a/tests/unit/test_backend_execution_dask_parallel.py +++ b/tests/unit/test_backend_execution_dask_parallel.py @@ -779,8 +779,11 @@ async def test_dask_args_kwargs_not_mutated_on_caller_task(): @pytest.mark.asyncio async def test_check_resources_satisfiable_uses_scheduler_identity(): """_check_resources_satisfiable must use scheduler.identity(), not the always- empty - scheduler_info() cache on an asynchronous client. Called with no n_workers= kwarg - since older `distributed` releases don't accept that parameter at all.""" + scheduler_info() cache on an asynchronous client. + + Called with no n_workers= kwarg + since older `distributed` releases don't accept that parameter at all. + """ try: from unittest.mock import AsyncMock from unittest.mock import MagicMock