Feature/refactor daskbackend

CHANGELOG.md

@@ -1,5 +1,76 @@
 # 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 <function ...>: 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()`, 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
+  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

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": {...}}`.

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.

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)

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)

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())

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.