diff --git a/components/src/dynamo/sglang/capacity.py b/components/src/dynamo/sglang/capacity.py index 1b41afc637a8..813a68a79d52 100644 --- a/components/src/dynamo/sglang/capacity.py +++ b/components/src/dynamo/sglang/capacity.py @@ -32,6 +32,35 @@ def local_dp_rank_bounds(server_args: Any) -> tuple[int, int]: return 0, 1 +def publishes_kv_events(server_args: Any) -> bool: + """Whether this node should advertise a KV-event source. + + The router keys KV sources by ``(worker_id, dp_rank)``, and non-leader nodes + publish under the leader's worker ID so the router-visible trees stay keyed + to one logical worker. That only yields a unique key per node while DP + attention gives each node a distinct rank slice. + + Without DP attention, ``local_dp_rank_bounds`` returns ``[0, 1)`` on every + node. Every node of a multinode gang would therefore advertise the same + ``(leader_worker_id, 0)`` source. The frontend marks that key ambiguous and + never activates the direct-ZMQ ingress. + + Only the leader owns the single logical rank in TP-only mode. SGLang emits + radix-cache events from the rank-0 scheduler, so non-leader sockets have + nothing distinct to contribute. + """ + dp_size = getattr(server_args, "dp_size", 1) or 1 + enable_dp_attention = getattr(server_args, "enable_dp_attention", False) + nnodes = getattr(server_args, "nnodes", 1) or 1 + node_rank = getattr(server_args, "node_rank", 0) or 0 + + # Mirrors the branch in local_dp_rank_bounds: per-node distinct slices. + if enable_dp_attention and dp_size > 1: + return True + + return not (nnodes > 1 and node_rank > 0) + + def model_card_dp_rank_bounds(server_args: Any) -> tuple[int, int]: dp_size = getattr(server_args, "dp_size", 1) or 1 return 0, dp_size diff --git a/components/src/dynamo/sglang/publisher.py b/components/src/dynamo/sglang/publisher.py index 0dd43be84dda..5cd6f4f69015 100644 --- a/components/src/dynamo/sglang/publisher.py +++ b/components/src/dynamo/sglang/publisher.py @@ -26,7 +26,11 @@ from dynamo.runtime import Endpoint from dynamo.sglang._disagg import SGLANG_WORKER_GROUP_ID_KEY, get_sglang_worker_group_id from dynamo.sglang.args import Config -from dynamo.sglang.capacity import kv_metrics_block_values, local_dp_rank_bounds +from dynamo.sglang.capacity import ( + kv_metrics_block_values, + local_dp_rank_bounds, + publishes_kv_events, +) def get_local_dp_rank_range(server_args) -> range: @@ -296,7 +300,14 @@ def init_kv_event_publish(self) -> List[KvEventPublisher]: List of KvEventPublisher instances if KV event publishing is enabled, empty list otherwise. """ - if self.dynamo_args.use_kv_events: + if self.dynamo_args.use_kv_events and not publishes_kv_events(self.server_args): + logging.info( + "Non-leader node (node_rank=%s) shares the leader's single KV " + "rank slice; skipping KV event publishing so the router sees " + "exactly one source per (worker_id, dp_rank).", + getattr(self.server_args, "node_rank", 0) or 0, + ) + elif self.dynamo_args.use_kv_events: kv_events = json.loads(self.server_args.kv_events_config) base_ep = kv_events.get("endpoint") if not base_ep: @@ -565,7 +576,9 @@ async def handle_non_leader_node( ) try: - if publisher.dynamo_args.use_kv_events: + if publisher.dynamo_args.use_kv_events and publishes_kv_events( + publisher.server_args + ): kv_worker_id = await _resolve_multinode_leader_worker_id( publisher.generate_endpoint, publisher.server_args, diff --git a/components/src/dynamo/sglang/tests/test_sglang_local_dp_ranks.py b/components/src/dynamo/sglang/tests/test_sglang_local_dp_ranks.py index 3f7b9bf60d97..5f865b3be0cd 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_local_dp_ranks.py +++ b/components/src/dynamo/sglang/tests/test_sglang_local_dp_ranks.py @@ -8,6 +8,12 @@ import pytest +from dynamo.sglang.capacity import ( + local_dp_rank_bounds, + model_card_dp_rank_bounds, + publishes_kv_events, +) + pytestmark = [ pytest.mark.unit, pytest.mark.sglang, @@ -22,8 +28,6 @@ def test_model_card_registration_keeps_global_dp_range(): - from dynamo.sglang.capacity import model_card_dp_rank_bounds - server_args = SimpleNamespace( dp_size=16, enable_dp_attention=True, @@ -32,3 +36,43 @@ def test_model_card_registration_keeps_global_dp_range(): ) assert model_card_dp_rank_bounds(server_args) == (0, 16) + + +def _args(**kwargs) -> SimpleNamespace: + base = dict(dp_size=1, enable_dp_attention=False, nnodes=1, node_rank=0) + base.update(kwargs) + return SimpleNamespace(**base) + + +def test_single_node_publishes_kv_events(): + assert publishes_kv_events(_args()) is True + + +def test_multinode_without_dp_attention_publishes_only_from_leader(): + """TP-only multinode must advertise one source per logical worker.""" + leader = _args(nnodes=2, node_rank=0) + follower = _args(nnodes=2, node_rank=1) + + # Precondition for the collision this guards against. + assert local_dp_rank_bounds(leader) == local_dp_rank_bounds(follower) == (0, 1) + + assert publishes_kv_events(leader) is True + assert publishes_kv_events(follower) is False + + +def test_dp_attention_publishes_from_every_node(): + """Each node owns a distinct slice when DP attention is enabled.""" + nodes = [ + _args(dp_size=4, enable_dp_attention=True, nnodes=2, node_rank=rank) + for rank in (0, 1) + ] + assert local_dp_rank_bounds(nodes[0]) != local_dp_rank_bounds(nodes[1]) + assert all(publishes_kv_events(node) is True for node in nodes) + + +def test_dp_size_one_with_dp_attention_still_leader_only(): + """DP size one keeps the shared [0, 1) slice even with the flag set.""" + assert ( + publishes_kv_events(_args(enable_dp_attention=True, nnodes=2, node_rank=1)) + is False + ) diff --git a/components/src/dynamo/sglang/tests/test_sglang_publisher.py b/components/src/dynamo/sglang/tests/test_sglang_publisher.py index abc066a6a8e9..087a51121699 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_publisher.py +++ b/components/src/dynamo/sglang/tests/test_sglang_publisher.py @@ -3,6 +3,7 @@ import asyncio from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock import pytest @@ -246,6 +247,8 @@ async def client(self): return FakeClient() server_args = SimpleNamespace( + dp_size=2, + enable_dp_attention=True, nnodes=2, node_rank=1, dist_timeout=5, @@ -289,6 +292,55 @@ def cleanup(self): assert metrics_task.cancelled() +@pytest.mark.asyncio +async def test_handle_non_leader_node_skips_tp_only_kv_event_setup(monkeypatch): + resolve_leader = AsyncMock(return_value=1234) + kv_event_publisher = Mock() + monkeypatch.setattr( + publisher_mod, + "_resolve_multinode_leader_worker_id", + resolve_leader, + ) + monkeypatch.setattr(publisher_mod, "KvEventPublisher", kv_event_publisher) + + server_args = SimpleNamespace( + dp_size=1, + enable_dp_attention=False, + nnodes=2, + node_rank=1, + ) + cleanup = Mock() + publisher = SimpleNamespace( + server_args=server_args, + dynamo_args=SimpleNamespace( + use_kv_events=True, + ), + generate_endpoint=SimpleNamespace(), + init_kv_event_publish=lambda: publisher_mod.KvEventPublisher(), + cleanup=cleanup, + ) + metrics_task = asyncio.create_task(asyncio.Event().wait()) + task = asyncio.create_task( + handle_non_leader_node( + SimpleNamespace(server_args=server_args), + publisher, + metrics_task, + ) + ) + + await asyncio.sleep(0) + + resolve_leader.assert_not_awaited() + kv_event_publisher.assert_not_called() + assert not task.done() + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + cleanup.assert_called_once_with() + assert metrics_task.cancelled() + + @pytest.mark.asyncio async def test_handle_non_leader_node_skips_kv_publish_without_resolved_worker( monkeypatch, diff --git a/lib/runtime/src/discovery/kube.rs b/lib/runtime/src/discovery/kube.rs index 40991161742d..eba14a46aaae 100644 --- a/lib/runtime/src/discovery/kube.rs +++ b/lib/runtime/src/discovery/kube.rs @@ -16,7 +16,7 @@ use utils::{KubeDiscoveryMode, PodInfo}; use crate::CancellationToken; use crate::discovery::{ Discovery, DiscoveryEvent, DiscoveryInstance, DiscoveryInstanceId, DiscoveryMetadata, - DiscoveryQuery, DiscoverySpec, DiscoveryStream, MetadataSnapshot, + DiscoveryQuery, DiscoverySpec, DiscoveryStream, MAX_JSON_SAFE_PUBLISHER_ID, MetadataSnapshot, }; use anyhow::Result; use async_trait::async_trait; @@ -26,8 +26,11 @@ use std::sync::Arc; use tokio::sync::RwLock; fn validate_kubernetes_publisher_id(publisher_id: u64) -> Result<()> { - if i64::try_from(publisher_id).is_err() { - anyhow::bail!("Kubernetes discovery publisher ID {publisher_id} exceeds i64::MAX"); + if publisher_id > MAX_JSON_SAFE_PUBLISHER_ID { + anyhow::bail!( + "Kubernetes discovery publisher ID {publisher_id} exceeds the JSON-safe maximum \ + {MAX_JSON_SAFE_PUBLISHER_ID}" + ); } Ok(()) @@ -524,8 +527,9 @@ mod tests { use super::*; #[test] - fn publisher_ids_must_fit_kubernetes_integer_range() { - assert!(validate_kubernetes_publisher_id(i64::MAX as u64).is_ok()); - assert!(validate_kubernetes_publisher_id((i64::MAX as u64) + 1).is_err()); + fn publisher_ids_must_fit_kubernetes_json_safe_range() { + assert!(validate_kubernetes_publisher_id(MAX_JSON_SAFE_PUBLISHER_ID).is_ok()); + assert!(validate_kubernetes_publisher_id(MAX_JSON_SAFE_PUBLISHER_ID + 1).is_err()); + assert!(validate_kubernetes_publisher_id(u64::MAX).is_err()); } } diff --git a/lib/runtime/src/discovery/kube/crd.rs b/lib/runtime/src/discovery/kube/crd.rs index 23d9f1426893..ed21188028b2 100644 --- a/lib/runtime/src/discovery/kube/crd.rs +++ b/lib/runtime/src/discovery/kube/crd.rs @@ -122,7 +122,9 @@ pub async fn apply_cr( #[cfg(test)] mod tests { use super::*; - use crate::discovery::{DiscoveryInstance, DiscoveryQuery, EventScope, EventSourceQuery}; + use crate::discovery::{ + DiscoveryInstance, DiscoveryQuery, EventScope, EventSourceQuery, MAX_JSON_SAFE_PUBLISHER_ID, + }; use crate::protocols::EndpointId; use kube::Resource; @@ -176,7 +178,7 @@ mod tests { endpoint: endpoint.clone(), }, topic: "kv-events".to_string(), - publisher_id: i64::MAX as u64, + publisher_id: MAX_JSON_SAFE_PUBLISHER_ID, metadata: serde_json::json!({"worker_id": 7, "dp_rank": 0}), }; metadata.register_event_source(source.clone()).unwrap(); @@ -187,7 +189,7 @@ mod tests { .and_then(|sources| sources.values().next()) .and_then(|source| source.get("publisher_id")) .expect("serialized event source publisher ID"); - assert!(publisher_id.is_i64()); + assert_eq!(publisher_id.as_u64(), Some(MAX_JSON_SAFE_PUBLISHER_ID)); let round_trip: DiscoveryMetadata = serde_json::from_value(cr.spec.data).unwrap(); diff --git a/lib/runtime/src/discovery/mod.rs b/lib/runtime/src/discovery/mod.rs index f63f7c1dd9ae..daf96d6a493f 100644 --- a/lib/runtime/src/discovery/mod.rs +++ b/lib/runtime/src/discovery/mod.rs @@ -25,6 +25,9 @@ pub mod utils; use crate::component::{DeviceType, TransportType}; pub use utils::watch_and_extract_field; +/// Largest publisher ID exactly representable by float64-backed JSON metadata. +pub(crate) const MAX_JSON_SAFE_PUBLISHER_ID: u64 = (1 << 53) - 1; + /// Transport kind for event plane - used for configuration and env var selection. /// /// This enum represents the *type* of transport without connection details. diff --git a/lib/runtime/src/transports/event_plane/mod.rs b/lib/runtime/src/transports/event_plane/mod.rs index fcc2ac53da80..29c655560555 100644 --- a/lib/runtime/src/transports/event_plane/mod.rs +++ b/lib/runtime/src/transports/event_plane/mod.rs @@ -43,6 +43,7 @@ use crate::DistributedRuntime; use crate::component::{Component, Endpoint, Namespace}; use crate::discovery::{ Discovery, DiscoveryInstance, DiscoveryQuery, DiscoverySpec, EventChannelQuery, EventTransport, + MAX_JSON_SAFE_PUBLISHER_ID, }; use crate::protocols::EndpointId; use crate::traits::DistributedRuntimeProvider; @@ -259,9 +260,9 @@ impl Stream for DeduplicatingStream { } } -/// Keep the shared wire, channel, and source publisher ID representable in Kubernetes metadata. +/// Keep publisher IDs exactly representable in float64-backed JSON metadata. fn discovery_safe_publisher_id(random_id: u64) -> u64 { - random_id & (i64::MAX as u64) + random_id & MAX_JSON_SAFE_PUBLISHER_ID } /// Event publisher for a specific topic. @@ -923,9 +924,24 @@ mod tests { use crate::config::environment_names::zmq_broker as broker_env; #[test] - fn publisher_ids_fit_kubernetes_discovery_integer_range() { - assert_eq!(discovery_safe_publisher_id(42), 42); - assert!(i64::try_from(discovery_safe_publisher_id(u64::MAX)).is_ok()); + fn publisher_ids_survive_a_json_number_round_trip() { + // This historical full-range ID is not JSON-safe. It was observed + // rounding to 13584172880116488000 after a discovery round trip. + let unsafe_id: u64 = 13_584_172_880_116_487_724; + assert!(unsafe_id > MAX_JSON_SAFE_PUBLISHER_ID); + assert_ne!(unsafe_id as f64 as u64, unsafe_id); + + for random_id in [0, 1, u64::MAX, unsafe_id, 6_633_287_539_119_378] { + let publisher_id = discovery_safe_publisher_id(random_id); + assert!( + publisher_id <= MAX_JSON_SAFE_PUBLISHER_ID, + "publisher ID {publisher_id} exceeds the JSON-safe integer range" + ); + assert_eq!( + publisher_id as f64 as u64, publisher_id, + "publisher ID {publisher_id} must survive an f64 round trip" + ); + } } #[test] @@ -1119,6 +1135,12 @@ mod tests { publisher_b.publisher_id(), ]) ); + for publisher_id in [publisher_a.publisher_id(), publisher_b.publisher_id()] { + assert!( + publisher_id <= MAX_JSON_SAFE_PUBLISHER_ID, + "publisher ID {publisher_id} exceeds the JSON-safe integer range" + ); + } }; tokio::time::timeout(std::time::Duration::from_secs(5), receive) .await