diff --git a/src/rhapsody/api/task.py b/src/rhapsody/api/task.py index 813fbaf..dd22a6d 100644 --- a/src/rhapsody/api/task.py +++ b/src/rhapsody/api/task.py @@ -244,10 +244,13 @@ def from_dict(cls, data: dict[str, Any]) -> BaseTask: Raises: TaskValidationError: If required fields are missing or invalid """ - # Determine task type based on fields present - if "prompt" in data: + # Determine task type based on fields present. Test for ``is not + # None`` (not truthiness) so a present-but-falsy value — e.g. an empty + # prompt list ``[]`` or empty string — still routes to the right class + # instead of being treated as absent. + if data.get("prompt") is not None: return AITask(**data) - elif "executable" in data or "function" in data: + elif data.get("executable") is not None or data.get("function") is not None: return ComputeTask(**data) else: raise TaskValidationError( diff --git a/src/rhapsody/backends/execution/__init__.py b/src/rhapsody/backends/execution/__init__.py index 50ec4db..cf84611 100644 --- a/src/rhapsody/backends/execution/__init__.py +++ b/src/rhapsody/backends/execution/__init__.py @@ -7,8 +7,9 @@ from __future__ import annotations from .concurrent import ConcurrentExecutionBackend # noqa: F401 +from .noop import NoopExecutionBackend # noqa: F401 -__all__ = ["ConcurrentExecutionBackend"] +__all__ = ["ConcurrentExecutionBackend", "NoopExecutionBackend"] # Try to import optional backends try: @@ -25,6 +26,13 @@ except ImportError: pass +try: + from .orbit import OrbitExecutionBackend # noqa: F401 + + __all__.append("OrbitExecutionBackend") +except ImportError: + pass + try: from .dragon import DragonExecutionBackendV1 # noqa: F401 from .dragon import DragonExecutionBackendV2 # noqa: F401 diff --git a/src/rhapsody/backends/execution/dragon.py b/src/rhapsody/backends/execution/dragon.py index 9098034..e3ac46c 100644 --- a/src/rhapsody/backends/execution/dragon.py +++ b/src/rhapsody/backends/execution/dragon.py @@ -1,4 +1,6 @@ import asyncio +import functools +import glob import logging import os import shlex @@ -363,6 +365,7 @@ def _aggregate_function_results( "exit_code": result.exit_code, "return_value": result.return_value, "exception": result.exception, + "traceback": result.traceback, "success": result.success, } else: @@ -389,6 +392,10 @@ def _aggregate_function_results( "exception": None if all_successful else "; ".join(str(r.exception) for r in results if not r.success), + "traceback": None + if all_successful + else "\n".join(r.traceback for r in results + if not r.success and r.traceback), "success": all_successful, } @@ -1656,6 +1663,38 @@ class TaskStateMapperV3: terminal_states = {DONE, FAILED, CANCELED} +def _v3_function_wrapper(user_func, stdout_path, stderr_path, *args, **kwargs): + """Run *user_func* with sys.stdout/stderr redirected to per-rank files. + + Dragon batch surfaces only a dragon-side ``DragonUserCodeError`` when a + ProcessGroup rank exits non-zero — the rank's own Python traceback is + dropped. Persisting it to ``.`` before re-raising + lets ``_deliver_batch`` read it post-mortem. + """ + import os + import sys + import traceback + + rank = (os.environ.get("DRAGON_RANK") + or os.environ.get("PMI_RANK") + or os.environ.get("OMPI_COMM_WORLD_RANK") + or os.environ.get("PALS_RANKID") + or os.environ.get("SLURM_PROCID") + or f"pid{os.getpid()}") + + with open(f"{stdout_path}.{rank}", "w") as out, \ + open(f"{stderr_path}.{rank}", "w") as err: + old_out, old_err = sys.stdout, sys.stderr + sys.stdout, sys.stderr = out, err + try: + return user_func(*args, **kwargs) + except BaseException: + traceback.print_exc(file=err) + raise + finally: + sys.stdout, sys.stderr = old_out, old_err + + # ============================================================================ # Main Backend Classes # V1 Integrates with Dragon HPC Native API @@ -3218,9 +3257,19 @@ def _monitor_loop(self): )[tuid] except KeyError: continue - self._monitored_batches.pop(tuid, None) - completed.append((uid, result, tb, raised, stdout, stderr)) + # Drain the per-rank function stdout/stderr files here, in + # the monitor thread, so the file I/O never blocks the + # asyncio loop that runs `_deliver_batch`. Executable tasks + # redirect via a shell script (`script_path` set) and have no + # per-rank files to drain. + fn_stdout = fn_stderr = None + info = self._task_registry.get(uid) + if info and not info.get("script_path"): + fn_stdout = self._drain_rank_files(info.get("stdout_path")) + fn_stderr = self._drain_rank_files(info.get("stderr_path")) + completed.append((uid, result, tb, raised, stdout, stderr, + fn_stdout, fn_stderr)) # One cross-thread wakeup for the entire sweep batch if completed: @@ -3235,14 +3284,46 @@ def _monitor_loop(self): self.logger.debug("Dragon batch monitor loop stopped") + @staticmethod + def _drain_rank_files(path: str) -> str: + """Read and unlink all per-rank ``.`` files written by + :func:`_v3_function_wrapper`, returning their concatenated content. + + Runs in the monitor thread (never the event loop), reads both the + success and failure output of every rank, and removes the files + afterwards so they do not accumulate in the work dir. + """ + if not path: + return "" + chunks = [] + for f in sorted(glob.glob(f"{path}.*")): + try: + # Worker output may contain non-UTF-8 bytes (C libs, progress + # bars); replace rather than drop the rank's whole output. + with open(f, encoding="utf-8", errors="replace") as fh: + content = fh.read() + except Exception: + content = "" + try: + os.unlink(f) + except OSError: + pass + if content: + rank = f.rsplit(".", 1)[-1] + chunks.append(f"[rank {rank}]\n{content}") + return "\n".join(chunks) + 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- - thread wakeups from O(tasks) to O(sweeps). + thread wakeups from O(tasks) to O(sweeps). Per-rank function output has + already been drained off-loop by the monitor thread and arrives as + ``fn_stdout`` / ``fn_stderr`` strings, so this method does no file I/O. """ - for uid, result, tb, raised, stdout, stderr in completions: + for (uid, result, tb, raised, stdout, stderr, + fn_stdout, fn_stderr) in completions: task_info = self._task_registry.pop(uid, None) if not task_info: continue @@ -3252,15 +3333,48 @@ def _deliver_batch(self, completions: list) -> None: task_desc = task_info["description"] stdout_path = task_info.get("stdout_path") stderr_path = task_info.get("stderr_path") + # Executable redirect writes one literal file per path (script_path + # set), so callers receive the path. Function tasks stream to + # per-rank files already folded into fn_stdout / fn_stderr. + is_exec_redirect = bool(task_info.get("script_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 "" + if is_exec_redirect: + task_desc["exception"] = result + task_desc["stderr"] = stderr_path or (tb if tb else str(result)) + task_desc["stdout"] = stdout_path or stdout or "" + else: + diagnostic = fn_stderr or "" + if diagnostic: + # Fold the rank's own traceback into the exception so a + # re-raise surfaces the real cause, not just Dragon's + # generic DragonUserCodeError. + try: + cls = type(result) + # A non-BaseException result (e.g. a string from a + # serialization fallback) would build fine but never + # be raised by session.py, silently masking the + # failure as success — coerce it to RuntimeError. + if not issubclass(cls, BaseException): + raise TypeError("result is not a BaseException") + augmented = cls(f"{result}\n--- worker output ---\n{diagnostic}") + except Exception: + augmented = RuntimeError( + f"{result}\n--- worker output ---\n{diagnostic}" + ) + task_desc["exception"] = augmented + else: + task_desc["exception"] = result + task_desc["stderr"] = diagnostic or stderr or (tb if tb else str(result)) + task_desc["stdout"] = fn_stdout or stdout or "" 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 "" + if is_exec_redirect: + task_desc["stdout"] = stdout_path or stdout or "" + task_desc["stderr"] = stderr_path or stderr or "" + else: + task_desc["stdout"] = fn_stdout or stdout or "" + task_desc["stderr"] = fn_stderr or stderr or "" self._callback_func(task_desc, "DONE") async def submit_tasks(self, tasks: list[dict]) -> None: @@ -3401,6 +3515,15 @@ def target(*a, **kw): target = "/bin/bash" task_args = (script_path,) + # Wrap every function target so each rank persists its own + # stdout/stderr; see _v3_function_wrapper for why. + if is_function: + stdout_path = os.path.join(self._work_dir, f"{uid}.stdout") + stderr_path = os.path.join(self._work_dir, f"{uid}.stderr") + target = functools.partial( + _v3_function_wrapper, target, stdout_path, stderr_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} diff --git a/src/rhapsody/backends/execution/noop.py b/src/rhapsody/backends/execution/noop.py new file mode 100644 index 0000000..91ed2b3 --- /dev/null +++ b/src/rhapsody/backends/execution/noop.py @@ -0,0 +1,136 @@ +"""No-op execution backend for performance benchmarking. + +Tasks are immediately marked as DONE without executing anything. +""" + +import asyncio +import logging +import uuid +from typing import Any +from typing import Callable + +from ..base import BaseBackend +from ..constants import BackendMainStates +from ..constants import StateMapper + + +def _get_logger() -> logging.Logger: + return logging.getLogger(__name__) + + +class NoopExecutionBackend(BaseBackend): + """Backend that completes every task instantly. + + Useful for measuring Orbit/broker/client overhead without any actual task execution cost. + """ + + def __init__(self, name: str = "noop"): + super().__init__(name=name) + self.logger = _get_logger() + self.tasks: dict[str, dict] = {} + self._futures: dict[str, asyncio.Task] = {} + self._callback_func: Callable = lambda t, s: None + self._initialized = False + self._backend_state = BackendMainStates.INITIALIZED + + def __await__(self): + return self._async_init().__await__() + + async def _async_init(self): + if not self._initialized: + StateMapper.register_backend_states_with_defaults(backend=self) + StateMapper.register_backend_tasks_states_with_defaults(backend=self) + self._backend_state = BackendMainStates.INITIALIZED + self._initialized = True + self.logger.info("Noop execution backend started") + return self + + def get_task_states_map(self): + return StateMapper(backend=self) + + async def submit_tasks(self, tasks: list[dict[str, Any]]) -> None: + if self._backend_state != BackendMainStates.RUNNING: + self._backend_state = BackendMainStates.RUNNING + + for task in tasks: + task.update( + { + "return_value": True, + "stdout": "", + "stderr": "", + "exit_code": 0, + } + ) + uid = task.setdefault("uid", f"noop.{uuid.uuid4().hex[:8]}") + self.tasks[uid] = task + # Track the completion future so it can be cancelled via + # cancel_task / cancel_all_tasks / shutdown. + self._futures[uid] = asyncio.create_task(self._complete(task)) + + async def _complete(self, task: dict) -> None: + uid = task["uid"] + try: + task["state"] = "DONE" + self._callback_func(task, "DONE") + finally: + # Drop terminal tasks so the backend does not retain every task it + # ever ran (this backend is used to benchmark millions of them). + self._futures.pop(uid, None) + self.tasks.pop(uid, None) + + async def cancel_task(self, uid: str) -> bool: + if uid not in self.tasks: + return False + + future = self._futures.pop(uid, None) + if future is not None and not future.done(): + future.cancel() + + # Remove the task (terminal); don't retain it. + task = self.tasks.pop(uid) + if task.get("state") not in ("DONE", "FAILED", "CANCELED"): + task["state"] = "CANCELED" + self._callback_func(task, "CANCELED") + return True + + async def cancel_all_tasks(self) -> int: + uids = list(self.tasks) + for uid in uids: + await self.cancel_task(uid) + return len(uids) + + async def shutdown(self) -> None: + await self.cancel_all_tasks() + self._backend_state = BackendMainStates.SHUTDOWN + self.tasks.clear() + self._futures.clear() + self.logger.info("Noop execution backend shutdown") + + def build_task(self, uid, task_desc, task_specific_kwargs): + pass + + 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: + return self._backend_state.value + + def task_state_cb(self, task, state): + pass + + 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, name: str = "noop") -> "NoopExecutionBackend": + """Alternative factory that returns an initialized backend.""" + backend = cls(name=name) + return await backend diff --git a/src/rhapsody/backends/execution/orbit.py b/src/rhapsody/backends/execution/orbit.py new file mode 100644 index 0000000..90e4fbb --- /dev/null +++ b/src/rhapsody/backends/execution/orbit.py @@ -0,0 +1,725 @@ +"""Orbit execution backend for remote task execution via ORBIT. + +This module provides a backend that submits tasks to a remote HPC node +through the ORBIT broker/endpoint infrastructure. The Endpoint node +runs a Rhapsody plugin with a local backend (e.g. Dragon V3) that +actually executes the work. + +Internally delegates to ``RhapsodyClient`` so all transport-level +optimizations (template compression, pipelined batching, event-based +wait, batch notifications) are inherited automatically. +""" + +from __future__ import annotations + +import asyncio +import logging +from typing import Any +from typing import Callable + +from ..base import BaseBackend +from ..constants import BackendMainStates +from ..constants import StateMapper + +# ``radical.orbit`` is imported at module level so that tests can patch +# ``EndpointRuntime`` here. When the package isn't installed we keep +# ``EndpointRuntime = None`` and remember the original ImportError; the +# actual chained re-raise happens in ``OrbitExecutionBackend.__init__`` +# so the user sees the real cause (e.g. a downstream import failure +# inside ``radical.orbit`` itself, not just "package missing"). +try: + from radical.orbit import EndpointRuntime + + _radical_orbit_import_error = None +except ImportError as exc: + EndpointRuntime = None + _radical_orbit_import_error = exc + +try: + import radical.prof as rprof +except ImportError: + rprof = None + + +class OrbitExecutionBackend(BaseBackend): + """Execution backend that delegates to a remote ORBIT node. + + Uses ``radical.orbit.EndpointRuntime`` and ``RhapsodyClient`` for all + communication — inheriting batching, template compression, + pipelined submission, and event-based notifications. + + When tasks are submitted individually (one at a time), a batching + layer collects them over a short time window (default 0.1 s) and + flushes them as a single bulk request, dramatically reducing + per-task HTTP round-trip overhead. + + Args: + broker_url: URL of the ORBIT broker + (e.g. ``"https://localhost:8000"``). If omitted, + falls back to the ``RADICAL_ORBIT_BROKER_URL`` env var + or ``~/.radical/orbit/broker.url`` (resolved by + ``EndpointRuntime`` itself; the bearer token is + resolved the same way from + ``RADICAL_ORBIT_BROKER_TOKEN`` / + ``~/.radical/orbit/broker.token``). + endpoint_name: Name of the endpoint to target. If omitted, the + backend auto-selects the first present endpoint whose + topology entry advertises a rhapsody plugin; the + broker and pure consumers (including this runtime + itself) are skipped. Raises ``RuntimeError`` from + ``await backend`` if no candidate is found. + backends: Backend names to request on the remote session + (default ``["dragon_v3"]``). + name: Backend name for Rhapsody registration + (default ``"orbit"``). + batch_window: Seconds to collect tasks before flushing + (default 0.25). Set to 0 to disable batching. + batch_limit: Max tasks per batch — triggers an immediate flush + when reached (default 1024). + start_timeout: Seconds to wait for broker registration and the + first topology frame (default 30). + init_timeout: Seconds to wait for the remote session to become + ready after registration (default 120). Together + with ``start_timeout`` this bounds the whole + ``await backend`` initialization. + """ + + _DEFAULT_BATCH_WINDOW = 0.25 + _PLUGIN_NAME = "rhapsody" + _TERMINAL_STATES = frozenset({"DONE", "FAILED", "CANCELED", "COMPLETED"}) + + def __init__( + self, + broker_url: str | None = None, + endpoint_name: str | None = None, + backends: list[str] | None = None, + name: str = "orbit", + plugin_name: str = _PLUGIN_NAME, + batch_window: float | None = None, + batch_limit: int = 1024, + start_timeout: float = 30.0, + init_timeout: float = 120.0, + ): + super().__init__(name=name) + + if EndpointRuntime is None: + raise ImportError( + f"OrbitExecutionBackend: cannot import radical.orbit: {_radical_orbit_import_error}" + ) from _radical_orbit_import_error + + self.logger = logging.getLogger(__name__) + self._broker_url = broker_url + self._endpoint_name = endpoint_name + self._plugin_name = plugin_name + self._remote_backends = backends or ["dragon_v3"] + self._start_timeout = start_timeout + self._init_timeout = init_timeout + self._endpoint_python_mm: tuple | None = None + self._endpoint_python_lookup_done = False + + self._runtime = None # EndpointRuntime + self._rh = None # RhapsodyClient (from get_rhapsody_handle) + self._tasks: dict[str, dict] = {} + # uids for which a terminal state callback has already fired, so the + # notification and local-cancel paths don't double-fire the same + # terminal state + self._terminal_fired: set[str] = set() + + self._callback_func: Callable = lambda t, s: None + self._initialized = False + # Serializes _async_init so concurrent first submissions (which now + # await inside init) don't run the initialization logic twice. + self._init_lock = asyncio.Lock() + self._backend_state = BackendMainStates.INITIALIZED + self._loop: asyncio.AbstractEventLoop | None = None + + # -- submission batching -- + self._batch_window = ( + batch_window if batch_window is not None else self._DEFAULT_BATCH_WINDOW + ) + self._batch_limit = batch_limit + self._batch_buffer: list[dict] = [] + # ``_batch_lock`` guards ``_batch_buffer`` / ``_flush_handle`` and is + # held only briefly (append, or detach-the-batch). ``_send_lock`` + # serializes the actual remote sends, so concurrent flushes neither + # overlap on the (not necessarily thread-safe) client nor reorder + # batches — while leaving the batch lock free for other submitters to + # keep buffering during a slow network flush. + self._batch_lock = asyncio.Lock() + self._send_lock = asyncio.Lock() + self._flush_handle: asyncio.TimerHandle | None = None + # Background (timer-triggered) flush tasks, awaited on shutdown so a + # batch that is mid-send is never dropped. + self._inflight_flushes: set[asyncio.Task] = set() + + # -- profiling -- + self._prof = rprof.Profiler("client.task", ns="radical.orbit") if rprof else None + + # ------------------------------------------------------------------ + # Async init + # ------------------------------------------------------------------ + + def __await__(self): + return self._async_init().__await__() + + async def _async_init(self): + if self._initialized: + return self + + async with self._init_lock: + # Re-check under the lock: a concurrent caller may have finished + # initialization while we waited to acquire it. + if self._initialized: + return self + + return await self._do_async_init() + + async def _do_async_init(self): + self._loop = asyncio.get_running_loop() + + # Register states + StateMapper.register_backend_states_with_defaults(backend=self) + StateMapper.register_backend_tasks_states_with_defaults(backend=self) + + # Runs blocking network I/O (EndpointRuntime.start, topology scan, + # remote session registration) so keep it off the event loop. Can + # block up to ``start_timeout + init_timeout`` — all connection and + # session waiting is paid here, never on the submit path. + self._rh = await asyncio.to_thread(self._get_rhapsody_handle) + + # Register persistent notification callback for task completions + self._rh.register_notification_callback(self._on_task_notification, topic="task_status") + self._rh.register_notification_callback( + self._on_task_notification, topic="task_status_batch" + ) + + self._initialized = True + self.logger.info( + "Orbit backend ready: %s/%s (session %s)", + self._endpoint_name, + self._plugin_name, + self._rh.sid, + ) + return self + + # ------------------------------------------------------------------ + # Notification handling + # ------------------------------------------------------------------ + + def _on_task_notification(self, endpoint, plugin, topic, data): + """Notification callback — runs on the runtime's callback thread. + + Marshal the whole update onto the event loop so that every read/write of + ``self._tasks`` and the task dicts happens single-threaded there, never + racing loop-side writers (submit / cancel / flush-error). + """ + loop = self._loop + if loop is not None and not loop.is_closed(): + loop.call_soon_threadsafe(self._handle_notification, topic, data) + + def _handle_notification(self, topic, data): + """Dispatch a notification on the event-loop thread.""" + if topic == "task_status_batch": + for t in data.get("tasks", []): + self._apply_task_update(t) + else: + self._apply_task_update(data) + + def _apply_task_update(self, body: dict): + """Apply a single task status update from a notification.""" + if not isinstance(body, dict): + return + + uid = body.get("uid") + state = body.get("state", "") + + if uid not in self._tasks: + return + + task = self._tasks[uid] + + # Decode base64-encoded return values (bytes results) + if body.get("_return_value_encoding") == "base64": + import base64 + + body["return_value"] = base64.b64decode(body["return_value"]) + del body["_return_value_encoding"] + + # A remote failure may serialize its exception as a string or dict; + # coerce it to a real exception so session.py raises it instead of + # silently treating the failed task as successful. + # ``.get(k, default)`` only falls back when the key is *absent*; a + # present-but-null "exception" (common in JSON) must still fall back to + # "error". + exc = body.get("exception") + if exc is None: + exc = body.get("error") + if exc is not None and not isinstance(exc, BaseException): + body["exception"] = RuntimeError(str(exc)) + + # Update local task dict with remote results + for key in ( + "state", + "stdout", + "stderr", + "exit_code", + "return_value", + "exception", + "traceback", + "error", + ): + if key in body: + task[key] = body[key] + + # Profile task completion on client side + if self._prof: + self._prof.prof("task_complete", uid=uid, state=state) + + # Already on the event-loop thread (dispatched from + # _on_task_notification via call_soon_threadsafe), so fire directly. + self._fire_callback(task, state) + + def _fire_callback(self, task: dict, state: str) -> None: + """Invoke the Rhapsody state callback, at most once per terminal state. + + Both the remote notification path and the local cancel path can report + the same terminal state for a task (e.g. a cancel fires the callback + locally *and* the endpoint emits a matching CANCELED notification). + This guard ensures consumers that resolve a future on the terminal + callback see it exactly once. + + Once a terminal state has fired, the task is dropped from ``_tasks`` + (and its uid from ``_terminal_fired``) to bound memory over a + long-running session. A later duplicate notification is then ignored + by ``_apply_task_update`` because the uid is no longer tracked. + + Must run on the event loop thread — ``_terminal_fired`` is accessed + without a lock, and every caller (notifications via + ``call_soon_threadsafe``, + cancel paths from coroutine bodies) already runs there. + """ + uid = task.get("uid") + terminal = bool(uid) and str(state).upper() in self._TERMINAL_STATES + if terminal: + if uid in self._terminal_fired: + return + self._terminal_fired.add(uid) + self._callback_func(task, state) + if terminal: + self._tasks.pop(uid, None) + self._terminal_fired.discard(uid) + + # ------------------------------------------------------------------ + # Endpoint auto-selection and Plugin retrieval + # ------------------------------------------------------------------ + + def _get_rhapsody_handle(self) -> Any: + """Connect the EndpointRuntime, pick an endpoint, return a RhapsodyClient. + + The endpoint is either the named one, or auto-selected as the first + present endpoint whose topology entry advertises a rhapsody plugin + (the broker and pure consumers — including this runtime itself — are + skipped). Raises ``RuntimeError`` if no candidate is found. + + Runs blocking network I/O; always called via ``asyncio.to_thread``. + """ + import uuid + + # A unique name suffix avoids the broker's name-in-use rejection when + # several rhapsody clients connect (or one restarts within the + # liveness grace window). + rt = EndpointRuntime( + broker_url=self._broker_url, + name=f"rhapsody.{uuid.uuid4().hex[:8]}", + ) + try: + rt.start(wait=True, timeout=self._start_timeout) + # start() returns silently on timeout — check registration + # explicitly, or the failure would surface further down as a + # misleading "no endpoint found". + if not rt.wait_registered(timeout=0): + raise RuntimeError( + f"failed to register with ORBIT broker {rt.broker_url!r} " + f"within {self._start_timeout}s" + ) + self._broker_url = rt.broker_url + + # find a suitable endpoint from the (local) topology snapshot + if not self._endpoint_name: + for eid, info in rt.topology().items(): + if eid == rt.name: + continue # this runtime itself + if info.get("role") in ("broker", "consumer"): + continue + if info.get("liveness") not in (None, "present"): + continue # suspect / lost + if self._plugin_name in (info.get("plugins") or {}): + self.logger.info( + "auto-selected endpoint %r (plugin %r)", eid, self._plugin_name + ) + self._endpoint_name = eid + break + + if not self._endpoint_name: + raise RuntimeError( + f"no endpoint advertises a {self._plugin_name!r} plugin " + f"on broker {self._broker_url}" + ) + + # get_plugin() registers the remote session and blocks until it + # is ready (bounded by init_timeout). + rh = rt.get_plugin( + self._endpoint_name, + self._plugin_name, + backends=self._remote_backends, + init_timeout=self._init_timeout, + ) + except Exception: + # Don't leak the runtime's daemon threads / WebSocket on a failed + # init — callers that never reach shutdown() would keep it alive. + rt.stop() + raise + + self._runtime = rt + return rh + + # ------------------------------------------------------------------ + # Python-version compatibility for cloudpickled function tasks + # ------------------------------------------------------------------ + + async def _check_python_compat(self, tasks: list) -> None: + """Raise if any cloudpickled task in *tasks* would hit a Python version mismatch on the + remote endpoint. No check is performed if the batch contains only executable or import-path + tasks, or if the endpoint's Python version could not be determined. + + Cloudpickle serializes function bytecode using ``CodeType``, + whose tuple shape changed between Python 3.10 and 3.11 — any + cross-minor-version skew fails deserialization on the endpoint. + """ + import sys + + def needs_compat(t: dict) -> bool: + fn = t.get("function") + # A raw Python callable (a ComputeTask function not yet serialized) + # takes the same cloudpickle path once submitted, so it must be + # checked too — not just already-``cloudpickle::``-encoded strings. + return ( + callable(fn) + or (isinstance(fn, str) and fn.startswith("cloudpickle::")) + or bool(t.get("_pickled_fields")) + ) + + if not any(needs_compat(t) for t in tasks): + return + + if not self._endpoint_python_lookup_done: + # Only the remote lookup is transient (network / plugin not yet + # ready) — on failure we leave the check armed so a later + # submission retries. A successful lookup settles the check for + # good, even if the reported version can't be parsed (that's + # deterministic and retrying wouldn't help). + info = None + + def _lookup(): + # get_plugin() auto-registers an ephemeral session even + # though host_role() needs none — close the client so failed + # retries don't accumulate sessions on the endpoint. + si = self._runtime.get_plugin(self._endpoint_name, "sysinfo") + try: + return si.host_role() + finally: + si.close() + + try: + # Blocking network I/O — run off the event loop. + info = await asyncio.to_thread(_lookup) + except Exception as exc: + self.logger.debug(f"pickle compat lookup failed, will retry: {exc}") + + if info is not None: + self._endpoint_python_lookup_done = True + pyver = info.get("python_version") or "" + parts = pyver.split(".") + try: + if len(parts) >= 2: + self._endpoint_python_mm = (int(parts[0]), int(parts[1])) + except ValueError: + pass + if not self._endpoint_python_mm: + self.logger.debug(f"skip pickle compat check ({pyver})") + + if not self._endpoint_python_mm: + return + + client_mm = (sys.version_info.major, sys.version_info.minor) + if self._endpoint_python_mm != client_mm: + raise RuntimeError( + f"function tasks cannot be submitted to endpoint " + f"{self._endpoint_name!r}: client Python " + f"{client_mm[0]}.{client_mm[1]} != endpoint Python " + f"{self._endpoint_python_mm[0]}.{self._endpoint_python_mm[1]}. " + f"cloudpickle is not portable across Python minor " + f"versions — align the venvs, or use executable / " + f"import-path tasks for this endpoint." + ) + + # ------------------------------------------------------------------ + # BaseBackend interface + # ------------------------------------------------------------------ + + async def submit_tasks(self, tasks: list[dict[str, Any]]) -> None: + """Submit tasks to the remote endpoint for execution. + + When batching is enabled (batch_window > 0), tasks are collected + in an internal buffer and flushed after the window expires or the + buffer reaches ``batch_limit``. This turns many small individual + submissions into fewer bulk HTTP requests. + + When batching is disabled (batch_window == 0), delegates directly + to ``RhapsodyClient.submit_tasks()``. + """ + if not self._initialized: + await self._async_init() + + if self._backend_state != BackendMainStates.RUNNING: + self._backend_state = BackendMainStates.RUNNING + + # Fail fast on cross-Python cloudpickle skew *before* registering or + # buffering, so the error always propagates to this caller. When + # batching defers the flush, the check would otherwise run inside an + # orphaned ``ensure_future`` task and never reach the awaiter. + await self._check_python_compat(tasks) + + # Assign UIDs, register locally, emit submit prof events — single pass + import uuid + + prof = self._prof + for task in tasks: + task.setdefault("uid", f"task.{uuid.uuid4().hex[:8]}") + self._tasks[task["uid"]] = task + if prof: + prof.prof("task_submit", uid=task["uid"]) + + # No batching — submit immediately + if self._batch_window <= 0: + task_dicts = [dict(t) for t in tasks] + await asyncio.to_thread(self._rh.submit_tasks, task_dicts) + return + + # Batching — collect and, if the buffer is now full, flush inline. + # The batch is detached *under* the lock but sent *outside* it, so a + # slow network send never blocks other submitters from buffering. + batch = None + async with self._batch_lock: + self._batch_buffer.extend(dict(t) for t in tasks) + + if len(self._batch_buffer) >= self._batch_limit: + # Buffer full — take the batch now, send it below. + batch = self._detach_batch() + elif self._flush_handle is None: + # Start the timer for the first task in this window + loop = asyncio.get_running_loop() + self._flush_handle = loop.call_later(self._batch_window, self._trigger_flush) + + if batch is not None: + await self._send_batch(batch) + + def _trigger_flush(self): + """Timer callback — schedule a background flush on the event loop. + + The task is tracked in ``_inflight_flushes`` so ``shutdown`` can await + an in-progress send instead of dropping the batch. + """ + task = asyncio.ensure_future(self._locked_flush()) + self._inflight_flushes.add(task) + task.add_done_callback(self._on_flush_done) + + def _on_flush_done(self, task: asyncio.Task) -> None: + """Untrack a finished background flush and retrieve its exception. + + ``_send_batch`` already logs the failure and marks tasks FAILED; we + retrieve the exception here only so asyncio does not emit a spurious + "Task exception was never retrieved" warning. + """ + self._inflight_flushes.discard(task) + if not task.cancelled(): + exc = task.exception() + if exc is not None: + self.logger.debug("background batch flush failed: %s", exc) + + async def _locked_flush(self): + """Detach the buffered batch under the batch lock, then send it unlocked.""" + async with self._batch_lock: + batch = self._detach_batch() + await self._send_batch(batch) + + def _detach_batch(self) -> list[dict]: + """Swap out the current buffer and cancel the pending flush timer. + + Must be called while holding ``_batch_lock``. Returns the detached + batch (possibly empty). + """ + batch = self._batch_buffer + self._batch_buffer = [] + + if self._flush_handle is not None: + self._flush_handle.cancel() + self._flush_handle = None + + return batch + + async def _send_batch(self, batch: list[dict]) -> None: + """Send an already-detached batch to the remote endpoint in one request. + + Runs *outside* ``_batch_lock`` (so other submitters can keep buffering + during the network I/O) but *under* ``_send_lock`` (so concurrent + flushes stay serialized and ordered, and never call the client + concurrently). + """ + if not batch: + return + + # NOTE: compat is validated up front in submit_tasks(); buffered tasks + # are already known-good by the time they reach the flush. + + prof = self._prof + if prof: + for t in batch: + prof.prof("task_batch_flush", uid=t.get("uid", "?")) + + self.logger.debug("Flushing batch of %d tasks", len(batch)) + async with self._send_lock: + try: + await asyncio.to_thread(self._rh.submit_tasks, batch) + except Exception as exc: + # A timer-triggered flush runs as a detached background task; an + # unhandled error here would leave the batch's tasks hung + # forever. Fail them explicitly so the failure reaches the + # client (and re-raise for the inline buffer-full caller). + self.logger.error("Failed to submit batch of %d tasks: %s", len(batch), exc) + for t in batch: + task = self._tasks.get(t.get("uid")) + if task is not None: + task["exception"] = exc + task["state"] = "FAILED" + self._fire_callback(task, "FAILED") + raise + + async def cancel_task(self, uid: str) -> bool: + """Cancel a single task on the remote endpoint.""" + if not self._initialized: + await self._async_init() + + if uid not in self._tasks: + return False + + # If the task is still buffered locally (not yet flushed), the endpoint + # never received it — drop it from the buffer and skip the remote + # cancel (which would target an unknown uid). + async with self._batch_lock: + buffered = any(t.get("uid") == uid for t in self._batch_buffer) + if buffered: + self._batch_buffer = [t for t in self._batch_buffer if t.get("uid") != uid] + + if not buffered: + await asyncio.to_thread(self._rh.cancel_task, uid) + + task = self._tasks[uid] + task["state"] = "CANCELED" + self._fire_callback(task, "CANCELED") + return True + + async def cancel_all_tasks(self) -> int: + """Cancel all non-terminal tasks on the remote endpoint. + + Mirrors ``cancel_task``: marks each non-terminal local task CANCELED + and fires its state callback, so consumers that resolve futures via + callbacks don't hang. ``_fire_callback`` dedups against the matching + CANCELED remote notification, so each task's terminal callback fires + once. + """ + if not self._initialized: + await self._async_init() + + # Drop anything still buffered locally so it is not flushed and executed + # after cancellation; the remote cancel_all only covers submitted tasks. + async with self._batch_lock: + self._batch_buffer = [] + if self._flush_handle is not None: + self._flush_handle.cancel() + self._flush_handle = None + + result = await asyncio.to_thread(self._rh.cancel_all_tasks) + + for task in list(self._tasks.values()): + if str(task.get("state", "")).upper() not in self._TERMINAL_STATES: + task["state"] = "CANCELED" + self._fire_callback(task, "CANCELED") + + return result.get("canceled", 0) + + async def shutdown(self) -> None: + """Flush pending tasks, close session and EndpointRuntime.""" + try: + await self._locked_flush() + except Exception as e: + # Never let a flush failure skip client cleanup below. + self.logger.warning("Failed to flush pending tasks during shutdown: %s", e) + + # Wait for any timer-triggered flush already in flight so its batch is + # sent before we tear down the clients. + if self._inflight_flushes: + await asyncio.gather(*list(self._inflight_flushes), return_exceptions=True) + + self._backend_state = BackendMainStates.SHUTDOWN + + # close() calls perform blocking network I/O — run them off the loop. + if self._rh: + try: + await asyncio.to_thread(self._rh.close) + except Exception as e: + self.logger.warning("Failed to close session: %s", e) + self._rh = None + + # stop() joins the runtime's threads and does blocking teardown calls. + if self._runtime: + try: + await asyncio.to_thread(self._runtime.stop) + except Exception as e: + self.logger.warning("Failed to stop endpoint runtime: %s", e) + self._runtime = None + + if self._prof: + self._prof.close() + + self.logger.info("Orbit execution backend shutdown complete") + + async def state(self) -> str: + return self._backend_state.value + + def task_state_cb(self, task, state): + pass + + def get_task_states_map(self): + return StateMapper(backend=self) + + def build_task(self, uid, task_desc, task_specific_kwargs): + pass + + 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 + + # ------------------------------------------------------------------ + # Context manager + # ------------------------------------------------------------------ + + async def __aenter__(self): + if not self._initialized: + await self._async_init() + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + await self.shutdown() diff --git a/tests/unit/test_backend_execution_dragon.py b/tests/unit/test_backend_execution_dragon.py index 5aa8551..b4a1d84 100644 --- a/tests/unit/test_backend_execution_dragon.py +++ b/tests/unit/test_backend_execution_dragon.py @@ -644,7 +644,7 @@ def test_v3_deliver_batch_success_stores_value_and_fires_done(backend_v3): task_desc = {"uid": uid} backend_v3._task_registry[uid] = {"uid": uid, "description": task_desc} - backend_v3._deliver_batch([(uid, 42, None, False, None, None)]) + backend_v3._deliver_batch([(uid, 42, None, False, None, None, None, None)]) assert task_desc["return_value"] == 42 assert task_desc["stdout"] == "" @@ -659,7 +659,7 @@ def test_v3_deliver_batch_propagates_stdout_stderr(backend_v3): task_desc = {"uid": uid} backend_v3._task_registry[uid] = {"uid": uid, "description": task_desc} - backend_v3._deliver_batch([(uid, "ok", None, False, "hello\n", "warn\n")]) + backend_v3._deliver_batch([(uid, "ok", None, False, "hello\n", "warn\n", None, None)]) assert task_desc["stdout"] == "hello\n" assert task_desc["stderr"] == "warn\n" @@ -673,7 +673,7 @@ def test_v3_deliver_batch_failure_stores_exc_and_fires_failed(backend_v3): 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_v3._deliver_batch([(uid, exc, None, True, None, None, None, None)]) assert task_desc["exception"] is exc assert "something went wrong" in task_desc["stderr"] @@ -689,11 +689,29 @@ def test_v3_deliver_batch_prefers_traceback_over_str_exc(backend_v3): exc = RuntimeError("boom") tb = "Traceback (most recent call last):\n File ...\nRuntimeError: boom" - backend_v3._deliver_batch([(uid, exc, tb, True, None, None)]) + backend_v3._deliver_batch([(uid, exc, tb, True, None, None, None, None)]) assert task_desc["stderr"] == tb +def test_v3_deliver_batch_non_exception_result_coerced_to_baseexception(backend_v3): + """A non-BaseException failure result must still surface as a raisable + exception, otherwise session.py would silently treat the failure as DONE.""" + uid = "task.unit-failed-str" + task_desc = {"uid": uid} + backend_v3._task_registry[uid] = {"uid": uid, "description": task_desc} + + # result is a plain string (e.g. a cross-node serialization fallback) and a + # per-rank diagnostic is present, so the augmentation path runs. + backend_v3._deliver_batch( + [(uid, "remote failure", None, True, None, None, None, "[rank 0]\ntb")] + ) + + assert isinstance(task_desc["exception"], BaseException) + assert "remote failure" in str(task_desc["exception"]) + backend_v3._callback_func.assert_called_once_with(task_desc, "FAILED") + + def test_v3_cancelled_task_skips_callback(backend_v3): """_deliver_batch is a no-op for UIDs in _cancelled_tasks.""" uid = "task.unit-cancelled" @@ -701,7 +719,7 @@ def test_v3_cancelled_task_skips_callback(backend_v3): backend_v3._task_registry[uid] = {"uid": uid, "description": task_desc} backend_v3._cancelled_tasks.add(uid) - backend_v3._deliver_batch([(uid, 99, None, False, None, None)]) + backend_v3._deliver_batch([(uid, 99, None, False, None, None, None, None)]) backend_v3._callback_func.assert_not_called() assert "return_value" not in task_desc @@ -1036,7 +1054,7 @@ def test_v3_deliver_batch_empty_dragon_stdout_written_as_empty_string(backend_v3 task_desc = {"uid": uid} backend_v3._task_registry[uid] = {"uid": uid, "description": task_desc} - backend_v3._deliver_batch([(uid, 0, None, False, "", "")]) + backend_v3._deliver_batch([(uid, 0, None, False, "", "", None, None)]) assert isinstance(task_desc["stdout"], str) assert isinstance(task_desc["stderr"], str) @@ -1053,9 +1071,13 @@ def test_v3_deliver_batch_stdout_path_takes_priority_over_empty_stdout(backend_v "description": task_desc, "stdout_path": "/work/uid.stdout", "stderr_path": "/work/uid.stderr", + # capture_stdio redirect writes a wrapper script alongside the literal + # stdout/stderr files; script_path marks this as an exec redirect (vs a + # function-wrap glob prefix), which is what makes the path take priority. + "script_path": "/work/uid.sh", } - backend_v3._deliver_batch([(uid, 0, None, False, "", "")]) + backend_v3._deliver_batch([(uid, 0, None, False, "", "", None, None)]) assert task_desc["stdout"] == "/work/uid.stdout" assert task_desc["stderr"] == "/work/uid.stderr" @@ -1068,6 +1090,6 @@ def test_v3_deliver_batch_failed_stdout_is_string(backend_v3): task_desc = {"uid": uid} backend_v3._task_registry[uid] = {"uid": uid, "description": task_desc} - backend_v3._deliver_batch([(uid, RuntimeError("boom"), None, True, "", "")]) + backend_v3._deliver_batch([(uid, RuntimeError("boom"), None, True, "", "", None, None)]) assert isinstance(task_desc["stdout"], str) diff --git a/tests/unit/test_backend_execution_orbit.py b/tests/unit/test_backend_execution_orbit.py new file mode 100644 index 0000000..1974cb9 --- /dev/null +++ b/tests/unit/test_backend_execution_orbit.py @@ -0,0 +1,895 @@ +"""Unit tests for OrbitExecutionBackend (refactored: delegates to RhapsodyClient).""" + +import asyncio +import sys +import threading +import time +from unittest.mock import AsyncMock +from unittest.mock import MagicMock +from unittest.mock import patch + +import pytest + +import rhapsody.backends.execution.orbit as orbit_mod +from rhapsody.backends.execution.orbit import OrbitExecutionBackend + + +@pytest.fixture(autouse=True) +def _stub_endpoint_runtime(): + """Allow the suite to run without ``radical.orbit`` installed (e.g. in CI). + + ``OrbitExecutionBackend.__init__`` raises ``ImportError`` when the optional + ``radical.orbit`` package is missing (module-level ``EndpointRuntime`` is + ``None``). Every test mocks the runtime/rhapsody chain anyway, so when the + real package is absent stub the symbol to a non-``None`` mock so the + backend can be constructed. When it is installed this is a no-op. + """ + if orbit_mod.EndpointRuntime is not None: + yield + return + with patch.object(orbit_mod, "EndpointRuntime", MagicMock()): + yield + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _mock_rhapsody_client(sid="session.abc123"): + """Return a mock RhapsodyClient (PluginClient).""" + rh = MagicMock() + rh.sid = sid + rh.submit_tasks = MagicMock(return_value=[{"uid": "t.001", "state": "SUBMITTED"}]) + rh.cancel_task = MagicMock(return_value={"uid": "t.001", "status": "canceled"}) + rh.cancel_all_tasks = MagicMock(return_value={"canceled": 5}) + rh.close = MagicMock() + rh.register_notification_callback = MagicMock() + return rh + + +def _mock_runtime(rh=None, topology=None): + """Return a mock EndpointRuntime whose ``get_plugin`` produces *rh*.""" + if rh is None: + rh = _mock_rhapsody_client() + rt = MagicMock() + # NOTE: ``MagicMock(name=...)`` sets the mock's repr, not the attribute — + # ``name`` must be assigned after construction. + rt.name = "rhapsody.test0000" + rt.broker_url = "http://localhost:8000" + rt.start = MagicMock(return_value=rt) + rt.wait_registered = MagicMock(return_value=True) + rt.topology = MagicMock(return_value=topology or {}) + rt.get_plugin = MagicMock(return_value=rh) + rt.stop = MagicMock() + return rt, rh + + +def _make_backend(**kwargs): + """Create an OrbitExecutionBackend (not yet initialised).""" + defaults = { + "broker_url": "http://localhost:8000", + "endpoint_name": "test_endpoint", + } + defaults.update(kwargs) + return OrbitExecutionBackend(**defaults) + + +async def _init_backend(**kwargs): + """Create and initialise a backend with a mocked EndpointRuntime chain.""" + topology = kwargs.pop("topology", None) + backend = _make_backend(**kwargs) + rt, rh = _mock_runtime(topology=topology) + with patch("rhapsody.backends.execution.orbit.EndpointRuntime", return_value=rt): + await backend._async_init() + # Expose mocks for assertions + backend._mock_rt = rt + backend._mock_rh = rh + return backend + + +# --------------------------------------------------------------------------- +# Construction +# --------------------------------------------------------------------------- + + +def test_endpoint_backend_construction(): + backend = _make_backend() + assert backend._broker_url == "http://localhost:8000" + assert backend._endpoint_name == "test_endpoint" + assert backend._plugin_name == "rhapsody" + assert backend._remote_backends == ["dragon_v3"] + assert backend._initialized is False + + +def test_endpoint_backend_custom_params(): + backend = _make_backend( + plugin_name="my_rhapsody", + backends=["dragon_v3"], + name="my_endpoint", + ) + assert backend._plugin_name == "my_rhapsody" + assert backend._remote_backends == ["dragon_v3"] + assert backend.name == "my_endpoint" + + +# --------------------------------------------------------------------------- +# Async init +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_async_init_creates_client_chain(): + backend = await _init_backend() + + assert backend._initialized is True + assert backend._runtime is not None + assert backend._rh is not None + + # runtime started and registration verified + backend._mock_rt.start.assert_called_once_with(wait=True, timeout=30.0) + # get_plugin called with endpoint + plugin name + session kwargs + backend._mock_rt.get_plugin.assert_called_once_with( + "test_endpoint", "rhapsody", backends=["dragon_v3"], init_timeout=120.0 + ) + + # Notification callbacks registered + calls = backend._mock_rh.register_notification_callback.call_args_list + topics = [c[1]["topic"] for c in calls] + assert "task_status" in topics + assert "task_status_batch" in topics + + +@pytest.mark.asyncio +async def test_async_init_idempotent(): + backend = await _init_backend() + rh = backend._mock_rh + + await backend._async_init() + # register_notification_callback should NOT be called again + assert rh.register_notification_callback.call_count == 2 # initial only + + +# --------------------------------------------------------------------------- +# submit_tasks +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_submit_tasks_delegates_to_rhapsody_client(): + backend = await _init_backend(batch_window=0) + + tasks = [{"uid": "t.001", "executable": "/bin/echo", "arguments": ["hi"]}] + await backend.submit_tasks(tasks) + + backend._mock_rh.submit_tasks.assert_called_once() + submitted = backend._mock_rh.submit_tasks.call_args[0][0] + assert len(submitted) == 1 + assert submitted[0]["uid"] == "t.001" + + # Task tracked locally + assert "t.001" in backend._tasks + + +@pytest.mark.asyncio +async def test_submit_tasks_sets_running_state(): + backend = await _init_backend() + assert await backend.state() == "INITIALIZED" + + await backend.submit_tasks([{"uid": "t.1", "executable": "/bin/true"}]) + assert await backend.state() == "RUNNING" + + +# --------------------------------------------------------------------------- +# cancel_task +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_cancel_task(): + backend = await _init_backend() + backend._tasks["t.001"] = {"uid": "t.001", "state": "RUNNING"} + cb = MagicMock() + backend.register_callback(cb) + + result = await backend.cancel_task("t.001") + assert result is True + # terminal task is delivered via callback, then pruned from the registry + task, state = cb.call_args[0] + assert task["uid"] == "t.001" + assert state == "CANCELED" + assert task["state"] == "CANCELED" + assert "t.001" not in backend._tasks + backend._mock_rh.cancel_task.assert_called_once_with("t.001") + + +@pytest.mark.asyncio +async def test_cancel_unknown_task(): + backend = await _init_backend() + result = await backend.cancel_task("no_such_task") + assert result is False + + +# --------------------------------------------------------------------------- +# cancel_all_tasks +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_cancel_all_tasks(): + backend = await _init_backend() + count = await backend.cancel_all_tasks() + assert count == 5 + backend._mock_rh.cancel_all_tasks.assert_called_once() + + +# --------------------------------------------------------------------------- +# shutdown +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_shutdown_closes_clients(): + backend = await _init_backend() + await backend.shutdown() + + backend._mock_rh.close.assert_called_once() + backend._mock_rt.stop.assert_called_once() + assert backend._rh is None + assert backend._runtime is None + assert await backend.state() == "SHUTDOWN" + + +# --------------------------------------------------------------------------- +# Notification handling +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_on_task_notification_single(): + # _on_task_notification marshals to the event loop; drive one tick to let + # the scheduled handler run before asserting. + backend = await _init_backend() + backend._tasks["t.001"] = {"uid": "t.001", "state": "SUBMITTED"} + cb = MagicMock() + backend.register_callback(cb) + + backend._on_task_notification( + endpoint="hpc1", + plugin="rhapsody", + topic="task_status", + data={"uid": "t.001", "state": "DONE", "stdout": "hello\n", "exit_code": 0}, + ) + await asyncio.sleep(0) + + task, state = cb.call_args[0] + assert state == "DONE" + assert task["state"] == "DONE" + assert task["stdout"] == "hello\n" + assert "t.001" not in backend._tasks # terminal -> pruned + + +@pytest.mark.asyncio +async def test_on_task_notification_batch(): + backend = await _init_backend() + backend._tasks["t.001"] = {"uid": "t.001", "state": "SUBMITTED"} + backend._tasks["t.002"] = {"uid": "t.002", "state": "SUBMITTED"} + cb = MagicMock() + backend.register_callback(cb) + + backend._on_task_notification( + endpoint="hpc1", + plugin="rhapsody", + topic="task_status_batch", + data={ + "tasks": [ + {"uid": "t.001", "state": "DONE"}, + {"uid": "t.002", "state": "FAILED", "error": "boom"}, + ] + }, + ) + await asyncio.sleep(0) + + fired = {c.args[0]["uid"]: c.args for c in cb.call_args_list} + assert fired["t.001"][1] == "DONE" + assert fired["t.002"][1] == "FAILED" + assert fired["t.002"][0]["error"] == "boom" + # both terminal -> pruned + assert "t.001" not in backend._tasks + assert "t.002" not in backend._tasks + + +@pytest.mark.asyncio +async def test_on_task_notification_ignores_unknown_task(): + backend = await _init_backend() + # No tasks registered — should not crash + backend._on_task_notification( + endpoint="hpc1", + plugin="rhapsody", + topic="task_status", + data={"uid": "unknown", "state": "DONE"}, + ) + await asyncio.sleep(0) + + +@pytest.mark.asyncio +async def test_on_task_notification_coerces_string_exception(): + """A string exception serialized over the wire must become a real BaseException so session.py + raises it instead of silently treating the failed task as successful.""" + backend = await _init_backend() + backend._tasks["t.001"] = {"uid": "t.001", "state": "SUBMITTED"} + cb = MagicMock() + backend.register_callback(cb) + + backend._on_task_notification( + endpoint="hpc1", + plugin="rhapsody", + topic="task_status", + data={"uid": "t.001", "state": "FAILED", "exception": "remote boom"}, + ) + await asyncio.sleep(0) + + task, _ = cb.call_args[0] + exc = task["exception"] + assert isinstance(exc, BaseException) + assert "remote boom" in str(exc) + + +# --------------------------------------------------------------------------- +# state / context manager +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_state(): + backend = await _init_backend() + assert await backend.state() == "INITIALIZED" + + +@pytest.mark.asyncio +async def test_context_manager(): + backend = _make_backend() + rt, rh = _mock_runtime() + with patch("rhapsody.backends.execution.orbit.EndpointRuntime", return_value=rt): + async with backend as b: + assert b._initialized is True + assert await b.state() == "SHUTDOWN" + + +# --------------------------------------------------------------------------- +# Endpoint auto-selection from the topology + registration failure +# --------------------------------------------------------------------------- + + +def _topo_entry(role, liveness="present", plugins=()): + return { + "role": role, + "liveness": liveness, + "plugins": {p: {"namespace": f"/{p}", "enabled": True} for p in plugins}, + } + + +@pytest.mark.asyncio +async def test_auto_select_endpoint_from_topology(): + """With no endpoint_name, the first present endpoint advertising a rhapsody plugin is picked.""" + topology = { + "ep0": _topo_entry("endpoint", plugins=["sysinfo"]), # no rhapsody + "ep1": _topo_entry("endpoint", plugins=["rhapsody", "sysinfo"]), + "ep2": _topo_entry("endpoint", plugins=["rhapsody"]), + } + backend = await _init_backend(endpoint_name=None, topology=topology) + + assert backend._endpoint_name == "ep1" + backend._mock_rt.get_plugin.assert_called_once_with( + "ep1", "rhapsody", backends=["dragon_v3"], init_timeout=120.0 + ) + + +@pytest.mark.asyncio +async def test_auto_select_skips_broker_consumer_self_and_lost(): + """The broker (even if it hosts a rhapsody plugin), pure consumers, this runtime's own topology + entry, and non-present endpoints must all be skipped.""" + topology = { + "broker": _topo_entry("broker", plugins=["rhapsody"]), + "rhapsody.test0000": _topo_entry("consumer", plugins=["rhapsody"]), # self + "watcher": _topo_entry("consumer", plugins=["rhapsody"]), + "ep.lost": _topo_entry("endpoint", liveness="lost", plugins=["rhapsody"]), + "ep.good": _topo_entry("endpoint", plugins=["rhapsody"]), + } + backend = await _init_backend(endpoint_name=None, topology=topology) + + assert backend._endpoint_name == "ep.good" + + +@pytest.mark.asyncio +async def test_auto_select_no_candidate_raises(): + """An empty/ineligible topology raises RuntimeError and stops the runtime.""" + topology = { + "broker": _topo_entry("broker", plugins=["rhapsody"]), + "ep0": _topo_entry("endpoint", plugins=["sysinfo"]), + } + backend = _make_backend(endpoint_name=None) + rt, _ = _mock_runtime(topology=topology) + with patch("rhapsody.backends.execution.orbit.EndpointRuntime", return_value=rt): + with pytest.raises(RuntimeError, match="no endpoint advertises"): + await backend._async_init() + + rt.stop.assert_called_once() # failed init must not leak the runtime + assert backend._initialized is False + + +@pytest.mark.asyncio +async def test_registration_timeout_raises(): + """A silent start() timeout must be detected as a missed registration, failing with a message + naming the broker — not a misleading 'no endpoint found'.""" + backend = _make_backend() + rt, _ = _mock_runtime() + rt.wait_registered = MagicMock(return_value=False) + with patch("rhapsody.backends.execution.orbit.EndpointRuntime", return_value=rt): + with pytest.raises(RuntimeError, match="failed to register with ORBIT broker"): + await backend._async_init() + + rt.stop.assert_called_once() + assert backend._initialized is False + + +# --------------------------------------------------------------------------- +# Python-version compat for cloudpickled function tasks +# +# The check fires only for tasks that carry a cloudpickled function or +# pickled fields. Executable / import-path function tasks bypass it. +# It queries sysinfo.host_role() once and caches the (major, minor) +# tuple on the backend instance. +# --------------------------------------------------------------------------- + + +def _runtime_with_sysinfo(rh, sysinfo_python_version): + """Runtime mock whose ``get_plugin(eid, 'sysinfo')`` returns a sysinfo plugin reporting the + given python_version, and whose ``get_plugin(eid, 'rhapsody', ...)`` returns *rh*.""" + sysinfo = MagicMock() + sysinfo.host_role = MagicMock( + return_value={ + "role": "compute", + "scheduler": "slurm", + "psij_executor": "slurm", + "job_id": "12345", + "python_version": sysinfo_python_version, + } + ) + sysinfo.close = MagicMock() + + rt, _ = _mock_runtime(rh=rh) + + def _get_plugin(eid, pname, **kwargs): + return sysinfo if pname == "sysinfo" else rh + + rt.get_plugin = MagicMock(side_effect=_get_plugin) + return rt, sysinfo + + +async def _init_backend_with_sysinfo(sysinfo_python_version, **kwargs): + """Like _init_backend but with a sysinfo plugin returning the given python_version on + host_role().""" + kwargs.setdefault("batch_window", 0) # immediate flush, no timer + backend = _make_backend(**kwargs) + rh = _mock_rhapsody_client() + rt, si = _runtime_with_sysinfo(rh, sysinfo_python_version) + with patch("rhapsody.backends.execution.orbit.EndpointRuntime", return_value=rt): + await backend._async_init() + backend._mock_rt = rt + backend._mock_rh = rh + backend._mock_sysinfo = si + return backend + + +@pytest.mark.asyncio +async def test_python_compat_skipped_for_executable_tasks(): + """No sysinfo lookup, no exception when batch contains only executable / import-path tasks.""" + client_mm = f"{sys.version_info.major}.{sys.version_info.minor}.0" + backend = await _init_backend_with_sysinfo(client_mm) + + await backend.submit_tasks( + [ + {"uid": "t.1", "executable": "/bin/true"}, + {"uid": "t.2", "function": "mod:func"}, # import-path, not pickled + ] + ) + + backend._mock_rh.submit_tasks.assert_called_once() + backend._mock_sysinfo.host_role.assert_not_called() + + +@pytest.mark.asyncio +async def test_python_compat_passes_for_matching_versions(): + """Cloudpickled task + matching endpoint Python -> submission proceeds.""" + client_mm = f"{sys.version_info.major}.{sys.version_info.minor}.0" + backend = await _init_backend_with_sysinfo(client_mm) + + await backend.submit_tasks( + [ + {"uid": "t.1", "function": "cloudpickle::ABCDEF"}, + ] + ) + + backend._mock_rh.submit_tasks.assert_called_once() + backend._mock_sysinfo.host_role.assert_called_once() + + +@pytest.mark.asyncio +async def test_python_compat_raises_on_mismatch(): + """Cloudpickled task + endpoint on a different Python minor -> RuntimeError; + rhapsody.submit_tasks is NOT called.""" + # Pick a minor version guaranteed to differ from the test runner's. + other_minor = sys.version_info.minor + 1 + endpoint_pyver = f"{sys.version_info.major}.{other_minor}.0" + backend = await _init_backend_with_sysinfo(endpoint_pyver) + + with pytest.raises(RuntimeError, match="cloudpickle is not portable"): + await backend.submit_tasks( + [ + {"uid": "t.1", "function": "cloudpickle::ABCDEF"}, + ] + ) + + backend._mock_rh.submit_tasks.assert_not_called() + + +@pytest.mark.asyncio +async def test_python_compat_caches_first_lookup(): + """Subsequent submits of cloudpickled tasks hit the sysinfo plugin only once (cached on the + backend instance).""" + client_mm = f"{sys.version_info.major}.{sys.version_info.minor}.0" + backend = await _init_backend_with_sysinfo(client_mm) + + for _ in range(3): + await backend.submit_tasks( + [ + {"uid": f"t.{_}", "function": "cloudpickle::ABCDEF"}, + ] + ) + + assert backend._mock_rh.submit_tasks.call_count == 3 + backend._mock_sysinfo.host_role.assert_called_once() + + +@pytest.mark.asyncio +async def test_python_compat_pickled_fields_marker_triggers_check(): + """A task with ``_pickled_fields`` (no cloudpickle:: prefix) still + triggers the check, because the deserialization path is the same.""" + other_minor = sys.version_info.minor + 1 + endpoint_pyver = f"{sys.version_info.major}.{other_minor}.0" + backend = await _init_backend_with_sysinfo(endpoint_pyver) + + with pytest.raises(RuntimeError, match="cloudpickle is not portable"): + await backend.submit_tasks( + [ + {"uid": "t.1", "function": "mod:func", "_pickled_fields": ["args"]}, + ] + ) + backend._mock_rh.submit_tasks.assert_not_called() + + +@pytest.mark.asyncio +async def test_python_compat_fails_fast_under_default_batching(): + """The compat check must propagate to the submit_tasks() caller even when batching is enabled + (the common case) — not vanish into a deferred flush. + + The offending task must not be registered locally either. + """ + other_minor = sys.version_info.minor + 1 + endpoint_pyver = f"{sys.version_info.major}.{other_minor}.0" + backend = await _init_backend_with_sysinfo(endpoint_pyver, batch_window=0.25) + + with pytest.raises(RuntimeError, match="cloudpickle is not portable"): + await backend.submit_tasks( + [ + {"uid": "t.1", "function": "cloudpickle::ABCDEF"}, + ] + ) + + backend._mock_rh.submit_tasks.assert_not_called() + assert "t.1" not in backend._tasks + + +@pytest.mark.asyncio +async def test_sysinfo_client_closed_after_compat_lookup(): + """The throwaway sysinfo client is closed after the lookup so its auto-registered ephemeral + session does not linger on the endpoint.""" + client_mm = f"{sys.version_info.major}.{sys.version_info.minor}.0" + backend = await _init_backend_with_sysinfo(client_mm) + + await backend.submit_tasks([{"uid": "t.1", "function": "cloudpickle::ABCDEF"}]) + + backend._mock_sysinfo.host_role.assert_called_once() + backend._mock_sysinfo.close.assert_called_once() + + +@pytest.mark.asyncio +async def test_python_compat_retries_after_lookup_failure(): + """A transient sysinfo lookup failure must not permanently disable the check — the next + submission retries the lookup.""" + client_mm = f"{sys.version_info.major}.{sys.version_info.minor}.0" + backend = await _init_backend_with_sysinfo(client_mm) + + # First lookup blows up (e.g. plugin not ready); second succeeds. + backend._mock_sysinfo.host_role.side_effect = [ + RuntimeError("endpoint not ready"), + {"role": "compute", "python_version": client_mm}, + ] + + # First submit: lookup fails -> check skipped, submission still proceeds, + # and the check stays armed for a retry. + await backend.submit_tasks([{"uid": "t.1", "function": "cloudpickle::X"}]) + assert backend._endpoint_python_lookup_done is False + + # Second submit: lookup retried and settles. + await backend.submit_tasks([{"uid": "t.2", "function": "cloudpickle::X"}]) + assert backend._endpoint_python_lookup_done is True + assert backend._mock_sysinfo.host_role.call_count == 2 + assert backend._mock_rh.submit_tasks.call_count == 2 + + +# --------------------------------------------------------------------------- +# cancel_all_tasks: local state + callbacks, and terminal-callback dedup +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_cancel_all_updates_local_state_and_callbacks(): + """cancel_all_tasks marks every non-terminal local task CANCELED and fires its callback, leaving + already-terminal tasks untouched.""" + backend = await _init_backend() + backend._tasks = { + "t.1": {"uid": "t.1", "state": "RUNNING"}, + "t.2": {"uid": "t.2", "state": "DONE"}, # terminal -> skip + "t.3": {"uid": "t.3", "state": "SUBMITTED"}, + } + cb = MagicMock() + backend.register_callback(cb) + + await backend.cancel_all_tasks() + + # cancelled tasks fire CANCELED then are pruned; the already-terminal one + # is left untouched. + assert "t.1" not in backend._tasks + assert "t.3" not in backend._tasks + assert backend._tasks["t.2"]["state"] == "DONE" + + fired = {call.args[0]["uid"] for call in cb.call_args_list} + assert fired == {"t.1", "t.3"} + assert all(call.args[1] == "CANCELED" for call in cb.call_args_list) + + +@pytest.mark.asyncio +async def test_terminal_callback_fires_once_across_cancel_and_sse(): + """A local cancel and the matching CANCELED SSE notification must not double-fire the terminal + callback.""" + backend = await _init_backend() + backend._tasks["t.1"] = {"uid": "t.1", "state": "RUNNING"} + cb = MagicMock() + backend.register_callback(cb) + + await backend.cancel_task("t.1") # local path fires once, prunes t.1 + # the matching CANCELED SSE notification arrives afterwards and is ignored + # because t.1 is no longer tracked + backend._on_task_notification( + endpoint="hpc1", + plugin="rhapsody", + topic="task_status", + data={"uid": "t.1", "state": "CANCELED"}, + ) + await asyncio.sleep(0) + + assert cb.call_count == 1 + assert "t.1" not in backend._tasks + + +# --------------------------------------------------------------------------- +# Batch flushing: lock is released during the network send, sends stay +# serialized/ordered, and buffered tasks survive shutdown +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_batch_send_releases_lock_during_send(): + """A slow batch send must not hold the batch lock — other submitters can keep buffering while + the network flush is in flight. + + With the send performed under the batch lock this would dead-time out. + """ + backend = await _init_backend(batch_window=10, batch_limit=2) + + send_started = threading.Event() + release_send = threading.Event() + + def slow_submit(batch): + send_started.set() + release_send.wait(timeout=5) + return [{"uid": batch[0]["uid"], "state": "SUBMITTED"}] + + backend._mock_rh.submit_tasks = MagicMock(side_effect=slow_submit) + + # Fill the buffer (limit=2) → inline flush that blocks inside the send. + first = asyncio.create_task( + backend.submit_tasks( + [ + {"uid": "t.1", "executable": "/bin/true"}, + {"uid": "t.2", "executable": "/bin/true"}, + ] + ) + ) + + # Wait until the send is actually in flight in the worker thread. + assert await asyncio.to_thread(send_started.wait, 5) + + # While the send is blocked, another submit must still complete (buffer) + # rather than hang on the batch lock. + await asyncio.wait_for( + backend.submit_tasks([{"uid": "t.3", "executable": "/bin/true"}]), + timeout=1.0, + ) + assert "t.3" in backend._tasks + + release_send.set() + await asyncio.wait_for(first, timeout=5) + await backend.shutdown() + + +@pytest.mark.asyncio +async def test_batch_sends_are_serialized(): + """Concurrent flushes must not call the client concurrently — the send lock serializes them, so + at most one send runs at a time.""" + backend = await _init_backend(batch_window=10, batch_limit=1) + + counter_lock = threading.Lock() + state = {"cur": 0, "max": 0} + order = [] + + def rec(batch): + with counter_lock: + state["cur"] += 1 + state["max"] = max(state["max"], state["cur"]) + time.sleep(0.02) + with counter_lock: + state["cur"] -= 1 + order.append(batch[0]["uid"]) + return [{"uid": batch[0]["uid"], "state": "SUBMITTED"}] + + backend._mock_rh.submit_tasks = MagicMock(side_effect=rec) + + # limit=1 → each submit flushes immediately; fire several concurrently. + await asyncio.gather( + *[backend.submit_tasks([{"uid": f"t.{i}", "executable": "/bin/true"}]) for i in range(5)] + ) + + assert len(order) == 5 + assert state["max"] == 1 + + +@pytest.mark.asyncio +async def test_shutdown_flushes_buffered_tasks(): + """Tasks still buffered when shutdown is called are flushed, not dropped.""" + backend = await _init_backend(batch_window=10, batch_limit=100) + + await backend.submit_tasks([{"uid": "t.1", "executable": "/bin/true"}]) + backend._mock_rh.submit_tasks.assert_not_called() # buffered, not sent yet + + await backend.shutdown() + + backend._mock_rh.submit_tasks.assert_called_once() + sent = backend._mock_rh.submit_tasks.call_args[0][0] + assert [t["uid"] for t in sent] == ["t.1"] + + +@pytest.mark.asyncio +async def test_sse_notification_applied_on_loop_thread(): + """A notification delivered from a non-loop (SSE) thread must be applied on the event-loop + thread — never on the caller's thread — so it can't race loop-side writers of self._tasks.""" + backend = await _init_backend() + backend._tasks["t.1"] = {"uid": "t.1", "state": "SUBMITTED"} + + applied = {} + orig_apply = backend._apply_task_update + + def spy(body): + applied["thread"] = threading.get_ident() + return orig_apply(body) + + backend._apply_task_update = spy + + sse = {} + + def deliver(): + sse["thread"] = threading.get_ident() + backend._on_task_notification( + endpoint="e", + plugin="rhapsody", + topic="task_status", + data={"uid": "t.1", "state": "RUNNING"}, # non-terminal: stays tracked + ) + + t = threading.Thread(target=deliver) + t.start() + t.join() + + # Scheduled on the loop, but not run yet (we haven't yielded). + assert "thread" not in applied + assert backend._tasks["t.1"]["state"] == "SUBMITTED" + + await asyncio.sleep(0) + + loop_thread = threading.get_ident() + assert applied["thread"] == loop_thread # applied on the loop thread + assert applied["thread"] != sse["thread"] # not the SSE (caller) thread + assert backend._tasks["t.1"]["state"] == "RUNNING" + + +# --------------------------------------------------------------------------- +# Memory bounding + cancel of still-buffered tasks +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_terminal_tasks_are_pruned(): + """Reaching a terminal state drops the task from _tasks and _terminal_fired so the backend does + not retain every task it ever ran.""" + backend = await _init_backend(batch_window=0) + await backend.submit_tasks([{"uid": "t.1", "executable": "/bin/true"}]) + assert "t.1" in backend._tasks + + backend._on_task_notification( + endpoint="e", + plugin="rhapsody", + topic="task_status", + data={"uid": "t.1", "state": "DONE"}, + ) + await asyncio.sleep(0) + + assert "t.1" not in backend._tasks + assert "t.1" not in backend._terminal_fired + + +@pytest.mark.asyncio +async def test_cancel_buffered_task_skips_remote_and_is_not_flushed(): + """Cancelling a task still in the batch buffer removes it from the buffer, skips the remote + cancel (endpoint never received it), and it is not sent on the next flush.""" + backend = await _init_backend(batch_window=10, batch_limit=100) # stays buffered + + await backend.submit_tasks([{"uid": "t.1", "executable": "/bin/true"}]) + assert any(t.get("uid") == "t.1" for t in backend._batch_buffer) + + ok = await backend.cancel_task("t.1") + assert ok is True + backend._mock_rh.cancel_task.assert_not_called() # not submitted yet + assert not any(t.get("uid") == "t.1" for t in backend._batch_buffer) # dropped + + # A subsequent shutdown flush must not send the cancelled task. + await backend.shutdown() + backend._mock_rh.submit_tasks.assert_not_called() + + +@pytest.mark.asyncio +async def test_cancel_submitted_task_calls_remote(): + """A task already submitted (not buffered) still goes to the remote cancel.""" + backend = await _init_backend(batch_window=0) # immediate submit, no buffering + await backend.submit_tasks([{"uid": "t.1", "executable": "/bin/true"}]) + assert not backend._batch_buffer + + ok = await backend.cancel_task("t.1") + assert ok is True + backend._mock_rh.cancel_task.assert_called_once_with("t.1") + + +@pytest.mark.asyncio +async def test_cancel_all_drops_buffered_tasks(): + """cancel_all_tasks clears the local buffer so buffered tasks are not flushed after + cancellation.""" + backend = await _init_backend(batch_window=10, batch_limit=100) + await backend.submit_tasks([{"uid": "t.1", "executable": "/bin/true"}]) + assert backend._batch_buffer + + await backend.cancel_all_tasks() + assert backend._batch_buffer == [] + + await backend.shutdown() + backend._mock_rh.submit_tasks.assert_not_called() diff --git a/tox.ini b/tox.ini index f8f61f2..ecaa71f 100644 --- a/tox.ini +++ b/tox.ini @@ -88,6 +88,15 @@ deps = commands = pytest tests/integration/ {posargs} +[testenv:unit] +deps = + pytest + pytest-asyncio + psycopg2-binary + .[dev,tracing,metrics,serialization] +commands = + pytest tests/unit/ {posargs} + ; [testenv:notebooks] ; deps = ; nbconvert