orbit backend: dispatcher pool mode

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=<name>`` (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.<uuid>` 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"])

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)