diff --git a/lib/llm/src/discovery/watcher.rs b/lib/llm/src/discovery/watcher.rs index 6f7859e72260..452d1da7e24e 100644 --- a/lib/llm/src/discovery/watcher.rs +++ b/lib/llm/src/discovery/watcher.rs @@ -1573,6 +1573,200 @@ mod tests { .expect("model did not reach the expected serving state") } + #[tokio::test] + async fn classify_aliases_share_endpoint_with_isolated_routing_and_removal() { + use dynamo_runtime::{ + discovery::{DiscoveryQuery, DiscoverySpec, EventTransportKind}, + distributed::{DiscoveryBackend, RequestPlaneMode}, + engine::AsyncEngineContextProvider, + pipeline::{ResponseStream, network::Ingress}, + storage::kv, + }; + + // The TCP server is process-global; isolate it from other test runtimes. + const TEST: &str = concat!( + module_path!(), + "::classify_aliases_share_endpoint_with_isolated_routing_and_removal" + ); + let test_name = TEST.split_once("::").unwrap().1; + if std::env::var("DYNAMO_ALIAS_TEST").as_deref() != Ok(test_name) { + let output = tokio::time::timeout( + Duration::from_secs(30), + tokio::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", test_name, "--nocapture"]) + .env("DYNAMO_ALIAS_TEST", test_name) + .env("DYN_TCP_RPC_HOST", "127.0.0.1") + .env("DYN_TCP_RPC_PORT", "0") + .env("DYN_TCP_RESPONSE_STREAM_HOST", "127.0.0.1") + .env("DYN_TCP_RESPONSE_STREAM_PORT", "0") + .kill_on_drop(true) + .output(), + ) + .await + .expect("classify subprocess must finish within its deadline") + .expect("classify subprocess must start"); + assert!( + output.status.success(), + "{}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + let stdout = String::from_utf8_lossy(&output.stdout); + assert!( + stdout + .lines() + .any(|line| line.starts_with("test result: ok. 1 passed; 0 failed;")), + "classify subprocess must run exactly one passing test: {stdout}" + ); + return; + } + + #[derive(Debug)] + struct ClassifyWorker(&'static str); + + #[async_trait] + impl + AsyncEngine< + SingleIn, + ManyOut>, + Error, + > for ClassifyWorker + { + async fn generate( + &self, + request: SingleIn, + ) -> Result>, Error> { + let mut response = NvCreateClassifyResponse::empty(); + response.model = self.0.to_string(); + Ok(ResponseStream::new( + Box::pin(futures::stream::iter([Annotated::from_data(response)])), + request.context(), + )) + } + } + + async fn check_response(manager: &ModelManager, alias: &str) { + let request = serde_json::from_value(serde_json::json!({ + "model": alias, "input": "test" + })) + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), async { + let mut response = manager + .get_classify_engine(alias) + .unwrap() + .generate(SingleIn::new(request)) + .await + .unwrap(); + assert_eq!(response.next().await.unwrap().data.unwrap().model, alias); + while response.next().await.is_some() {} + }) + .await + .expect("classify request must complete"); + } + + let runtime = Runtime::from_current().unwrap(); + let store = tempfile::tempdir().unwrap(); + let config = || DistributedConfig { + discovery_backend: DiscoveryBackend::KvStore(kv::Selector::File(store.path().into())), + nats_config: None, + request_plane: RequestPlaneMode::Tcp, + response_plane: None, + event_transport_kind: EventTransportKind::Zmq, + }; + let frontend = DistributedRuntime::new(runtime.clone(), config()) + .await + .unwrap(); + let manager = Arc::new(ModelManager::new()); + let cancellation = CancellationToken::new(); + let stream = frontend + .discovery() + .list_and_watch(DiscoveryQuery::AllModels, Some(cancellation.clone())) + .await + .unwrap(); + let watcher = Arc::new(ModelWatcher::new( + frontend, + manager.clone(), + RouterConfig { + router_mode: RouterMode::RoundRobin, + ..Default::default() + }, + 0, + None, + None, + None, + Arc::new(Metrics::new()), + )); + let task = tokio::spawn(watcher.watch(stream, NamespaceFilter::Global)); + let mut workers = Vec::new(); + for alias in ["alias-a", "alias-b"] { + let drt = DistributedRuntime::new(runtime.clone(), config()) + .await + .unwrap(); + let endpoint = drt + .namespace("classify-aliases") + .unwrap() + .component("workers") + .unwrap() + .endpoint("generate"); + let serving = endpoint + .endpoint_builder() + .handler(Ingress::for_engine(Arc::new(ClassifyWorker(alias))).unwrap()) + .start_with_registration() + .await + .unwrap(); + let mut card = ModelDeploymentCard::with_name_only(alias); + card.model_input = ModelInput::Text; + card.model_type = ModelType::Classify; + card.worker_type = Some(WorkerType::Aggregated); + card.source_path = Some("org/shared-classifier".to_string()); + let registration = drt + .discovery() + .register( + DiscoverySpec::from_model( + "classify-aliases".into(), + "workers".into(), + "generate".into(), + &card, + ) + .unwrap(), + ) + .await + .unwrap(); + workers.push((drt, serving, registration)); + } + for alias in ["alias-a", "alias-b"] { + wait_for_model(&manager, alias, Model::is_ready_to_serve).await; + } + let mut names = manager.list_classify_models(); + names.sort(); + assert_eq!(names, ["alias-a", "alias-b"]); + // Repeated round-robin calls expose an unfiltered shared endpoint pool. + for _ in 0..4 { + check_response(&manager, "alias-a").await; + check_response(&manager, "alias-b").await; + } + let (drt, serving, registration) = workers.remove(0); + drt.discovery().unregister(registration).await.unwrap(); + serving.shutdown().await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), async { + while manager.get_committed_model("alias-a").is_some() { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("removed alias must leave the catalog"); + assert_eq!(manager.list_classify_models(), ["alias-b"]); + assert!(manager.get_classify_engine("alias-a").is_err()); + check_response(&manager, "alias-b").await; + for (drt, serving, registration) in workers { + drt.discovery().unregister(registration).await.unwrap(); + serving.shutdown().await.unwrap(); + } + cancellation.cancel(); + task.await.unwrap(); + runtime.shutdown(); + } + #[tokio::test] async fn incompatible_advertisements_preserve_the_incumbent_catalog_and_readiness() { let runtime = Runtime::from_current().unwrap(); diff --git a/lib/runtime/src/discovery/mock.rs b/lib/runtime/src/discovery/mock.rs index f47818108d80..86652a3d1e20 100644 --- a/lib/runtime/src/discovery/mock.rs +++ b/lib/runtime/src/discovery/mock.rs @@ -609,40 +609,131 @@ mod tests { } #[tokio::test] - async fn register_rejects_distinct_base_cards_with_same_source_path_on_same_endpoint() { - let registry = SharedMockRegistry::new(); - let discovery1 = MockDiscovery::new(Some(1), registry.clone()); - let discovery2 = MockDiscovery::new(Some(2), registry); - let spec = |display_name: &str| DiscoverySpec::Model { - namespace: "ns".to_string(), - component: "comp".to_string(), - endpoint: "generate".to_string(), - card_json: serde_json::json!({ - "display_name": display_name, - "source_path": "org/base-model", - }), - model_suffix: None, - }; - - discovery1.register(spec("public-name-a")).await.unwrap(); - let err = discovery2 - .register(spec("public-name-b")) - .await - .unwrap_err(); - - assert!(err.to_string().contains( - "Cannot register model 'public-name-b' on endpoint 'ns/comp/generate': a different model 'public-name-a' is already registered there" - )); - - let instances = discovery1 - .list(DiscoveryQuery::EndpointModels { + async fn register_non_lora_alias_compatibility() { + for (name, source_a, source_b, compatible) in [ + ("alias-b", Some("org/base"), Some("org/base"), true), + ("alias-a", Some("/mount/a"), Some("/mount/b"), true), + ("alias-b", Some("org/base"), Some("org/other"), false), + ("alias-b", Some("org/base"), None, false), + ("alias-b", None, Some("org/base"), false), + ("alias-b", Some(""), Some(""), false), + ] { + let registry = SharedMockRegistry::new(); + let discovery1 = MockDiscovery::new(Some(1), registry.clone()); + let discovery2 = MockDiscovery::new(Some(2), registry); + let spec = |display_name: &str, source_path: Option<&str>| DiscoverySpec::Model { namespace: "ns".to_string(), component: "comp".to_string(), endpoint: "generate".to_string(), - }) - .await - .unwrap(); - assert_eq!(instances.len(), 1); + card_json: serde_json::json!({ + "display_name": display_name, + "source_path": source_path, + }), + model_suffix: None, + }; + discovery1 + .register(spec("alias-a", source_a)) + .await + .unwrap(); + let result = discovery2.register(spec(name, source_b)).await; + assert_eq!( + result.is_ok(), + compatible, + "{name}: {source_a:?}, {source_b:?}: {result:?}" + ); + if let Err(err) = result { + assert!( + err.to_string() + .contains("a different model 'alias-a' is already registered there") + ); + } + let instances = discovery1 + .list(DiscoveryQuery::EndpointModels { + namespace: "ns".to_string(), + component: "comp".to_string(), + endpoint: "generate".to_string(), + }) + .await + .unwrap(); + assert_eq!(instances.len(), if compatible { 2 } else { 1 }); + } + } + + #[tokio::test] + async fn register_shared_source_requires_disjoint_served_names() { + for (name_b, aliases_a, aliases_b, compatible) in [ + ("b", vec!["b"], vec![], false), + ("b", vec![], vec!["a"], false), + ("b", vec!["shared"], vec!["shared"], false), + ("b", vec!["a", "extra-a"], vec!["b", "extra-b"], true), + ("a", vec!["shared"], vec!["shared"], true), + ] { + let registry = SharedMockRegistry::new(); + let first = MockDiscovery::new(Some(1), registry.clone()); + let second = MockDiscovery::new(Some(2), registry); + let spec = |name: &str, aliases: Vec<&str>| DiscoverySpec::Model { + namespace: "ns".into(), + component: "comp".into(), + endpoint: "generate".into(), + card_json: serde_json::json!({ + "display_name": name, + "aliases": aliases, + "source_path": "org/base", + }), + model_suffix: None, + }; + let incumbent = first.register(spec("a", aliases_a)).await.unwrap(); + let result = second.register(spec(name_b, aliases_b)).await; + assert_eq!(result.is_ok(), compatible, "{result:?}"); + let instances = first + .list(DiscoveryQuery::EndpointModels { + namespace: "ns".into(), + component: "comp".into(), + endpoint: "generate".into(), + }) + .await + .unwrap(); + assert!(instances.contains(&incumbent)); + assert_eq!(instances.len(), if compatible { 2 } else { 1 }); + } + } + + #[tokio::test] + async fn register_checks_every_existing_model() { + // B is compatible with A by name and C by source, but A and C conflict. + for order in [[0, 1, 2], [2, 1, 0]] { + let registry = SharedMockRegistry::new(); + let cards = [("a", "/mount/a"), ("a", "/mount/b"), ("b", "/mount/b")]; + let mut accepted = Vec::new(); + for (position, index) in order.into_iter().enumerate() { + let discovery = MockDiscovery::new(Some(index as u64 + 1), registry.clone()); + let (name, source) = cards[index]; + let result = discovery + .register(DiscoverySpec::Model { + namespace: "ns".into(), + component: "comp".into(), + endpoint: "generate".into(), + card_json: serde_json::json!({"display_name": name, "source_path": source}), + model_suffix: None, + }) + .await; + if position < 2 { + accepted.push(result.unwrap()); + } else { + assert!(result.is_err()); + let instances = discovery + .list(DiscoveryQuery::EndpointModels { + namespace: "ns".into(), + component: "comp".into(), + endpoint: "generate".into(), + }) + .await + .unwrap(); + assert_eq!(instances.len(), accepted.len()); + assert!(accepted.iter().all(|instance| instances.contains(instance))); + } + } + } } #[tokio::test] @@ -725,6 +816,58 @@ mod tests { .unwrap(); } + #[tokio::test] + async fn register_base_and_lora_require_distinct_served_names() { + for (adapter_name, compatible) in [("base", false), ("alias", false), ("adapter", true)] { + for adapter_first in [false, true] { + let registry = SharedMockRegistry::new(); + let first = MockDiscovery::new(Some(1), registry.clone()); + let second = MockDiscovery::new(Some(2), registry); + let base = DiscoverySpec::Model { + namespace: "ns".into(), + component: "comp".into(), + endpoint: "generate".into(), + card_json: serde_json::json!({ + "display_name": "base", + "aliases": ["alias"], + "source_path": "org/base", + }), + model_suffix: None, + }; + let adapter = lora_model_spec( + "ns", + "comp", + "generate", + adapter_name, + "org/base", + adapter_name, + ); + let (incumbent, newcomer) = if adapter_first { + (adapter, base) + } else { + (base, adapter) + }; + let incumbent = first.register(incumbent).await.unwrap(); + let result = second.register(newcomer).await; + assert_eq!( + result.is_ok(), + compatible, + "{adapter_name}, adapter_first={adapter_first}: {result:?}" + ); + let instances = first + .list(DiscoveryQuery::EndpointModels { + namespace: "ns".into(), + component: "comp".into(), + endpoint: "generate".into(), + }) + .await + .unwrap(); + assert!(instances.contains(&incumbent)); + assert_eq!(instances.len(), if compatible { 2 } else { 1 }); + } + } + } + #[tokio::test] async fn register_rejects_lora_adapter_for_different_base_model() { let registry = SharedMockRegistry::new(); diff --git a/lib/runtime/src/discovery/mod.rs b/lib/runtime/src/discovery/mod.rs index 959fca42570f..b88a53c7efe7 100644 --- a/lib/runtime/src/discovery/mod.rs +++ b/lib/runtime/src/discovery/mod.rs @@ -1208,6 +1208,7 @@ pub type DiscoveryStream = Pin> + S #[derive(Clone, Debug, PartialEq, Eq)] struct ModelRegistrationIdentity { display_name: String, + aliases: Vec, source_path: Option, is_lora: bool, } @@ -1218,10 +1219,30 @@ impl ModelRegistrationIdentity { } fn is_compatible_with(&self, other: &Self) -> bool { - if self.is_lora || other.is_lora { + if self.is_lora != other.is_lora { + let (adapter, base) = if self.is_lora { + (self, other) + } else { + (other, self) + }; + adapter.base_identity() == base.base_identity() + && adapter.display_name != base.display_name + && !base.aliases.contains(&adapter.display_name) + } else if self.is_lora { self.base_identity() == other.base_identity() } else { + // Preserve existing same-name registration compatibility across local model paths. self.display_name == other.display_name + || self.source_path.as_deref().is_some_and(|source| { + !source.is_empty() + && other.source_path.as_deref() == Some(source) + && !self.aliases.contains(&other.display_name) + && !other.aliases.contains(&self.display_name) + && !self + .aliases + .iter() + .any(|alias| other.aliases.contains(alias)) + }) } } } @@ -1241,11 +1262,20 @@ fn extract_model_registration_identity( .get("source_path") .and_then(serde_json::Value::as_str) .map(str::to_owned); + let aliases = card_json + .get("aliases") + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .filter_map(serde_json::Value::as_str) + .map(str::to_owned) + .collect(); let is_lora = model_suffix.is_some() || card_json.get("lora").is_some_and(|value| !value.is_null()); Ok(ModelRegistrationIdentity { display_name, + aliases, source_path, is_lora, })