diff --git a/src/rhapsody/backends/execution/orbit.py b/src/rhapsody/backends/execution/orbit.py index 90e4fbb..0a8f3ba 100644 --- a/src/rhapsody/backends/execution/orbit.py +++ b/src/rhapsody/backends/execution/orbit.py @@ -72,6 +72,22 @@ class OrbitExecutionBackend(BaseBackend): (default ``["dragon_v3"]``). name: Backend name for Rhapsody registration (default ``"orbit"``). + pool: Name of a dispatcher-managed pool to run tasks in. + Setting it targets a task-dispatcher-style plugin + instead of a per-endpoint rhapsody session: every + submitted task is stamped ``pool=`` (a task + already carrying one keeps it), endpoint + auto-selection is skipped in favour of the broker + participant hosting the dispatcher, and the + python-version handshake for cloudpickled tasks + resolves against the pool's *executing* endpoint + rather than the submission target. Combine with + ``plugin_name="task_dispatcher"``. + session_kwargs: Extra keyword arguments for the remote session + registration (``get_plugin``) -- e.g. ``sid`` to + join an existing dispatcher session whose pools + this backend should see, or ``pools`` to declare + them. 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 @@ -94,7 +110,10 @@ def __init__( endpoint_name: str | None = None, backends: list[str] | None = None, name: str = "orbit", + participant_name: str | None = None, plugin_name: str = _PLUGIN_NAME, + pool: str | None = None, + session_kwargs: dict | None = None, batch_window: float | None = None, batch_limit: int = 1024, start_timeout: float = 30.0, @@ -109,8 +128,11 @@ def __init__( self.logger = logging.getLogger(__name__) self._broker_url = broker_url + self._participant_name = participant_name self._endpoint_name = endpoint_name self._plugin_name = plugin_name + self._pool = pool + self._session_kwargs = dict(session_kwargs or {}) self._remote_backends = backends or ["dragon_v3"] self._start_timeout = start_timeout self._init_timeout = init_timeout @@ -325,12 +347,19 @@ def _get_rhapsody_handle(self) -> Any: """ 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). + # The broker shows every runtime as a participant, so an anonymous + # `rhapsody.` reads as noise in a topology view. A caller + # which knows what this backend is *for* can say so via + # ``participant_name`` -- uniqueness is then the caller's contract. + # Without one, a unique 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]}", + name=self._participant_name or f"rhapsody.{uuid.uuid4().hex[:8]}", + # the advertised role defaults to 'consumer', which says nothing. + # This participant is the workflow engine's hand-off into ORBIT. + role="engine", ) try: rt.start(wait=True, timeout=self._start_timeout) @@ -344,6 +373,12 @@ def _get_rhapsody_handle(self) -> Any: ) self._broker_url = rt.broker_url + # pool mode: the dispatcher lives on the broker participant, so + # there is no endpoint to select -- the pool's pilots pick the + # executing endpoints, not this backend + if self._pool and not self._endpoint_name: + self._endpoint_name = "broker" + # find a suitable endpoint from the (local) topology snapshot if not self._endpoint_name: for eid, info in rt.topology().items(): @@ -373,6 +408,7 @@ def _get_rhapsody_handle(self) -> Any: self._plugin_name, backends=self._remote_backends, init_timeout=self._init_timeout, + **self._session_kwargs, ) except Exception: # Don't leak the runtime's daemon threads / WebSocket on a failed @@ -421,10 +457,16 @@ def needs_compat(t: dict) -> bool: info = None def _lookup(): + # Pool mode: the submission target (the dispatcher's host) + # does not execute anything — the pool's endpoint does, so + # that is what the cloudpickle handshake compares against. + target = self._endpoint_name + if self._pool and hasattr(self._rh, "pool_detail"): + target = self._rh.pool_detail(self._pool).get("endpoint_name") or target # 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") + si = self._runtime.get_plugin(target, "sysinfo") try: return si.host_role() finally: @@ -496,6 +538,10 @@ async def submit_tasks(self, tasks: list[dict[str, Any]]) -> None: prof = self._prof for task in tasks: task.setdefault("uid", f"task.{uuid.uuid4().hex[:8]}") + # pool mode: route every task into this backend's pool; a task + # that already names one keeps it + if self._pool: + task.setdefault("pool", self._pool) self._tasks[task["uid"]] = task if prof: prof.prof("task_submit", uid=task["uid"]) diff --git a/tests/unit/test_backend_execution_orbit.py b/tests/unit/test_backend_execution_orbit.py index 1974cb9..9b21927 100644 --- a/tests/unit/test_backend_execution_orbit.py +++ b/tests/unit/test_backend_execution_orbit.py @@ -893,3 +893,77 @@ async def test_cancel_all_drops_buffered_tasks(): await backend.shutdown() backend._mock_rh.submit_tasks.assert_not_called() + + +# --------------------------------------------------------------------------- +# Pool mode (dispatcher-managed pools) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_pool_mode_targets_broker_and_forwards_session_kwargs(): + """With a pool and no endpoint, the broker participant (hosting the dispatcher) is the target -- + no endpoint auto-selection -- and session kwargs ride into get_plugin so the backend can join a + session by sid.""" + backend = await _init_backend( + endpoint_name=None, + pool="insitu", + plugin_name="task_dispatcher", + session_kwargs={"sid": "session.abc"}, + ) + + assert backend._endpoint_name == "broker" + backend._mock_rt.get_plugin.assert_called_once_with( + "broker", + "task_dispatcher", + backends=["dragon_v3"], + init_timeout=120.0, + sid="session.abc", + ) + + +@pytest.mark.asyncio +async def test_pool_mode_stamps_tasks(): + """Every submitted task carries the backend's pool; an explicit per-task pool wins.""" + backend = await _init_backend( + endpoint_name=None, pool="insitu", plugin_name="task_dispatcher", batch_window=0 + ) + + await backend.submit_tasks( + [ + {"uid": "t.001", "executable": "/bin/echo"}, + {"uid": "t.002", "executable": "/bin/echo", "pool": "special"}, + ] + ) + + submitted = backend._mock_rh.submit_tasks.call_args[0][0] + assert submitted[0]["pool"] == "insitu" + assert submitted[1]["pool"] == "special" + + +@pytest.mark.asyncio +async def test_pool_mode_compat_resolves_executing_endpoint(): + """The cloudpickle python handshake compares against the pool's executing endpoint (from the + dispatcher's pool detail), not the dispatcher's host.""" + import sys + + backend = await _init_backend( + endpoint_name=None, pool="insitu", plugin_name="task_dispatcher", batch_window=0 + ) + backend._mock_rh.pool_detail = MagicMock( + return_value={"name": "insitu", "endpoint_name": "compute_ep"} + ) + + si = MagicMock() + mm = sys.version_info + si.host_role = MagicMock(return_value={"python_version": f"{mm.major}.{mm.minor}.0"}) + si.close = MagicMock() + backend._mock_rt.get_plugin = MagicMock(return_value=si) + + def fn(): + return 42 + + await backend.submit_tasks([{"uid": "t.003", "function": fn}]) + + backend._mock_rt.get_plugin.assert_called_once_with("compute_ep", "sysinfo") + assert backend._endpoint_python_mm == (mm.major, mm.minor)