[model-gateway] Use UUIDs for router-managed worker resources (#15540)

This commit is contained in:
shuwenn
2025-12-21 07:40:03 -10:00
committed by GitHub
parent ef001fb811
commit 1d9ba2ce4d
6 changed files with 189 additions and 64 deletions
@@ -121,6 +121,26 @@ impl WorkerRegistry {
worker_id
}
/// Reserve (or retrieve) a stable UUID for a worker URL.
pub fn reserve_id_for_url(&self, url: &str) -> WorkerId {
if let Some(existing_id) = self.url_to_id.get(url) {
return existing_id.clone();
}
let worker_id = WorkerId::new();
self.url_to_id.insert(url.to_string(), worker_id.clone());
worker_id
}
/// Best-effort lookup of the URL for a given worker ID.
pub fn get_url_by_id(&self, worker_id: &WorkerId) -> Option<String> {
if let Some(worker) = self.get(worker_id) {
return Some(worker.url().to_string());
}
self.url_to_id
.iter()
.find_map(|entry| (entry.value() == worker_id).then(|| entry.key().clone()))
}
/// Remove a worker by ID
pub fn remove(&self, worker_id: &WorkerId) -> Option<Arc<dyn Worker>> {
if let Some((_, worker)) = self.workers.remove(worker_id) {
+122 -40
View File
@@ -19,6 +19,7 @@ use serde::Deserialize;
use serde_json::{json, Value};
use tokio::{signal, spawn};
use tracing::{debug, error, info, warn, Level};
use uuid::Uuid;
use crate::{
app_context::AppContext,
@@ -484,6 +485,10 @@ async fn create_worker(
// Submit job for async processing
let worker_url = config.url.clone();
let worker_id = state
.context
.worker_registry
.reserve_id_for_url(&worker_url);
let job = Job::AddWorker {
config: Box::new(config),
};
@@ -495,12 +500,20 @@ async fn create_worker(
.expect("JobQueue not initialized");
match job_queue.submit(job).await {
Ok(_) => {
let location = format!("/workers/{}", worker_id.as_str());
let response = json!({
"status": "accepted",
"worker_id": worker_url,
"worker_id": worker_id.as_str(),
"url": worker_url,
"location": location,
"message": "Worker addition queued for background processing"
});
(StatusCode::ACCEPTED, Json(response)).into_response()
(
StatusCode::ACCEPTED,
[(http::header::LOCATION, location)],
Json(response),
)
.into_response()
}
Err(error) => {
let error_response = WorkerErrorResponse {
@@ -513,8 +526,15 @@ async fn create_worker(
}
async fn list_workers_rest(State(state): State<Arc<AppState>>) -> Response {
let workers = state.context.worker_registry.get_all();
let worker_infos: Vec<WorkerInfo> = workers.iter().map(worker_to_info).collect();
let workers = state.context.worker_registry.get_all_with_ids();
let worker_infos: Vec<WorkerInfo> = workers
.iter()
.map(|(worker_id, worker)| {
let mut info = worker_to_info(worker);
info.id = worker_id.as_str().to_string();
info
})
.collect();
let response = json!({
"workers": worker_infos,
@@ -528,57 +548,102 @@ async fn list_workers_rest(State(state): State<Arc<AppState>>) -> Response {
Json(response).into_response()
}
async fn get_worker(State(state): State<Arc<AppState>>, Path(url): Path<String>) -> Response {
fn parse_worker_id(raw: &str) -> Result<crate::core::WorkerId, String> {
Uuid::parse_str(raw)
.map_err(|e| format!("Invalid worker_id '{raw}' (expected UUID). Error: {e}"))?;
Ok(crate::core::WorkerId::from_string(raw.to_string()))
}
async fn get_worker(
State(state): State<Arc<AppState>>,
Path(worker_id_raw): Path<String>,
) -> Response {
let worker_id = match parse_worker_id(&worker_id_raw) {
Ok(id) => id,
Err(msg) => {
let error = WorkerErrorResponse {
error: msg,
code: "BAD_REQUEST".to_string(),
};
return (StatusCode::BAD_REQUEST, Json(error)).into_response();
}
};
let job_queue = state
.context
.worker_job_queue
.get()
.expect("JobQueue not initialized");
if let Some(worker) = state.context.worker_registry.get_by_url(&url) {
if let Some(worker) = state.context.worker_registry.get(&worker_id) {
// Worker exists in registry, get its full info and attach job status if any
let worker_url = worker.url().to_string();
let mut worker_info = worker_to_info(&worker);
if let Some(status) = job_queue.get_status(&url) {
worker_info.id = worker_id.as_str().to_string();
if let Some(status) = job_queue.get_status(&worker_url) {
worker_info.job_status = Some(status);
}
return Json(worker_info).into_response();
}
// Worker not in registry, check job queue for its status
if let Some(status) = job_queue.get_status(&url) {
// Create a partial WorkerInfo to report the job status
let worker_info = WorkerInfo {
id: url.clone(),
url: url.clone(),
model_id: "unknown".to_string(),
priority: 0,
cost: 1.0,
worker_type: "unknown".to_string(),
is_healthy: false,
load: 0,
connection_mode: "unknown".to_string(),
runtime_type: None,
tokenizer_path: None,
reasoning_parser: None,
tool_parser: None,
chat_template: None,
bootstrap_port: None,
metadata: HashMap::new(),
job_status: Some(status),
};
return Json(worker_info).into_response();
// Worker not in registry yet. If we can map id -> url (reserved IDs), return job status.
if let Some(worker_url) = state.context.worker_registry.get_url_by_id(&worker_id) {
if let Some(status) = job_queue.get_status(&worker_url) {
// Create a partial WorkerInfo to report the job status
let worker_info = WorkerInfo {
id: worker_id.as_str().to_string(),
url: worker_url.clone(),
model_id: "unknown".to_string(),
priority: 0,
cost: 1.0,
worker_type: "unknown".to_string(),
is_healthy: false,
load: 0,
connection_mode: "unknown".to_string(),
runtime_type: None,
tokenizer_path: None,
reasoning_parser: None,
tool_parser: None,
chat_template: None,
bootstrap_port: None,
metadata: HashMap::new(),
job_status: Some(status),
};
return Json(worker_info).into_response();
}
}
// Worker not found in registry or job queue
let error = WorkerErrorResponse {
error: format!("Worker {url} not found"),
error: format!("Worker {worker_id_raw} not found"),
code: "WORKER_NOT_FOUND".to_string(),
};
(StatusCode::NOT_FOUND, Json(error)).into_response()
}
async fn delete_worker(State(state): State<Arc<AppState>>, Path(url): Path<String>) -> Response {
let worker_id = url.clone();
async fn delete_worker(
State(state): State<Arc<AppState>>,
Path(worker_id_raw): Path<String>,
) -> Response {
let worker_id = match parse_worker_id(&worker_id_raw) {
Ok(id) => id,
Err(msg) => {
let error = WorkerErrorResponse {
error: msg,
code: "BAD_REQUEST".to_string(),
};
return (StatusCode::BAD_REQUEST, Json(error)).into_response();
}
};
let Some(url) = state.context.worker_registry.get_url_by_id(&worker_id) else {
let error = WorkerErrorResponse {
error: format!("Worker {worker_id_raw} not found"),
code: "WORKER_NOT_FOUND".to_string(),
};
return (StatusCode::NOT_FOUND, Json(error)).into_response();
};
let job = Job::RemoveWorker { url };
let job_queue = state
@@ -590,7 +655,7 @@ async fn delete_worker(State(state): State<Arc<AppState>>, Path(url): Path<Strin
Ok(_) => {
let response = json!({
"status": "accepted",
"worker_id": worker_id,
"worker_id": worker_id.as_str(),
"message": "Worker removal queued for background processing"
});
(StatusCode::ACCEPTED, Json(response)).into_response()
@@ -607,10 +672,28 @@ async fn delete_worker(State(state): State<Arc<AppState>>, Path(url): Path<Strin
async fn update_worker(
State(state): State<Arc<AppState>>,
Path(url): Path<String>,
Path(worker_id_raw): Path<String>,
Json(update): Json<WorkerUpdateRequest>,
) -> Response {
let worker_id = url.clone();
let worker_id = match parse_worker_id(&worker_id_raw) {
Ok(id) => id,
Err(msg) => {
let error = WorkerErrorResponse {
error: msg,
code: "BAD_REQUEST".to_string(),
};
return (StatusCode::BAD_REQUEST, Json(error)).into_response();
}
};
let Some(url) = state.context.worker_registry.get_url_by_id(&worker_id) else {
let error = WorkerErrorResponse {
error: format!("Worker {worker_id_raw} not found"),
code: "WORKER_NOT_FOUND".to_string(),
};
return (StatusCode::NOT_FOUND, Json(error)).into_response();
};
let job = Job::UpdateWorker {
url,
update: Box::new(update),
@@ -625,7 +708,7 @@ async fn update_worker(
Ok(_) => {
let response = json!({
"status": "accepted",
"worker_id": worker_id,
"worker_id": worker_id.as_str(),
"message": "Worker update queued for background processing"
});
(StatusCode::ACCEPTED, Json(response)).into_response()
@@ -734,10 +817,9 @@ pub fn build_app(
));
let worker_routes = Router::new()
.route("/workers", post(create_worker))
.route("/workers", get(list_workers_rest))
.route("/workers", post(create_worker).get(list_workers_rest))
.route(
"/workers/{url}",
"/workers/{worker_id}",
get(get_worker).put(update_worker).delete(delete_worker),
)
.route_layer(axum::middleware::from_fn_with_state(