rob thijssen
e42e8ee81f
Replace NodeConfig (static vram_mb, pinned) with NeuronEndpoint.
Hardware discovery and model pinning now come from neuron API and
models.toml catalogue respectively.
- config.rs: nodes -> neurons, add models_config path
- catalogue.rs: ModelProfile with pinned_on, ModelCatalogue
- poller.rs: poll neuron GET /models (ModelInfo format)
- router.rs: resolve inference endpoint via neuron GET /models/{id}/endpoint
- evictor.rs: call neuron POST /models/unload
- node.rs: remove vram_mb, pinned fields (come from discovery/catalogue)
- All 22 gateway tests updated to mock neuron API
- Remove MistralModelsResponse, ModelLifecycleRequest (no longer needed)
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
112 lines
3.1 KiB
Rust
112 lines
3.1 KiB
Rust
use anyhow::Result;
|
|
use clap::{Parser, Subcommand};
|
|
use cortex_core::config::GatewayConfig;
|
|
use tracing_subscriber::EnvFilter;
|
|
|
|
#[derive(Parser)]
|
|
#[command(name = "cortex")]
|
|
#[command(about = "Unified inference gateway for multi-node mistral.rs clusters")]
|
|
#[command(version)]
|
|
struct Cli {
|
|
#[command(subcommand)]
|
|
command: Commands,
|
|
}
|
|
|
|
#[derive(Subcommand)]
|
|
enum Commands {
|
|
/// Start the gateway server.
|
|
Serve {
|
|
/// Path to the gateway config file.
|
|
#[arg(short, long, default_value = "cortex.toml")]
|
|
config: String,
|
|
},
|
|
/// Print the fleet status (models, nodes, health).
|
|
Status {
|
|
/// Gateway API endpoint to query.
|
|
#[arg(short, long, default_value = "http://localhost:8000")]
|
|
endpoint: String,
|
|
},
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<()> {
|
|
// Initialize tracing with env filter (e.g. RUST_LOG=cortex_gateway=debug).
|
|
tracing_subscriber::fmt()
|
|
.with_env_filter(
|
|
EnvFilter::try_from_default_env()
|
|
.unwrap_or_else(|_| EnvFilter::new("info,cortex_gateway=debug")),
|
|
)
|
|
.init();
|
|
|
|
let cli = Cli::parse();
|
|
|
|
match cli.command {
|
|
Commands::Serve { config } => {
|
|
let cfg = GatewayConfig::load(&config)
|
|
.map_err(|e| anyhow::anyhow!("failed to load config from '{config}': {e}"))?;
|
|
|
|
tracing::info!(
|
|
neurons = cfg.neurons.len(),
|
|
listen = %cfg.gateway.listen,
|
|
"starting cortex"
|
|
);
|
|
|
|
// Install Prometheus metrics exporter on a separate port.
|
|
cortex_gateway::metrics::install(&cfg.gateway.metrics_listen)?;
|
|
|
|
cortex_gateway::run(cfg).await?;
|
|
}
|
|
Commands::Status { endpoint } => {
|
|
print_status(&endpoint).await?;
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
async fn print_status(endpoint: &str) -> Result<()> {
|
|
let client = reqwest::Client::new();
|
|
|
|
// Fetch health.
|
|
let health: serde_json::Value = client
|
|
.get(format!("{endpoint}/health"))
|
|
.send()
|
|
.await?
|
|
.json()
|
|
.await?;
|
|
|
|
println!("Fleet health: {}", serde_json::to_string_pretty(&health)?);
|
|
|
|
// Fetch models.
|
|
let models: serde_json::Value = client
|
|
.get(format!("{endpoint}/v1/models"))
|
|
.send()
|
|
.await?
|
|
.json()
|
|
.await?;
|
|
|
|
println!("\nModels:");
|
|
if let Some(data) = models.get("data").and_then(|d| d.as_array()) {
|
|
for model in data {
|
|
let id = model.get("id").and_then(|v| v.as_str()).unwrap_or("?");
|
|
let locations = model
|
|
.get("locations")
|
|
.and_then(|v| v.as_array())
|
|
.map(|arr| {
|
|
arr.iter()
|
|
.filter_map(|l| {
|
|
let node = l.get("node")?.as_str()?;
|
|
let status = l.get("status")?.as_str()?;
|
|
Some(format!("{node}({status})"))
|
|
})
|
|
.collect::<Vec<_>>()
|
|
.join(", ")
|
|
})
|
|
.unwrap_or_default();
|
|
println!(" {id:40} {locations}");
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|