From 4972c7d1e7fa6ca04424e0c2715a2628028eafa0 Mon Sep 17 00:00:00 2001 From: rob thijssen Date: Thu, 4 Jun 2026 13:49:54 +0300 Subject: [PATCH] =?UTF-8?q?feat(cortex-gateway):=20C3=20=E2=80=94=20propag?= =?UTF-8?q?ate=20vision=20capabilities=20through=20/v1/models?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ModelEntry and CortexModelEntry gain a `capabilities: Vec` field (serde-default for back-compat). The poller copies it verbatim from each neuron's ModelInfo.capabilities; list_models computes the union across every node where a model is loaded so a checkpoint loaded text-only on one neuron and text+vision on another reports both to the fleet. Catalogue-only and mid-prewarm entries default to empty until the catalogue gains a capabilities declaration. Aliases inherit their target's capability union. New gateway test mocks two nodes with differing capability arrays and asserts the unioned /v1/models response. Closes part of #16 (Stage C3). Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/cortex-core/src/node.rs | 12 ++++ crates/cortex-gateway/src/handlers.rs | 16 +++++ crates/cortex-gateway/src/poller.rs | 2 + crates/cortex-gateway/src/router.rs | 1 + crates/cortex-gateway/tests/aliases.rs | 3 + crates/cortex-gateway/tests/common/mod.rs | 1 + crates/cortex-gateway/tests/eviction.rs | 3 + crates/cortex-gateway/tests/poller.rs | 83 +++++++++++++++++++++++ 8 files changed, 121 insertions(+) diff --git a/crates/cortex-core/src/node.rs b/crates/cortex-core/src/node.rs index b008577..88e8d61 100644 --- a/crates/cortex-core/src/node.rs +++ b/crates/cortex-core/src/node.rs @@ -37,6 +37,12 @@ pub struct ModelEntry { pub last_accessed: Option>, /// Estimated VRAM usage in MB when loaded. pub vram_estimate_mb: Option, + /// Modalities the loaded model advertises (e.g. `["text", "vision"]`), + /// copied verbatim from the neuron's `ModelInfo.capabilities` at poll + /// time. Empty when the neuron reports none. `#[serde(default)]` keeps + /// older persisted/serialised entries deserialisable. + #[serde(default)] + pub capabilities: Vec, } /// Model lifecycle status. @@ -85,6 +91,12 @@ pub struct CortexModelEntry { /// disjoint from) `feasible_on` depending on whether the catalogue /// covers this model. pub locations: Vec, + /// Union of the modalities advertised by every neuron that has this + /// model loaded (e.g. `["text", "vision"]`). Empty for catalogue-only + /// entries with no loaded location — the catalogue profile doesn't + /// declare capabilities yet (tracked separately from C3). + #[serde(default)] + pub capabilities: Vec, } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/crates/cortex-gateway/src/handlers.rs b/crates/cortex-gateway/src/handlers.rs index fe6e294..bac408c 100644 --- a/crates/cortex-gateway/src/handlers.rs +++ b/crates/cortex-gateway/src/handlers.rs @@ -414,6 +414,9 @@ async fn list_models(State(fleet): State>) -> Json { loaded: false, feasible_on, locations: Vec::new(), + // Catalogue profiles don't declare capabilities yet; + // the union is filled in Pass 2 from loaded locations. + capabilities: Vec::new(), }, ); } @@ -438,6 +441,14 @@ async fn list_models(State(fleet): State>) -> Json { if was_loaded { e.loaded = true; } + // Union the per-node capabilities so a model loaded + // on several neurons reports every modality any of + // them advertises. + for cap in &entry.capabilities { + if !e.capabilities.contains(cap) { + e.capabilities.push(cap.clone()); + } + } }) .or_insert_with(|| CortexModelEntry { id: model_id.clone(), @@ -449,6 +460,7 @@ async fn list_models(State(fleet): State>) -> Json { // feasibility; leave empty. feasible_on: Vec::new(), locations: vec![location], + capabilities: entry.capabilities.clone(), }); } } @@ -498,6 +510,9 @@ async fn list_models(State(fleet): State>) -> Json { loaded: false, feasible_on: Vec::new(), locations: vec![location], + // A model that's only mid-prewarm has no loaded + // location to read capabilities from yet. + capabilities: Vec::new(), }); } } @@ -527,6 +542,7 @@ async fn list_models(State(fleet): State>) -> Json { loaded: target_entry.loaded, feasible_on: target_entry.feasible_on, locations: target_entry.locations, + capabilities: target_entry.capabilities, }, ); } diff --git a/crates/cortex-gateway/src/poller.rs b/crates/cortex-gateway/src/poller.rs index 367c814..5c3e9f6 100644 --- a/crates/cortex-gateway/src/poller.rs +++ b/crates/cortex-gateway/src/poller.rs @@ -107,12 +107,14 @@ async fn poll_neuron(fleet: &CortexState, name: &str, endpoint: &str) { .and_modify(|e| { e.status = status; e.vram_estimate_mb = upstream.vram_used_mb; + e.capabilities = upstream.capabilities.clone(); }) .or_insert_with(|| ModelEntry { id: upstream.id.clone(), status, last_accessed: None, vram_estimate_mb: upstream.vram_used_mb, + capabilities: upstream.capabilities.clone(), }); } diff --git a/crates/cortex-gateway/src/router.rs b/crates/cortex-gateway/src/router.rs index eeb20b4..5f23d5a 100644 --- a/crates/cortex-gateway/src/router.rs +++ b/crates/cortex-gateway/src/router.rs @@ -244,6 +244,7 @@ async fn cold_load( status: ModelStatus::Loaded, last_accessed: Some(chrono::Utc::now()), vram_estimate_mb: profile.vram_mb, + capabilities: Vec::new(), }, ); } diff --git a/crates/cortex-gateway/tests/aliases.rs b/crates/cortex-gateway/tests/aliases.rs index 4bf5a40..ef59373 100644 --- a/crates/cortex-gateway/tests/aliases.rs +++ b/crates/cortex-gateway/tests/aliases.rs @@ -74,6 +74,7 @@ async fn test_alias_resolves_in_chat_completions() { status: ModelStatus::Loaded, last_accessed: None, vram_estimate_mb: None, + capabilities: Vec::new(), }, ); } @@ -154,6 +155,7 @@ async fn test_aliases_surface_in_v1_models() { status: ModelStatus::Loaded, last_accessed: None, vram_estimate_mb: Some(2000), + capabilities: Vec::new(), }, ); } @@ -235,6 +237,7 @@ async fn test_alias_falls_through_for_unmapped_model() { status: ModelStatus::Loaded, last_accessed: None, vram_estimate_mb: None, + capabilities: Vec::new(), }, ); } diff --git a/crates/cortex-gateway/tests/common/mod.rs b/crates/cortex-gateway/tests/common/mod.rs index 5a12c9b..294474b 100644 --- a/crates/cortex-gateway/tests/common/mod.rs +++ b/crates/cortex-gateway/tests/common/mod.rs @@ -305,6 +305,7 @@ pub async fn spawn_gateway_with_state(mock_url: &str) -> (Arc, Stri status: ModelStatus::Loaded, last_accessed: None, vram_estimate_mb: Some(8000), + capabilities: Vec::new(), }, ); } diff --git a/crates/cortex-gateway/tests/eviction.rs b/crates/cortex-gateway/tests/eviction.rs index b2b9ab8..71c8637 100644 --- a/crates/cortex-gateway/tests/eviction.rs +++ b/crates/cortex-gateway/tests/eviction.rs @@ -91,6 +91,7 @@ async fn test_evict_lru_model() { status: ModelStatus::Loaded, last_accessed: Some(Utc::now() - chrono::Duration::hours(2)), vram_estimate_mb: Some(8000), + capabilities: Vec::new(), }, ); node.models.insert( @@ -100,6 +101,7 @@ async fn test_evict_lru_model() { status: ModelStatus::Loaded, last_accessed: Some(Utc::now()), vram_estimate_mb: Some(8000), + capabilities: Vec::new(), }, ); } @@ -163,6 +165,7 @@ async fn test_eviction_increments_lifecycle_cycles() { status: ModelStatus::Loaded, last_accessed: None, vram_estimate_mb: None, + capabilities: Vec::new(), }, ); } diff --git a/crates/cortex-gateway/tests/poller.rs b/crates/cortex-gateway/tests/poller.rs index fe76ada..91ab42f 100644 --- a/crates/cortex-gateway/tests/poller.rs +++ b/crates/cortex-gateway/tests/poller.rs @@ -118,6 +118,87 @@ async fn test_poller_updates_gateway_models_endpoint() { } } +#[tokio::test] +async fn test_models_endpoint_unions_capabilities_across_nodes() { + // C3: two neurons each have the same model loaded but advertise + // different capability sets. The gateway's /v1/models must report + // the union — a model loaded text-only on one node and + // text+vision on another is vision-capable to the fleet. + let node_a = common::spawn_mock_neuron_with_models(json!([ + {"id": "shared-model", "harness": "candle", "status": "loaded", "devices": [0], "vram_used_mb": null, "capabilities": ["text"]} + ])) + .await; + let node_b = common::spawn_mock_neuron_with_models(json!([ + {"id": "shared-model", "harness": "candle", "status": "loaded", "devices": [1], "vram_used_mb": null, "capabilities": ["text", "vision"]} + ])) + .await; + + let config = GatewayConfig { + gateway: GatewaySettings { + listen: "127.0.0.1:0".into(), + metrics_listen: "127.0.0.1:0".into(), + }, + eviction: EvictionSettings { + strategy: EvictionStrategy::Lru, + defrag_after_cycles: 0, + }, + neurons: vec![ + NeuronEndpoint { + name: "node-a".into(), + endpoint: node_a, + }, + NeuronEndpoint { + name: "node-b".into(), + endpoint: node_b, + }, + ], + models_config: "/dev/null".into(), + }; + + let fleet = Arc::new(CortexState::from_config(&config)); + cortex_gateway::poller::poll_once(&fleet).await; + + let app = cortex_gateway::build_app(Arc::clone(&fleet)); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + + let client = reqwest::Client::new(); + let body: serde_json::Value = client + .get(format!("http://{addr}/v1/models")) + .send() + .await + .expect("request should succeed") + .json() + .await + .unwrap(); + + let model = body["data"] + .as_array() + .expect("data array") + .iter() + .find(|m| m["id"] == "shared-model") + .expect("shared-model should be present"); + + let caps: Vec<&str> = model["capabilities"] + .as_array() + .expect("capabilities array") + .iter() + .filter_map(|c| c.as_str()) + .collect(); + assert!(caps.contains(&"text"), "union must include text: {caps:?}"); + assert!( + caps.contains(&"vision"), + "union must include vision: {caps:?}" + ); + assert_eq!(caps.len(), 2, "union must not duplicate text: {caps:?}"); + + // Both nodes hold the model, so two locations regardless of caps. + assert_eq!(model["locations"].as_array().unwrap().len(), 2); +} + #[tokio::test] async fn test_poller_marks_unreachable_node_unhealthy() { let config = GatewayConfig { @@ -216,6 +297,7 @@ async fn test_poller_removes_stale_models() { status: ModelStatus::Loaded, last_accessed: None, vram_estimate_mb: None, + capabilities: Vec::new(), }, ); node.models.insert( @@ -225,6 +307,7 @@ async fn test_poller_removes_stale_models() { status: ModelStatus::Loaded, last_accessed: None, vram_estimate_mb: None, + capabilities: Vec::new(), }, ); }