diff --git a/README.md b/README.md index e41a5f3..0680cf3 100644 --- a/README.md +++ b/README.md @@ -23,11 +23,11 @@ A unified runtime for executing **AI and HPC workloads** on supercomputing infra ```python import asyncio from rhapsody.api import Session, ComputeTask, AITask -from rhapsody.backends import DragonExecutionBackendV3, DragonVllmInferenceBackend +from rhapsody.backends import DragonExecutionBackend, DragonVllmInferenceBackend async def main(): # Initialize backends - hpc_backend = await DragonExecutionBackendV3(name="hpc") + hpc_backend = await DragonExecutionBackend(name="hpc") ai_backend = await DragonVllmInferenceBackend(name="vllm", model="Qwen2.5-7B") # Create session with multiple backends diff --git a/docs/getting-started/advanced-usage.md b/docs/getting-started/advanced-usage.md index ce503b8..9dbe1c4 100644 --- a/docs/getting-started/advanced-usage.md +++ b/docs/getting-started/advanced-usage.md @@ -216,8 +216,11 @@ named after the task UID (`task.000001.stdout`, `task.000001.stderr`). !!! note "Supported backends" `capture_stdio` is supported by `ConcurrentExecutionBackend`, - `DaskExecutionBackend`, and `DragonExecutionBackendV3`. Function tasks - (`function=`) are unaffected regardless of the flag value. + `DaskExecutionBackend`, and `DragonExecutionBackend`. On + `ConcurrentExecutionBackend` and `DaskExecutionBackend`, function tasks + (`function=`) are unaffected by this flag. On `DragonExecutionBackend`, + `capture_stdio` applies to **all** task types including function tasks — see + [Dragon: Task Output and File Paths](#dragon-task-output-and-file-paths). !!! note "Session work_dir" Files are placed in `{session.work_dir}/{session.uid}/`. Pass `work_dir=` @@ -240,7 +243,7 @@ import asyncio import rhapsody from rhapsody.api import Session from rhapsody.api import ComputeTask -from rhapsody.backends import ConcurrentExecutionBackend, DragonExecutionBackendV3 +from rhapsody.backends import ConcurrentExecutionBackend, DragonExecutionBackend async def compute_function(): @@ -250,7 +253,7 @@ async def compute_function(): async def multi_backend_demo(): # Initialize two different backends be_local = ConcurrentExecutionBackend(name="local_cpu") - be_remote = DragonExecutionBackendV3(name="hpc_gpu") + be_remote = DragonExecutionBackend(name="hpc_gpu") await asyncio.gather(be_local, be_remote) @@ -302,12 +305,12 @@ For large-scale HPC deployments, the Dragon backend provides native integration import asyncio import rhapsody from rhapsody.api import Session, ComputeTask -from rhapsody.backends import DragonExecutionBackendV3 +from rhapsody.backends import DragonExecutionBackend async def dragon_demo(): # Initialize the Dragon backend (optimized for HPC) # V3 uses native Dragon Batch for high-performance wait/callbacks - dragon_be = await DragonExecutionBackendV3( + dragon_be = await DragonExecutionBackend( name="dragon_hpc" ) @@ -336,7 +339,7 @@ if __name__ == "__main__": ## Dragon Process Templates -When using `DragonExecutionBackendV3`, you can pass Dragon-native process configuration to `ComputeTask` via the `task_backend_specific_kwargs` parameter. This exposes the `process_template` and `process_templates` options, which map directly to Dragon's [ProcessTemplate](https://dragonhpc.github.io/dragon/doc/_build/html/ref/native/dragon.native.process.html#dragon.native.process.ProcessTemplate). +When using `DragonExecutionBackend`, you can pass Dragon-native process configuration to `ComputeTask` via the `task_backend_specific_kwargs` parameter. This exposes the `process_template` and `process_templates` options, which map directly to Dragon's [ProcessTemplate](https://dragonhpc.github.io/dragon/doc/_build/html/ref/native/dragon.native.process.html#dragon.native.process.ProcessTemplate). The `ProcessTemplate` accepts the following parameters: @@ -351,7 +354,7 @@ The `ProcessTemplate` accepts the following parameters: | `options` | `ProcessOptions`, optional | Process options, such as allowing the process to connect to the infrastructure | !!! warning "Excluded Parameters" - `DragonExecutionBackendV3` manages `target` (the binary or Python callable), `args`, and `kwargs` internally through the `ComputeTask` interface. These three parameters are **not** available in `process_template` / `process_templates` — use `ComputeTask.executable`, `ComputeTask.arguments`, and `ComputeTask.function` instead. + `DragonExecutionBackend` manages `target` (the binary or Python callable), `args`, and `kwargs` internally through the `ComputeTask` interface. These three parameters are **not** available in `process_template` / `process_templates` — use `ComputeTask.executable`, `ComputeTask.arguments`, and `ComputeTask.function` instead. ### Single-Process Template @@ -409,7 +412,7 @@ from dragon.infrastructure.policy import Policy import rhapsody from rhapsody.api import ComputeTask, Session -from rhapsody.backends import DragonExecutionBackendV3 +from rhapsody.backends import DragonExecutionBackend from dragon.native.machine import System, Node @@ -446,7 +449,7 @@ def make_policies(all_gpus, nprocs=32): async def main(): - backend = await DragonExecutionBackendV3() + backend = await DragonExecutionBackend() session = Session(backends=[backend]) all_gpus = find_gpus() # e.g. [("node-0", 0), ("node-0", 1), ("node-1", 0), ...] @@ -497,11 +500,11 @@ The following example demonstrates mixing native functions, single-process tasks import asyncio import rhapsody from rhapsody.api import ComputeTask, Session -from rhapsody.backends import DragonExecutionBackendV3 +from rhapsody.backends import DragonExecutionBackend async def main(): - backend = await DragonExecutionBackendV3() + backend = await DragonExecutionBackend() session = Session(backends=[backend]) async def single_function(): @@ -573,6 +576,25 @@ if __name__ == "__main__": dragon my_heterogeneous_workload.py ``` +## Dragon: Task Output and File Paths + +After a task completes, `task.stdout` and `task.stderr` work as follows: + +- **`capture_stdio=False` (default)** — content is an inline string (whatever the task printed). Read it directly: `task.stdout`. +- **`capture_stdio=True`** — content is a file path under the Rhapsody session directory (`rhapsody.session.{uid}/task.000001.stdout`). Open it to read: `open(task.stdout).read()`. + +!!! note "Performance at scale" + At 100K+ tasks on shared filesystems, the default file-read overhead (~100–400 s on Lustre) can be avoided by passing `batch_kwargs={"task_logs": False}` to `DragonExecutionBackend`. With this option, `task.stdout` and `task.stderr` will be **empty** for all `capture_stdio=False` tasks — output goes to the console only. A warning is logged when this option is active. + +!!! warning "Functions launched via process_template or process_templates" + If you use `function=` together with `process_template` or `process_templates`, + Dragon runs the callable as a subprocess — **not** in Dragon's native function pool. + In this mode `task.return_value` is the **exit code** (`0` on success), not the + Python return value of your function. Use a bare `ComputeTask(function=my_fn)` (no + templates) when you need the actual return value. + +--- + ## Mixed HPC and AI Workloads One of RHAPSODY's most powerful features is the ability to orchestrate traditional binaries alongside AI inference services (like vLLM). diff --git a/docs/getting-started/configuration.md b/docs/getting-started/configuration.md index e7f5cf8..92f7b7f 100644 --- a/docs/getting-started/configuration.md +++ b/docs/getting-started/configuration.md @@ -62,9 +62,9 @@ ComputeTask( High-performance execution using the Dragon runtime. ```python -from rhapsody.backends import DragonExecutionBackendV3 +from rhapsody.backends import DragonExecutionBackend -backend = DragonExecutionBackendV3( +backend = DragonExecutionBackend( name="dragon", num_nodes=4, # Total nodes for the worker pool (optional) pool_nodes=2, # Nodes per worker pool (optional) @@ -80,12 +80,12 @@ backend = DragonExecutionBackendV3( | `disable_telemetry` | `bool` | `False` | Disable Dragon internal telemetry | !!! note "Streaming pipeline" - `DragonExecutionBackendV3` uses Dragon's streaming batch pipeline — tasks submitted via + `DragonExecutionBackend` uses Dragon's streaming batch pipeline — tasks submitted via `session.submit_tasks()` are dispatched individually by a continuously running background thread. There is no compile or start step; tasks begin executing as soon as they are submitted. !!! note "Dragon Versions" - While `DragonExecutionBackendV3` is recommended for most users and will be always maintained, V1 and V2 are also available for legacy compatibility. + `DragonExecutionBackend` is the only supported Dragon backend as of RHAPSODY 0.5.0. ### RADICAL-Pilot Backend Large-scale HPC execution using RADICAL-Pilot. diff --git a/docs/getting-started/resource-specification.md b/docs/getting-started/resource-specification.md index 36e4665..cf9e37b 100644 --- a/docs/getting-started/resource-specification.md +++ b/docs/getting-started/resource-specification.md @@ -13,7 +13,7 @@ This design keeps the task API thin and makes resource semantics explicit per ba |---------|-------------|-----|-----|-------| | `ConcurrentExecutionBackend` | not supported | `task_backend_specific_kwargs={"cwd": "..."}` | `task_backend_specific_kwargs={"env": {...}}` | `task_backend_specific_kwargs={"shell": True}` | | `DaskExecutionBackend` | `task_backend_specific_kwargs={"resources": {"GPU": 1}}` | `task_backend_specific_kwargs={"cwd": "..."}` | `task_backend_specific_kwargs={"env": {...}}` | `task_backend_specific_kwargs={"shell": True}` | -| `DragonExecutionBackendV3` | `task_backend_specific_kwargs={"process_template": {"policy": ...}}` | `task_backend_specific_kwargs={"process_template": {"cwd": "..."}}` | `task_backend_specific_kwargs={"process_template": {"env": {...}}}` | not applicable | +| `DragonExecutionBackend` | `task_backend_specific_kwargs={"process_template": {"policy": ...}}` | `task_backend_specific_kwargs={"process_template": {"cwd": "..."}}` | `task_backend_specific_kwargs={"process_template": {"env": {...}}}` | not applicable | --- @@ -140,7 +140,7 @@ backend = await DaskExecutionBackend(cluster=cluster) ## Dragon Backend (V3) -`DragonExecutionBackendV3` submits tasks to the Dragon runtime. All per-task resource +`DragonExecutionBackend` submits tasks to the Dragon runtime. All per-task resource and placement settings flow through Dragon's [ProcessTemplate](https://dragonhpc.github.io/dragon/doc/_build/html/ref/native/dragon.native.process.html#dragon.native.process.ProcessTemplate). There is **no backend-level `working_directory`** in V3. Every process @@ -152,7 +152,7 @@ Pass a `batch_kwargs` dict when constructing the backend to tune the underlying `dragon.workflows.batch.Batch` instance: ```python -backend = await DragonExecutionBackendV3(batch_kwargs={ +backend = await DragonExecutionBackend(batch_kwargs={ "num_nodes": 8, "results_ddict_mem": 4 * 1024**3, # 4 GiB results DDict "scheduler_workers": 16, diff --git a/docs/hpc-machines/anvil.md b/docs/hpc-machines/anvil.md index c78e86c..d9df557 100644 --- a/docs/hpc-machines/anvil.md +++ b/docs/hpc-machines/anvil.md @@ -63,7 +63,7 @@ dragon ## Example -This example runs 32 MPI jobs concurrently using RHAPSODY's `Session` API with `DragonExecutionBackendV3`. Each job gets a random number of ranks (between 2 and 32), all scheduled across the 2 allocated nodes. Worker count is determined automatically by Dragon from the allocation. +This example runs 32 MPI jobs concurrently using RHAPSODY's `Session` API with `DragonExecutionBackend`. Each job gets a random number of ranks (between 2 and 32), all scheduled across the 2 allocated nodes. Worker count is determined automatically by Dragon from the allocation. ```python title="rhapsody-mpi.py" import asyncio @@ -73,7 +73,7 @@ import time import rhapsody from rhapsody.api import ComputeTask, Session -from rhapsody.backends import DragonExecutionBackendV3 +from rhapsody.backends import DragonExecutionBackend rhapsody.enable_logging(level=logging.DEBUG) @@ -100,7 +100,7 @@ async def main(): sleepsecs = 2 maxranks = 32 - backend = await DragonExecutionBackendV3() + backend = await DragonExecutionBackend() session = Session(backends=[backend]) print("--- Submitting Tasks ---") diff --git a/docs/hpc-machines/bridges2.md b/docs/hpc-machines/bridges2.md index f38c814..c421bb2 100644 --- a/docs/hpc-machines/bridges2.md +++ b/docs/hpc-machines/bridges2.md @@ -84,7 +84,7 @@ dragon -w ssh --network-config slurm.yaml -t tcp 01-workload-async-gather.py ## Example This example submits a heterogeneous workload mixing function tasks and -executable tasks using RHAPSODY's `Session` API with `DragonExecutionBackendV3`. +executable tasks using RHAPSODY's `Session` API with `DragonExecutionBackend`. ```python title="01-workload-async-gather.py" import asyncio @@ -92,7 +92,7 @@ import logging import rhapsody from rhapsody.api import ComputeTask, Session -from rhapsody.backends import DragonExecutionBackendV3 +from rhapsody.backends import DragonExecutionBackend rhapsody.enable_logging(level=logging.DEBUG) @@ -102,7 +102,7 @@ def compute(x: int) -> int: async def main(): - backend = await DragonExecutionBackendV3() + backend = await DragonExecutionBackend() session = Session(backends=[backend]) tasks = [ComputeTask(function=compute, args=(i,)) for i in range(16)] diff --git a/docs/hpc-machines/index.md b/docs/hpc-machines/index.md index 12031c5..7701308 100644 --- a/docs/hpc-machines/index.md +++ b/docs/hpc-machines/index.md @@ -6,7 +6,7 @@ This section provides machine-specific setup guides and verified examples for ru **`ConcurrentExecutionBackend (designed for debug and tests single node only)`** works out of the box on every machine listed here — no machine-specific configuration is required. Install RHAPSODY and run. -**`DragonExecutionBackendV3 (designed for large scale multi-node)`** requires machine-specific setup: each HPC system has its own software stack, MPI library, interconnect, and job launcher. The exact steps — module loading, virtual environment creation, transport configuration, and launch command — differ per machine and are documented in each section below. +**`DragonExecutionBackend (designed for large scale multi-node)`** requires machine-specific setup: each HPC system has its own software stack, MPI library, interconnect, and job launcher. The exact steps — module loading, virtual environment creation, transport configuration, and launch command — differ per machine and are documented in each section below. Each guide covers environment setup, module loading, job allocation, and working examples. diff --git a/docs/hpc-machines/ncsa-delta.md b/docs/hpc-machines/ncsa-delta.md index 4461c31..39e03c4 100644 --- a/docs/hpc-machines/ncsa-delta.md +++ b/docs/hpc-machines/ncsa-delta.md @@ -60,7 +60,7 @@ dragon ## Example -This example runs 32 MPI jobs concurrently using RHAPSODY's `Session` API with the `DragonExecutionBackendV3`. Each job gets a random number of ranks (between 2 and 32), all scheduled across the 2 allocated nodes. Worker count is determined automatically by Dragon from the allocation. +This example runs 32 MPI jobs concurrently using RHAPSODY's `Session` API with the `DragonExecutionBackend`. Each job gets a random number of ranks (between 2 and 32), all scheduled across the 2 allocated nodes. Worker count is determined automatically by Dragon from the allocation. ```python title="mpi-rhapsody.py" import asyncio @@ -70,7 +70,7 @@ import time import rhapsody from rhapsody.api import ComputeTask, Session -from rhapsody.backends import DragonExecutionBackendV3 +from rhapsody.backends import DragonExecutionBackend rhapsody.enable_logging(level=logging.DEBUG) @@ -97,7 +97,7 @@ async def main(): sleepsecs = 2 maxranks = 32 - backend = await DragonExecutionBackendV3() + backend = await DragonExecutionBackend() session = Session(backends=[backend]) print("--- Submitting Tasks ---") @@ -140,7 +140,7 @@ dragon mpi-rhapsody.py 2026-03-25 23:07:36,938 | DEBUG | [api_setup] | got handshake 2026-03-25 23:07:36,938 | INFO | [api_setup] | debug entry hooked 2026-03-25 23:07:36,942 | DEBUG | [dragon.native.queue] | Created queue {self!r} - 2026-03-25 23:07:37,648 | INFO | [rhapsody.backends.execution.dragon] | DragonExecutionBackendV3: 128 workers, 2 managers + 2026-03-25 23:07:37,648 | INFO | [rhapsody.backends.execution.dragon] | DragonExecutionBackend: 128 workers, 2 managers 2026-03-25 23:07:37,648 | DEBUG | [rhapsody.backends.execution.dragon] | Starting Dragon backend V3 async initialization... 2026-03-25 23:07:37,648 | DEBUG | [rhapsody.backends.execution.dragon] | Registering backend states... 2026-03-25 23:07:37,649 | DEBUG | [rhapsody.backends.execution.dragon] | Registering task states... diff --git a/docs/hpc-machines/perlmutter.md b/docs/hpc-machines/perlmutter.md index b3a0d6f..ee3519a 100644 --- a/docs/hpc-machines/perlmutter.md +++ b/docs/hpc-machines/perlmutter.md @@ -79,7 +79,7 @@ dragon This example demonstrates a two-phase HPC-AI workflow on Perlmutter: 1. **Simulation phase** — run a batch of scientific simulations using LLM inference (`DragonVllmInferenceBackend`). Each `AITask` sends a domain-specific prompt and collects the generated output as synthetic simulation data. -2. **Training phase** — fine-tune a small model on the collected simulation outputs using a `ComputeTask` dispatched through `DragonExecutionBackendV3`. +2. **Training phase** — fine-tune a small model on the collected simulation outputs using a `ComputeTask` dispatched through `DragonExecutionBackend`. Both phases run within the same RHAPSODY `Session`, sharing the same Dragon runtime across GPU and CPU resources. @@ -90,7 +90,7 @@ import multiprocessing as mp import rhapsody from rhapsody.api import AITask, ComputeTask, Session -from rhapsody.backends import DragonExecutionBackendV3, DragonVllmInferenceBackend +from rhapsody.backends import DragonExecutionBackend, DragonVllmInferenceBackend rhapsody.enable_logging(level=logging.INFO) @@ -160,7 +160,7 @@ def fine_tune(simulation_outputs: list): async def main(): mp.set_start_method("dragon") - execution_backend = await DragonExecutionBackendV3() + execution_backend = await DragonExecutionBackend() inference_backend = await DragonVllmInferenceBackend( config_file="config.yaml", diff --git a/docs/index.md b/docs/index.md index 67adfbe..6861816 100644 --- a/docs/index.md +++ b/docs/index.md @@ -19,11 +19,11 @@ RHAPSODY is a high-performance runtime system designed for orchestrating complex ```python import asyncio from rhapsody import Session, ComputeTask -from rhapsody.backends import DragonExecutionBackendV3 +from rhapsody.backends import DragonExecutionBackend async def main(): # 1. Initialize session with a backend - backend = await DragonExecutionBackendV3() + backend = await DragonExecutionBackend() async with Session(backends=[backend]) as session: # 2. Define a task diff --git a/docs/integrations.md b/docs/integrations.md index 5c575ae..f007fa6 100644 --- a/docs/integrations.md +++ b/docs/integrations.md @@ -43,14 +43,14 @@ import multiprocessing as mp from rhapsody import Session from rhapsody.api import AITask, ComputeTask -from rhapsody.backends import DragonExecutionBackendV3, DragonVllmInferenceBackend +from rhapsody.backends import DragonExecutionBackend, DragonVllmInferenceBackend async def main(): # Set Dragon as the multiprocessing start method mp.set_start_method("dragon") # Initialize execution backend for compute tasks - execution_backend = await DragonExecutionBackendV3() + execution_backend = await DragonExecutionBackend() # Initialize inference backend for AI tasks inference_backend = DragonVllmInferenceBackend( @@ -231,14 +231,14 @@ import asyncio import multiprocessing as mp from radical.asyncflow import WorkflowEngine -from rhapsody.backends import DragonExecutionBackendV3 +from rhapsody.backends import DragonExecutionBackend async def main(): # Set Dragon as the multiprocessing start method mp.set_start_method("dragon") # Initialize RHAPSODY backend - backend = await DragonExecutionBackendV3() + backend = await DragonExecutionBackend() # Create AsyncFlow workflow engine with RHAPSODY backend flow = await WorkflowEngine.create(backend=backend) @@ -347,8 +347,8 @@ from rhapsody.backends import DaskExecutionBackend backend = await DaskExecutionBackend() # Dragon HPC -from rhapsody.backends import DragonExecutionBackendV3 -backend = await DragonExecutionBackendV3() +from rhapsody.backends import DragonExecutionBackend +backend = await DragonExecutionBackend() # Create workflow with chosen backend flow = await WorkflowEngine.create(backend=backend) diff --git a/docs/telemetry/reference.md b/docs/telemetry/reference.md index 51b90d5..9e32613 100644 --- a/docs/telemetry/reference.md +++ b/docs/telemetry/reference.md @@ -24,7 +24,7 @@ All telemetry originates as **frozen Python dataclasses** (`BaseEvent` subclasse | `event_time` | `float` | Wall-clock time the event occurred (`time.time()`) | | `emit_time` | `float` | Wall-clock time of queue insertion (`time.time()`) | | `session_id` | `str` | Owning session identifier | -| `backend` | `str` | Backend name (e.g. `"dragon_v3"`, `"dask"`, `"concurrent"`) | +| `backend` | `str` | Backend name (e.g. `"dragon"`, `"dask"`, `"concurrent"`) | | `task_id` | `str \| None` | Task UID — `None` for session-level events | | `node_id` | `str \| None` | Hostname or worker address — `None` when not applicable | | `attributes` | `dict` | Optional metadata (see per-event notes below) | @@ -418,7 +418,7 @@ task span ├── parent: session span ├── attributes: │ session_id = "session.0001" -│ backend = "dragon_v3" +│ backend = "dragon" │ task_id = "task.0042" │ executable = "/usr/bin/python" │ task_type = "compute" diff --git a/examples/01-workload-async-gather.py b/examples/01-workload-async-gather.py index 9c8f420..95d572b 100644 --- a/examples/01-workload-async-gather.py +++ b/examples/01-workload-async-gather.py @@ -6,7 +6,7 @@ import rhapsody from rhapsody.api import ComputeTask from rhapsody.api import Session -from rhapsody.backends import DragonExecutionBackendV3 +from rhapsody.backends import DragonExecutionBackend rhapsody.enable_logging(level=logging.DEBUG) @@ -17,7 +17,7 @@ def func_task(): async def main(): # Initialize session with Concurrent backend - backend = await DragonExecutionBackendV3() + backend = await DragonExecutionBackend() session = Session(backends=[backend]) print("--- Submitting Tasks ---") diff --git a/examples/02-workload-heterogeneous.py b/examples/02-workload-heterogeneous.py index 29aeda1..e7a7d5a 100644 --- a/examples/02-workload-heterogeneous.py +++ b/examples/02-workload-heterogeneous.py @@ -4,14 +4,14 @@ import rhapsody from rhapsody.api import ComputeTask from rhapsody.api import Session -from rhapsody.backends import DragonExecutionBackendV3 +from rhapsody.backends import DragonExecutionBackend rhapsody.enable_logging(level=logging.DEBUG) async def main(): # Initialize backend - backend = await DragonExecutionBackendV3() + backend = await DragonExecutionBackend() session = Session(backends=[backend]) async def single_function(): diff --git a/examples/03-workload-ai-hpc.py b/examples/03-workload-ai-hpc.py index ee64862..4c50fe7 100644 --- a/examples/03-workload-ai-hpc.py +++ b/examples/03-workload-ai-hpc.py @@ -6,7 +6,7 @@ from rhapsody.api import AITask from rhapsody.api import ComputeTask from rhapsody.api import Session -from rhapsody.backends import DragonExecutionBackendV3 +from rhapsody.backends import DragonExecutionBackend from rhapsody.backends import DragonVllmInferenceBackend rhapsody.enable_logging(level=logging.DEBUG) @@ -17,7 +17,7 @@ async def main(): mp.set_start_method("dragon") - execution_backend = await DragonExecutionBackendV3() + execution_backend = await DragonExecutionBackend() inference_backend = DragonVllmInferenceBackend( config_file="config.yaml", diff --git a/examples/04-integration-asyncflow.py b/examples/04-integration-asyncflow.py index edb05f1..daf0412 100644 --- a/examples/04-integration-asyncflow.py +++ b/examples/04-integration-asyncflow.py @@ -4,13 +4,13 @@ from radical.asyncflow import WorkflowEngine -from rhapsody.backends import DragonExecutionBackendV3 +from rhapsody.backends import DragonExecutionBackend async def main(): mp.set_start_method("dragon") - backend = await DragonExecutionBackendV3() + backend = await DragonExecutionBackend() flow = await WorkflowEngine.create(backend=backend) @flow.function_task diff --git a/examples/README.md b/examples/README.md index 8a6293b..c30dcbc 100644 --- a/examples/README.md +++ b/examples/README.md @@ -9,7 +9,7 @@ Different examples require different launchers depending on the backend they use | Backend | Launcher | Command | |---------|----------|---------| | `ConcurrentExecutionBackend` | Standard Python | `python example.py` | -| `DragonExecutionBackendV3` | Dragon runtime | `dragon example.py` | +| `DragonExecutionBackend` | Dragon runtime | `dragon example.py` | | `DragonVllmInferenceBackend` | Dragon runtime + GPU | `dragon example.py` | > **Why `dragon` instead of `python`?** @@ -37,7 +37,7 @@ The simplest possible RHAPSODY example. Uses `ConcurrentExecutionBackend` (Pytho dragon 01-workload-async-gather.py ``` -Submits **1024 function tasks** to `DragonExecutionBackendV3` and awaits all of them with `asyncio.gather()`. Demonstrates RHAPSODY's batch submission throughput and the `async with session` context manager pattern. +Submits **1024 function tasks** to `DragonExecutionBackend` and awaits all of them with `asyncio.gather()`. Demonstrates RHAPSODY's batch submission throughput and the `async with session` context manager pattern. **What you'll learn:** Large-scale batch submission, `asyncio.gather(*futures)` for collecting results, Dragon backend initialization. @@ -69,7 +69,7 @@ Runs **five different execution modes** in a single session: dragon 03-workload-ai-hpc.py ``` -Combines `DragonExecutionBackendV3` (HPC compute) with `DragonVllmInferenceBackend` (LLM inference) in a single session. Submits a mix of `AITask` and `ComputeTask` objects — RHAPSODY routes each to the correct backend automatically. +Combines `DragonExecutionBackend` (HPC compute) with `DragonVllmInferenceBackend` (LLM inference) in a single session. Submits a mix of `AITask` and `ComputeTask` objects — RHAPSODY routes each to the correct backend automatically. **Requires:** GPU access, vLLM installed, and a `config.yaml` file. You must create this config file for your environment based on the sample at [vllm-dragonhpc/config.sample](https://github.com/radical-cybertools/vllm-dragonhpc/blob/main/config.sample). diff --git a/pyproject.toml b/pyproject.toml index d9d6acf..479cdea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,7 +48,9 @@ telemetry = ["opentelemetry-sdk>=1.20.0", "nvidia-ml-py"] # Backend-specific dependencies dask = ["dask[distributed]>=2023.0.0"] radical_pilot = ["radical.pilot>=1.30.0"] -dragon = ["dragonhpc==0.14.0"] # Requires Python >= 3.10, supports up to 3.13 + +dragon = ["dragonhpc"] # Requires Python >= 3.10, supports up to 3.13 + vllm-dragon = [ "aiohttp>=3.8.0", "pyyaml>=6.0", diff --git a/src/rhapsody/backends/__init__.py b/src/rhapsody/backends/__init__.py index b544c18..bc37895 100644 --- a/src/rhapsody/backends/__init__.py +++ b/src/rhapsody/backends/__init__.py @@ -45,17 +45,9 @@ pass try: - from .execution import DragonExecutionBackendV1 # noqa: F401 - from .execution import DragonExecutionBackendV2 # noqa: F401 - from .execution import DragonExecutionBackendV3 # noqa: F401 + from .execution import DragonExecutionBackend # noqa: F401 - __all__.extend( - [ - "DragonExecutionBackendV1", - "DragonExecutionBackendV2", - "DragonExecutionBackendV3", - ] - ) + __all__.append("DragonExecutionBackend") except ImportError: pass diff --git a/src/rhapsody/backends/constants.py b/src/rhapsody/backends/constants.py index c998f62..8ae7d5d 100644 --- a/src/rhapsody/backends/constants.py +++ b/src/rhapsody/backends/constants.py @@ -383,7 +383,7 @@ def to_main_state(self, backend_state: Any) -> TasksMainStates: Example: :: - mapper = StateMapper('dragon_v3') + mapper = StateMapper('dragon') main_state = mapper.to_main_state('COMPLETED') # TasksMainStates.DONE """ try: diff --git a/src/rhapsody/backends/discovery.py b/src/rhapsody/backends/discovery.py index 35f05ab..2d0853f 100644 --- a/src/rhapsody/backends/discovery.py +++ b/src/rhapsody/backends/discovery.py @@ -61,7 +61,7 @@ def _discover_backends(cls) -> None: # ConcurrentExecutionBackend -> concurrent # DaskExecutionBackend -> dask # RadicalExecutionBackend -> radical_pilot - # DragonExecutionBackendV1 -> dragon_v1 + # DragonExecutionBackend -> dragon backend_name = cls._derive_backend_name(class_name) cls._backends[backend_name] = backend_class @@ -69,6 +69,10 @@ def _discover_backends(cls) -> None: # Backend class not available despite being in __all__ pass + # Deprecated name alias — dragon_v3 resolves to dragon for backward compatibility + if "dragon" in cls._backends: + cls._backends.setdefault("dragon_v3", cls._backends["dragon"]) + cls._initialized = True @classmethod @@ -80,7 +84,7 @@ def _derive_backend_name(cls, class_name: str) -> str: - ConcurrentExecutionBackend -> concurrent - DaskExecutionBackend -> dask - RadicalExecutionBackend -> radical_pilot - - DragonExecutionBackendV1 -> dragon_v1 + - DragonExecutionBackend -> dragon Args: class_name: The class name (e.g., "DaskExecutionBackend") diff --git a/src/rhapsody/backends/execution/__init__.py b/src/rhapsody/backends/execution/__init__.py index 50ec4db..b3521d4 100644 --- a/src/rhapsody/backends/execution/__init__.py +++ b/src/rhapsody/backends/execution/__init__.py @@ -26,16 +26,8 @@ pass try: - from .dragon import DragonExecutionBackendV1 # noqa: F401 - from .dragon import DragonExecutionBackendV2 # noqa: F401 - from .dragon import DragonExecutionBackendV3 # noqa: F401 - - __all__.extend( - [ - "DragonExecutionBackendV1", - "DragonExecutionBackendV2", - "DragonExecutionBackendV3", - ] - ) + from .dragon import DragonExecutionBackend # noqa: F401 + + __all__.append("DragonExecutionBackend") except ImportError: pass diff --git a/src/rhapsody/backends/execution/dask_parallel.py b/src/rhapsody/backends/execution/dask_parallel.py index c45fbb2..87552d6 100644 --- a/src/rhapsody/backends/execution/dask_parallel.py +++ b/src/rhapsody/backends/execution/dask_parallel.py @@ -219,7 +219,8 @@ async def cancel_task(self, uid: str) -> bool: task = self.tasks[uid] future = task.get("future") if future: - return await future.cancel() + await future.cancel() + return True return False async def submit_tasks(self, tasks: list[dict[str, Any]]) -> None: diff --git a/src/rhapsody/backends/execution/dragon.py b/src/rhapsody/backends/execution/dragon.py index 9098034..d2c5b29 100644 --- a/src/rhapsody/backends/execution/dragon.py +++ b/src/rhapsody/backends/execution/dragon.py @@ -1,3085 +1,76 @@ import asyncio import logging import os -import shlex -import sys import threading -import time -import uuid -from dataclasses import dataclass -from dataclasses import field -from enum import Enum from typing import Any from typing import Callable from typing import Optional -import typeguard - -from ..base import BaseBackend -from ..constants import BackendMainStates -from ..constants import StateMapper - -try: - import multiprocessing as mp - - import dragon - from dragon.data.ddict.ddict import DDict - from dragon.infrastructure.policy import Policy - - # Node Telemetry only - from dragon.native.event import Event - from dragon.native.machine import System - from dragon.native.process import Popen - from dragon.native.process import Process - from dragon.native.process import ProcessTemplate - from dragon.native.process_group import DragonUserCodeError - from dragon.native.process_group import ProcessGroup - from dragon.native.queue import Queue - from dragon.workflows.batch import Batch - -except ImportError: # pragma: no cover - environment without Dragon - dragon = None - Process = None - ProcessTemplate = None - ProcessGroup = None - Popen = None - Queue = None - DDict = None - System = None - Policy = None - Batch = None - Event = None - - -def _get_logger() -> logging.Logger: - """Get logger for dragon backend module. - - This function provides lazy logger evaluation, ensuring the logger is created after the user has - configured logging, not at module import time. - """ - return logging.getLogger(__name__) - - -DRAGON_DEFAULT_REF_THRESHOLD = int(os.environ.get("DRAGON_DEFAULT_REF_THRESHOLD", 1024 * 1024)) - -# ============================================================================ -# V1 Integration Helper classes -# ============================================================================ - - -class TaskTypeV1(Enum): - """Enumeration of supported task types.""" - - SINGLE_FUNCTION = "single_function" - SINGLE_EXECUTABLE = "single_executable" - MULTI_FUNCTION = "multi_function" - MULTI_EXECUTABLE = "multi_executable" - MPI_FUNCTION = "mpi_function" - MPI_EXECUTABLE = "mpi_executable" - - -@dataclass -class TaskInfoV1: - """Container for task runtime information.""" - - task_type: TaskTypeV1 - ranks: int - start_time: float - canceled: bool = False - process: Optional[Process] = None - group: Optional[ProcessGroup] = None - - -@dataclass -class ExecutableTaskCompletionV1: - """Task completion data from executable process.""" - - task_uid: str - rank: int - process_id: int - stdout: str - stderr: str - exit_code: int - timestamp: float - - def to_result_dict(self) -> dict: - """Convert to task result dictionary.""" - return { - "stdout": self.stdout, - "stderr": self.stderr, - "exit_code": self.exit_code, - "return_value": None, - "exception": None - if self.exit_code == 0 - else f"Process exited with code {self.exit_code}", - } - - -@dataclass -class FunctionTaskCompletionV1: - """Task completion data from function process - sent via queue.""" - - task_uid: str - rank: int - process_id: int - stdout: str - stderr: str - exit_code: int - timestamp: float - success: bool - exception: Optional[str] - traceback: Optional[str] - return_value: Any # Actual value or DataReference - stored_in_ddict: bool # True if return_value is in DDict - - def to_result_dict(self) -> dict: - """Convert to task result dictionary.""" - return { - "stdout": self.stdout, - "stderr": self.stderr, - "exit_code": self.exit_code, - "return_value": self.return_value, - "exception": self.exception, - "success": self.success, - } - - -class DataReferenceV1: - """Zero-copy reference to function results stored in DDict. - - Points directly to result keys written by function wrapper. User controls when to resolve and - fetch data from DDict. - """ - - def __init__(self, task_uid: str, ranks: int, ddict: DDict, backend_id: str): - self._task_uid = task_uid - self._ranks = ranks - self._ddict = ddict - self._backend_id = backend_id - - @property - def task_uid(self) -> str: - return self._task_uid - - @property - def ranks(self) -> int: - return self._ranks - - @property - def backend_id(self) -> str: - return self._backend_id - - def resolve(self) -> Any: - """Resolve reference to actual return values from DDict. - - Fetches data directly from function-wrapper-written keys in DDict: - - Single rank (ranks=1): Returns single return_value - - Multi-rank (ranks>1): Returns list of return_values indexed by rank - """ - if self._ranks == 1: - # Single rank - return single value - result_key = f"return_{self._task_uid}_rank_0" - if result_key not in self._ddict: - raise KeyError(f"Result data not found for task: {self._task_uid}") - - return self._ddict[result_key] - else: - # Multi-rank - return list of values - return_values = [] - for rank in range(self._ranks): - result_key = f"return_{self._task_uid}_rank_{rank}" - if result_key not in self._ddict: - raise KeyError(f"Result data not found for task {self._task_uid} rank {rank}") - - return_values.append(self._ddict[result_key]) - - return return_values - - def __repr__(self) -> str: - return f"DataReferenceV1(task_uid='{self._task_uid}', ranks={self._ranks}, backend_id='{self._backend_id}')" - - -class SharedMemoryManagerV1: - """Manages optional DDict storage for large function results.""" - - def __init__(self, ddict: DDict, system: System, logger: logging.Logger): - self.ddict = ddict - self.system = system - self.logger = logger - self.backend_id = f"dragon_{uuid.uuid4().hex[:8]}" - - async def initialize(self): - """Initialize the storage manager.""" - self.logger.debug("SharedMemoryManagerV1 initialized with optional DDict storage") - - def create_data_reference(self, task_uid: str, ranks: int) -> DataReferenceV1: - """Create a zero-copy reference to existing result keys in DDict. - - Creates a reference object that points to - keys: return_{task_uid}_rank_{i} - """ - return DataReferenceV1(task_uid, ranks, self.ddict, self.backend_id) - - def cleanup_reference(self, ref: DataReferenceV1): - """Clean up reference data from DDict.""" - try: - for rank in range(ref.ranks): - result_key = f"return_{ref.task_uid}_rank_{rank}" - if result_key in self.ddict: - del self.ddict[result_key] - except Exception as e: - self.logger.warning(f"Error cleaning up reference {ref.task_uid}: {e}") - - # DDict operations for direct task data sharing - def store_task_data(self, key: str, data: Any) -> None: - """Store data in shared DDict.""" - self.ddict.pput(key, data) - - def get_task_data(self, key: str, default=None) -> Any: - """Retrieve data from shared DDict.""" - try: - if key in self.ddict: - return self.ddict[key] - return default - except Exception: - return default - - def list_task_data_keys(self) -> list: - """List all keys in the shared DDict.""" - try: - return list(self.ddict.keys()) - except Exception: - return [] - - def clear_task_data(self, key: str = None) -> None: - """Clear specific key or all data from shared DDict.""" - try: - if key: - if key in self.ddict: - del self.ddict[key] - else: - self.ddict.clear() - except Exception as e: - self.logger.warning(f"Error clearing DDict data: {e}") - - -class ResultCollectorV1: - """Unified queue-based result collection for both executable and function tasks.""" - - def __init__( - self, - shared_memory_manager: SharedMemoryManagerV1, - result_queue: Queue, - logger: logging.Logger, - ): - self.shared_memory = shared_memory_manager - self.ddict = shared_memory_manager.ddict - self.result_queue = result_queue - self.logger = logger - - # Track completion for all task types - self.completion_counts: dict[str, int] = {} # task_uid -> received_count - self.expected_counts: dict[str, int] = {} # task_uid -> expected_count - self.aggregated_results: dict[str, list] = {} # task_uid -> [completions] - - def register_task(self, task_uid: str, ranks: int): - """Register any task (executable or function) for result tracking.""" - self.expected_counts[task_uid] = ranks - self.completion_counts[task_uid] = 0 - self.aggregated_results[task_uid] = [] - - def try_consume_result(self) -> Optional[str]: - """Try to consume one result from queue. - - Returns task_uid if task completed, None otherwise. - """ - try: - result = self.result_queue.get(block=False) - if result == "SHUTDOWN": - return None - - if isinstance(result, (ExecutableTaskCompletionV1, FunctionTaskCompletionV1)): - return self._process_completion(result) - - except Exception: # noqa: S110 - # Queue empty - pass - return None - - def _process_completion(self, completion) -> Optional[str]: - """Process completion and return task_uid if task is now complete.""" - task_uid = completion.task_uid - - if task_uid not in self.expected_counts: - self.logger.warning(f"Received completion for unregistered task {task_uid}") - return None - - # Increment completion count - self.completion_counts[task_uid] += 1 - self.aggregated_results[task_uid].append(completion) - - expected = self.expected_counts[task_uid] - received = self.completion_counts[task_uid] - - if received >= expected: - return task_uid - - return None - - def get_completed_task(self, task_uid: str) -> Optional[dict]: - """Get completed task data and clean up tracking.""" - if task_uid not in self.expected_counts: - return None - - results = self.aggregated_results.get(task_uid, []) - if not results: - return None - - # Sort by rank - results.sort(key=lambda r: r.rank) - - # Determine if this is function or executable task - is_function_task = isinstance(results[0], FunctionTaskCompletionV1) - - if is_function_task: - result_data = self._aggregate_function_results(task_uid, results) - else: - result_data = self._aggregate_executable_results(results) - - # Clean up tracking - self.cleanup_task(task_uid) - return result_data - - def _aggregate_function_results( - self, task_uid: str, results: list[FunctionTaskCompletionV1] - ) -> dict: - """Aggregate function task results.""" - if len(results) == 1: - # Single rank - return as-is - result = results[0] - return { - "stdout": result.stdout, - "stderr": result.stderr, - "exit_code": result.exit_code, - "return_value": result.return_value, - "exception": result.exception, - "success": result.success, - } - else: - # Multi-rank - aggregate - stdout_parts = [f"Rank {r.rank}: {r.stdout}" for r in results] - stderr_parts = [f"Rank {r.rank}: {r.stderr}" for r in results] - max_exit_code = max(r.exit_code for r in results) - all_successful = all(r.success for r in results) - - # Collect return values - return_values = [r.return_value for r in results] - - # If any stored in DDict, create single reference for all ranks - if any(r.stored_in_ddict for r in results): - return_value = self.shared_memory.create_data_reference(task_uid, len(results)) - else: - return_value = return_values - - return { - "stdout": "\n".join(stdout_parts), - "stderr": "\n".join(stderr_parts), - "exit_code": max_exit_code, - "return_value": return_value, - "exception": None - if all_successful - else "; ".join(str(r.exception) for r in results if not r.success), - "success": all_successful, - } - - def _aggregate_executable_results(self, results: list[ExecutableTaskCompletionV1]) -> dict: - """Aggregate executable task results.""" - if len(results) == 1: - return results[0].to_result_dict() - else: - stdout_parts = [f"Rank {r.rank}: {r.stdout}" for r in results] - stderr_parts = [f"Rank {r.rank}: {r.stderr}" for r in results] - max_exit_code = max(r.exit_code for r in results) - - return { - "stdout": "\n".join(stdout_parts), - "stderr": "\n".join(stderr_parts), - "exit_code": max_exit_code, - "return_value": None, - "exception": None if max_exit_code == 0 else "One or more processes failed", - } - - def cleanup_task(self, task_uid: str): - """Clean up task tracking data.""" - self.completion_counts.pop(task_uid, None) - self.expected_counts.pop(task_uid, None) - self.aggregated_results.pop(task_uid, None) - - def is_task_complete(self, task_uid: str) -> bool: - """Check if task is complete based on completion counts.""" - if task_uid not in self.expected_counts: - return False - - expected = self.expected_counts[task_uid] - received = self.completion_counts.get(task_uid, 0) - return received >= expected - - -class TaskLauncherV1: - """Unified task launching for all task types.""" - - def __init__(self, ddict: DDict, result_queue: Queue, working_dir: str, logger: logging.Logger): - self.ddict = ddict - self.result_queue = result_queue - self.working_dir = working_dir - self.logger = logger - - async def launch_task(self, task: dict) -> TaskInfoV1: - """Launch any type of task and return TaskInfo.""" - task_type = self._determine_task_type(task) - - if task_type.name.startswith("SINGLE_"): - return await self._launch_single_task(task, task_type) - else: - return await self._launch_group_task(task, task_type) - - def _determine_task_type(self, task: dict) -> TaskTypeV1: - """Determine task type based on task configuration.""" - is_function = bool(task.get("function")) - backend_kwargs = task.get("task_backend_specific_kwargs", {}) - ranks = int(backend_kwargs.get("ranks", 1)) - mpi = backend_kwargs.get("pmi", None) - - if mpi: - if ranks < 2: - raise ValueError("MPI tasks must have ranks > 1") - return TaskTypeV1.MPI_FUNCTION if is_function else TaskTypeV1.MPI_EXECUTABLE - if ranks == 1: - return TaskTypeV1.SINGLE_FUNCTION if is_function else TaskTypeV1.SINGLE_EXECUTABLE - else: # ranks > 1 and not MPI - return TaskTypeV1.MULTI_FUNCTION if is_function else TaskTypeV1.MULTI_EXECUTABLE - - async def _launch_single_task(self, task: dict, task_type: TaskTypeV1) -> TaskInfoV1: - """Launch single-rank task.""" - uid = task["uid"] - - if task_type == TaskTypeV1.SINGLE_FUNCTION: - process = await self._create_function_process(task, 0) - else: - process = self._create_executable_process(task, 0) - - process.start() - self.logger.debug(f"Started single-rank Dragon process for task {uid}") - - return TaskInfoV1(task_type=task_type, ranks=1, start_time=time.time(), process=process) - - async def _launch_group_task(self, task: dict, task_type: TaskTypeV1) -> TaskInfoV1: - """Launch multi-rank or MPI task.""" - uid = task["uid"] - backend_kwargs = task.get("task_backend_specific_kwargs", {}) - ranks = int(backend_kwargs.get("ranks", 2)) - - if task_type.name.startswith("MPI_"): - mpi = backend_kwargs.get("pmi", None) - if mpi is None: - raise ValueError("Missing required 'pmi' value in backend_kwargs.") - - group = ProcessGroup(restart=False, policy=None, pmi=mpi) - self.logger.debug(f"Started MPI group task {uid} ({mpi}) with {ranks} ranks") - else: - group = ProcessGroup(restart=False, policy=None) - self.logger.debug(f"Started group task {uid} with {ranks} ranks") - - if task_type.name.endswith("_FUNCTION"): - await self._add_function_processes_to_group(group, task, ranks) - else: - self._add_executable_processes_to_group(group, task, ranks) - - group.init() - group.start() - - self.logger.debug(f"Started group task {uid} with {ranks} ranks") - - return TaskInfoV1(task_type=task_type, ranks=ranks, start_time=time.time(), group=group) - - async def _create_function_process(self, task: dict, rank: int) -> Process: - """Create a single function process.""" - backend_kwargs = task.get("task_backend_specific_kwargs", {}) - use_ddict_storage = backend_kwargs.get("use_ddict_storage", False) - - return Process( - target=_function_wrapper_v1, - args=(self.result_queue, self.ddict, task, rank, use_ddict_storage), - ) - - def _create_executable_process(self, task: dict, rank: int) -> Process: - """Create a single executable process.""" - executable = task["executable"] - args = list(task.get("arguments", [])) - uid = task["uid"] - backend_kwargs = task.get("task_backend_specific_kwargs", {}) - execute_in_shell = backend_kwargs.get("shell", False) - - return Process( - target=_executable_wrapper_v1, - args=( - self.result_queue, - executable, - args, - uid, - rank, - self.working_dir, - execute_in_shell, - ), - ) - - async def _add_function_processes_to_group( - self, group: ProcessGroup, task: dict, ranks: int - ) -> None: - """Add function processes to process group.""" - task["uid"] - backend_kwargs = task.get("task_backend_specific_kwargs", {}) - use_ddict_storage = backend_kwargs.get("use_ddict_storage", False) - - for rank in range(ranks): - env = os.environ.copy() - env["DRAGON_RANK"] = str(rank) - - template = ProcessTemplate( - target=_function_wrapper_v1, - args=(self.result_queue, self.ddict, task, rank, use_ddict_storage), - env=env, - cwd=self.working_dir, - stdout=Popen.PIPE, - stderr=Popen.PIPE, - stdin=Popen.DEVNULL, - ) - group.add_process(nproc=1, template=template) - - def _add_executable_processes_to_group( - self, group: ProcessGroup, task: dict, ranks: int - ) -> None: - """Add executable processes to process group.""" - executable = task["executable"] - args = list(task.get("arguments", [])) - uid = task["uid"] - backend_kwargs = task.get("task_backend_specific_kwargs", {}) - execute_in_shell = backend_kwargs.get("shell", False) - - for rank in range(ranks): - env = os.environ.copy() - env["DRAGON_RANK"] = str(rank) - - template = ProcessTemplate( - target=_executable_wrapper_v1, - args=( - self.result_queue, - executable, - args, - uid, - rank, - self.working_dir, - execute_in_shell, - ), - env=env, - cwd=self.working_dir, - ) - group.add_process(nproc=1, template=template) - - -def _executable_wrapper_v1( - result_queue: Queue, - executable: str, - args: list, - task_uid: str, - rank: int, - working_dir: str, - execute_in_shell: bool = False, -): - """Wrapper function that executes executable and pushes completion to queue.""" - import subprocess - import time - - try: - # Execute the process and capture output - if execute_in_shell: - # Shell mode: join executable and arguments into single command string - cmd = " ".join([executable] + args) - result = subprocess.run( # noqa: S602 - cmd, - cwd=working_dir, - capture_output=True, - shell=True, - text=True, - timeout=3600, # 1 hour timeout - ) - else: - # Exec mode: pass executable and arguments separately (no shell) - result = subprocess.run( - [executable] + args, - cwd=working_dir, - capture_output=True, - shell=False, - text=True, - timeout=3600, # 1 hour timeout - ) - - # Create completion object - completion = ExecutableTaskCompletionV1( - task_uid=task_uid, - rank=rank, - process_id=os.getpid(), - stdout=result.stdout, - stderr=result.stderr, - exit_code=result.returncode, - timestamp=time.time(), - ) - - # Push completion directly to queue - result_queue.put(completion, block=True) - - except subprocess.TimeoutExpired: - completion = ExecutableTaskCompletionV1( - task_uid=task_uid, - rank=rank, - process_id=os.getpid(), - stdout="", - stderr="Process timed out", - exit_code=124, - timestamp=time.time(), - ) - result_queue.put(completion, block=True) - - except Exception as e: - completion = ExecutableTaskCompletionV1( - task_uid=task_uid, - rank=rank, - process_id=os.getpid(), - stdout="", - stderr=f"Execution error: {str(e)}", - exit_code=1, - timestamp=time.time(), - ) - result_queue.put(completion, block=True) - - -def _function_wrapper_v1( - result_queue: Queue, ddict: DDict, task: dict, rank: int, use_ddict_storage: bool -): - """Wrapper function that executes user functions and pushes completion to queue. - - DDict is only used when user explicitly sets use_ddict_storage=True. Otherwise, return value is - sent directly via queue. - """ - import io - import traceback - - task_uid = task["uid"] - os.environ["DRAGON_RANK"] = str(rank) - - # Capture stdout/stderr - old_out, old_err = sys.stdout, sys.stderr - out_buf, err_buf = io.StringIO(), io.StringIO() - - function = task["function"] - args = task.get("args", ()) - kwargs = task.get("kwargs", {}) - - try: - sys.stdout, sys.stderr = out_buf, err_buf - - # Execute function - if asyncio.iscoroutinefunction(function): - result = asyncio.run(function(*args, **kwargs)) - else: - raise RuntimeError("Sync functions are not supported, please define it as async") - - # Store in DDict only if user requested - stored_in_ddict = False - return_value_for_queue = result - - if use_ddict_storage: - result_key = f"return_{task_uid}_rank_{rank}" - ddict.pput(result_key, result) - stored_in_ddict = True - return_value_for_queue = None - - # Create completion and send via queue - completion = FunctionTaskCompletionV1( - task_uid=task_uid, - rank=rank, - process_id=os.getpid(), - stdout=out_buf.getvalue(), - stderr=err_buf.getvalue(), - exit_code=0, - timestamp=time.time(), - success=True, - exception=None, - traceback=None, - return_value=return_value_for_queue, - stored_in_ddict=stored_in_ddict, - ) - - except Exception as e: - # Error case - completion = FunctionTaskCompletionV1( - task_uid=task_uid, - rank=rank, - process_id=os.getpid(), - stdout=out_buf.getvalue(), - stderr=err_buf.getvalue(), - exit_code=1, - timestamp=time.time(), - success=False, - exception=str(e), - traceback=traceback.format_exc(), - return_value=None, - stored_in_ddict=False, - ) - - finally: - sys.stdout, sys.stderr = old_out, old_err - - # Send completion to queue - try: - result_queue.put(completion, block=True, timeout=30) - except Exception as queue_error: - print(f"Failed to send completion to queue: {queue_error}", file=sys.stderr) - - # Detach from DDict if used - try: - if stored_in_ddict: - ddict.detach() - except Exception: # noqa: S110 - pass - - -# ============================================================================ -# V2 Integration Helper Classes -# ============================================================================ - - -class TaskTypeV2(Enum): - """Enumeration of supported task types.""" - - FUNCTION = "function" - EXECUTABLE = "executable" - - -class WorkerPinningPolicyV2(Enum): - """Worker pinning policy for task assignment.""" - - STRICT = "strict" # Wait indefinitely for hinted worker - SOFT = "soft" # Wait N seconds, then fallback to any worker - AFFINITY = "affinity" # Prefer hinted worker, use others if not immediately available - EXCLUSIVE = "exclusive" # Only hinted worker can run, reject if insufficient capacity - - -class WorkerTypeV2(Enum): - """Worker type enumeration.""" - - COMPUTE = "compute" - TRAINING = "training" - - -@dataclass -class WorkerRequestV2: - """Request sent to worker pool.""" - - task_uid: str - task_type: TaskTypeV2 - rank: int - total_ranks: int - gpu_ids: list[int] = field(default_factory=list) - use_ddict_storage: bool = False - function: Optional[Callable] = None - args: tuple = () - kwargs: dict = None - executable: Optional[str] = None - exec_args: list = None - working_dir: str = "." - execute_in_shell: bool = False - - def __post_init__(self): - if self.kwargs is None: - self.kwargs = {} - if self.exec_args is None: - self.exec_args = [] - - -@dataclass -class WorkerResponseV2: - """Response from worker.""" - - task_uid: str - rank: int - success: bool - worker_name: str = "" - return_value: Any = None - stored_ref_key: Optional[str] = None - is_reference: bool = False - stdout: str = "" - stderr: str = "" - exception: Optional[str] = None - exit_code: int = 0 - timestamp: float = 0.0 - - -@dataclass -class TaskInfoV2: - """Container for task runtime information.""" - - task_type: TaskTypeV2 - ranks: int - start_time: float - worker_name: str = "" - gpu_allocations: dict[int, list[int]] = field(default_factory=dict) - canceled: bool = False - completed_ranks: int = 0 - - -@dataclass -class PolicyConfigV2: - """Configuration for a single policy. - - For COMPUTE workers: nprocs MUST be specified - For TRAINING workers: nprocs MUST be None (omitted) - """ - - policy: Optional[Policy] = None - nprocs: Optional[int] = None - ngpus: int = 0 - - -@dataclass -class WorkerGroupConfigV2: - """Unified configuration for worker groups. - - Attributes: - name: Worker group name - worker_type: COMPUTE or TRAINING - policies: List of PolicyConfig objects - kwargs: Additional config (for training: passed to configure_training_group) - - Examples: - # Compute worker - WorkerGroupConfigV2( - name="compute", - worker_type=WorkerTypeV2.COMPUTE, - policies=[PolicyConfigV2(policy=p, nprocs=128, ngpus=2)] - ) - - # Training worker - WorkerGroupConfigV2( - name="training", - worker_type=WorkerTypeV2.TRAINING, - policies=[PolicyConfigV2(policy=p1), PolicyConfigV2(policy=p2)], - kwargs={'nprocs': 2, 'ppn': 32} - ) - """ - - name: str - worker_type: WorkerTypeV2 = WorkerTypeV2.COMPUTE - policies: list[PolicyConfigV2] = field(default_factory=list) - kwargs: dict[str, Any] = field(default_factory=dict) - - def __post_init__(self): - """Validate configuration.""" - if self.worker_type == WorkerTypeV2.TRAINING: - for i, pc in enumerate(self.policies): - if pc.nprocs is not None: - raise ValueError( - f"Training worker '{self.name}': PolicyConfig[{i}] must NOT specify nprocs" - ) - else: # COMPUTE - for i, pc in enumerate(self.policies): - if pc.nprocs is None: - raise ValueError( - f"Compute worker '{self.name}': PolicyConfig[{i}] must specify nprocs" - ) - - def total_slots(self) -> int: - """Total CPU slots.""" - if self.worker_type == WorkerTypeV2.TRAINING: - if "nprocs" in self.kwargs: - return int(self.kwargs["nprocs"]) - return len(self.policies) - else: - return sum(p.nprocs for p in self.policies) - - def total_gpus(self) -> int: - """Total GPUs.""" - if self.worker_type == WorkerTypeV2.TRAINING: - return 0 # Training workers manage GPUs via policies - else: - return sum(p.ngpus for p in self.policies) - - -class DataReferenceV2: - """Reference to data stored in Cross Node Distributed Dict.""" - - def __init__( - self, ref_id: str, backend_id: str, ddict: DDict, rank_info: Optional[dict] = None - ): - self._ref_id = ref_id - self._backend_id = backend_id - self._ddict = ddict - self._rank_info = rank_info - - @property - def ref_id(self) -> str: - return self._ref_id - - @property - def backend_id(self) -> str: - return self._backend_id - - def resolve(self) -> Any: - """Resolve reference to actual data.""" - if self._rank_info is None: - data_key = f"data_{self._ref_id}" - if data_key not in self._ddict: - raise KeyError(f"Reference data not found: {self._ref_id}") - return self._ddict[data_key] - else: - rank_keys = self._rank_info["rank_keys"] - results = [] - for ref_key in rank_keys: - if ref_key is None: - results.append(None) - else: - data_key = f"data_{ref_key}" - if data_key in self._ddict: - results.append(self._ddict[data_key]) - else: - results.append(None) - return results - - def __repr__(self) -> str: - if self._rank_info: - return f"DataReferenceV2(ref_id='{self._ref_id}', backend_id='{self._backend_id}', ranks={len(self._rank_info['rank_keys'])})" - return f"DataReferenceV2(ref_id='{self._ref_id}', backend_id='{self._backend_id}')" - - -class SharedMemoryManagerV2: - """Manages data storage using DDict.""" - - def __init__( - self, - ddict: DDict, - system: System, - logger: logging.Logger, - reference_threshold: int = DRAGON_DEFAULT_REF_THRESHOLD, - ): - self.ddict = ddict - self.system = system - self.logger = logger - self.backend_id = f"dragon_{uuid.uuid4().hex[:8]}" - self.reference_threshold = reference_threshold - - async def initialize(self): - """Initialize storage manager.""" - self.logger.debug( - f"SharedMemoryManagerV2 initialized (threshold: {self.reference_threshold} bytes)" - ) - - async def store_data(self, data: Any, node_id: int = 0) -> DataReferenceV2: - """Store data in DDict and return reference.""" - ref_id = f"ref_{uuid.uuid4().hex}" - try: - self._store_in_ddict(ref_id, data) - self.logger.debug(f"Stored data {ref_id} in DDict") - except Exception as e: - self.logger.error(f"Failed to store data {ref_id}: {e}") - raise - return DataReferenceV2(ref_id, self.backend_id, self.ddict) - - def _store_in_ddict(self, ref_id: str, data: Any): - """Store data in DDict.""" - self.ddict.pput(f"data_{ref_id}", data) - self.ddict.pput(f"meta_{ref_id}", {"backend_id": self.backend_id, "stored_at": time.time()}) - - def cleanup_reference(self, ref: DataReferenceV2): - """Clean up reference data.""" - try: - for key in [f"meta_{ref.ref_id}", f"data_{ref.ref_id}"]: - if key in self.ddict: - del self.ddict[key] - except Exception as e: - self.logger.warning(f"Error cleaning up reference {ref.ref_id}: {e}") - - -def _worker_loop_v2( - worker_id: int, - worker_name: str, - input_queue: Queue, - output_queue: Queue, - ddict: DDict, - working_dir: str, -) -> None: - """Persistent worker that processes tasks from queue. - - Handles both function and executable tasks: - - Functions: Execute Python callables - - Executables: Run external processes via subprocess - """ - import io - import subprocess - import traceback - - os.environ["DRAGON_WORKER_ID"] = str(worker_id) - os.environ["DRAGON_WORKER_NAME"] = worker_name - - backend_id = os.environ.get("DRAGON_BACKEND_ID", f"dragon_{uuid.uuid4().hex[:8]}") - dragon_cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", None) - - # Map Dragon training env vars to PyTorch standard names - if "DRAGON_PG_RANK" in os.environ: - os.environ["RANK"] = os.environ["DRAGON_PG_RANK"] - if "DRAGON_PG_LOCAL_RANK" in os.environ: - os.environ["LOCAL_RANK"] = os.environ["DRAGON_PG_LOCAL_RANK"] - if "DRAGON_PG_WORLD_SIZE" in os.environ: - os.environ["WORLD_SIZE"] = os.environ["DRAGON_PG_WORLD_SIZE"] - if "DRAGON_PG_MASTER_ADDR" in os.environ and "MASTER_ADDR" not in os.environ: - os.environ["MASTER_ADDR"] = os.environ["DRAGON_PG_MASTER_ADDR"] - if "DRAGON_PG_MASTER_PORT" in os.environ and "MASTER_PORT" not in os.environ: - os.environ["MASTER_PORT"] = os.environ["DRAGON_PG_MASTER_PORT"] - - try: - while True: - # Get task from queue (blocking) - request = input_queue.get() - - # Shutdown signal - if request is None: - break - - if not isinstance(request, WorkerRequestV2): - continue - - # Set rank environment variable for the task - os.environ["DRAGON_RANK"] = str(request.rank) - - # Set GPU visibility for this rank - if request.gpu_ids: - os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, request.gpu_ids)) - elif dragon_cuda_visible is not None: - os.environ["CUDA_VISIBLE_DEVICES"] = dragon_cuda_visible - - response = WorkerResponseV2( - task_uid=request.task_uid, - rank=request.rank, - worker_name=worker_name, - success=False, - timestamp=time.time(), - ) - - try: - if request.task_type == TaskTypeV2.FUNCTION: - # Execute Python function - old_out, old_err = sys.stdout, sys.stderr - out_buf, err_buf = io.StringIO(), io.StringIO() - - try: - sys.stdout, sys.stderr = out_buf, err_buf - - # Handle async functions - if asyncio.iscoroutinefunction(request.function): - result = asyncio.run(request.function(*request.args, **request.kwargs)) - else: - result = request.function(*request.args, **request.kwargs) - - if result is not None and request.use_ddict_storage: - ref_key = f"return_{request.task_uid}_rank_{request.rank}" - try: - ddict.pput(f"data_{ref_key}", result) - ddict.pput( - f"meta_{ref_key}", - {"backend_id": backend_id, "stored_at": time.time()}, - ) - response.stored_ref_key = ref_key - response.is_reference = True - response.return_value = None - except Exception as store_error: - raise RuntimeError( - f"Failed to store return value in DDict: {store_error}" - ) - else: - response.return_value = result - response.is_reference = False - - response.success = True - response.stdout = out_buf.getvalue() - response.stderr = err_buf.getvalue() - response.exit_code = 0 - - except Exception as e: - response.success = False - response.exception = str(e) - response.stderr = err_buf.getvalue() + "\n" + traceback.format_exc() - response.exit_code = 1 - - finally: - sys.stdout, sys.stderr = old_out, old_err - - elif request.task_type == TaskTypeV2.EXECUTABLE: - # Execute external process - if request.execute_in_shell: - # Shell mode: join executable and arguments into single command string - cmd = " ".join([request.executable] + request.exec_args) - result = subprocess.run( # noqa: S602 - cmd, - cwd=request.working_dir, - capture_output=True, - text=True, - shell=True, - timeout=3600, # FIXME: should be user defined - ) - else: - # Exec mode: pass executable and arguments separately (no shell) - result = subprocess.run( - [request.executable] + request.exec_args, - cwd=request.working_dir, - capture_output=True, - text=True, - shell=False, - timeout=3600, # FIXME: should be user defined - ) - - response.success = result.returncode == 0 - response.stdout = result.stdout - response.stderr = result.stderr - response.exit_code = result.returncode - if result.returncode != 0: - response.exception = f"Process exited with code {result.returncode}" - - except subprocess.TimeoutExpired: - response.exception = "Process timed out" - response.exit_code = 124 - response.stderr = "Process timed out" - - except Exception as e: - response.exception = f"Worker error: {str(e)}" - response.exit_code = 1 - response.stderr = traceback.format_exc() - - # Send response back - output_queue.put(response) - - except Exception as e: - print(f"Worker {worker_id} ({worker_name}) fatal error: {e}") - raise - finally: - try: - ddict.detach() - except Exception: # noqa: S110 - pass - - -class WorkerPoolV2: - """Manages persistent worker pool with per-worker queues and slot reservation. - - Key features: - - Each worker group has its own input queue - - Slot tracking per worker for load balancing - - Tasks are assigned to workers with sufficient free slots - """ - - def __init__( - self, - worker_configs: list[WorkerGroupConfigV2], - ddict: DDict, - working_dir: str, - logger: logging.Logger, - system: System, - backend_id: str, - ): - self.worker_configs = worker_configs - self.ddict = ddict - self.working_dir = working_dir - self.logger = logger - self.system = system - self.backend_id = backend_id - - self.output_queue: Optional[Queue] = None - self.worker_queues: dict[str, Queue] = {} - # CPU slot tracking - self.worker_slots: dict[str, int] = {} - self.worker_free_slots: dict[str, int] = {} - # GPU tracking - self.worker_gpus: dict[str, int] = {} - self.worker_free_gpus: dict[str, list[int]] = {} - self.worker_types: dict[str, WorkerTypeV2] = {} - - self.process_groups: list[ProcessGroup] = [] - self.total_workers = len(worker_configs) - self.total_slots = sum(cfg.total_slots() for cfg in worker_configs) - self.total_gpus = sum(cfg.total_gpus() for cfg in worker_configs) - self.initialized = False - - # Thread-safe resource management - self._resource_lock = asyncio.Lock() - - async def initialize(self): - """Initialize worker pool.""" - if self.initialized: - return - - try: - # Single shared output queue - self.output_queue = Queue() - - worker_id = 0 - - for worker_config in self.worker_configs: - worker_name = worker_config.name - worker_type = worker_config.worker_type - - # Create dedicated input queue for this worker - input_queue = Queue() - self.worker_queues[worker_name] = input_queue - self.worker_types[worker_name] = worker_type - - # Track CPU slots - total_worker_slots = worker_config.total_slots() - self.worker_slots[worker_name] = total_worker_slots - self.worker_free_slots[worker_name] = total_worker_slots - - # Track GPUs - assign sequential IDs - total_worker_gpus = worker_config.total_gpus() - self.worker_gpus[worker_name] = total_worker_gpus - self.worker_free_gpus[worker_name] = list(range(total_worker_gpus)) - - if worker_type == WorkerTypeV2.TRAINING: - # Training worker: extract policies and pass kwargs to configure_training_group - config_kwargs = { - "training_fn": _worker_loop_v2, - "training_args": ( - worker_id, - worker_name, - input_queue, - self.output_queue, - self.ddict, - self.working_dir, - ), - "policies": [pc.policy for pc in worker_config.policies], - } - config_kwargs.update(worker_config.kwargs) - - self.logger.info( - f"Creating training worker '{worker_name}' with {total_worker_slots} slots" - ) - - process_group = ProcessGroup.configure_training_group(**config_kwargs) - process_group.init() - process_group.start() - self.process_groups.append(process_group) - - else: - # Compute worker: use PolicyConfig.nprocs - process_group = ProcessGroup(restart=False) - - for policy_config in worker_config.policies: - env = os.environ.copy() - env["DRAGON_WORKER_ID"] = str(worker_id) - env["DRAGON_WORKER_NAME"] = worker_name - env["DRAGON_BACKEND_ID"] = self.backend_id - - template = ProcessTemplate( - target=_worker_loop_v2, - args=( - worker_id, - worker_name, - input_queue, - self.output_queue, - self.ddict, - self.working_dir, - ), - env=env, - cwd=self.working_dir, - policy=policy_config.policy, - ) - - process_group.add_process(nproc=policy_config.nprocs, template=template) - - process_group.init() - process_group.start() - self.process_groups.append(process_group) - - worker_id += 1 - - self.initialized = True - - config_summary = [] - for cfg in self.worker_configs: - config_summary.append( - f"{cfg.name} ({cfg.worker_type.value}): {cfg.total_slots()} slots, {cfg.total_gpus()} GPUs" - ) - - self.logger.info( - f"Worker pool initialized: {self.total_workers} workers, " - f"{self.total_slots} total slots, {self.total_gpus} total GPUs. " - f"Config: {'; '.join(config_summary)}" - ) - - except Exception as e: - self.logger.exception(f"Failed to initialize worker pool: {e}") - raise - - def find_worker_for_task( - self, - ranks: int, - gpus_per_rank: int = 0, - preferred_worker: Optional[str] = None, - worker_type_hint: Optional[str] = None, - ) -> Optional[str]: - """Find worker with sufficient capacity using least-loaded strategy.""" - total_gpus_needed = ranks * gpus_per_rank - - # Check preferred worker first - if preferred_worker and preferred_worker in self.worker_free_slots: - worker_type = self.worker_types.get(preferred_worker) - if worker_type == WorkerTypeV2.TRAINING: - if self.worker_free_slots[preferred_worker] >= ranks: - return preferred_worker - else: - if ( - self.worker_free_slots[preferred_worker] >= ranks - and len(self.worker_free_gpus[preferred_worker]) >= total_gpus_needed - ): - return preferred_worker - - # Find all eligible workers and pick the least loaded one - eligible_workers = [] - - for worker_name in self.worker_free_slots.keys(): - # Filter by worker type if specified - if worker_type_hint: - try: - target_type = WorkerTypeV2(worker_type_hint.lower()) - if self.worker_types.get(worker_name) != target_type: - continue - except ValueError: - pass - - worker_type = self.worker_types.get(worker_name) - - # Check if worker has sufficient capacity - if worker_type == WorkerTypeV2.TRAINING: - if self.worker_free_slots[worker_name] >= ranks: - eligible_workers.append((worker_name, self.worker_free_slots[worker_name])) - else: - if ( - self.worker_free_slots[worker_name] >= ranks - and len(self.worker_free_gpus[worker_name]) >= total_gpus_needed - ): - eligible_workers.append((worker_name, self.worker_free_slots[worker_name])) - - # Return the worker with the most free slots (least loaded) - if eligible_workers: - # Sort by free slots (descending) to get least loaded worker - eligible_workers.sort(key=lambda x: x[1], reverse=True) - return eligible_workers[0][0] - - return None - - def worker_has_capacity(self, worker_name: str, ranks: int, gpus_per_rank: int = 0) -> bool: - """Check if worker has capacity.""" - if worker_name not in self.worker_free_slots: - return False - - worker_type = self.worker_types.get(worker_name) - if worker_type == WorkerTypeV2.TRAINING: - return self.worker_free_slots[worker_name] >= ranks - - total_gpus_needed = ranks * gpus_per_rank - return ( - self.worker_free_slots[worker_name] >= ranks - and len(self.worker_free_gpus[worker_name]) >= total_gpus_needed - ) - - def worker_exists(self, worker_name: str) -> bool: - """Check if worker exists.""" - return worker_name in self.worker_slots - - async def reserve_resources( - self, worker_name: str, ranks: int, gpus_per_rank: int = 0 - ) -> tuple[bool, dict[int, list[int]]]: - """Reserve resources (thread-safe).""" - async with self._resource_lock: - if worker_name not in self.worker_free_slots: - return False, {} - - worker_type = self.worker_types.get(worker_name) - - if worker_type == WorkerTypeV2.TRAINING: - if self.worker_free_slots[worker_name] < ranks: - return False, {} - - self.worker_free_slots[worker_name] -= ranks - return True, {} - - total_gpus_needed = ranks * gpus_per_rank - if ( - self.worker_free_slots[worker_name] < ranks - or len(self.worker_free_gpus[worker_name]) < total_gpus_needed - ): - return False, {} - - self.worker_free_slots[worker_name] -= ranks - - gpu_allocations = {} - if gpus_per_rank > 0: - for rank in range(ranks): - allocated_gpus = [] - for _ in range(gpus_per_rank): - gpu_id = self.worker_free_gpus[worker_name].pop(0) - allocated_gpus.append(gpu_id) - gpu_allocations[rank] = allocated_gpus - - return True, gpu_allocations - - async def release_resources( - self, worker_name: str, ranks: int, gpu_allocations: dict[int, list[int]] - ): - """Release resources (thread-safe).""" - async with self._resource_lock: - if worker_name in self.worker_free_slots: - self.worker_free_slots[worker_name] += ranks - - worker_type = self.worker_types.get(worker_name) - if worker_type != WorkerTypeV2.TRAINING: - for rank_gpus in gpu_allocations.values(): - self.worker_free_gpus[worker_name].extend(rank_gpus) - self.worker_free_gpus[worker_name].sort() - - def submit_request(self, worker_name: str, request: WorkerRequestV2): - """Submit task request to worker.""" - if not self.initialized: - raise RuntimeError("Worker pool not initialized") - - if worker_name not in self.worker_queues: - raise ValueError(f"Unknown worker: {worker_name}") - - try: - self.worker_queues[worker_name].put(request, timeout=10) - except Exception as e: - self.logger.error(f"Failed to submit request to {worker_name}: {e}") - raise - - def try_get_response(self) -> Optional[WorkerResponseV2]: - """Try to get response from output queue.""" - try: - return self.output_queue.get(block=False) - except Exception: - return None - - async def shutdown(self): - """Shutdown worker pool gracefully.""" - if not self.initialized: - return - - try: - self.logger.info("Initiating worker pool shutdown...") - - # Send shutdown signal (None) to each worker process - # Each worker has multiple processes, so we need to send one None per process - for worker_name, input_queue in self.worker_queues.items(): - worker_slots = self.worker_slots[worker_name] - self.logger.debug(f"Sending {worker_slots} shutdown signals to {worker_name}") - for _ in range(worker_slots): - try: - input_queue.put(None, timeout=1.0) - except Exception as e: - self.logger.warning(f"Failed to send shutdown signal to {worker_name}: {e}") - - # Give processes time to finish current work and exit cleanly - await asyncio.sleep(0.5) - - # Join and stop all process groups - for idx, process_group in enumerate(self.process_groups): - try: - self.logger.debug(f"Stopping ProcessGroup {idx}") - - # Join first to wait for processes to exit - try: - process_group.join(timeout=5.0) - self.logger.debug(f"ProcessGroup {idx} joined successfully") - except Exception as e: - self.logger.warning(f"ProcessGroup {idx} join timeout or error: {e}") - - # Then stop and close - try: - process_group.stop() - self.logger.debug(f"ProcessGroup {idx} stopped") - except Exception as e: - self.logger.debug( - f"ProcessGroup {idx} stop error (may already be stopped): {e}" - ) - - try: - process_group.close() - self.logger.debug(f"ProcessGroup {idx} closed") - except Exception as e: - self.logger.debug(f"ProcessGroup {idx} close error: {e}") - - except DragonUserCodeError as e: - self.logger.debug(f"ProcessGroup {idx} user code error during shutdown: {e}") - except Exception as e: - self.logger.warning(f"Error stopping ProcessGroup {idx}: {e}") - - self.logger.info("Worker pool shutdown complete") - - except Exception as e: - self.logger.exception(f"Error shutting down worker pool: {e}") - finally: - self.initialized = False - self.process_groups.clear() - self.worker_queues.clear() - - -class ResultCollectorV2: - """Collects and aggregates results from worker pool.""" - - def __init__(self, shared_memory_manager: SharedMemoryManagerV2, logger: logging.Logger): - self.shared_memory = shared_memory_manager - self.logger = logger - - # Track task completions - self.task_responses: dict[str, list[WorkerResponseV2]] = {} - self.task_expected: dict[str, int] = {} - - def register_task(self, task_uid: str, ranks: int): - """Register a task for result tracking.""" - self.task_expected[task_uid] = ranks - self.task_responses[task_uid] = [] - - def process_response(self, response: WorkerResponseV2) -> Optional[str]: - """Process worker response.""" - task_uid = response.task_uid - - if task_uid not in self.task_expected: - self.logger.warning(f"Received response for unregistered task {task_uid}") - return None - - self.task_responses[task_uid].append(response) - - if len(self.task_responses[task_uid]) >= self.task_expected[task_uid]: - return task_uid - - return None - - async def get_task_result(self, task_uid: str) -> Optional[dict]: - """Get aggregated task result.""" - if task_uid not in self.task_responses: - return None - - responses = self.task_responses[task_uid] - responses.sort(key=lambda r: r.rank) - - if len(responses) == 1: - r = responses[0] - - if r.is_reference: - final_return = DataReferenceV2( - ref_id=r.stored_ref_key, - backend_id=self.shared_memory.backend_id, - ddict=self.shared_memory.ddict, - ) - else: - final_return = r.return_value - - result = { - "stdout": r.stdout, - "stderr": r.stderr, - "exit_code": r.exit_code, - "return_value": final_return, - "exception": r.exception, - } - else: - stdout_parts = [f"Rank {r.rank}: {r.stdout}" for r in responses] - stderr_parts = [f"Rank {r.rank}: {r.stderr}" for r in responses] - max_exit_code = max(r.exit_code for r in responses) - all_successful = all(r.success for r in responses) - - has_any_reference = any(r.is_reference for r in responses if r.success) - - if has_any_reference: - rank_keys = [] - for r in responses: - if r.success and r.is_reference: - rank_keys.append(r.stored_ref_key) - else: - rank_keys.append(None) - - final_return = DataReferenceV2( - ref_id=f"unified_{task_uid}", - backend_id=self.shared_memory.backend_id, - ddict=self.shared_memory.ddict, - rank_info={"task_uid": task_uid, "rank_keys": rank_keys}, - ) - else: - return_values = [r.return_value for r in responses] - final_return = return_values[0] if len(return_values) == 1 else return_values - - result = { - "stdout": "\n".join(stdout_parts), - "stderr": "\n".join(stderr_parts), - "exit_code": max_exit_code, - "return_value": final_return, - "exception": None - if all_successful - else "; ".join(r.exception for r in responses if r.exception), - } - - # Cleanup - self.cleanup_task(task_uid) - return result - - def cleanup_task(self, task_uid: str): - """Clean up task tracking.""" - self.task_responses.pop(task_uid, None) - self.task_expected.pop(task_uid, None) - - -# ============================================================================ -# V3 Integration Helper Classes -# ============================================================================ - - -class TaskStateMapperV3: - PENDING = "PENDING" - RUNNING = "RUNNING" - DONE = "DONE" - FAILED = "FAILED" - CANCELED = "CANCELED" - terminal_states = {DONE, FAILED, CANCELED} - - -# ============================================================================ -# Main Backend Classes -# V1 Integrates with Dragon HPC Native API -# V2 Integrates with Dragon API and builds Compute and Training Workers -# V3 Integrates with Dragon.Batch API (Most performant) -# ============================================================================ - - -class DragonExecutionBackendV1(BaseBackend): - """Dragon execution backend with unified queue-based architecture. - - ┌────────────────────────────┐ - │ DRAGON EXECUTION BACKEND │ - └────────────────────────────┘ - | - ┌──────────────┐ - │ TaskLauncher │ - └──────┬───────┘ - ┌────────────-──┴──────────────┐ - ▼ ▼ - ┌────────────────────┐ ┌────────────────────┐ - │ Executable Tasks │ │ Function Tasks │ - │ _executable_wrapper │ │ _function_wrapper │ - └──────────┬─────────┘ └──────────┬─────────┘ - ▼ ▼ - ┌──────────────────────────────────────────┐ - │ DRAGON QUEUE (Unified) │ - │ ExecutableCompletion | FunctionCompletion │ - └──────────────────┬───────────────────────┘ - ▼ - ┌───────────────────┐ - │ ResultCollector │ - │ (Unified Logic) │ - └─────────┬─────────┘ - ▼ - ┌────────────┴────────────┐ - ▼ ▼ - ┌──────────────┐ ┌──────────────────┐ - │ Small Results│ │ Large Results │ - │ (via Queue) │ │ (DDict optional) │ - └──────────────┘ └──────────────────┘ - ▼ - ┌──────────────┐ - │ DataReference│ - │ .resolve() │ - └──────────────┘ - - Characteristics: - - Single unified queue for all tasks - - DDict only for large results or user-requested storage - - Consistent result collection pattern - - Lower latency for small function results - """ - - @typeguard.typechecked - def __init__( - self, - resources: Optional[dict] = None, - ddict: Optional[DDict] = None, - name: Optional[str] = "dragon", - ): - if dragon is None: - raise ImportError("Dragon is required for DragonExecutionBackendV1.") - if DDict is None: - raise ImportError("Dragon DDict is required for this backend version.") - if System is None: - raise ImportError("Dragon System is required for this backend version.") - - super().__init__(name=name) - - self.logger = _get_logger() - self.tasks: dict[str, dict[str, Any]] = {} - self._callback_func: Callable = lambda t, s: None - self._resources = resources or {} - self._initialized = False - self._backend_state = BackendMainStates.INITIALIZED - - # Resource management - self._slots: int = int(self._resources.get("slots", mp.cpu_count() or 1)) - self._free_slots: int = self._slots - self._working_dir: str = self._resources.get("working_dir", os.getcwd()) - - # Task tracking - self._running_tasks: dict[str, TaskInfoV1] = {} - - # Dragon components - self._ddict: Optional[DDict] = ddict - self._system_alloc: Optional[System] = None - self._result_queue: Optional[Queue] = None - - # Shared memory manager - self._shared_memory: Optional[SharedMemoryManagerV1] = None - - # Utilities - self._result_collector: Optional[ResultCollectorV1] = None - self._task_launcher: Optional[TaskLauncherV1] = None - - # Async management - self._monitor_task: Optional[asyncio.Task] = None - self._shutdown_event = asyncio.Event() - - # --------------------------- Lifecycle --------------------------- - def __await__(self): - return self._async_init().__await__() - - async def _async_init(self): - """Unified async initialization with backend and task state registration. - - Pattern: - 1. Register backend states first - 2. Register task states - 3. Set backend state to INITIALIZED - 4. Initialize backend components - """ - if not self._initialized: - try: - self.logger.debug("Starting Dragon backend V1 async initialization...") - - # Step 1: Register backend states - self.logger.debug("Registering backend states...") - StateMapper.register_backend_states_with_defaults(backend=self) - - # Step 2: Register task states - self.logger.debug("Registering task states...") - StateMapper.register_backend_tasks_states_with_defaults(backend=self) - - # Step 3: Set backend state to INITIALIZED - self._backend_state = BackendMainStates.INITIALIZED - self.logger.debug(f"Backend state set to: {self._backend_state.value}") - - # Step 4: Initialize backend components - await self._initialize() - self._initialized = True - self.logger.info("Dragon backend V1 fully initialized and ready") - - except Exception as e: - self.logger.exception(f"Dragon backend V1 initialization failed: {e}") - self._initialized = False - raise - return self - - async def _initialize(self) -> None: - try: - self.logger.debug("Initializing Dragon backend V1 with unified queue architecture...") - await self._initialize_dragon() - - # Initialize system allocation - self.logger.debug("Creating System allocation...") - self._system_alloc = System() - nnodes = self._system_alloc.nnodes - self.logger.debug(f"System allocation created with {nnodes} nodes") - - # Initialize DDict (optional - only for large results) - self.logger.debug("Creating DDict") - if not self._ddict: - self._ddict = DDict( - n_nodes=nnodes, - total_mem=nnodes * int(4 * 1024 * 1024 * 1024), # 4GB per node - wait_for_keys=True, - working_set_size=4, - timeout=200, - ) - self.logger.debug("DDict created successfully") - - # Initialize global result queue (unified for all task types) - self.logger.debug("Creating unified result queue...") - self._result_queue = Queue() - self.logger.debug("Result queue created successfully") - - # Initialize shared memory manager - self._shared_memory = SharedMemoryManagerV1( - self._ddict, self._system_alloc, self.logger - ) - await self._shared_memory.initialize() - - # Initialize utilities - self._result_collector = ResultCollectorV1( - self._shared_memory, self._result_queue, self.logger - ) - self._task_launcher = TaskLauncherV1( - self._ddict, self._result_queue, self._working_dir, self.logger - ) - - # Start task monitoring - self.logger.debug("Starting unified task monitoring...") - self._monitor_task = asyncio.create_task(self._monitor_tasks()) - await asyncio.sleep(0.1) - - self.logger.info( - f"Dragon backend V1 initialized with {self._slots} slots, " - f"unified queue architecture" - ) - except Exception as e: - self.logger.exception(f"Failed to initialize Dragon backend V1: {str(e)}") - raise - - async def _initialize_dragon(self): - """Ensure start method is 'dragon' and proceed.""" - try: - current_method = mp.get_start_method() - self.logger.debug(f"Current multiprocessing start method: {current_method}") - if current_method != "dragon": - mp.set_start_method("dragon", force=True) - except RuntimeError: - pass - self.logger.debug("Dragon backend V1 active with unified queue-based architecture.") - - def get_task_states_map(self): - return StateMapper(backend=self) - - async def submit_tasks(self, tasks: list[dict[str, Any]]) -> None: - self._ensure_initialized() - - # Set backend state to RUNNING when tasks are submitted - 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: - # Validate task - is_valid, error_msg = self._validate_task(task) - if not is_valid: - task["exception"] = ValueError(error_msg) - self._callback_func(task, "FAILED") - continue - - self.tasks[task["uid"]] = task - - try: - await self._submit_task(task) - except Exception as e: - task["exception"] = e - self._callback_func(task, "FAILED") - - async def _submit_task(self, task: dict[str, Any]) -> None: - """Submit a single task for execution.""" - uid = task["uid"] - backend_kwargs = task.get("task_backend_specific_kwargs", {}) - ranks = int(backend_kwargs.get("ranks", 1)) - - # Wait for available slots - while self._free_slots < ranks: - self.logger.debug(f"Waiting for {ranks} slots for task {uid}, {self._free_slots} free") - await asyncio.sleep(0.1) - - self._free_slots -= ranks - - try: - # Register all tasks with unified result collector - self._result_collector.register_task(uid, ranks) - - # Launch task using unified launcher - task_info = await self._task_launcher.launch_task(task) - self._running_tasks[uid] = task_info - self._callback_func(task, "RUNNING") - - except Exception: - self._free_slots += ranks - raise - - async def _monitor_tasks(self) -> None: - """Monitor running tasks with unified queue consumption.""" - self.logger.debug("Monitor task started") - while not self._shutdown_event.is_set(): - try: - # Batch consume queue results (unified for all task types) - completed_tasks = [] - - for _ in range(1000): # Process up to 1000 results per iteration - completed_task_uid = self._result_collector.try_consume_result() - if completed_task_uid: - completed_tasks.append(completed_task_uid) - else: - break - - # Process all completed tasks - for uid in completed_tasks: - if uid in self._running_tasks: - task_info = self._running_tasks[uid] - task = self.tasks.get(uid) - - if task: - # Get aggregated results from collector - result_data = self._result_collector.get_completed_task(uid) - if result_data: - task.update(result_data) - - # Determine task status and notify callback - if task.get("canceled", False): - self._callback_func(task, "CANCELED") - elif task.get("exception") or task.get("exit_code", 0) != 0: - self._callback_func(task, "FAILED") - else: - self._callback_func(task, "DONE") - - # Free up slots - self._free_slots += task_info.ranks - - # Remove from running tasks - self._running_tasks.pop(uid, None) - - await asyncio.sleep(0.01) # Short sleep for responsiveness - - except Exception as e: - self.logger.exception(f"Error in task monitoring: {e}") - await asyncio.sleep(1) - - async def cancel_task(self, uid: str) -> bool: - """Cancel a specific running task with proper cleanup.""" - self._ensure_initialized() - - task_info = self._running_tasks.get(uid) - if not task_info: - return False - - try: - success = await self._cancel_task_by_info(task_info) - - if success: - task_info.canceled = True - - # Clean up result collector tracking - self._result_collector.cleanup_task(uid) - - # Clean up data references if stored in DDict - try: - task_result = self.tasks.get(uid, {}).get("return_value") - if isinstance(task_result, DataReferenceV1): - self._shared_memory.cleanup_reference(task_result) - except Exception as e: - self.logger.warning(f"Error cleaning up references for task {uid}: {e}") - - return success - - except Exception as e: - self.logger.exception(f"Error cancelling task {uid}: {e}") - return False - - async def _cancel_task_by_info(self, task_info: TaskInfoV1) -> bool: - """Cancel task based on TaskInfo.""" - proc, group = task_info.process, task_info.group - - if proc: - if proc.is_alive: - proc.terminate() - proc.join(2.0) - if proc.is_alive: - proc.kill() - proc.join(1.0) - return True - - if group and not group.inactive_puids: - try: - group.stop() - group.close() - except DragonUserCodeError: - pass - return True - return False - - async def cancel_all_tasks(self) -> int: - """Cancel all running tasks.""" - self._ensure_initialized() - canceled = 0 - for task_uid in list(self._running_tasks.keys()): - try: - if await self.cancel_task(task_uid): - canceled += 1 - except Exception: # noqa: S110 - pass - return canceled - - def _validate_task(self, task: dict) -> tuple[bool, str]: - """Validate task configuration before submission.""" - uid = task.get("uid") - if not uid: - return False, "Task must have a 'uid' field" - - function = task.get("function") - executable = task.get("executable") - - if not function and not executable: - return False, "Task must specify either 'function' or 'executable'" - - if function and executable: - return False, "Task cannot specify both 'function' and 'executable'" - - if function and not callable(function): - return False, "Task 'function' must be callable" - - backend_kwargs = task.get("task_backend_specific_kwargs", {}) - ranks = backend_kwargs.get("ranks", 1) - - try: - ranks = int(ranks) - if ranks < 1: - return False, "Task 'ranks' must be >= 1" - except (ValueError, TypeError): - return False, "Task 'ranks' must be a valid integer" - - return True, "" - - def _ensure_initialized(self): - """Ensure backend is properly initialized.""" - if not self._initialized: - raise RuntimeError( - "DragonExecutionBackendV1 must be awaited before use. " - "Use: backend = await DragonExecutionBackendV1(resources)" - ) - - def get_ddict(self) -> DDict: - """Get the shared DDict for cross-task data sharing.""" - self._ensure_initialized() - return self._ddict - - def get_result_queue(self) -> Queue: - """Get the global result queue.""" - self._ensure_initialized() - return self._result_queue - - # Data management methods delegated to SharedMemoryManager - def store_task_data(self, key: str, data: Any) -> None: - """Store data in shared DDict for cross-task access.""" - self._ensure_initialized() - self._shared_memory.store_task_data(key, data) - - def get_task_data(self, key: str, default=None) -> Any: - """Retrieve data from shared DDict.""" - self._ensure_initialized() - return self._shared_memory.get_task_data(key, default) - - def list_task_data_keys(self) -> list: - """List all keys in the shared DDict.""" - self._ensure_initialized() - return self._shared_memory.list_task_data_keys() - - def clear_task_data(self, key: str = None) -> None: - """Clear specific key or all data from shared DDict.""" - self._ensure_initialized() - self._shared_memory.clear_task_data(key) - - def link_explicit_data_deps(self, src_task=None, dst_task=None, file_name=None, file_path=None): - """Link explicit data dependencies between tasks.""" - pass - - def link_implicit_data_deps(self, src_task, dst_task): - """Link implicit data dependencies between tasks.""" - pass - - async def state(self) -> str: - """Get backend state. - - Returns: - str: Current backend state (INITIALIZED, RUNNING, SHUTDOWN) - """ - return self._backend_state.value - - async def task_state_cb(self, task: dict, state: str) -> None: - """Task state callback.""" - pass - - async def build_task(self, task: dict) -> None: - """Build task.""" - pass - - async def shutdown(self) -> None: - """Shutdown with proper cleanup.""" - if not self._initialized: - return - - try: - # Set backend state to SHUTDOWN - self._backend_state = BackendMainStates.SHUTDOWN - self.logger.debug(f"Backend state set to: {self._backend_state.value}") - - self._shutdown_event.set() - await self.cancel_all_tasks() - - # Signal result collector to stop - try: - if self._result_queue: - self._result_queue.put("SHUTDOWN", block=False) - except Exception as e: - self.logger.warning(f"Error signaling queue shutdown: {e}") - - # Stop monitoring task - if self._monitor_task and not self._monitor_task.done(): - try: - await asyncio.wait_for(self._monitor_task, timeout=5.0) - except asyncio.TimeoutError: - self._monitor_task.cancel() - - # Clean up result queue - try: - if self._result_queue: - while True: - try: - self._result_queue.get(block=False) - except Exception: - break - self._result_queue = None - self.logger.debug("Result queue cleaned up") - except Exception as e: - self.logger.warning(f"Error cleaning up result queue: {e}") - - # Clean up DDict - try: - if self._ddict: - self._ddict.clear() - self._ddict.destroy() - self._ddict = None - self.logger.debug("DDict cleaned up and destroyed") - except Exception as e: - self.logger.warning(f"Error cleaning up DDict: {e}") - - # Clean up system allocation - try: - if self._system_alloc: - self._system_alloc = None - self.logger.debug("System allocation cleaned up") - except Exception as e: - self.logger.warning(f"Error cleaning up system allocation: {e}") - - self.logger.info("Dragon execution backend V1 shutdown complete") - - except Exception as e: - self.logger.exception(f"Error during shutdown: {e}") - finally: - self.tasks.clear() - self._running_tasks.clear() - self._initialized = False - - async def __aenter__(self): - if not self._initialized: - await self._async_init() - return self - - async def __aexit__(self, exc_type, exc_val, exc_tb): - await self.shutdown() - - @classmethod - async def create(cls, resources: Optional[dict] = None): - """Create and initialize a DragonExecutionBackendV1.""" - backend = cls(resources) - return await backend - - -class DragonExecutionBackendV2(BaseBackend): - """Dragon execution backend with parallel scheduling. - - Features: - - Per-worker queue architecture for true load balancing - - Slot reservation ensures tasks go to workers with capacity - - Tasks with same rank requirement run in parallel on different workers - - High performance with minimal coordination overhead - - Support for both compute and training workers - - Example configuration: - from radical.asyncflow.backends.execution.dragon import ( - WorkerGroupConfigV2, PolicyConfigV2, WorkerTypeV2 - ) - from dragon.infrastructure.policy import Policy - - # Define policies for node placement - policy_n0 = Policy(host_id=0, distribution=Policy.Distribution.BLOCK) - policy_n1 = Policy(host_id=1, distribution=Policy.Distribution.BLOCK) - - # Compute worker - compute_worker = WorkerGroupConfigV2( - name="cpu_worker", - worker_type=WorkerTypeV2.COMPUTE, - policies=[ - PolicyConfigV2(policy=policy_n0, nprocs=128), - PolicyConfigV2(policy=policy_n1, nprocs=128) - ] - ) - - # GPU compute worker - gpu_worker = WorkerGroupConfigV2( - name="gpu_worker", - worker_type=WorkerTypeV2.COMPUTE, - policies=[ - PolicyConfigV2(policy=policy_n0, nprocs=128, ngpus=2), - PolicyConfigV2(policy=policy_n1, nprocs=128, ngpus=2) - ] - ) - - # Training worker (for distributed training with DDP/NCCL) - import socket - hostname = socket.gethostname() - - policy_rank0 = Policy( - placement=Policy.Placement.HOST_NAME, - host_name=hostname, - gpu_affinity=[0] - ) - policy_rank1 = Policy( - placement=Policy.Placement.HOST_NAME, - host_name=hostname, - gpu_affinity=[1] - ) - - training_worker = WorkerGroupConfigV2( - name="training_worker", - worker_type=WorkerTypeV2.TRAINING, - policies=[ - PolicyConfigV2(policy=policy_rank0), - PolicyConfigV2(policy=policy_rank1) - ], - kwargs={'nprocs': 2, 'ppn': 32, 'port': 29500} - ) - - # Create backend - resources = {"workers": [compute_worker, gpu_worker, training_worker]} - backend = await DragonExecutionBackendV2(resources) - """ - - @typeguard.typechecked - def __init__( - self, - resources: Optional[dict] = None, - ddict: Optional[DDict] = None, - name: Optional[str] = "dragon", - ): - if dragon is None: - raise ImportError("Dragon is required for DragonExecutionBackendV2.") - - super().__init__(name=name) - - self.logger = _get_logger() - self.tasks: dict[str, dict[str, Any]] = {} - self._callback_func: Callable = lambda t, s: None - self._resources = resources or {} - self._initialized = False - self._backend_state = BackendMainStates.INITIALIZED - self._canceled_tasks = set() - - # Parse worker configuration - self._worker_configs = self._parse_worker_config(self._resources) - self._total_slots = sum(cfg.total_slots() for cfg in self._worker_configs) - self._total_gpus = sum(cfg.total_gpus() for cfg in self._worker_configs) - - # Other resources - self._working_dir: str = self._resources.get("working_dir", os.getcwd()) - self._reference_threshold: int = int( - self._resources.get("reference_threshold", DRAGON_DEFAULT_REF_THRESHOLD) - ) - - # Task tracking - self._running_tasks: dict[str, TaskInfoV2] = {} - self._pending_tasks: asyncio.Queue = asyncio.Queue() - - # Dragon components - self._ddict: Optional[DDict] = ddict - self._system_alloc: Optional[System] = None - self._shared_memory: Optional[SharedMemoryManagerV2] = None - self._result_collector: Optional[ResultCollectorV2] = None - self._worker_pool: Optional[WorkerPoolV2] = None - self._backend_id: str = f"dragon_{uuid.uuid4().hex[:8]}" - - # Async management - self._monitor_task: Optional[asyncio.Task] = None - self._scheduler_task: Optional[asyncio.Task] = None - self._shutdown_event = asyncio.Event() - - def _parse_worker_config(self, resources: dict) -> list[WorkerGroupConfigV2]: - """Parse worker configuration from resources. - - If no workers specified, creates a default compute worker. - """ - if "workers" in resources: - workers = resources["workers"] - if not isinstance(workers, list): - raise TypeError( - "resources['workers'] must be a list of WorkerGroupConfigV2 objects" - ) - - if not workers: - raise ValueError("resources['workers'] cannot be empty") - - for idx, worker in enumerate(workers): - if not isinstance(worker, WorkerGroupConfigV2): - raise TypeError( - f"Worker at index {idx} must be a WorkerGroupConfigV2 object. " - f"Got {type(worker).__name__} instead." - ) - - return workers - else: - # No workers specified - create default compute worker - slots = int(resources.get("slots", mp.cpu_count() or 1)) - self.logger.info( - f"No workers specified, creating default compute worker with {slots} slots" - ) - return [ - WorkerGroupConfigV2( - name="default_worker", - worker_type=WorkerTypeV2.COMPUTE, - policies=[PolicyConfigV2(nprocs=slots, policy=None, ngpus=0)], - ) - ] - - def __await__(self): - return self._async_init().__await__() - - async def _async_init(self): - """Unified async initialization with backend and task state registration. - - Pattern: - 1. Register backend states first - 2. Register task states - 3. Set backend state to INITIALIZED - 4. Initialize backend components - """ - if not self._initialized: - try: - self.logger.debug("Starting Dragon backend V2 async initialization...") - - # Step 1: Register backend states - self.logger.debug("Registering backend states...") - StateMapper.register_backend_states_with_defaults(backend=self) - - # Step 2: Register task states - self.logger.debug("Registering task states...") - StateMapper.register_backend_tasks_states_with_defaults(backend=self) - - # Step 3: Set backend state to INITIALIZED - self._backend_state = BackendMainStates.INITIALIZED - self.logger.debug(f"Backend state set to: {self._backend_state.value}") - - # Step 4: Initialize backend components - await self._initialize() - self._initialized = True - self.logger.info("Dragon backend V2 fully initialized and ready") - - except Exception as e: - self.logger.exception(f"Dragon backend V2 initialization failed: {e}") - self._initialized = False - raise - return self - - async def _initialize(self) -> None: - try: - # Set multiprocessing method - try: - if mp.get_start_method() != "dragon": - mp.set_start_method("dragon", force=True) - except RuntimeError: - pass - - # Initialize system - self._system_alloc = System() - nnodes = self._system_alloc.nnodes - self.logger.debug(f"System allocation created with {nnodes} nodes") - - # Initialize DDict - if not self._ddict: - self._ddict = DDict( - n_nodes=nnodes, - total_mem=nnodes * int(4 * 1024 * 1024 * 1024), - wait_for_keys=True, - working_set_size=4, - timeout=200, - ) - self.logger.debug("DDict initialized") - - # Initialize shared memory manager - self._shared_memory = SharedMemoryManagerV2( - self._ddict, self._system_alloc, self.logger, self._reference_threshold - ) - await self._shared_memory.initialize() - - # Initialize result collector - self._result_collector = ResultCollectorV2(self._shared_memory, self.logger) - - # Initialize worker pool with per-worker queues - self._worker_pool = WorkerPoolV2( - self._worker_configs, - self._ddict, - self._working_dir, - self.logger, - self._system_alloc, - self._backend_id, - ) - await self._worker_pool.initialize() - - # Start monitoring and scheduling - self._monitor_task = asyncio.create_task(self._monitor_tasks()) - self._scheduler_task = asyncio.create_task(self._schedule_tasks()) - - self.logger.info( - f"Dragon backend V2 initialized: {len(self._worker_configs)} workers, " - f"{self._total_slots} total slots, {self._total_gpus} total GPUs, " - f"reference threshold: {self._reference_threshold} bytes" - ) - - except Exception as e: - self.logger.exception(f"Failed to initialize Dragon backend V2: {e}") - raise - - def get_task_states_map(self): - return StateMapper(backend=self) - - async def submit_tasks(self, tasks: list[dict[str, Any]]) -> None: - self._ensure_initialized() - - # Set backend state to RUNNING when tasks are submitted - 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: - is_valid, error_msg = self._validate_task(task) - if not is_valid: - task["exception"] = ValueError(error_msg) - self._callback_func(task, "FAILED") - continue - - self.tasks[task["uid"]] = task - - try: - # Add to pending queue for scheduler - await self._pending_tasks.put(task) - except Exception as e: - task["exception"] = e - self._callback_func(task, "FAILED") - - async def _schedule_tasks(self) -> None: - """Parallel scheduler with multiple worker coroutines.""" - - num_scheduler_workers = min(32, max(4, self._total_slots // 4)) - - self.logger.info( - f"Starting parallel scheduler with {num_scheduler_workers} scheduler workers" - ) - - async def scheduler_worker(worker_id: int): - """Individual scheduler worker.""" - tasks_processed = 0 - - while not self._shutdown_event.is_set(): - try: - try: - task = await asyncio.wait_for(self._pending_tasks.get(), timeout=0.1) - except asyncio.TimeoutError: - continue - - task["uid"] - backend_kwargs = task.get("task_backend_specific_kwargs", {}) - ranks = int(backend_kwargs.get("ranks", 1)) - gpus_per_rank = int(backend_kwargs.get("gpus_per_rank", 0)) - worker_hint = backend_kwargs.get("worker_hint") - worker_type_hint = backend_kwargs.get("worker_type") - - pinning_policy_str = backend_kwargs.get("pinning_policy", "").lower() - try: - pinning_policy = ( - WorkerPinningPolicyV2(pinning_policy_str) - if pinning_policy_str - else None - ) - except ValueError: - pinning_policy = None - - pinning_timeout = float(backend_kwargs.get("pinning_timeout", 30.0)) - - worker_name = await self._apply_pinning_policy( - task, - ranks, - gpus_per_rank, - worker_hint, - worker_type_hint, - pinning_policy, - pinning_timeout, - ) - - if not worker_name: - continue - - max_retries = 3 - retry_count = 0 - success = False - gpu_allocations = {} - - while retry_count < max_retries and not self._shutdown_event.is_set(): - success, gpu_allocations = await self._worker_pool.reserve_resources( - worker_name, ranks, gpus_per_rank - ) - - if success: - break - - retry_count += 1 - await asyncio.sleep(0.001 * retry_count) - - if not success: - await self._pending_tasks.put(task) - await asyncio.sleep(0.005) - continue - - try: - await self._submit_task_to_worker(task, worker_name, ranks, gpu_allocations) - tasks_processed += 1 - - except Exception as e: - await self._worker_pool.release_resources( - worker_name, ranks, gpu_allocations - ) - task["exception"] = e - self._callback_func(task, "FAILED") - - except Exception as e: - self.logger.exception(f"Scheduler worker {worker_id}: Unexpected error: {e}") - await asyncio.sleep(0.1) - - self.logger.debug( - f"Scheduler worker {worker_id} shutting down (processed {tasks_processed} tasks)" - ) - - scheduler_tasks = [] - for worker_id in range(num_scheduler_workers): - task = asyncio.create_task(scheduler_worker(worker_id)) - scheduler_tasks.append(task) - - self.logger.info(f"Scheduler worker pool active with {num_scheduler_workers} workers") - - await self._shutdown_event.wait() - - self.logger.info("Shutdown signal received, stopping scheduler workers...") - - try: - await asyncio.wait_for( - asyncio.gather(*scheduler_tasks, return_exceptions=True), timeout=10.0 - ) - except asyncio.TimeoutError: - self.logger.warning("Scheduler workers did not complete within timeout, canceling...") - for task in scheduler_tasks: - if not task.done(): - task.cancel() - - await asyncio.gather(*scheduler_tasks, return_exceptions=True) - - self.logger.info("All scheduler workers stopped") - - async def _apply_pinning_policy( - self, - task: dict, - ranks: int, - gpus_per_rank: int, - worker_hint: Optional[str], - worker_type_hint: Optional[str], - pinning_policy: Optional[WorkerPinningPolicyV2], - timeout: float, - ) -> Optional[str]: - """Apply worker pinning policy to find appropriate worker.""" - - if not worker_hint or not pinning_policy: - worker_name = self._worker_pool.find_worker_for_task( - ranks, gpus_per_rank, worker_type_hint=worker_type_hint - ) - while not worker_name and not self._shutdown_event.is_set(): - await asyncio.sleep(0.01) - worker_name = self._worker_pool.find_worker_for_task( - ranks, gpus_per_rank, worker_type_hint=worker_type_hint - ) - return worker_name - - if not self._worker_pool.worker_exists(worker_hint): - error_msg = f"Worker hint '{worker_hint}' does not exist. Available workers: {list(self._worker_pool.worker_slots.keys())}" - self.logger.error(error_msg) - task["exception"] = ValueError(error_msg) - self._callback_func(task, "FAILED") - return None - - if pinning_policy == WorkerPinningPolicyV2.AFFINITY: - if self._worker_pool.worker_has_capacity(worker_hint, ranks, gpus_per_rank): - self.logger.debug( - f"Task {task['uid']}: AFFINITY policy - using preferred worker {worker_hint}" - ) - return worker_hint - else: - worker_name = self._worker_pool.find_worker_for_task( - ranks, gpus_per_rank, worker_type_hint=worker_type_hint - ) - if worker_name: - self.logger.debug( - f"Task {task['uid']}: AFFINITY policy - fallback to {worker_name}" - ) - return worker_name - while not worker_name and not self._shutdown_event.is_set(): - await asyncio.sleep(0.01) - worker_name = self._worker_pool.find_worker_for_task( - ranks, gpus_per_rank, worker_type_hint=worker_type_hint - ) - return worker_name - - elif pinning_policy == WorkerPinningPolicyV2.STRICT: - self.logger.debug( - f"Task {task['uid']}: STRICT policy - waiting for worker {worker_hint}" - ) - while not self._shutdown_event.is_set(): - if self._worker_pool.worker_has_capacity(worker_hint, ranks, gpus_per_rank): - self.logger.debug( - f"Task {task['uid']}: STRICT policy - worker {worker_hint} now available" - ) - return worker_hint - await asyncio.sleep(0.01) - return None - - elif pinning_policy == WorkerPinningPolicyV2.SOFT: - self.logger.debug( - f"Task {task['uid']}: SOFT policy - waiting {timeout}s for worker {worker_hint}" - ) - start_time = time.time() - - while time.time() - start_time < timeout: - if self._worker_pool.worker_has_capacity(worker_hint, ranks, gpus_per_rank): - self.logger.debug( - f"Task {task['uid']}: SOFT policy - worker {worker_hint} available" - ) - return worker_hint - await asyncio.sleep(0.01) - if self._shutdown_event.is_set(): - return None - - self.logger.debug(f"Task {task['uid']}: SOFT policy - timeout reached, using fallback") - worker_name = self._worker_pool.find_worker_for_task( - ranks, gpus_per_rank, worker_type_hint=worker_type_hint - ) - while not worker_name and not self._shutdown_event.is_set(): - await asyncio.sleep(0.01) - worker_name = self._worker_pool.find_worker_for_task( - ranks, gpus_per_rank, worker_type_hint=worker_type_hint - ) - - if worker_name: - self.logger.debug(f"Task {task['uid']}: SOFT policy - fallback to {worker_name}") - return worker_name - - elif pinning_policy == WorkerPinningPolicyV2.EXCLUSIVE: - if self._worker_pool.worker_has_capacity(worker_hint, ranks, gpus_per_rank): - self.logger.debug( - f"Task {task['uid']}: EXCLUSIVE policy - using worker {worker_hint}" - ) - return worker_hint - else: - total_capacity = self._worker_pool.worker_slots.get(worker_hint, 0) - total_gpu_capacity = self._worker_pool.worker_gpus.get(worker_hint, 0) - total_gpus_needed = ranks * gpus_per_rank - - if ranks > total_capacity or total_gpus_needed > total_gpu_capacity: - error_msg = ( - f"Task {task['uid']}: EXCLUSIVE policy - worker '{worker_hint}' " - f"has insufficient total capacity ({total_capacity} slots, {total_gpu_capacity} GPUs) " - f"for {ranks} ranks × {gpus_per_rank} GPUs/rank" - ) - else: - error_msg = ( - f"Task {task['uid']}: EXCLUSIVE policy - worker '{worker_hint}' " - f"currently has insufficient free resources" - ) - - self.logger.error(error_msg) - task["exception"] = ValueError(error_msg) - self._callback_func(task, "FAILED") - return None - - return None - - async def _submit_task_to_worker( - self, - task: dict[str, Any], - worker_name: str, - ranks: int, - gpu_allocations: dict[int, list[int]], - ) -> None: - """Submit task to specific worker.""" - uid = task["uid"] - - try: - # Register task with result collector - self._result_collector.register_task(uid, ranks) - - # Determine task type - is_function = bool(task.get("function")) - task_type = TaskTypeV2.FUNCTION if is_function else TaskTypeV2.EXECUTABLE - - backend_kwargs = task.get("task_backend_specific_kwargs", {}) - use_ddict_storage = backend_kwargs.get("use_ddict_storage", False) - execute_in_shell = backend_kwargs.get("shell", False) - - for rank in range(ranks): - gpu_ids = gpu_allocations.get(rank, []) - - if is_function: - request = WorkerRequestV2( - task_uid=uid, - task_type=TaskTypeV2.FUNCTION, - rank=rank, - total_ranks=ranks, - gpu_ids=gpu_ids, - use_ddict_storage=use_ddict_storage, - function=task["function"], - args=task.get("args", ()), - kwargs=task.get("kwargs", {}), - ) - else: - request = WorkerRequestV2( - task_uid=uid, - task_type=TaskTypeV2.EXECUTABLE, - rank=rank, - total_ranks=ranks, - gpu_ids=gpu_ids, - use_ddict_storage=False, - executable=task["executable"], - exec_args=list(task.get("arguments", [])), - working_dir=self._working_dir, - execute_in_shell=execute_in_shell, - ) - - # Submit to specific worker's queue - self._worker_pool.submit_request(worker_name, request) - - # Track task - self._running_tasks[uid] = TaskInfoV2( - task_type=task_type, - ranks=ranks, - worker_name=worker_name, - gpu_allocations=gpu_allocations, - start_time=time.time(), - ) - - self._callback_func(task, "RUNNING") - - except Exception: - raise - - async def _monitor_tasks(self) -> None: - """Monitor tasks by consuming responses from worker pool.""" - while not self._shutdown_event.is_set(): - try: - completed_tasks = [] - # Consume responses from worker pool (batch processing) - for _ in range(1000): - response = self._worker_pool.try_get_response() - if response: - completed_uid = self._result_collector.process_response(response) - if completed_uid and completed_uid not in completed_tasks: - completed_tasks.append(completed_uid) - else: - break - - for uid in completed_tasks: - # Check if task was canceled - if uid in self._canceled_tasks: - self.logger.debug(f"Ignoring response for canceled task {uid}") - # Now clean it up from result collector - self._result_collector.cleanup_task(uid) - self._canceled_tasks.discard(uid) - continue - - # Check if task is still being tracked - task_info = self._running_tasks.get(uid) - if not task_info: - # Already processed or never tracked - continue - - if task_info.canceled: - # Redundant check, but safe - self.logger.debug(f"Skipping already-canceled task {uid}") - self._result_collector.cleanup_task(uid) - self._running_tasks.pop(uid, None) - continue - - # Normal task completion processing - task = self.tasks.get(uid) - if task: - result = await self._result_collector.get_task_result(uid) - if result: - task.update(result) - - await self._worker_pool.release_resources( - task_info.worker_name, task_info.ranks, task_info.gpu_allocations - ) - - # Send appropriate callback - if task.get("exception") or task.get("exit_code", 0) != 0: - self._callback_func(task, "FAILED") - else: - self._callback_func(task, "DONE") - - self._running_tasks.pop(uid, None) - - await asyncio.sleep(0.001) - except Exception as e: - self.logger.exception(f"Error in task monitoring: {e}") - await asyncio.sleep(1) - - async def cancel_task(self, uid: str) -> bool: - """Cancel a running task. - - Note: With worker pool architecture, cancellation is best-effort. - Workers that already picked up the task will complete it. - """ - task_info = self._running_tasks.get(uid) - if not task_info: - return False - - # Mark as canceled FIRST - task_info.canceled = True - self._canceled_tasks.add(uid) # Track it - - # Update task dict with canceled flag - if uid in self.tasks: - self.tasks[uid]["canceled"] = True - - # Send immediate CANCELED callback - task = self.tasks.get(uid) - if task: - self._callback_func(task, "CANCELED") - - # Release resources - await self._worker_pool.release_resources( - task_info.worker_name, task_info.ranks, task_info.gpu_allocations - ) - - # Remove from running tasks - self._running_tasks.pop(uid, None) - - return True - - async def cancel_all_tasks(self) -> int: - """Cancel all running tasks.""" - canceled = 0 - for uid in list(self._running_tasks.keys()): - if await self.cancel_task(uid): - canceled += 1 - return canceled - - def _validate_task(self, task: dict) -> tuple[bool, str]: - """Validate task configuration.""" - uid = task.get("uid") - if not uid: - return False, "Task must have a 'uid' field" - - function = task.get("function") - executable = task.get("executable") - - if not function and not executable: - return False, "Task must specify either 'function' or 'executable'" - - if function and executable: - return False, "Task cannot specify both 'function' and 'executable'" - - if function and not callable(function): - return False, "Task 'function' must be callable" - - backend_kwargs = task.get("task_backend_specific_kwargs", {}) - ranks = backend_kwargs.get("ranks", 1) - - try: - ranks = int(ranks) - if ranks < 1: - return False, "Task 'ranks' must be >= 1" - except (ValueError, TypeError): - return False, "Task 'ranks' must be a valid integer" - - gpus_per_rank = backend_kwargs.get("gpus_per_rank", 0) - try: - gpus_per_rank = int(gpus_per_rank) - if gpus_per_rank < 0: - return False, "Task 'gpus_per_rank' must be >= 0" - except (ValueError, TypeError): - return False, "Task 'gpus_per_rank' must be a valid integer" - - pinning_policy = backend_kwargs.get("pinning_policy", "").lower() - if pinning_policy: - try: - WorkerPinningPolicyV2(pinning_policy) - except ValueError: - valid_policies = [p.value for p in WorkerPinningPolicyV2] - return ( - False, - f"Invalid pinning_policy '{pinning_policy}'. Must be one of: {valid_policies}", - ) - - worker_hint = backend_kwargs.get("worker_hint") - if worker_hint and not isinstance(worker_hint, str): - return False, "worker_hint must be a string" - - timeout = backend_kwargs.get("pinning_timeout") - if timeout is not None: - try: - float(timeout) - except (ValueError, TypeError): - return False, "pinning_timeout must be a number" - - return True, "" - - def _ensure_initialized(self): - """Ensure backend is initialized.""" - if not self._initialized: - raise RuntimeError( - "DragonExecutionBackendV2 must be awaited before use. " - "Use: backend = await DragonExecutionBackendV2(resources)" - ) - - def get_ddict(self) -> DDict: - """Get shared DDict for cross-task data sharing.""" - self._ensure_initialized() - return self._ddict - - def link_explicit_data_deps(self, src_task=None, dst_task=None, file_name=None, file_path=None): - pass - - def link_implicit_data_deps(self, src_task, dst_task): - pass - - async def state(self) -> str: - """Get backend state. - - Returns: - str: Current backend state (INITIALIZED, RUNNING, SHUTDOWN) - """ - return self._backend_state.value - - async def task_state_cb(self, task: dict, state: str) -> None: - pass - - async def build_task(self, task: dict) -> None: - pass +from ..base import BaseBackend +from ..constants import BackendMainStates +from ..constants import StateMapper - async def shutdown(self) -> None: - """Shutdown backend gracefully.""" - if not self._initialized: - return +DRAGON_BATCH_INIT_ERROR = None - try: - # Set backend state to SHUTDOWN - self._backend_state = BackendMainStates.SHUTDOWN - self.logger.debug(f"Backend state set to: {self._backend_state.value}") - self.logger.info("Starting Dragon backend V2 shutdown...") - self._shutdown_event.set() +try: + import dragon + from dragon.infrastructure.policy import Policy + from dragon.native.process import ProcessTemplate + from dragon.workflows.batch import Batch + from dragon.workflows.batch import BatchError + from dragon.workflows.batch import TaskCancelledError + from dragon.workflows.batch import TaskNotReadyError - # Cancel all running tasks - canceled = await self.cancel_all_tasks() - if canceled > 0: - self.logger.info(f"Canceled {canceled} running tasks") +except ImportError as e: # pragma: no cover - environment without Dragon + dragon = None + ProcessTemplate = None + Policy = None + Batch = None + BatchError = None + TaskCancelledError = None + TaskNotReadyError = None + DRAGON_BATCH_INIT_ERROR = e - # Stop scheduler - if self._scheduler_task and not self._scheduler_task.done(): - self.logger.debug("Stopping scheduler task...") - try: - await asyncio.wait_for(self._scheduler_task, timeout=5.0) - except asyncio.TimeoutError: - self.logger.warning("Scheduler task timeout, canceling...") - self._scheduler_task.cancel() - try: - await self._scheduler_task - except asyncio.CancelledError: - pass - # Stop monitoring - if self._monitor_task and not self._monitor_task.done(): - self.logger.debug("Stopping monitor task...") - try: - await asyncio.wait_for(self._monitor_task, timeout=5.0) - except asyncio.TimeoutError: - self.logger.warning("Monitor task timeout, canceling...") - self._monitor_task.cancel() - try: - await self._monitor_task - except asyncio.CancelledError: - pass +def _get_logger() -> logging.Logger: + """Get logger for dragon backend module. - # Shutdown worker pool (this waits for worker processes to exit) - if self._worker_pool: - await self._worker_pool.shutdown() + This function provides lazy logger evaluation, ensuring the logger is created after the user has + configured logging, not at module import time. + """ + return logging.getLogger(__name__) - # Clean up DDict - if self._ddict: - try: - self.logger.debug("Cleaning up DDict...") - self._ddict.clear() - self._ddict.destroy() - self.logger.debug("DDict cleanup complete") - except Exception as e: - self.logger.warning(f"Error cleaning up DDict: {e}") - self.logger.info("Dragon backend V2 shutdown complete") +# ============================================================================ +# Backend Helper Classes +# ============================================================================ - except Exception as e: - self.logger.exception(f"Error during shutdown: {e}") - finally: - self.tasks.clear() - self._running_tasks.clear() - self._initialized = False - async def __aenter__(self): - if not self._initialized: - await self._async_init() - return self +class TaskStateMapper: + PENDING = "PENDING" + RUNNING = "RUNNING" + DONE = "DONE" + FAILED = "FAILED" + CANCELED = "CANCELED" + terminal_states = {DONE, FAILED, CANCELED} - async def __aexit__(self, exc_type, exc_val, exc_tb): - await self.shutdown() - @classmethod - async def create(cls, resources: Optional[dict] = None): - """Create and initialize backend.""" - backend = cls(resources) - return await backend +# ============================================================================ +# Main Backend Class — Dragon.Batch API (event-driven, HPC-scale) +# ============================================================================ -class DragonExecutionBackendV3(BaseBackend): +class DragonExecutionBackend(BaseBackend): """Dragon Batch backend using the streaming pipeline model. Tasks submitted via batch.function()/process()/job() are auto-dispatched by the Batch - background thread. A single monitor thread polls each task's manager-specific DDict shard - and fires callbacks when results become available. + background thread. A single monitor thread uses ``Batch.poll()`` for event-driven result + delivery — zero busy-wait, zero private DDict access. Args: batch_kwargs: Forwarded verbatim to ``dragon.workflows.batch.Batch()``. - Supported keys (Dragon 0.14.0): + Supported keys (Dragon 0.14.1+): - ``num_nodes`` *(int, optional)* — nodes to use; defaults to the full allocation. @@ -3091,10 +82,14 @@ class DragonExecutionBackendV3(BaseBackend): - ``results_ddict_mem`` *(int, optional)* — bytes to allocate for the results DDict (default: 1 GiB × num_nodes). Increase for workloads that return large arrays or submit millions of tasks. - - ``pool_nodes`` — **no-op in Dragon 0.14.0**, kept for API compatibility. + - ``results_ddict_managers_per_pool`` *(int, optional)* — DDict shard count + per worker pool (default: 4). At large scale (64K+ tasks), increasing this + reduces DDict write contention. Valid range: ``[1, workers_per_node]``; + Dragon clamps automatically. + - ``pool_nodes`` — **no-op in Dragon 0.14.0+**, kept for API compatibility. Note on working directory: - DragonExecutionBackendV3 does not support a backend-level working directory. + DragonExecutionBackend does not support a backend-level working directory. Set it per task via ``task_backend_specific_kwargs`` with ``process_template`` (single process) or ``process_templates`` (MPI job):: @@ -3120,17 +115,23 @@ def __init__( name: Optional[str] = "dragon", ): if not Batch: - raise RuntimeError("Dragon Batch not available") + raise RuntimeError(DRAGON_BATCH_INIT_ERROR) super().__init__(name=name) self.logger = _get_logger() - self.batch = Batch(**(batch_kwargs or {})) + effective_kwargs = {"task_logs": True, **(batch_kwargs or {})} + if not effective_kwargs.get("task_logs", True): + self.logger.warning( + "task_logs is disabled: task['stdout']/task['stderr'] will be empty " + "for capture_stdio=False tasks; output goes directly to the console." + ) + self.batch = Batch(**effective_kwargs) self._backend_state = BackendMainStates.INITIALIZED self._callback_func: Callable = lambda t, s: None self._task_registry: dict[str, Any] = {} - self._task_states = TaskStateMapperV3() + self._task_states = TaskStateMapper() self._initialized = False self._cancelled_tasks: set[str] = set() self._monitored_batches = {} @@ -3140,7 +141,7 @@ def __init__( self._shutdown_event = threading.Event() self.logger.info( - f"DragonExecutionBackendV3: {self.batch.num_workers} workers, " + f"DragonExecutionBackend: {self.batch.num_workers} workers, " f"{self.batch.num_managers} managers" ) @@ -3158,7 +159,7 @@ async def _async_init(self): """ if not self._initialized: try: - self.logger.debug("Starting Dragon backend V3 async initialization...") + self.logger.debug("Starting Dragon backend async initialization...") # Step 1: Register backend states self.logger.debug("Registering backend states...") @@ -3169,69 +170,94 @@ async def _async_init(self): StateMapper.register_backend_tasks_states_with_defaults(backend=self) self._initialized = True - self.logger.info("Dragon backend V3 fully initialized and ready") + self.logger.info("Dragon backend fully initialized and ready") except Exception as e: - self.logger.exception(f"Dragon backend V3 initialization failed: {e}") + self.logger.exception(f"Dragon backend initialization failed: {e}") self._initialized = False raise return self - def _monitor_loop(self): - """Single thread to monitor all active tasks. + def _monitor_loop(self) -> None: + """Single thread to monitor all active tasks via event-driven Batch.poll(). - Tasks are auto-dispatched by the Batch background thread the moment they are created. This - loop polls each registered task non-blocking and fires callbacks when results arrive. + Blocks on poll() at the OS queue level — zero CPU when idle. Drains all simultaneously + available completions before waking the asyncio event loop, preserving O(sweeps) cross- + thread wakeups instead of O(tasks). """ - self.logger.debug("Starting Dragon batch monitor loop (polling mode)") + self.logger.debug("Starting Dragon batch monitor loop (event-driven)") while not self._shutdown_event.is_set() or self._monitored_batches: try: - # Iterate over a copy of keys to allow modification during iteration - batch_tuids = list(self._monitored_batches.keys()) - - if not batch_tuids: - # No active tasks, sleep briefly to avoid high CPU - time.sleep(0.01) + tuid = self.batch.poll(timeout=0.05) + if tuid is None: continue - # Collect all completed tasks in this sweep, then deliver in one batch + # Drain all already-queued completions in one pass + tuids = [tuid] + while True: + nxt = self.batch.poll(timeout=0) + if nxt is None: + break + tuids.append(nxt) + completed = [] - for tuid in batch_tuids: - entry = self._monitored_batches.get(tuid) + for t in tuids: + entry = self._monitored_batches.pop(t, None) if entry is None: - continue + continue # user-cancelled: already delivered as CANCELED batch_task, uid = entry - # Wait until the Batch client compiler has placed this task on a manager. - # manager_idx is None in the window between submission and compile. - manager_idx = batch_task.core.manager_idx - if manager_idx is None: - continue - - # SCHEDULER_MANAGER_IDX == -1 maps to shard 0; subnode managers map 1-to-1. - ddict_shard_idx = 0 if manager_idx == -1 else manager_idx - try: - result, tb, raised, stdout, stderr = self.batch.results_ddict.manager( - ddict_shard_idx - )[tuid] - except KeyError: + # Block until the result is fully committed to the DDict. + # poll() and the DDict write are not atomic; block=True here + # ensures get_stdout(block=False) is safe afterward. + result = batch_task.get(block=True) + raised, tb = False, None + except TaskCancelledError: + # Dragon-side timeout / upstream cancel (not a user cancel_task call) + task_info = self._task_registry.pop(uid, None) + if task_info: + self._callback_func(task_info["description"], "CANCELED") continue + except Exception as exc: + result, raised, tb = exc, True, batch_task.traceback - self._monitored_batches.pop(tuid, None) + task_reg = self._task_registry.get(uid) + is_native_function = ( + task_reg.get("is_native_function", False) if task_reg else False + ) + if not is_native_function: + # Dragon returns exit codes as int/list[int] for process/job tasks. + # Normalize non-zero to raised=True so Rhapsody delivers FAILED. + if isinstance(result, int) and result != 0: + result = SystemExit(result) + raised = True + elif isinstance(result, list) and any( + c != 0 for c in result if isinstance(c, int) + ): + result = SystemExit(result) + raised = True + + capture_stdio = ( + task_reg.get("description", {}).get("capture_stdio") if task_reg else False + ) + if capture_stdio: + stdout = batch_task.stdout_path or "" + stderr = batch_task.stderr_path or "" + else: + try: + stdout = batch_task.get_stdout(block=False) or "" + stderr = batch_task.get_stderr(block=False) or "" + except OSError: + stdout, stderr = "", "" completed.append((uid, result, tb, raised, stdout, stderr)) - # One cross-thread wakeup for the entire sweep batch if completed: self._loop.call_soon_threadsafe(self._deliver_batch, completed) - # Small sleep after each full sweep to prevent tight loop - time.sleep(0.005) - except Exception as e: self.logger.exception(f"Critical error in monitor loop: {e}") - time.sleep(0.1) self.logger.debug("Dragon batch monitor loop stopped") @@ -3239,7 +265,7 @@ def _deliver_batch(self, completions: list) -> None: """Deliver a batch of completed tasks. Runs on the asyncio event loop (via call_soon_threadsafe). - Called once per monitor sweep with all tasks that completed in that sweep, reducing cross- + Called once per monitor drain with all tasks that completed in that sweep, reducing cross- thread wakeups from O(tasks) to O(sweeps). """ for uid, result, tb, raised, stdout, stderr in completions: @@ -3250,17 +276,15 @@ def _deliver_batch(self, completions: list) -> None: self._cancelled_tasks.discard(uid) continue task_desc = task_info["description"] - stdout_path = task_info.get("stdout_path") - stderr_path = task_info.get("stderr_path") if raised: task_desc["exception"] = result - task_desc["stderr"] = stderr_path if stderr_path else (tb if tb else str(result)) - task_desc["stdout"] = stdout_path or stdout or "" + task_desc["stderr"] = stderr if stderr else (tb if tb else str(result)) + task_desc["stdout"] = stdout self._callback_func(task_desc, "FAILED") else: task_desc["return_value"] = result - task_desc["stdout"] = stdout_path or stdout or "" - task_desc["stderr"] = stderr_path or stderr or "" + task_desc["stdout"] = stdout + task_desc["stderr"] = stderr self._callback_func(task_desc, "DONE") async def submit_tasks(self, tasks: list[dict]) -> None: @@ -3316,18 +340,16 @@ async def build_task(self, task: dict): """Translate AsyncFlow task to Dragon Batch task. Translation Priority (in order): - 1. If process_templates (list) provided → Job mode (ignore type='mpi', ignore ranks) [function/executable] + 1. If process_templates (list) provided → Job mode [function/executable] 2. If process_template (single) provided → Process mode [function/executable] - 3. If type='mpi' AND ranks provided (no templates) → Job mode (auto-build) [function/executable] - 4. If is_function (no templates, no MPI) → Function mode (native) [function only] - 5. If is_executable (no templates, no MPI) → Process mode (auto-build) [executable only] + 3. If is_function (no templates) → Function mode (native) [function only] + 4. If is_executable (no templates) → Process mode (auto-build) [executable only] Execution Modes: - Function Native: batch.function() - direct Python function call - Function Process: batch.process() - function wrapped in ProcessTemplate - - Function Job: batch.job() - function in MPI job with multiple ranks - Executable Process: batch.process() - single executable process - - Executable Job: batch.job() - executable in MPI job with multiple ranks + - Executable Job: batch.job() - executable in MPI job via process_templates Setting cwd (working directory): Pass ``cwd`` inside ``process_template`` or each entry of ``process_templates`` @@ -3368,108 +390,68 @@ def target(*a, **kw): # Get template configs once process_templates_config = backend_kwargs.get("process_templates") process_template_config = backend_kwargs.get("process_template") + is_native_function = ( + is_function and process_templates_config is None and process_template_config is None + ) - # When capture_stdio is requested for subprocess tasks, write a shell script that - # redirects both stdout and stderr to files, then invoke it as `bash script.sh`. - # Dragon's PIPE for the bash wrapper sees an empty stream → immediate EOF → no hang. - redirect = task.get("capture_stdio") and not is_function - - # Dragon 0.13.2 unconditionally set template.stdout = Popen.PIPE for every - # process task, so stdout was always captured. Dragon 0.14.0 removed that — - # needs_output is only True when the template explicitly carries stdout=Popen.PIPE. - # Restore that behavior for subprocess tasks: - # • no capture_stdio redirect → stdout=Popen.PIPE (direct capture via DDict) - # • capture_stdio redirect → stdout=None (bash script writes to file; - # PIPE would just see EOF) - # Functions capture stdout via io.StringIO internally; no PIPE needed. - stdout_pipe = None if (redirect or is_function) else Popen.PIPE - stdout_path = stderr_path = script_path = None - if redirect: + # Compute per-task stdio paths when capture is requested. + # Dragon creates parent directories automatically. + # None = forward output to console (no file capture). + stdout_path = stderr_path = None + if task.get("capture_stdio"): stdout_path = os.path.join(self._work_dir, f"{uid}.stdout") stderr_path = os.path.join(self._work_dir, f"{uid}.stderr") - script_path = os.path.join(self._work_dir, f"{uid}.sh") - cmd_line = " ".join( - [shlex.quote(str(target))] + [shlex.quote(str(a)) for a in task_args] - ) - with open(script_path, "w") as _f: - _f.write( - f"#!/usr/bin/bash\n" - f"{cmd_line}" - f" 1>{shlex.quote(stdout_path)}" - f" 2>{shlex.quote(stderr_path)}\n" - ) - target = "/bin/bash" - task_args = (script_path,) def _build_process_template_kwargs(template_cfg: dict[str, Any]) -> dict[str, Any]: - """Build ProcessTemplate kwargs: stdout_pipe as default; user config overrides.""" - return {"stdout": stdout_pipe, **template_cfg, "args": task_args, "kwargs": task_kwargs} + """Build ProcessTemplate kwargs from user config.""" + return {**template_cfg, "args": task_args, "kwargs": task_kwargs} + + # Build task options once and reuse across all submission paths. + # Dragon 0.14.1-rc requires metadata (name, timeout, stdio paths) to go through + # Batch.options() rather than as direct kwargs to process/job/function. + task_opts = self.batch.options( + name=name, + timeout=timeout, + stdout=stdout_path, + stderr=stderr_path, + ) # Single decision tree - no redundant checks if process_templates_config is not None: # Priority 1: Job with user templates - # stdout_pipe is the default; user's tc dict can override it. process_templates = [ - ( - nranks, - ProcessTemplate(target, **_build_process_template_kwargs(tc)), - ) + (nranks, ProcessTemplate(target, **_build_process_template_kwargs(tc))) for nranks, tc in process_templates_config ] - batch_task = self.batch.job(process_templates, name=name, timeout=timeout) - execution_mode = "job" + batch_task = task_opts.job(process_templates) elif process_template_config is not None: # Priority 2: Process with user template - # stdout_pipe is the default; user's process_template_config can override it. - batch_task = self.batch.process( - ProcessTemplate(target, **_build_process_template_kwargs(process_template_config)), - name=name, - timeout=timeout, - ) - execution_mode = "process" - - elif backend_kwargs.get("type") == "mpi": - # Priority 3: Job auto-build - batch_task = self.batch.job( - [ - ( - backend_kwargs.get("ranks", 1), - ProcessTemplate(target, **_build_process_template_kwargs({})), - ) - ], - name=name, - timeout=timeout, + batch_task = task_opts.process( + ProcessTemplate(target, **_build_process_template_kwargs(process_template_config)) ) - execution_mode = "job" elif is_function: - # Priority 4: Function native - batch_task = self.batch.function( - target, *task_args, name=name, timeout=timeout, **task_kwargs - ) - execution_mode = "function" + # Priority 3: Function native (batch.function — returns Python value, not exit code) + batch_task = task_opts.function(target, *task_args, **task_kwargs) else: - # Priority 5: Executable process auto-build - batch_task = self.batch.process( - ProcessTemplate(target, **_build_process_template_kwargs({})), - name=name, - timeout=timeout, + # Priority 4: Executable process auto-build + batch_task = task_opts.process( + ProcessTemplate(target, **_build_process_template_kwargs({})) ) - execution_mode = "process" # Register and return self._task_registry[uid] = { "uid": uid, "description": task, "batch_task": batch_task, - "stdout_path": stdout_path, - "stderr_path": stderr_path, - "script_path": script_path, + "is_native_function": is_native_function, } - self.logger.debug(f"Created {execution_mode} task: {uid}") + self.logger.debug( + f"Created {'function' if is_native_function else 'process/job'} task: {uid}" + ) return batch_task @@ -3512,9 +494,8 @@ async def cancel_task(self, uid: str) -> bool: # Task completed naturally between the executor call and here — not cancelled. return False task_desc = registry_entry["description"] - # Stop the monitor loop from polling this task's DDict entry. Without this, - # the blocking DDict IPC for a cancelled task can stall the monitor thread - # and delay result delivery for subsequently submitted tasks. + # Eagerly pop from _monitored_batches so when the cancelled tuid arrives from + # poll(), the None entry is skipped without delivering a spurious result. self._monitored_batches.pop(batch_task.uid, None) self._callback_func(task_desc, "CANCELED") @@ -3528,7 +509,7 @@ async def shutdown(self) -> None: # Set backend state to SHUTDOWN self._backend_state = BackendMainStates.SHUTDOWN self.logger.debug(f"Backend state set to: {self._backend_state.value}") - self.logger.info("Shutting down V3 backend") + self.logger.info("Shutting down Dragon backend") self._shutdown_event.set() # Wait for monitor thread to finish if it exists @@ -3559,11 +540,21 @@ async def shutdown(self) -> None: self._task_registry.clear() self._state = "idle" - self.logger.info("Shutdown V3 complete") + self.logger.info("Dragon backend shutdown complete") + + def wait(self, timeout: Optional[float] = None) -> None: + """Wait for all submitted tasks to complete. - # Batch features - def fence(self): - self.batch.fence() + Delegates to Dragon's Batch.fence(). Tasks submitted before this call + will all complete before any tasks submitted after it are started. + + Args: + timeout: Maximum seconds to wait. Uses Dragon's default if not set. + """ + if timeout is not None: + self.batch.fence(timeout=timeout) + else: + self.batch.fence() def create_ddict(self, *args, **kwargs): from dragon.data.ddict.ddict import DDict @@ -3575,6 +566,20 @@ async def create( cls, batch_kwargs: Optional[dict] = None, ): - """Create and initialize a DragonExecutionBackendV3.""" + """Create and initialize a DragonExecutionBackend.""" backend = cls(batch_kwargs=batch_kwargs) return await backend + + +def __getattr__(name: str): + import warnings + + if name == "DragonExecutionBackendV3": + warnings.warn( + "DragonExecutionBackendV3 is deprecated and will be removed in a future release. " + "Use DragonExecutionBackend instead.", + DeprecationWarning, + stacklevel=2, + ) + return DragonExecutionBackend + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/rhapsody/backends/execution/radical_pilot.py b/src/rhapsody/backends/execution/radical_pilot.py index 1244086..38be412 100644 --- a/src/rhapsody/backends/execution/radical_pilot.py +++ b/src/rhapsody/backends/execution/radical_pilot.py @@ -427,7 +427,7 @@ def backend_callback(rp_task, state) -> None: try: if rp_task.mode == rp.TASK_SERVICE and state == rp.AGENT_EXECUTING: - service_callback = service_ready_callback + service_callback = service_ready_callback # noqa: F841 # Get the original Task object original_task = self.tasks.get(rp_task.uid) @@ -457,7 +457,7 @@ def backend_callback(rp_task, state) -> None: original_task["exception"] = rp_task_dict["exception"] # Call the registered callback with the original Task object - func(original_task, state, service_callback=service_callback) + func(original_task, state) except Exception: self.logger.exception( f"Backend callback failed for task {getattr(rp_task, 'uid', None)} " diff --git a/src/rhapsody/logger.py b/src/rhapsody/logger.py index 5efe73f..ae8c6d2 100644 --- a/src/rhapsody/logger.py +++ b/src/rhapsody/logger.py @@ -50,7 +50,7 @@ def enable_logging( import logging import rhapsody rhapsody.enable_logging(logging.DEBUG) - backend = await rhapsody.get_backend("dragon_v3") + backend = await rhapsody.get_backend("dragon") """ if format_string is None: format_string = "%(asctime)s | %(levelname)-8s | [%(name)s] | %(message)s" diff --git a/tests/integration/test-hpc/dragon/conftest.py b/tests/integration/test-hpc/dragon/conftest.py index 2ceeaee..04e559a 100644 --- a/tests/integration/test-hpc/dragon/conftest.py +++ b/tests/integration/test-hpc/dragon/conftest.py @@ -106,20 +106,20 @@ def _apply_topology_skips(request, topology, gpu_nodes): @pytest.fixture(scope="module") async def dragon_backend(request): - """Initialise a DragonExecutionBackendV3 for the test module. + """Initialise a DragonExecutionBackend for the test module. Scale tests can override batch_kwargs via indirect parametrisation or by setting the RHAPSODY_DDICT_MEM_GB env var (default: 2 GiB/node). """ from dragon.native.machine import System - from rhapsody.backends import DragonExecutionBackendV3 + from rhapsody.backends import DragonExecutionBackend num_nodes = len(list(System().nodes)) ddict_mem_gb = int(os.environ.get("RHAPSODY_DDICT_MEM_GB", 2)) ddict_mem = ddict_mem_gb * num_nodes * (1024**3) - backend = await DragonExecutionBackendV3(batch_kwargs={"results_ddict_mem": ddict_mem}) + backend = await DragonExecutionBackend(batch_kwargs={"results_ddict_mem": ddict_mem}) yield backend await backend.shutdown() diff --git a/tests/integration/test-hpc/dragon/test_scale.py b/tests/integration/test-hpc/dragon/test_scale.py index daa7207..28a6d8f 100644 --- a/tests/integration/test-hpc/dragon/test_scale.py +++ b/tests/integration/test-hpc/dragon/test_scale.py @@ -62,16 +62,16 @@ def _scale_skip(): @pytest.fixture(scope="module") async def scale_backend(): - """DragonExecutionBackendV3 with enlarged results DDict for scale tests.""" + """DragonExecutionBackend with enlarged results DDict for scale tests.""" from dragon.native.machine import System - from rhapsody.backends import DragonExecutionBackendV3 + from rhapsody.backends import DragonExecutionBackend num_nodes = len(list(System().nodes)) ddict_gb = int(os.environ.get("RHAPSODY_DDICT_MEM_GB", 4)) ddict_mem = ddict_gb * num_nodes * (1024**3) - backend = await DragonExecutionBackendV3( + backend = await DragonExecutionBackend( batch_kwargs={ "results_ddict_mem": ddict_mem, "scheduler_workers": num_nodes * 4, diff --git a/tests/integration/test-hpc/dragon/test_throughput.py b/tests/integration/test-hpc/dragon/test_throughput.py new file mode 100644 index 0000000..2b2ec3d --- /dev/null +++ b/tests/integration/test-hpc/dragon/test_throughput.py @@ -0,0 +1,76 @@ +import argparse +import asyncio +import time + +import rhapsody +from rhapsody.api import ComputeTask +from rhapsody.api import Session +from rhapsody.backends import DragonExecutionBackend + +width_num_tasks = 10 +width_throughput = 20 + + +def no_op(): + pass + + +async def run_tasks(session: Session, num_tasks: int, warmup: bool = False, backend=None) -> None: + start_time = time.time() + + tasks = [ComputeTask(function=no_op) for _ in range(num_tasks)] + futures = await session.submit_tasks(tasks) + + backend.wait() # high performance wait loop + + end_time = time.time() + + if not warmup: + runtime = end_time - start_time + throughput = num_tasks / runtime + print( + f"{num_tasks:<{width_num_tasks}} {throughput:<{width_throughput}}", + flush=True, + ) + + backend.batch.clear_results() + + +async def run_bench(session: Session, min_tasks: int, max_tasks: int, backend=None) -> None: + print("task throughput benchmark", flush=True) + print("-------------------------", flush=True) + + await run_tasks(session, min_tasks, warmup=True, backend=backend) + + num_tasks_str = "num tasks" + throughput_str = "throughput [tasks/s]" + print( + f"{num_tasks_str.ljust(width_num_tasks)} {throughput_str.ljust(width_throughput)}", + flush=True, + ) + + num_tasks = min_tasks + + while num_tasks <= max_tasks: + await run_tasks(session, num_tasks, backend=backend) + num_tasks *= 2 + + print("", flush=True) + + +async def main(min_tasks: int, max_tasks: int) -> None: + backend = await DragonExecutionBackend(batch_kwargs={"task_logs": False}) + async with Session(backends=[backend]) as session: + await run_bench(session, min_tasks, max_tasks, backend=backend) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Task throughput benchmark") + parser.add_argument("--min_tasks", type=int, default=4, help="minimum number of tasks to run") + parser.add_argument("--max_tasks", type=int, default=128, help="maximum number of tasks to run") + args = parser.parse_args() + + if args.max_tasks < args.min_tasks: + args.max_tasks = args.min_tasks + + asyncio.run(main(args.min_tasks, args.max_tasks)) diff --git a/tests/integration/test_backend_functionality.py b/tests/integration/test_backend_functionality.py index e8cc307..b59e3d7 100644 --- a/tests/integration/test_backend_functionality.py +++ b/tests/integration/test_backend_functionality.py @@ -6,10 +6,10 @@ Test Mode Control: ------------------ Set RHAPSODY_TEST_MODE environment variable to control which backends are tested: -- RHAPSODY_TEST_MODE=regular (default): Tests non-Dragon backends (concurrent, dask, radical_pilot) +- RHAPSODY_TEST_MODE=regular (default): Tests local backends (concurrent, dask) Run with: pytest tests/integration/test_backend_functionality.py -- RHAPSODY_TEST_MODE=dragon: Tests Dragon backends only (dragon_v1, dragon_v2, dragon_v3) +- RHAPSODY_TEST_MODE=dragon: Tests Dragon backends only (dragon) Run with: dragon pytest tests/integration/test_backend_functionality.py - RHAPSODY_TEST_MODE=all: Tests all backends @@ -47,15 +47,59 @@ def get_available_backends_for_mode() -> list[str]: if mode == "dragon": # Only Dragon backends - return [name for name in available if name.startswith("dragon_")] + return [name for name in available if name == "dragon" or name.startswith("dragon_")] elif mode == "regular": - # Exclude Dragon backends - return [name for name in available if not name.startswith("dragon_")] + # Exclude Dragon and radical_pilot (both require dedicated HPC infrastructure) + return [ + name + for name in available + if name not in {"dragon", "radical_pilot"} and not name.startswith("dragon_") + ] else: # mode == "all" # All available backends return available +def _attach_state_collector(backend) -> dict[str, str]: + """Register a terminal-state callback on the backend and return the shared dict. + + Use this when the callback must be registered BEFORE submit_tasks() is called (e.g. when the + backend fires FAILED synchronously during build_task, as Dragon does when the executable path is + invalid). Call _wait_for_states() afterward. + """ + final_states: dict[str, str] = {} + + def on_state(task_desc: dict, state: str) -> None: + if state in {"DONE", "FAILED", "CANCELED"}: + final_states[task_desc["uid"]] = state + + backend.register_callback(on_state) + return final_states + + +async def _wait_for_states( + final_states: dict[str, str], task_uids: list[str], timeout: float = 10.0 +) -> dict[str, str]: + """Wait until all uids appear in final_states or timeout expires.""" + deadline = asyncio.get_event_loop().time() + timeout + while len(final_states) < len(task_uids) and asyncio.get_event_loop().time() < deadline: + await asyncio.sleep(0.1) + return final_states + + +async def _collect_terminal_states( + backend, task_uids: list[str], timeout: float = 10.0 +) -> dict[str, str]: + """Register a capturing callback, submit wait, return uid→state map. + + Convenience wrapper for the common case where callback registration happens AFTER + submit_tasks(). For backends that deliver FAILED synchronously during submission use + _attach_state_collector() + _wait_for_states() instead. + """ + final_states = _attach_state_collector(backend) + return await _wait_for_states(final_states, task_uids, timeout) + + async def setup_test_backend(backend_name=None): """Helper to set up a backend with proper initialization and callback. @@ -72,7 +116,7 @@ async def setup_test_backend(backend_name=None): try: # Initialize backend with proper resources - backend = rhapsody.get_backend(backend_name) + backend = await rhapsody.get_backend(backend_name) # Handle async initialization for backends that need it if hasattr(backend, "__await__"): @@ -95,238 +139,148 @@ class TestBackendFunctionality: """Test suite for backend functionality required by AsyncFlow.""" async def test_backend_task_cancellation(self): - """Test task cancellation functionality.""" + """Submit a long-running task, cancel it, and assert the CANCELED state is delivered.""" backend = await setup_test_backend("dask" if get_test_mode() == "regular" else None) + if not hasattr(backend, "cancel_task"): + await backend.shutdown() + pytest.skip(f"{type(backend).__name__} does not support cancel_task") + + task = rhapsody.ComputeTask(executable="/bin/sleep", arguments=["30"]) + final_states: dict[str, str] = {} + + def on_state(task_desc: dict, state: str) -> None: + if state in {"DONE", "FAILED", "CANCELED"}: + final_states[task_desc["uid"]] = state + + backend.register_callback(on_state) + try: - # Create a long-running task - tasks = [ - rhapsody.ComputeTask( - executable="/bin/sleep", - arguments=["30"], # 30 second sleep - ) - ] + await backend.submit_tasks([task]) - # Submit task (Task objects accepted directly) - await backend.submit_tasks(tasks) + # Allow the task to reach RUNNING state before cancelling + await asyncio.sleep(0.5) - # For most backends, cancellation may not be implemented - # Just verify the backend responds to cancel requests - try: - # Try to cancel the task (may not be implemented) - if hasattr(backend, "cancel_tasks"): - # Cancellation may not be supported - canceled = {} - assert isinstance(canceled, (dict, list, bool)) - except (NotImplementedError, AttributeError): - # Cancellation not implemented - that's okay - pass - - # Get current state (may not track individual states) - if hasattr(backend, "get_task_states"): - # Check if backend supports state tracking - - if hasattr(backend, "get_task_states_map"): - states = backend.get_task_states_map() - - else: - states = {} - assert isinstance(states, dict) + cancelled = await backend.cancel_task(task.uid) + assert cancelled is True, "cancel_task must return True when cancellation succeeds" + + # Wait up to 5 s for the CANCELED callback + for _ in range(50): + if task.uid in final_states: + break + await asyncio.sleep(0.1) + + assert task.uid in final_states, "No terminal-state callback received within 5 s" + assert final_states[task.uid] == "CANCELED", ( + f"Expected CANCELED, got {final_states[task.uid]}" + ) finally: await backend.shutdown() async def test_backend_resource_management(self): - """Test backend resource allocation and management.""" - # Test different backend configurations based on mode + """Each available backend can submit and complete a task.""" available_names = get_available_backends_for_mode() if not available_names: - mode = get_test_mode() - pytest.skip(f"No backends available for test mode: {mode}") + pytest.skip(f"No backends available for test mode: {get_test_mode()}") - backends = [] + backend_name = available_names[0] + backend = await rhapsody.get_backend(backend_name) + task = rhapsody.ComputeTask( + executable="/bin/echo", arguments=[f"resource-test {backend_name}"] + ) try: - # Test each available backend - for _i, backend_name in enumerate(available_names): - try: - backend = rhapsody.get_backend(backend_name) - backends.append(backend) - - # Submit a simple task to verify backend works - tasks = [ - rhapsody.ComputeTask( - executable="/bin/echo", - arguments=[f"Testing backend {backend_name}"], - ) - ] - - await backend.submit_tasks(tasks) - - # Quick wait to see if it processes - try: - # Just verify the task was submitted successfully - await asyncio.sleep(0.1) # Brief wait for processing - except Exception: # noqa: S110 - pass # Any issues are okay for this test - - except Exception as e: - # Skip backends that can't be initialized - continue - + await backend.submit_tasks([task]) + states = await _collect_terminal_states(backend, [task.uid], timeout=10.0) + assert states.get(task.uid) == "DONE", ( + f"Backend {backend_name}: expected DONE, got {states.get(task.uid)!r}" + ) finally: - # Cleanup all backends - for backend in backends: - try: - await backend.shutdown() - except Exception: # noqa: S110 - pass + await backend.shutdown() async def test_backend_state_transitions(self): - """Test task state transitions through backend lifecycle.""" + """Task progresses from submission to DONE state.""" backend = await setup_test_backend() + task = rhapsody.ComputeTask(executable="/bin/echo", arguments=["state-transition"]) try: - tasks = [ - rhapsody.ComputeTask( - executable="/bin/echo", - arguments=["State transition test"], - ) - ] - - # Submit tasks - await backend.submit_tasks(tasks) - - # For this test, just verify submission completed - # State tracking is complex and varies by backend - assert True # Task submission successful + await backend.submit_tasks([task]) + states = await _collect_terminal_states(backend, [task.uid], timeout=10.0) + assert task.uid in states, "Task never reached a terminal state within 10 s" + assert states[task.uid] == "DONE", f"Expected DONE, got {states[task.uid]!r}" finally: await backend.shutdown() async def test_backend_error_recovery(self): - """Test backend behavior with failing tasks.""" + """A failing task reaches FAILED while a good concurrent task reaches DONE.""" backend = await setup_test_backend() + good_task = rhapsody.ComputeTask(executable="/bin/echo", arguments=["ok"]) + bad_task = rhapsody.ComputeTask(executable="/nonexistent/command", arguments=[]) + task_uids = [good_task.uid, bad_task.uid] - try: - # Mix of good and bad tasks - tasks = [ - rhapsody.ComputeTask( - executable="/bin/echo", - arguments=["This should work"], - ), - rhapsody.ComputeTask( - executable="/nonexistent/command", - arguments=["This will fail"], - ), - ] - - # Submit tasks - await backend.submit_tasks(tasks) - - # Wait for completion - # Backend submission completed - assert True + # Register the callback BEFORE submit_tasks: some backends (Dragon) validate + # the executable at build_task time and fire FAILED synchronously during submission. + final_states = _attach_state_collector(backend) - # Task submission completed successfully - - # Error recovery test completed successfully - assert True + try: + await backend.submit_tasks([good_task, bad_task]) + await _wait_for_states(final_states, task_uids, timeout=10.0) + assert final_states.get(good_task.uid) == "DONE", ( + f"Good task: expected DONE, got {final_states.get(good_task.uid)!r}" + ) + assert final_states.get(bad_task.uid) == "FAILED", ( + f"Bad task: expected FAILED, got {final_states.get(bad_task.uid)!r}" + ) finally: await backend.shutdown() async def test_backend_batch_operations(self): - """Test backend with large batches of tasks.""" + """All tasks in a 10-task batch reach DONE.""" backend = await setup_test_backend() + tasks = [ + rhapsody.ComputeTask(executable="/bin/echo", arguments=[f"batch-{i}"]) + for i in range(10) + ] + task_uids = [t.uid for t in tasks] try: - # Create batch of tasks - batch_size = 10 - tasks = [ - rhapsody.ComputeTask( - executable="/bin/echo", - arguments=[f"Batch task {i}"], - ) - for i in range(batch_size) - ] - - # Submit batch await backend.submit_tasks(tasks) + states = await _collect_terminal_states(backend, task_uids, timeout=30.0) - # Wait for all to complete - task_uids = [task.uid for task in tasks] - # Backend submission completed - assert True - - # Batch processing completed successfully - - # Batch processing completed successfully - assert True - + not_done = [uid for uid in task_uids if states.get(uid) != "DONE"] + assert not not_done, f"{len(not_done)}/10 tasks did not reach DONE: {not_done}" finally: await backend.shutdown() async def test_backend_async_patterns(self): - """Test that backends work properly with asyncio patterns.""" - # Test only dask backend if in regular mode, skip if in dragon mode + """Three tasks submitted concurrently via asyncio.gather all reach DONE.""" mode = get_test_mode() - if mode == "dragon": - pytest.skip("Async pattern test not applicable for Dragon mode") + pytest.skip("Async pattern test uses dask only") - available_backends = rhapsody.discover_backends() - if not available_backends.get("dask", False): + try: + backend = await rhapsody.get_backend("dask") + except (ImportError, Exception): pytest.skip("Dask backend not available for async pattern testing") - backends = [] - try: - # Only test with dask backend to avoid concurrency issues with radical_pilot - backend = rhapsody.get_backend("dask") + tasks = [ + rhapsody.ComputeTask(executable="/bin/echo", arguments=[f"async-{i}"]) + for i in range(3) + ] - # Handle async initialization for backends that need it - try: - backend = await backend # type: ignore[misc] - except TypeError: - # Backend doesn't support await, that's fine - pass - - # Register a callback function - def task_callback(task: dict, state: str) -> None: - # Simple callback that does nothing but satisfies the interface - pass - - backend.register_callback(task_callback) - backends.append(backend) - - # Submit tasks to backend - async def submit_to_backend(backend, backend_idx): - tasks = [ - rhapsody.ComputeTask( - executable="/bin/echo", - arguments=[f"Async test {backend_idx}"], - ) - ] - - await backend.submit_tasks(tasks) - return True # Task submitted successfully - - # Execute multiple submissions concurrently to the same backend - results = await asyncio.gather(*[submit_to_backend(backends[0], i) for i in range(3)]) - - # All should complete without interference - assert len(results) == 3 - for result in results: - # Results are boolean True indicating successful submission - assert result is True + # Submit all three concurrently + await asyncio.gather(*[backend.submit_tasks([t]) for t in tasks]) + states = await _collect_terminal_states(backend, [t.uid for t in tasks], timeout=30.0) + not_done = [t.uid for t in tasks if states.get(t.uid) != "DONE"] + assert not not_done, f"Tasks did not reach DONE: {not_done}" finally: - # Cleanup all backends - await asyncio.gather( - *[backend.shutdown() for backend in backends], return_exceptions=True - ) + await backend.shutdown() @pytest.mark.asyncio @@ -334,79 +288,68 @@ class TestBackendCompatibility: """Test compatibility between different backends.""" async def test_backend_interface_consistency(self): - """Test that all backends implement the same interface.""" + """All instantiable backends expose the required interface AND can execute a task.""" available_names = get_available_backends_for_mode() - backend_instances = {} + tested = 0 - try: - # Create instances of all available backends for current mode - for name in available_names: - try: - backend = await rhapsody.get_backend(name) - backend_instances[name] = backend - except Exception as e: - print(f"Could not create {name} backend: {e}") - - # Test that all backends have the core interface - for _name, backend in backend_instances.items(): - # Check required methods exist - assert hasattr(backend, "submit_tasks") - assert hasattr(backend, "shutdown") - - # Test they're callable - assert callable(backend.submit_tasks) - assert callable(backend.shutdown) - - # Optional methods may or may not be present - # Different backends have different capabilities + for name in available_names: + try: + backend = await rhapsody.get_backend(name) + except (ImportError, Exception): + continue # skip backends whose runtime deps are not installed - finally: - # Cleanup - for backend in backend_instances.values(): - try: - await backend.shutdown() - except Exception: # noqa: S110 - pass + # Structural check + assert hasattr(backend, "submit_tasks") and callable(backend.submit_tasks), ( + f"{name}: submit_tasks not callable" + ) + assert hasattr(backend, "shutdown") and callable(backend.shutdown), ( + f"{name}: shutdown not callable" + ) - async def test_backend_switching(self): - """Test switching between different backends in same workflow.""" - # Create a simple task that should work on any backend - test_task_template = { - "executable": "/bin/echo", - "arguments": ["Backend switching test"], - } + # Behavioural check — actually invoke submit_tasks and assert result + task = rhapsody.ComputeTask( + executable="/bin/echo", arguments=[f"interface-check-{name}"] + ) + try: + await backend.submit_tasks([task]) + states = await _collect_terminal_states(backend, [task.uid], timeout=10.0) + assert states.get(task.uid) == "DONE", ( + f"Backend {name}: expected DONE, got {states.get(task.uid)!r}" + ) + tested += 1 + finally: + await backend.shutdown() + if tested == 0: + pytest.skip("No backends could be instantiated in this environment") + + async def test_backend_switching(self): + """Each instantiable backend can independently submit and complete an echo task.""" available_names = get_available_backends_for_mode() - results = {} + if not available_names: + pytest.skip(f"No backends available for test mode: {get_test_mode()}") + + tested = 0 for name in available_names: try: - # Create backend backend = await rhapsody.get_backend(name) + except (ImportError, Exception): + continue # skip backends whose runtime deps are not installed - # Create a fresh task for this backend - test_task = rhapsody.ComputeTask(**test_task_template) - - # Submit task - await backend.submit_tasks([test_task]) - - # Wait for completion - # Backend submission completed - assert True - results[name] = True - - # Cleanup + task = rhapsody.ComputeTask(executable="/bin/echo", arguments=[f"switch-test-{name}"]) + try: + await backend.submit_tasks([task]) + states = await _collect_terminal_states(backend, [task.uid], timeout=10.0) + assert states.get(task.uid) == "DONE", ( + f"Backend {name}: expected DONE, got {states.get(task.uid)!r}" + ) + tested += 1 + finally: await backend.shutdown() - except Exception as e: - print(f"Backend {name} failed: {e}") - results[name] = None - - # Should have at least one backend working - working_backends = get_available_backends_for_mode() - assert len(working_backends) > 0 - assert working_backends[0] in results - print(f"Backend switching results: {results}") + if tested == 0: + pytest.skip("No backends could be instantiated in this environment") async def main(): @@ -422,7 +365,7 @@ async def main(): return backend_name = available_names[0] - backend = rhapsody.get_backend(backend_name) + backend = await rhapsody.get_backend(backend_name) tasks = [ rhapsody.ComputeTask( executable="/bin/echo", diff --git a/tests/unit/telemetry/test_adapters_dragon.py b/tests/unit/telemetry/test_adapters_dragon.py index 6317ad8..4725144 100644 --- a/tests/unit/telemetry/test_adapters_dragon.py +++ b/tests/unit/telemetry/test_adapters_dragon.py @@ -85,7 +85,7 @@ def test_noop_when_telemetry_level_zero(): loop = asyncio.new_event_loop() adapter = DragonTelemetryAdapter( session_id="test-session", - backend_name="dragon_v3", + backend_name="dragon", interval=0.5, ) manager = _FakeManager() @@ -130,7 +130,7 @@ def test_adapter_emits_resource_update(): adapter = DragonTelemetryAdapter( session_id="test-dragon-adapter", - backend_name="dragon_v3", + backend_name="dragon", interval=1.0, ) manager = _FakeManager() @@ -157,7 +157,7 @@ async def _run_adapter(): ev = updates[0] assert ev.session_id == "test-dragon-adapter" - assert ev.backend == "dragon_v3" + assert ev.backend == "dragon" assert ev.node_id is not None and ev.node_id != "" # cpu_percent and memory_percent must be collected at level=1 @@ -191,7 +191,7 @@ def test_adapter_stops_cleanly(): adapter = DragonTelemetryAdapter( session_id="test-stop", - backend_name="dragon_v3", + backend_name="dragon", interval=1.0, ) manager = _FakeManager() diff --git a/tests/unit/telemetry/test_otel_contract.py b/tests/unit/telemetry/test_otel_contract.py index 3e14d4e..52a5377 100644 --- a/tests/unit/telemetry/test_otel_contract.py +++ b/tests/unit/telemetry/test_otel_contract.py @@ -123,7 +123,7 @@ def _make_dragon_adapter(session_id: str = "contract-session") -> Any: return DragonTelemetryAdapter( session_id=session_id, - backend_name="dragon_v3", + backend_name="dragon", interval=1.0, ) diff --git a/tests/unit/test_backend_execution_dragon.py b/tests/unit/test_backend_execution_dragon.py index 5aa8551..89a71d5 100644 --- a/tests/unit/test_backend_execution_dragon.py +++ b/tests/unit/test_backend_execution_dragon.py @@ -1,19 +1,15 @@ -"""Tests for Dragon execution backends (V1, V2, V3) using Session. +"""Tests for DragonExecutionBackend using Session. Structure --------- Shared behavior tests (1-15) - Parametrized across dragon_v1 / dragon_v2 / dragon_v3. Verify observable task - execution semantics that every Dragon backend must satisfy. + Parametrized against dragon. Verify observable task execution semantics. -V3 integration tests — process_template / process_templates routing - Use a dedicated ``session_v3`` fixture so only one V3 session is created, - rather than creating three parametrized sessions and skipping two. +Dragon integration tests — process_template / process_templates routing + Use a dedicated ``session_dragon`` fixture so only one V3 session is created. -V3 unit tests — constructor and internal methods - Verify the API surface introduced by the batch.py migration (new constructor - parameters, _deliver_batch, fence delegation) with Batch - mocked out. No Dragon cluster is required for these tests. +Dragon unit tests — constructor and internal methods + Verify the API surface with Batch mocked out. No Dragon cluster required. Run with: dragon python -m pytest tests/unit/test_backend_execution_dragon.py -v @@ -39,9 +35,9 @@ # ============================================================================ -@pytest.fixture(scope="module", params=["dragon_v1", "dragon_v2", "dragon_v3"]) +@pytest.fixture(scope="module", params=["dragon"]) def backend_name(request): - """Parametrize shared behavior tests across all Dragon backend versions.""" + """Backend name for shared behavior tests (Dragon backend only).""" return request.param @@ -55,32 +51,32 @@ async def session(backend_name): @pytest_asyncio.fixture(scope="module", loop_scope="module") -async def session_v3(): - """Session backed exclusively by DragonExecutionBackendV3. +async def session_dragon(): + """Session backed exclusively by DragonExecutionBackend. - Used by V3-specific tests so only one session is created instead of three. + Used by Dragon-specific tests so only one session is created instead of three. """ - backend_instance = await get_backend("dragon_v3") + backend_instance = await get_backend("dragon") session_instance = Session(backends=[backend_instance]) yield session_instance await session_instance.close() @pytest.fixture -def backend_v3(): - """DragonExecutionBackendV3 with Batch fully mocked — no Dragon cluster required. +def backend_dragon(): + """DragonExecutionBackend with Batch fully mocked — no Dragon cluster required. Suitable for unit tests that verify constructor wiring, internal callback helpers, and method delegation without running actual Dragon workers. """ - from rhapsody.backends.execution.dragon import DragonExecutionBackendV3 + from rhapsody.backends.execution.dragon import DragonExecutionBackend mock_batch = MagicMock() mock_batch.num_workers = 16 mock_batch.num_managers = 2 with patch("rhapsody.backends.execution.dragon.Batch", return_value=mock_batch): - backend = DragonExecutionBackendV3() + backend = DragonExecutionBackend() backend._callback_func = MagicMock() backend._loop = asyncio.new_event_loop() @@ -314,15 +310,22 @@ async def test_task_cancellation(session): @pytest.mark.asyncio(loop_scope="module") async def test_backend_state(session): - """Test backend state is queryable and non-null.""" + """Backend state is RUNNING while tasks are being processed and after they complete.""" backend = next(iter(session.backends.values())) - state = await backend.state() - assert state is not None - task = ComputeTask(executable="echo", arguments=["test"]) + task = ComputeTask(executable="echo", arguments=["state-test"]) await session.submit_tasks([task]) + + # state() must return "RUNNING" from the moment the first task was submitted + # (earlier tests in this module-scoped session have already triggered RUNNING) + state_mid = await backend.state() + assert state_mid == "RUNNING", f"Expected RUNNING during processing, got {state_mid!r}" + await session.wait_tasks([task]) + state_after = await backend.state() + assert state_after == "RUNNING", f"Expected RUNNING after completion, got {state_after!r}" + # ============================================================================ # Test 12: Multiple Submissions (Sequential) @@ -402,13 +405,13 @@ async def test_task_uid_uniqueness(session): # ============================================================================ -# V3 integration tests — per-task cwd and process_template routing +# Dragon integration tests — per-task cwd and process_template routing # ============================================================================ @pytest.mark.asyncio(loop_scope="module") -async def test_executable_with_cwd_via_process_template(session_v3): - """Test that cwd is honoured when set via process_template (V3 only).""" +async def test_executable_with_cwd_via_process_template(session_dragon): + """Test that cwd is honoured when set via process_template (Dragon backend only).""" import sys task = ComputeTask( @@ -417,114 +420,100 @@ async def test_executable_with_cwd_via_process_template(session_v3): task_backend_specific_kwargs={"process_template": {"cwd": "/tmp"}}, ) - await session_v3.submit_tasks([task]) - results = await session_v3.wait_tasks([task]) + await session_dragon.submit_tasks([task]) + results = await session_dragon.wait_tasks([task]) assert results[0].state == "DONE" assert "/tmp" in results[0].get("stdout", "") -@pytest.mark.asyncio(loop_scope="module") -async def test_process_template_cwd_built_and_passed(session_v3): - """Test A: process_template with cwd produces a ProcessTemplate with correct cwd.""" - from dragon.native.process import ProcessTemplate +@pytest.mark.asyncio +async def test_process_template_cwd_built_and_passed(backend_dragon): + """Test A: process_template with cwd produces a ProcessTemplate with correct cwd. - backend = session_v3.backends["dragon"] - captured = [] + Uses backend_dragon (mocked Batch) — no Dragon cluster required. The code calls + batch.options(...).process(pt); we capture pt via a mock proxy on batch.options. + """ + from dragon.native.process import ProcessTemplate - def capture(pt, **kw): - captured.append(pt) - return MagicMock() + captured_pt = [] + mock_proxy = MagicMock() + mock_proxy.process.side_effect = lambda pt: captured_pt.append(pt) or MagicMock() task = _make_task_dict(lambda: None, backend_specific={"process_template": {"cwd": "/tmp"}}) - with patch.object(backend.batch, "process", side_effect=capture): - await backend.build_task(task) + with patch.object(backend_dragon.batch, "options", return_value=mock_proxy): + await backend_dragon.build_task(task) - assert len(captured) == 1 - pt = captured[0] + assert len(captured_pt) == 1 + pt = captured_pt[0] assert isinstance(pt, ProcessTemplate) assert pt.cwd == "/tmp" -@pytest.mark.asyncio(loop_scope="module") -async def test_process_template_policy_gpu_affinity_built_and_passed(session_v3): +@pytest.mark.asyncio +async def test_process_template_policy_gpu_affinity_built_and_passed(backend_dragon): """Test B: process_template with policy(gpu_affinity) produces ProcessTemplate with correct policy.""" from dragon.infrastructure.policy import Policy from dragon.native.process import ProcessTemplate - backend = session_v3.backends["dragon"] - captured = [] - - def capture(pt, **kw): - captured.append(pt) - return MagicMock() + captured_pt = [] + mock_proxy = MagicMock() + mock_proxy.process.side_effect = lambda pt: captured_pt.append(pt) or MagicMock() policy = Policy(gpu_affinity=[0, 1, 2, 3]) task = _make_task_dict(lambda: None, backend_specific={"process_template": {"policy": policy}}) - with patch.object(backend.batch, "process", side_effect=capture): - await backend.build_task(task) + with patch.object(backend_dragon.batch, "options", return_value=mock_proxy): + await backend_dragon.build_task(task) - assert len(captured) == 1 - pt = captured[0] + assert len(captured_pt) == 1 + pt = captured_pt[0] assert isinstance(pt, ProcessTemplate) assert pt.policy is policy assert pt.policy.gpu_affinity == [0, 1, 2, 3] -@pytest.mark.asyncio(loop_scope="module") -async def test_process_template_empty_dict_uses_process_mode(session_v3): - """Test C: process_template={} still routes to batch.process(), not batch.function(). +@pytest.mark.asyncio +async def test_process_template_empty_dict_uses_process_mode(backend_dragon): + """Test C: process_template={} still routes to proxy.process(), not proxy.function(). - Regression test: a truthiness check on the dict silently falls through on an - empty dict; the ``is not None`` guard in build_task prevents this. + Regression: a truthiness check on the dict falls through on an empty dict; the + ``is not None`` guard in build_task prevents this. """ - backend = session_v3.backends["dragon"] process_calls = [] function_calls = [] - - def capture_process(pt, **kw): - process_calls.append(pt) - return MagicMock() - - def capture_function(target, *args, **kw): - function_calls.append(target) - return MagicMock() + mock_proxy = MagicMock() + mock_proxy.process.side_effect = lambda pt: process_calls.append(pt) or MagicMock() + mock_proxy.function.side_effect = lambda fn, *a, **kw: function_calls.append(fn) or MagicMock() task = _make_task_dict(lambda: None, backend_specific={"process_template": {}}) - with ( - patch.object(backend.batch, "process", side_effect=capture_process), - patch.object(backend.batch, "function", side_effect=capture_function), - ): - await backend.build_task(task) + with patch.object(backend_dragon.batch, "options", return_value=mock_proxy): + await backend_dragon.build_task(task) - assert len(process_calls) == 1, "batch.process() should have been called (Priority 2)" + assert len(process_calls) == 1, "options().process() should have been called (Priority 2)" assert len(function_calls) == 0, ( - "batch.function() must NOT be called when process_template is provided" + "options().function() must NOT be called when process_template is provided" ) -@pytest.mark.asyncio(loop_scope="module") -async def test_process_templates_list_built_and_passed_to_job(session_v3): - """Test D: process_templates list produces correct (nranks, ProcessTemplate) tuples for batch.job().""" +@pytest.mark.asyncio +async def test_process_templates_list_built_and_passed_to_job(backend_dragon): + """Test D: process_templates list produces correct (nranks, ProcessTemplate) tuples for job().""" from dragon.native.process import ProcessTemplate - backend = session_v3.backends["dragon"] captured_args = [] - - def capture_job(templates, **kw): - captured_args.append(templates) - return MagicMock() + mock_proxy = MagicMock() + mock_proxy.job.side_effect = lambda templates: captured_args.append(templates) or MagicMock() task = _make_task_dict( lambda: None, backend_specific={"process_templates": [(2, {"cwd": "/tmp"})]}, ) - with patch.object(backend.batch, "job", side_effect=capture_job): - await backend.build_task(task) + with patch.object(backend_dragon.batch, "options", return_value=mock_proxy): + await backend_dragon.build_task(task) assert len(captured_args) == 1 templates = captured_args[0] @@ -535,19 +524,16 @@ def capture_job(templates, **kw): assert pt.cwd == "/tmp" -@pytest.mark.asyncio(loop_scope="module") -async def test_process_template_combined_spec_policy_cwd_args(session_v3): - """Test E: process_template with policy + cwd all land on the ProcessTemplate correctly.""" +@pytest.mark.asyncio +async def test_process_template_combined_spec_policy_cwd_args(backend_dragon): + """Test E: process_template with policy + cwd + args all land on ProcessTemplate correctly.""" import cloudpickle from dragon.infrastructure.policy import Policy from dragon.native.process import ProcessTemplate - backend = session_v3.backends["dragon"] - captured = [] - - def capture(pt, **kw): - captured.append(pt) - return MagicMock() + captured_pt = [] + mock_proxy = MagicMock() + mock_proxy.process.side_effect = lambda pt: captured_pt.append(pt) or MagicMock() policy = Policy(gpu_affinity=[0]) task = _make_task_dict( @@ -557,11 +543,11 @@ def capture(pt, **kw): backend_specific={"process_template": {"policy": policy, "cwd": "/tmp"}}, ) - with patch.object(backend.batch, "process", side_effect=capture): - await backend.build_task(task) + with patch.object(backend_dragon.batch, "options", return_value=mock_proxy): + await backend_dragon.build_task(task) - assert len(captured) == 1 - pt = captured[0] + assert len(captured_pt) == 1 + pt = captured_pt[0] assert isinstance(pt, ProcessTemplate) assert pt.policy is policy assert pt.cwd == "/tmp" @@ -572,13 +558,13 @@ def capture(pt, **kw): # ============================================================================ -# V3 unit tests — constructor and internal methods (no Dragon cluster required) +# Dragon unit tests — constructor and internal methods (no Dragon cluster required) # ============================================================================ -def test_v3_constructor_batch_kwargs_forwarded_verbatim(): +def test_constructor_batch_kwargs_forwarded_verbatim(): """batch_kwargs contents are splatted into Batch() unchanged.""" - from rhapsody.backends.execution.dragon import DragonExecutionBackendV3 + from rhapsody.backends.execution.dragon import DragonExecutionBackend mock_batch = MagicMock() mock_batch.num_workers = 8 @@ -589,15 +575,15 @@ def test_v3_constructor_batch_kwargs_forwarded_verbatim(): with patch( "rhapsody.backends.execution.dragon.Batch", return_value=mock_batch ) as mock_batch_cls: - backend = DragonExecutionBackendV3(batch_kwargs=kwargs) + backend = DragonExecutionBackend(batch_kwargs=kwargs) - mock_batch_cls.assert_called_once_with(**kwargs) + mock_batch_cls.assert_called_once_with(task_logs=True, **kwargs) assert backend.batch is mock_batch -def test_v3_constructor_no_batch_kwargs_calls_batch_with_no_args(): - """DragonExecutionBackendV3() with no args calls Batch() with no args.""" - from rhapsody.backends.execution.dragon import DragonExecutionBackendV3 +def test_constructor_no_batch_kwargs_calls_batch_with_no_args(): + """DragonExecutionBackend() with no args calls Batch() with no args.""" + from rhapsody.backends.execution.dragon import DragonExecutionBackend mock_batch = MagicMock() mock_batch.num_workers = 4 @@ -606,18 +592,18 @@ def test_v3_constructor_no_batch_kwargs_calls_batch_with_no_args(): with patch( "rhapsody.backends.execution.dragon.Batch", return_value=mock_batch ) as mock_batch_cls: - DragonExecutionBackendV3() + DragonExecutionBackend() - mock_batch_cls.assert_called_once_with() + mock_batch_cls.assert_called_once_with(task_logs=True) -def test_v3_constructor_rejects_bare_batch_params(): +def test_constructor_rejects_bare_batch_params(): """Batch params passed directly (not via batch_kwargs) raise TypeError. num_nodes, pool_nodes, disable_telemetry, and other old top-level params are no longer accepted as direct constructor arguments — they must go through batch_kwargs. """ - from rhapsody.backends.execution.dragon import DragonExecutionBackendV3 + from rhapsody.backends.execution.dragon import DragonExecutionBackend mock_batch = MagicMock() mock_batch.num_workers = 8 @@ -625,193 +611,261 @@ def test_v3_constructor_rejects_bare_batch_params(): with patch("rhapsody.backends.execution.dragon.Batch", return_value=mock_batch): with pytest.raises(TypeError): - DragonExecutionBackendV3(num_nodes=4) + DragonExecutionBackend(num_nodes=4) with pytest.raises(TypeError): - DragonExecutionBackendV3(pool_nodes=2) + DragonExecutionBackend(pool_nodes=2) with pytest.raises(TypeError): - DragonExecutionBackendV3(disable_telemetry=True) + DragonExecutionBackend(disable_telemetry=True) with pytest.raises(TypeError): - DragonExecutionBackendV3(num_workers=4) + DragonExecutionBackend(num_workers=4) with pytest.raises(TypeError): - DragonExecutionBackendV3(disable_background_batching=True) + DragonExecutionBackend(disable_background_batching=True) with pytest.raises(TypeError): - DragonExecutionBackendV3(disable_batch_submission=True) + DragonExecutionBackend(disable_batch_submission=True) + + +@pytest.mark.asyncio +async def test_initial_state_is_initialized(): + """A freshly created backend reports INITIALIZED before any tasks are submitted.""" + from rhapsody.backends.execution.dragon import DragonExecutionBackend + + mock_batch = MagicMock() + mock_batch.num_workers = 4 + mock_batch.num_managers = 1 + + with patch("rhapsody.backends.execution.dragon.Batch", return_value=mock_batch): + backend = DragonExecutionBackend() + + assert await backend.state() == "INITIALIZED" -def test_v3_deliver_batch_success_stores_value_and_fires_done(backend_v3): +def test_deliver_batch_success_stores_value_and_fires_done(backend_dragon): """_deliver_batch stores return_value on the task dict and fires the DONE callback.""" uid = "task.unit-done" task_desc = {"uid": uid} - backend_v3._task_registry[uid] = {"uid": uid, "description": task_desc} + backend_dragon._task_registry[uid] = {"uid": uid, "description": task_desc} - backend_v3._deliver_batch([(uid, 42, None, False, None, None)]) + backend_dragon._deliver_batch([(uid, 42, None, False, "", "")]) assert task_desc["return_value"] == 42 assert task_desc["stdout"] == "" assert task_desc["stderr"] == "" - backend_v3._callback_func.assert_called_once_with(task_desc, "DONE") - assert uid not in backend_v3._task_registry + backend_dragon._callback_func.assert_called_once_with(task_desc, "DONE") + assert uid not in backend_dragon._task_registry -def test_v3_deliver_batch_propagates_stdout_stderr(backend_v3): +def test_deliver_batch_propagates_stdout_stderr(backend_dragon): """_deliver_batch stores stdout/stderr on task_desc when non-empty.""" uid = "task.unit-done-out" task_desc = {"uid": uid} - backend_v3._task_registry[uid] = {"uid": uid, "description": task_desc} + backend_dragon._task_registry[uid] = {"uid": uid, "description": task_desc} - backend_v3._deliver_batch([(uid, "ok", None, False, "hello\n", "warn\n")]) + backend_dragon._deliver_batch([(uid, "ok", None, False, "hello\n", "warn\n")]) assert task_desc["stdout"] == "hello\n" assert task_desc["stderr"] == "warn\n" -def test_v3_deliver_batch_failure_stores_exc_and_fires_failed(backend_v3): +def test_deliver_batch_failure_stores_exc_and_fires_failed(backend_dragon): """_deliver_batch stores the exception and stderr string, fires the FAILED callback.""" uid = "task.unit-failed" task_desc = {"uid": uid} - backend_v3._task_registry[uid] = {"uid": uid, "description": task_desc} + backend_dragon._task_registry[uid] = {"uid": uid, "description": task_desc} exc = RuntimeError("something went wrong") # raised=True, tb=None: stderr falls back to str(exc) - backend_v3._deliver_batch([(uid, exc, None, True, None, None)]) + backend_dragon._deliver_batch([(uid, exc, None, True, "", "")]) assert task_desc["exception"] is exc assert "something went wrong" in task_desc["stderr"] - backend_v3._callback_func.assert_called_once_with(task_desc, "FAILED") - assert uid not in backend_v3._task_registry + backend_dragon._callback_func.assert_called_once_with(task_desc, "FAILED") + assert uid not in backend_dragon._task_registry -def test_v3_deliver_batch_prefers_traceback_over_str_exc(backend_v3): +def test_deliver_batch_prefers_traceback_over_str_exc(backend_dragon): """_deliver_batch uses Dragon's traceback string when available.""" uid = "task.unit-failed-tb" task_desc = {"uid": uid} - backend_v3._task_registry[uid] = {"uid": uid, "description": task_desc} + backend_dragon._task_registry[uid] = {"uid": uid, "description": task_desc} exc = RuntimeError("boom") tb = "Traceback (most recent call last):\n File ...\nRuntimeError: boom" - backend_v3._deliver_batch([(uid, exc, tb, True, None, None)]) + backend_dragon._deliver_batch([(uid, exc, tb, True, "", "")]) assert task_desc["stderr"] == tb -def test_v3_cancelled_task_skips_callback(backend_v3): +def test_cancelled_task_skips_callback(backend_dragon): """_deliver_batch is a no-op for UIDs in _cancelled_tasks.""" uid = "task.unit-cancelled" task_desc = {"uid": uid} - backend_v3._task_registry[uid] = {"uid": uid, "description": task_desc} - backend_v3._cancelled_tasks.add(uid) + backend_dragon._task_registry[uid] = {"uid": uid, "description": task_desc} + backend_dragon._cancelled_tasks.add(uid) - backend_v3._deliver_batch([(uid, 99, None, False, None, None)]) + backend_dragon._deliver_batch([(uid, 99, None, False, "", "")]) - backend_v3._callback_func.assert_not_called() + backend_dragon._callback_func.assert_not_called() assert "return_value" not in task_desc - assert uid not in backend_v3._cancelled_tasks + assert uid not in backend_dragon._cancelled_tasks -def test_v3_fence_delegates_to_batch(backend_v3): +def test_fence_delegates_to_batch(backend_dragon): """backend.fence() calls batch.fence() exactly once.""" - backend_v3.fence() - backend_v3.batch.fence.assert_called_once() + backend_dragon.fence() + backend_dragon.batch.fence.assert_called_once() + + +def _run_monitor_one_cycle(backend, poll_tuid, monitored_entry=None, task_registry_entry=None): + """Drive _monitor_loop for one completion cycle and return the captured mock_loop. + + - poll call 1 returns `poll_tuid` (the tuid to deliver) + - poll call ≥2 sets _shutdown_event and returns None (draining + outer loop → exits) + - backend._loop is replaced with a MagicMock so call_soon_threadsafe is capturable + - If `monitored_entry` is given, it is inserted into _monitored_batches under poll_tuid + - If `task_registry_entry` is given, it is inserted into _task_registry + """ + mock_loop = MagicMock() + backend._loop = mock_loop + + if monitored_entry is not None: + backend._monitored_batches[poll_tuid] = monitored_entry + if task_registry_entry is not None: + uid = task_registry_entry["uid"] + backend._task_registry[uid] = task_registry_entry + + calls = [0] + def poll_side_effect(timeout=0): + calls[0] += 1 + if calls[0] == 1: + return poll_tuid + backend._shutdown_event.set() + return None -def test_v3_monitor_loop_skips_task_with_no_manager_idx(backend_v3): - """Monitor loop skips a task whose manager_idx is not yet set (None).""" + backend.batch.poll.side_effect = poll_side_effect + backend._monitor_loop() + return mock_loop + + +def test_monitor_loop_get_called_after_poll(backend_dragon): + """_monitor_loop calls batch_task.get(block=False) after poll() returns its tuid.""" + uid = "task.poll-get" + tuid = "dragon-tuid-poll-get" mock_task = MagicMock() - mock_task.core.manager_idx = None - uid = "task.monitor-no-idx" - backend_v3._monitored_batches[uid] = (mock_task, "flow-uid") - - # Run one sweep manually - backend_v3.batch.results_ddict.manager.side_effect = AssertionError("should not be called") - batch_tuids = list(backend_v3._monitored_batches.keys()) - completed = [] - for tuid in batch_tuids: - if tuid not in backend_v3._monitored_batches: - continue - batch_task, flow_uid = backend_v3._monitored_batches[tuid] - manager_idx = batch_task.core.manager_idx - if manager_idx is None: - continue - ddict_shard_idx = 0 if manager_idx == -1 else manager_idx - try: - result, tb, raised, stdout, stderr = backend_v3.batch.results_ddict.manager( - ddict_shard_idx - )[tuid] - except KeyError: - continue - backend_v3._monitored_batches.pop(tuid) - completed.append((flow_uid, result, tb, raised, stdout, stderr)) - - assert completed == [] - assert uid in backend_v3._monitored_batches # not consumed - - -def test_v3_monitor_loop_routes_to_correct_shard(backend_v3): - """Monitor loop calls results_ddict.manager(shard_idx) with the task's manager_idx.""" + mock_task.get.return_value = "done-result" + mock_task.traceback = None + mock_task.stdout_path = None + mock_task.stderr_path = None + mock_task.get_stdout.return_value = None + mock_task.get_stderr.return_value = None + + mock_loop = _run_monitor_one_cycle( + backend_dragon, + poll_tuid=tuid, + monitored_entry=(mock_task, uid), + task_registry_entry={"uid": uid, "description": {"uid": uid}, "is_native_function": True}, + ) + + mock_task.get.assert_called_once_with(block=False) + mock_loop.call_soon_threadsafe.assert_called_once() + _, completions = mock_loop.call_soon_threadsafe.call_args.args + assert completions == [(uid, "done-result", None, False, "", "")] + + +def test_monitor_loop_skips_cancelled_tuid(backend_dragon): + """When poll() returns a tuid absent from _monitored_batches, _deliver_batch is not called. + + This is the user-cancel path: cancel_task() eagerly pops from _monitored_batches before + the monitor loop sees the tuid via poll(). The loop must silently skip it. + """ + tuid = "dragon-cancelled-tuid" + # Intentionally not inserting into _monitored_batches — simulates cancel_task already popped it. + mock_loop = _run_monitor_one_cycle(backend_dragon, poll_tuid=tuid) + + mock_loop.call_soon_threadsafe.assert_not_called() + + +def test_monitor_loop_reads_paths_from_batch_task(backend_dragon): + """Stdout/stderr paths are taken from batch_task.stdout_path/.stderr_path after poll(). + + The monitor loop must NOT look up paths from _task_registry — they live on the + batch_task object as set at submission time via Batch.options(stdout=..., stderr=...). + """ + uid = "task.poll-paths" + tuid = "dragon-tuid-paths" mock_task = MagicMock() - mock_task.core.manager_idx = 2 - uid = "task.monitor-shard" - backend_v3._monitored_batches[uid] = (mock_task, "flow-uid-2") - - mock_shard = MagicMock() - mock_shard.__getitem__ = MagicMock(return_value=("result", None, False, "", "")) - backend_v3.batch.results_ddict.manager.return_value = mock_shard - - batch_tuids = list(backend_v3._monitored_batches.keys()) - completed = [] - for tuid in batch_tuids: - if tuid not in backend_v3._monitored_batches: - continue - batch_task, flow_uid = backend_v3._monitored_batches[tuid] - manager_idx = batch_task.core.manager_idx - if manager_idx is None: - continue - ddict_shard_idx = 0 if manager_idx == -1 else manager_idx - try: - result, tb, raised, stdout, stderr = backend_v3.batch.results_ddict.manager( - ddict_shard_idx - )[tuid] - except KeyError: - continue - backend_v3._monitored_batches.pop(tuid) - completed.append((flow_uid, result, tb, raised, stdout, stderr)) - - backend_v3.batch.results_ddict.manager.assert_called_once_with(2) - assert len(completed) == 1 - assert completed[0][0] == "flow-uid-2" - assert completed[0][1] == "result" - - -def test_v3_monitor_loop_scheduler_manager_maps_to_shard_0(backend_v3): - """SCHEDULER_MANAGER_IDX (-1) maps to DDict shard 0.""" + mock_task.get.return_value = "result" + mock_task.traceback = None + mock_task.stdout_path = "/work/uid.stdout" + mock_task.stderr_path = "/work/uid.stderr" + + mock_loop = _run_monitor_one_cycle( + backend_dragon, + poll_tuid=tuid, + monitored_entry=(mock_task, uid), + task_registry_entry={ + "uid": uid, + "description": {"uid": uid, "capture_stdio": True}, + "is_native_function": False, + }, + ) + + mock_loop.call_soon_threadsafe.assert_called_once() + _, completions = mock_loop.call_soon_threadsafe.call_args.args + assert len(completions) == 1 + assert completions[0][4] == "/work/uid.stdout" + assert completions[0][5] == "/work/uid.stderr" + + +def test_monitor_loop_nonzero_exit_delivers_failed(backend_dragon): + """Process task: get() returning int 1 must produce raised=True + SystemExit(1).""" + uid, tuid = "task.exit-1", "tuid-exit-1" mock_task = MagicMock() - mock_task.core.manager_idx = -1 # SCHEDULER_MANAGER_IDX - uid = "task.monitor-scheduler" - backend_v3._monitored_batches[uid] = (mock_task, "flow-sched") - - mock_shard = MagicMock() - mock_shard.__getitem__ = MagicMock(return_value=("sched-result", None, False, "", "")) - backend_v3.batch.results_ddict.manager.return_value = mock_shard - - batch_tuids = list(backend_v3._monitored_batches.keys()) - for tuid in batch_tuids: - if tuid not in backend_v3._monitored_batches: - continue - batch_task, _ = backend_v3._monitored_batches[tuid] - manager_idx = batch_task.core.manager_idx - if manager_idx is None: - continue - ddict_shard_idx = 0 if manager_idx == -1 else manager_idx - try: - backend_v3.batch.results_ddict.manager(ddict_shard_idx)[tuid] - except KeyError: - continue - - backend_v3.batch.results_ddict.manager.assert_called_once_with(0) + mock_task.get.return_value = 1 + mock_task.traceback = None + mock_task.get_stdout.return_value = None + mock_task.get_stderr.return_value = None + + mock_loop = _run_monitor_one_cycle( + backend_dragon, + poll_tuid=tuid, + monitored_entry=(mock_task, uid), + task_registry_entry={"uid": uid, "description": {"uid": uid}, "is_native_function": False}, + ) + + _, completions = mock_loop.call_soon_threadsafe.call_args.args + uid_c, result_c, _tb, raised_c, *_ = completions[0] + assert uid_c == uid + assert raised_c is True + assert isinstance(result_c, SystemExit) + assert result_c.code == 1 + + +def test_monitor_loop_function_int_return_delivers_done(backend_dragon): + """Function task returning int 42 must deliver DONE — exit-code check must not fire.""" + uid, tuid = "task.func-42", "tuid-func-42" + mock_task = MagicMock() + mock_task.get.return_value = 42 + mock_task.traceback = None + mock_task.get_stdout.return_value = None + mock_task.get_stderr.return_value = None + + mock_loop = _run_monitor_one_cycle( + backend_dragon, + poll_tuid=tuid, + monitored_entry=(mock_task, uid), + task_registry_entry={"uid": uid, "description": {"uid": uid}, "is_native_function": True}, + ) + + _, completions = mock_loop.call_soon_threadsafe.call_args.args + uid_c, result_c, _tb, raised_c, *_ = completions[0] + assert uid_c == uid + assert raised_c is False + assert result_c == 42 @pytest.mark.asyncio -async def test_v3_cancel_task_calls_dragon_cancel(backend_v3): +async def test_cancel_task_calls_dragon_cancel(backend_dragon): """cancel_task delegates to batch_task.cancel() and fires CANCELED callback on success. The callback must be called synchronously — no asyncio.sleep(0) needed to observe it. This @@ -825,26 +879,26 @@ async def test_v3_cancel_task_calls_dragon_cancel(backend_v3): mock_batch_task = MagicMock() mock_batch_task.cancel.return_value = True task_desc = {"uid": uid} - backend_v3._task_registry[uid] = { + backend_dragon._task_registry[uid] = { "uid": uid, "description": task_desc, "batch_task": mock_batch_task, } - backend_v3._monitored_batches[mock_batch_task.uid] = (mock_batch_task, uid) + backend_dragon._monitored_batches[mock_batch_task.uid] = (mock_batch_task, uid) - result = await backend_v3.cancel_task(uid) + result = await backend_dragon.cancel_task(uid) # Callback and state must be visible immediately — no yield required assert result is True mock_batch_task.cancel.assert_called_once() - backend_v3._callback_func.assert_called_once_with(task_desc, "CANCELED") + backend_dragon._callback_func.assert_called_once_with(task_desc, "CANCELED") # Task must be eagerly removed so the monitor loop stops polling - assert uid not in backend_v3._task_registry - assert mock_batch_task.uid not in backend_v3._monitored_batches + assert uid not in backend_dragon._task_registry + assert mock_batch_task.uid not in backend_dragon._monitored_batches @pytest.mark.asyncio -async def test_v3_cancel_task_returns_false_when_task_completed_before_cancel_lands(backend_v3): +async def test_cancel_task_returns_false_when_task_completed_before_cancel_lands(backend_dragon): """cancel_task returns False when the task was already delivered (registry entry gone). Race: the event loop yields inside run_in_executor; _deliver_batch could pop the task @@ -858,7 +912,7 @@ async def test_v3_cancel_task_returns_false_when_task_completed_before_cancel_la mock_batch_task = MagicMock() mock_batch_task.cancel.return_value = True task_desc = {"uid": uid} - backend_v3._task_registry[uid] = { + backend_dragon._task_registry[uid] = { "uid": uid, "description": task_desc, "batch_task": mock_batch_task, @@ -869,68 +923,68 @@ async def test_v3_cancel_task_returns_false_when_task_completed_before_cancel_la def cancel_and_deliver(): result = original_cancel() - backend_v3._task_registry.pop(uid, None) + backend_dragon._task_registry.pop(uid, None) return result mock_batch_task.cancel = cancel_and_deliver - result = await backend_v3.cancel_task(uid) + result = await backend_dragon.cancel_task(uid) assert result is False - backend_v3._callback_func.assert_not_called() - assert uid not in backend_v3._task_registry + backend_dragon._callback_func.assert_not_called() + assert uid not in backend_dragon._task_registry @pytest.mark.asyncio -async def test_v3_cancel_task_dragon_returns_false_no_callback(backend_v3): +async def test_cancel_task_dragon_returns_false_no_callback(backend_dragon): """cancel_task does not fire callback when Dragon reports the task cannot be cancelled.""" uid = "task.unit-cancel-false" mock_batch_task = MagicMock() mock_batch_task.cancel.return_value = False task_desc = {"uid": uid} - backend_v3._task_registry[uid] = { + backend_dragon._task_registry[uid] = { "uid": uid, "description": task_desc, "batch_task": mock_batch_task, } - result = await backend_v3.cancel_task(uid) + result = await backend_dragon.cancel_task(uid) assert result is False - backend_v3._callback_func.assert_not_called() - assert uid not in backend_v3._cancelled_tasks + backend_dragon._callback_func.assert_not_called() + assert uid not in backend_dragon._cancelled_tasks @pytest.mark.asyncio -async def test_v3_cancel_task_dragon_raises_falls_back_to_soft_cancel(backend_v3): +async def test_cancel_task_dragon_raises_falls_back_to_soft_cancel(backend_dragon): """cancel_task falls back to soft-cancel (callback fired, cancelled=True) if Dragon raises.""" uid = "task.unit-cancel-exc" mock_batch_task = MagicMock() mock_batch_task.cancel.side_effect = RuntimeError("dragon scheduler unreachable") task_desc = {"uid": uid} - backend_v3._task_registry[uid] = { + backend_dragon._task_registry[uid] = { "uid": uid, "description": task_desc, "batch_task": mock_batch_task, } - backend_v3._monitored_batches[mock_batch_task.uid] = (mock_batch_task, uid) + backend_dragon._monitored_batches[mock_batch_task.uid] = (mock_batch_task, uid) - result = await backend_v3.cancel_task(uid) + result = await backend_dragon.cancel_task(uid) assert result is True - backend_v3._callback_func.assert_called_once_with(task_desc, "CANCELED") - assert uid not in backend_v3._task_registry - assert mock_batch_task.uid not in backend_v3._monitored_batches + backend_dragon._callback_func.assert_called_once_with(task_desc, "CANCELED") + assert uid not in backend_dragon._task_registry + assert mock_batch_task.uid not in backend_dragon._monitored_batches @pytest.mark.asyncio -async def test_v3_shutdown_calls_join_and_destroy_not_close(backend_v3): +async def test_shutdown_calls_join_and_destroy_not_close(backend_dragon): """Shutdown() calls batch.join() then batch.destroy(); batch.close() must NOT be called.""" - await backend_v3.shutdown() + await backend_dragon.shutdown() - backend_v3.batch.join.assert_called_once() - backend_v3.batch.destroy.assert_called_once() - backend_v3.batch.close.assert_not_called() + backend_dragon.batch.join.assert_called_once() + backend_dragon.batch.destroy.assert_called_once() + backend_dragon.batch.close.assert_not_called() # ============================================================================ @@ -939,9 +993,14 @@ async def test_v3_shutdown_calls_join_and_destroy_not_close(backend_v3): @pytest.mark.asyncio -async def test_v3_capture_stdio_build_task_creates_script(backend_v3, tmp_path): - """capture_stdio=True writes a shell script and populates registry paths.""" - backend_v3._work_dir = str(tmp_path) +async def test_capture_stdio_true_passes_paths_to_batch(backend_dragon, tmp_path): + """capture_stdio=True passes stdout/stderr file paths to Batch.options() — no shell script. + + Dragon 0.14.1-rc requires metadata (name, timeout, stdio paths) in Batch.options(), not as + direct kwargs to process()/function(). The old patch on batch.process is dead code. Assert via + batch.options call_args instead. + """ + backend_dragon._work_dir = str(tmp_path) task = { "uid": "task.capture-test-001", @@ -955,25 +1014,22 @@ async def test_v3_capture_stdio_build_task_creates_script(backend_v3, tmp_path): "capture_stdio": True, } - await backend_v3.build_task(task) - - uid = task["uid"] - reg = backend_v3._task_registry[uid] - - assert reg["script_path"] is not None - assert reg["script_path"].endswith(".sh") - assert reg["stdout_path"].endswith(".stdout") - assert reg["stderr_path"].endswith(".stderr") + await backend_dragon.build_task(task) - script = open(reg["script_path"]).read() - assert "1>" in script - assert "2>" in script + options_kwargs = backend_dragon.batch.options.call_args.kwargs + assert options_kwargs["stdout"].endswith(".stdout") + assert options_kwargs["stderr"].endswith(".stderr") + # No shell script written — Dragon handles redirection natively + assert not list(tmp_path.glob("*.sh")) + # Registry no longer stores paths — they live on the batch_task object + reg = backend_dragon._task_registry[task["uid"]] + assert "script_path" not in reg @pytest.mark.asyncio -async def test_v3_capture_stdio_false_no_script(backend_v3, tmp_path): - """capture_stdio=False (default) creates no shell script.""" - backend_v3._work_dir = str(tmp_path) +async def test_capture_stdio_false_passes_none_to_batch(backend_dragon, tmp_path): + """capture_stdio=False (default) passes stdout=None, stderr=None to Batch.options().""" + backend_dragon._work_dir = str(tmp_path) task = { "uid": "task.capture-test-002", @@ -987,21 +1043,18 @@ async def test_v3_capture_stdio_false_no_script(backend_v3, tmp_path): "capture_stdio": False, } - await backend_v3.build_task(task) - - uid = task["uid"] - reg = backend_v3._task_registry[uid] + await backend_dragon.build_task(task) - assert reg["script_path"] is None - assert reg["stdout_path"] is None - assert reg["stderr_path"] is None + options_kwargs = backend_dragon.batch.options.call_args.kwargs + assert options_kwargs["stdout"] is None + assert options_kwargs["stderr"] is None assert not list(tmp_path.glob("*.sh")) @pytest.mark.asyncio -async def test_v3_capture_stdio_function_task_no_script(backend_v3, tmp_path): - """capture_stdio=True on a function task is ignored — no shell script written.""" - backend_v3._work_dir = str(tmp_path) +async def test_capture_stdio_function_task_passes_paths_to_batch_function(backend_dragon, tmp_path): + """capture_stdio=True on a function task passes stdout/stderr paths to Batch.options().""" + backend_dragon._work_dir = str(tmp_path) task = { "uid": "task.capture-test-003", @@ -1011,16 +1064,15 @@ async def test_v3_capture_stdio_function_task_no_script(backend_v3, tmp_path): "args": [], "kwargs": {}, "name": "func-capture-test", - "task_backend_specific_kwargs": {"process_template": {}}, + "task_backend_specific_kwargs": {}, "capture_stdio": True, } - await backend_v3.build_task(task) + await backend_dragon.build_task(task) - uid = task["uid"] - reg = backend_v3._task_registry[uid] - - assert reg["script_path"] is None + options_kwargs = backend_dragon.batch.options.call_args.kwargs + assert options_kwargs["stdout"].endswith(".stdout") + assert options_kwargs["stderr"].endswith(".stderr") assert not list(tmp_path.glob("*.sh")) @@ -1030,44 +1082,43 @@ async def test_v3_capture_stdio_function_task_no_script(backend_v3, tmp_path): @pytest.mark.result_contract -def test_v3_deliver_batch_empty_dragon_stdout_written_as_empty_string(backend_v3): +def test_deliver_batch_empty_dragon_stdout_written_as_empty_string(backend_dragon): """Regression: stdout='' from Dragon must land as '' not None or absent key.""" uid = "task.contract-empty-stdout" task_desc = {"uid": uid} - backend_v3._task_registry[uid] = {"uid": uid, "description": task_desc} + backend_dragon._task_registry[uid] = {"uid": uid, "description": task_desc} - backend_v3._deliver_batch([(uid, 0, None, False, "", "")]) + backend_dragon._deliver_batch([(uid, 0, None, False, "", "")]) assert isinstance(task_desc["stdout"], str) assert isinstance(task_desc["stderr"], str) - backend_v3._callback_func.assert_called_once_with(task_desc, "DONE") + backend_dragon._callback_func.assert_called_once_with(task_desc, "DONE") @pytest.mark.result_contract -def test_v3_deliver_batch_stdout_path_takes_priority_over_empty_stdout(backend_v3): - """stdout_path wins when capture_stdio redirect is active.""" +def test_deliver_batch_path_from_tuple_stored_on_task(backend_dragon): + """Stdout/stderr paths in the completion tuple are stored directly on task_desc. + + The monitor loop passes batch_task.stdout_path as the 5th element; _deliver_batch stores it + verbatim — no registry lookup needed. + """ uid = "task.contract-redirect" task_desc = {"uid": uid} - backend_v3._task_registry[uid] = { - "uid": uid, - "description": task_desc, - "stdout_path": "/work/uid.stdout", - "stderr_path": "/work/uid.stderr", - } + backend_dragon._task_registry[uid] = {"uid": uid, "description": task_desc} - backend_v3._deliver_batch([(uid, 0, None, False, "", "")]) + backend_dragon._deliver_batch([(uid, 0, None, False, "/work/uid.stdout", "/work/uid.stderr")]) assert task_desc["stdout"] == "/work/uid.stdout" assert task_desc["stderr"] == "/work/uid.stderr" @pytest.mark.result_contract -def test_v3_deliver_batch_failed_stdout_is_string(backend_v3): +def test_deliver_batch_failed_stdout_is_string(backend_dragon): """Stdout must be a str in the FAILED path.""" uid = "task.contract-failed-stdout" task_desc = {"uid": uid} - backend_v3._task_registry[uid] = {"uid": uid, "description": task_desc} + backend_dragon._task_registry[uid] = {"uid": uid, "description": task_desc} - backend_v3._deliver_batch([(uid, RuntimeError("boom"), None, True, "", "")]) + backend_dragon._deliver_batch([(uid, RuntimeError("boom"), None, True, "", "")]) assert isinstance(task_desc["stdout"], str) diff --git a/tutorials/README.md b/tutorials/README.md index ac9d2d4..e554d0a 100644 --- a/tutorials/README.md +++ b/tutorials/README.md @@ -14,7 +14,7 @@ jupyter notebook ### With Dragon Backend (dragon-resource-placement.ipynb, rhapsody-tutorial.ipynb) -Tutorials that use `DragonExecutionBackendV3` require the Dragon runtime. You **must** launch the notebook server using `dragon-jupyter`: +Tutorials that use `DragonExecutionBackend` require the Dragon runtime. You **must** launch the notebook server using `dragon-jupyter`: ```bash dragon-jupyter