feat(select-svc): Expose session affinity replication for external hosts

lib/bindings/python/rust/lib.rs

@@ -396,6 +396,8 @@ fn register_core(m: &Bound<'_, PyModule>) -> PyResult<()> {
     #[cfg(feature = "select-service")]
     m.add_class::<llm::kv::SelectionService>()?;
     #[cfg(feature = "select-service")]
+    m.add_class::<llm::kv::SessionBindingStream>()?;
+    #[cfg(feature = "select-service")]
     m.add_class::<llm::kv::SelectionCacheConfig>()?;
     m.add_class::<llm::kv::WorkerMetricsPublisher>()?;
     m.add_class::<llm::kv::MultimodalEmbeddingCachePublisher>()?;

lib/bindings/python/rust/llm/kv.rs

@@ -686,12 +686,36 @@ pub(crate) struct SelectionService {
     inner: Arc<RustSelectionService>,
 }
 
+#[cfg(feature = "select-service")]
+#[pyclass]
+pub(crate) struct SessionBindingStream {
+    updates: Arc<selection::SessionBindingUpdates>,
+}
+
+#[cfg(feature = "select-service")]
+#[pymethods]
+impl SessionBindingStream {
+    fn __aiter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        slf
+    }
+
+    fn __anext__<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
+        let updates = self.updates.clone();
+        crate::future_into_py(py, async move {
+            let events = updates.recv_many().await.ok_or_else(|| {
+                pyo3::exceptions::PyStopAsyncIteration::new_err("selection service shut down")
+            })?;
+            Python::with_gil(|py| pythonize(py, &events).map(|o| o.unbind()).map_err(to_pyerr))
+        })
+    }
+}
+
 #[cfg(feature = "select-service")]
 #[pymethods]
 impl SelectionService {
     /// Create a selection service. `indexer_threads` sizes the KV indexer pool.
     #[new]
-    #[pyo3(signature = (*, indexer_threads = 4, indexer_peers = None, replica_sync_port = None, replica_sync_peers = None, selection_cache = None, session_affinity_ttl_secs = None))]
+    #[pyo3(signature = (*, indexer_threads = 4, indexer_peers = None, replica_sync_port = None, replica_sync_peers = None, selection_cache = None, session_affinity_ttl_secs = None, session_affinity_mode = "hard", session_replication_writer_id = None))]
     #[allow(clippy::too_many_arguments)]
     fn new(
         py: Python<'_>,
@@ -701,7 +725,12 @@ impl SelectionService {
         replica_sync_peers: Option<Vec<String>>,
         selection_cache: Option<SelectionCacheConfig>,
         session_affinity_ttl_secs: Option<f64>,
+        session_affinity_mode: &str,
+        session_replication_writer_id: Option<u64>,
     ) -> PyResult<Self> {
+        let session_affinity_mode = session_affinity_mode
+            .parse::<RsSessionAffinityMode>()
+            .map_err(PyValueError::new_err)?;
         let replica_sync_peers = replica_sync_peers.unwrap_or_default();
         if replica_sync_port.is_none() && !replica_sync_peers.is_empty() {
             return Err(PyValueError::new_err(
@@ -719,7 +748,8 @@ impl SelectionService {
         )
         .indexer_threads(indexer_threads)
         .indexer_peers(indexer_peers.unwrap_or_default())
-        .selection_cache(selection_cache.unwrap_or_default().inner);
+        .selection_cache(selection_cache.unwrap_or_default().inner)
+        .session_affinity_mode(session_affinity_mode);
         if let Some(port) = replica_sync_port {
             builder = builder.replica_sync(port, replica_sync_peers);
         }
@@ -728,6 +758,9 @@ impl SelectionService {
                 session_affinity_ttl_from_secs(ttl).map_err(PyValueError::new_err)?,
             );
         }
+        if let Some(writer_id) = session_replication_writer_id {
+            builder = builder.session_binding_replication(writer_id);
+        }
         let inner = py
             .allow_threads(|| crate::bridge_runtime().block_on(builder.build()))
             .map_err(to_pyerr)?;
@@ -746,6 +779,26 @@ impl SelectionService {
         py.allow_threads(|| crate::bridge_runtime().block_on(service.shutdown()));
     }
 
+    /// Session bindings this replica commits or refreshes, for the host to
+    /// broadcast to its peers. Ends when the service shuts down.
+    fn session_binding_updates(&self) -> PyResult<SessionBindingStream> {
+        Ok(SessionBindingStream {
+            updates: Arc::new(
+                self.inner
+                    .session_binding_updates()
+                    .map_err(selection_to_pyerr)?,
+            ),
+        })
+    }
+
+    fn apply_session_binding(&self, py: Python<'_>, update: PyObject) -> PyResult<()> {
+        let event =
+            depythonize(update.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?;
+        self.inner
+            .apply_session_binding(event)
+            .map_err(selection_to_pyerr)
+    }
+
     /// Await service shutdown without blocking the Python event loop.
     fn shutdown_async<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
         let service = Arc::clone(&self.inner);
@@ -1020,9 +1073,10 @@ mod selection_service_lifecycle_tests {
     #[test]
     fn idempotent_shutdown() {
         pyo3::prepare_freethreaded_python();
-        let service =
-            Python::with_gil(|py| SelectionService::new(py, 1, None, None, None, None, None))
-                .unwrap();
+        let service = Python::with_gil(|py| {
+            SelectionService::new(py, 1, None, None, None, None, None, "hard", None)
+        })
+        .unwrap();
         Python::with_gil(|py| {
             service.shutdown(py);
             service.shutdown(py);

lib/bindings/python/src/dynamo/_core.pyi

@@ -783,11 +783,17 @@ class SelectionService:
         replica_sync_peers: Optional[list[str]] = None,
         selection_cache: Optional[SelectionCacheConfig] = None,
         session_affinity_ttl_secs: Optional[float] = None,
+        session_affinity_mode: Literal["hard", "soft"] = "hard",
+        session_replication_writer_id: Optional[int] = None,
     ) -> None:
         """Create a selection service. `indexer_threads` sizes the KV indexer pool.
 
         `session_affinity_ttl_secs` enables session affinity with an idle TTL
         between 1 and 31536000 seconds.
+        `session_affinity_mode`: hard rejects dispatch away from a live binding;
+        soft rebinds when the bound worker is excluded from the request's candidates.
+        `session_replication_writer_id` exports affinity updates for a host transport;
+        use an id unique to this replica incarnation.
         """
         ...
 
@@ -801,6 +807,27 @@ class SelectionService:
         """
         ...
 
+    async def shutdown_async(self) -> None:
+        """Await service shutdown without blocking the Python event loop."""
+        ...
+
+    def session_binding_updates(self) -> AsyncIterator[List[JsonLike]]:
+        """Stream session bindings this service commits or refreshes, for the
+        host to broadcast to its peers.
+
+        Each item is every binding queued since the last one (at least one).
+        Requires session_replication_writer_id. Ends on shutdown.
+        """
+        ...
+
+    def apply_session_binding(self, update: JsonLike) -> None:
+        """Merge a versioned peer binding without echoing it.
+
+        A binding this replica cannot place yet is dropped; the owner's next
+        refresh repairs it.
+        """
+        ...
+
     async def upsert_worker(self, worker: JsonLike) -> JsonLike:
         """Upsert a worker and subscribe to its live KV events; returns its catalog record."""
         ...

lib/bindings/python/tests/test_kv_bindings.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 
+import inspect
 import json
 import subprocess
 import sys
@@ -291,20 +292,153 @@ async def test_selection_service_selects_registered_worker(monkeypatch):
         await service.shutdown_async()
 
 
+@pytest.fixture
+async def session_binding_services(monkeypatch):
+    monkeypatch.setenv("DYN_ROUTER_USE_KV_EVENTS", "false")
+    source = SelectionService(
+        indexer_threads=1,
+        session_affinity_ttl_secs=10,
+        session_replication_writer_id=11,
+    )
+    peer = SelectionService(
+        indexer_threads=1,
+        session_affinity_ttl_secs=10,
+        session_affinity_mode="soft",
+        session_replication_writer_id=22,
+    )
+    try:
+        for service in (source, peer):
+            for worker_id in (1, 2):
+                await service.upsert_worker(
+                    {
+                        "worker_id": worker_id,
+                        "model_name": "model",
+                        "routing_group": "host",
+                        "endpoint": f"http://worker-{worker_id}",
+                        "block_size": 4,
+                        "max_num_batched_tokens": 1024,
+                    }
+                )
+        yield source, peer
+    finally:
+        await source.shutdown_async()
+        await peer.shutdown_async()
+
+
 @pytest.mark.skipif(
     SelectionService is None,
     reason="SelectionService requires the select-service Cargo feature",
 )
 @pytest.mark.timeout(5)
 @pytest.mark.asyncio
-async def test_selection_service_malformed_payload_raises_value_error(monkeypatch):
+async def test_selection_service_session_binding_stream(session_binding_services):
+    source, peer = session_binding_services
+    stream = source.session_binding_updates()
+    peer_stream = peer.session_binding_updates()
+    with pytest.raises(SelectionServiceError):
+        source.session_binding_updates()
+    partition = {"model_name": "model", "routing_group": "host"}
+    selected = await source.select_and_reserve(
+        {
+            **partition,
+            "token_ids": [1, 2, 3, 4],
+            "selection_id": "origin",
+            "session_id": "s",
+            "allowed_worker_ids": [1],
+        }
+    )
+    assert "session_binding" not in selected
+    (binding,) = json.loads(json.dumps(await anext(stream)))
+    assert binding["writer_id"] == 11
+    peer.apply_session_binding(binding)
+    await source.free_reservation("origin")
+    assert await anext(stream) == [binding]  # A repeated version refreshes TTL.
+    selected = await peer.select_and_reserve(
+        {
+            **partition,
+            "token_ids": [5, 6, 7, 8],
+            "selection_id": "next",
+            "session_id": "s",
+        }
+    )
+    assert selected["worker_id"] == 1
+    assert await anext(peer_stream) == [
+        binding
+    ]  # Local commits preserve imported versions.
+    # The default hard mode rejects an excluded binding; soft mode rebinds.
+    request = {
+        **partition,
+        "token_ids": [9, 10, 11, 12],
+        "selection_id": "excluded",
+        "session_id": "s",
+        "allowed_worker_ids": [2],
+    }
+    with pytest.raises(ValueError, match="session s is bound to worker 1, not 2"):
+        await source.select_and_reserve(request)
+    assert (await peer.select_and_reserve(request))["worker_id"] == 2
+    assert (await anext(peer_stream))[0]["worker_id"] == 2
+    await source.shutdown_async()
+    with pytest.raises(StopAsyncIteration):
+        await anext(stream)
+    with pytest.raises(SelectionServiceError):
+        source.apply_session_binding(binding)
+
+
+@pytest.mark.skipif(
+    SelectionService is None,
+    reason="SelectionService requires the select-service Cargo feature",
+)
+@pytest.mark.timeout(5)
+@pytest.mark.asyncio
+async def test_selection_service_session_binding_stream_drains_queue(
+    session_binding_services,
+):
+    source, _ = session_binding_services
+    stream = source.session_binding_updates()
+    for session_id in ("a", "b", "c"):
+        await source.select_and_reserve(
+            {
+                "model_name": "model",
+                "routing_group": "host",
+                "token_ids": [1, 2, 3, 4],
+                "selection_id": session_id,
+                "session_id": session_id,
+            }
+        )
+    # Bindings queued while the consumer was away arrive as one item.
+    batch = await anext(stream)
+    assert [binding["session_id"] for binding in batch] == ["a", "b", "c"]
+
+
+@pytest.mark.skipif(
+    SelectionService is None,
+    reason="SelectionService requires the select-service Cargo feature",
+)
+@pytest.mark.parametrize(
+    "method,payload",
+    [
+        pytest.param("upsert_worker", {"model_name": "model"}, id="worker"),
+        pytest.param(
+            "apply_session_binding",
+            {"session_id": "s"},
+            id="session-binding",
+        ),
+    ],
+)
+@pytest.mark.timeout(5)
+@pytest.mark.asyncio
+async def test_selection_service_malformed_payload_raises_value_error(
+    monkeypatch, method, payload
+):
     monkeypatch.setenv("DYN_ROUTER_USE_KV_EVENTS", "false")
     service = SelectionService(indexer_threads=1)
 
     try:
         # Malformed input surfaces as a plain ValueError, not a SelectionServiceError.
         with pytest.raises(ValueError):
-            await service.upsert_worker({"model_name": "model"})
+            result = getattr(service, method)(payload)
+            if inspect.isawaitable(result):
+                await result
         # Prompt booleans are not coerced by truthiness.
         with pytest.raises(ValueError, match="boolean"):
             await service.select({"token_ids": [1, 2, 3, 4], "is_eagle": "false"})