[model-gateway] Refactor worker steps and add update workflow (#15085)

sgl-model-gateway/src/core/job_queue.rs

@@ -22,7 +22,7 @@ use crate::{
     },
     mcp::McpConfig,
     observability::metrics::RouterMetrics,
-    protocols::worker_spec::{JobStatus, WorkerConfigRequest},
+    protocols::worker_spec::{JobStatus, WorkerConfigRequest, WorkerUpdateRequest},
     workflow::{WorkflowContext, WorkflowEngine, WorkflowId, WorkflowInstanceId, WorkflowStatus},
 };
 
@@ -32,6 +32,10 @@ pub enum Job {
     AddWorker {
         config: Box<WorkerConfigRequest>,
     },
+    UpdateWorker {
+        url: String,
+        update: Box<WorkerUpdateRequest>,
+    },
     RemoveWorker {
         url: String,
     },
@@ -57,6 +61,7 @@ impl Job {
     pub fn job_type(&self) -> &str {
         match self {
             Job::AddWorker { .. } => "AddWorker",
+            Job::UpdateWorker { .. } => "UpdateWorker",
             Job::RemoveWorker { .. } => "RemoveWorker",
             Job::InitializeWorkersFromConfig { .. } => "InitializeWorkersFromConfig",
             Job::InitializeMcpServers { .. } => "InitializeMcpServers",
@@ -70,6 +75,7 @@ impl Job {
     pub fn worker_url(&self) -> &str {
         match self {
             Job::AddWorker { config } => &config.url,
+            Job::UpdateWorker { url, .. } => url,
             Job::RemoveWorker { url } => url,
             Job::InitializeWorkersFromConfig { .. } => "startup",
             Job::InitializeMcpServers { .. } => "startup",
@@ -342,6 +348,24 @@ impl JobQueue {
                 )
                 .await
             }
+            Job::UpdateWorker { url, update } => {
+                let engine = context
+                    .workflow_engine
+                    .get()
+                    .ok_or_else(|| "Workflow engine not initialized".to_string())?;
+
+                let instance_id =
+                    Self::start_worker_update_workflow(engine, url, update, context).await?;
+
+                debug!(
+                    "Started worker update workflow for {} (instance: {})",
+                    url, instance_id
+                );
+
+                let timeout_duration = Duration::from_secs(30);
+
+                Self::wait_for_workflow_completion(engine, instance_id, url, timeout_duration).await
+            }
             Job::RemoveWorker { url } => {
                 let engine = context
                     .workflow_engine
@@ -684,6 +708,26 @@ impl JobQueue {
             .map_err(|e| format!("Failed to start worker removal workflow: {:?}", e))
     }
 
+    /// Start worker update workflow
+    async fn start_worker_update_workflow(
+        engine: &Arc<WorkflowEngine>,
+        url: &str,
+        update: &WorkerUpdateRequest,
+        context: &Arc<AppContext>,
+    ) -> Result<WorkflowInstanceId, String> {
+        let mut workflow_context = WorkflowContext::new(WorkflowInstanceId::new());
+        // Pass URL and dp_aware separately, workflow step handles the rest
+        workflow_context.set("worker_url", url.to_string());
+        workflow_context.set("dp_aware", context.router_config.dp_aware);
+        workflow_context.set("update_request", update.clone());
+        workflow_context.set_arc("app_context", Arc::clone(context));
+
+        engine
+            .start_workflow(WorkflowId::new("worker_update"), workflow_context)
+            .await
+            .map_err(|e| format!("Failed to start worker update workflow: {:?}", e))
+    }
+
     /// Start MCP server registration workflow
     async fn start_mcp_registration_workflow(
         engine: &Arc<WorkflowEngine>,

sgl-model-gateway/src/core/steps/mod.rs

@@ -1,7 +1,7 @@
 //! Workflow step implementations
 //!
 //! This module contains concrete step implementations for various workflows:
-//! - Worker management (registration, removal, future: updates)
+//! - Worker management (registration, removal, updates)
 //! - MCP server registration
 //! - WASM module registration and removal
 //! - Future: Tokenizer fetching, LoRA updates, etc.
@@ -22,6 +22,7 @@ pub use worker::{
     create_external_worker_workflow,
     create_local_worker_workflow,
     create_worker_removal_workflow,
+    create_worker_update_workflow,
     // Utility functions
     group_models_into_cards,
     infer_model_type_from_id,
@@ -36,6 +37,8 @@ pub use worker::{
     DiscoverMetadataStep,
     DiscoverModelsStep,
     DpInfo,
+    // Update steps
+    FindWorkerToUpdateStep,
     // Removal steps
     FindWorkersToRemoveStep,
     ModelInfo,
@@ -43,8 +46,10 @@ pub use worker::{
     RegisterWorkersStep,
     RemoveFromPolicyRegistryStep,
     RemoveFromWorkerRegistryStep,
+    UpdatePoliciesForWorkerStep,
     UpdatePoliciesStep,
     UpdateRemainingPoliciesStep,
+    UpdateWorkerPropertiesStep,
     WorkerList,
     WorkerRemovalRequest,
 };

sgl-model-gateway/src/core/steps/worker/local/find_worker_to_update.rs

@@ -0,0 +1,64 @@
+//! Step to find a worker to update based on URL.
+
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use tracing::debug;
+
+use super::find_workers_by_url;
+use crate::{
+    app_context::AppContext,
+    workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
+};
+
+/// Step to find workers to update based on URL.
+///
+/// For DP-aware workers, finds all workers with matching URL prefix.
+/// For regular workers, finds the single worker with exact URL match.
+///
+/// Expects the following context values:
+/// - "worker_url": String - the URL of the worker to update
+/// - "dp_aware": bool - whether to find all DP-aware workers with matching prefix
+/// - "app_context": Arc<AppContext>
+///
+/// Sets the following context values:
+/// - "workers_to_update": Vec<Arc<dyn Worker>>
+pub struct FindWorkerToUpdateStep;
+
+#[async_trait]
+impl StepExecutor for FindWorkerToUpdateStep {
+    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+        let worker_url: Arc<String> = context.get_or_err("worker_url")?;
+        let dp_aware: Arc<bool> = context.get_or_err("dp_aware")?;
+        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+
+        let workers_to_update =
+            find_workers_by_url(&app_context.worker_registry, &worker_url, *dp_aware);
+
+        if workers_to_update.is_empty() {
+            let error_msg = if *dp_aware {
+                format!("No workers found with prefix {}@", *worker_url)
+            } else {
+                format!("Worker {} not found", *worker_url)
+            };
+            return Err(WorkflowError::StepFailed {
+                step_id: StepId::new("find_worker_to_update"),
+                message: error_msg,
+            });
+        }
+
+        debug!(
+            "Found {} worker(s) to update for {}",
+            workers_to_update.len(),
+            *worker_url
+        );
+
+        context.set("workers_to_update", workers_to_update);
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false
+    }
+}

sgl-model-gateway/src/core/steps/worker/local/find_workers_to_remove.rs

@@ -0,0 +1,75 @@
+//! Step to find workers to remove based on URL.
+
+use std::{collections::HashSet, sync::Arc};
+
+use async_trait::async_trait;
+use tracing::debug;
+
+use super::find_workers_by_url;
+use crate::{
+    app_context::AppContext,
+    workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
+};
+
+/// Request structure for worker removal.
+#[derive(Debug, Clone)]
+pub struct WorkerRemovalRequest {
+    pub url: String,
+    pub dp_aware: bool,
+}
+
+/// Step to find workers to remove based on URL.
+///
+/// For DP-aware workers, finds all workers with matching URL prefix.
+/// For regular workers, finds the single worker with exact URL match.
+pub struct FindWorkersToRemoveStep;
+
+#[async_trait]
+impl StepExecutor for FindWorkersToRemoveStep {
+    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+        let request: Arc<WorkerRemovalRequest> = context.get_or_err("removal_request")?;
+        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+
+        let workers_to_remove =
+            find_workers_by_url(&app_context.worker_registry, &request.url, request.dp_aware);
+
+        if workers_to_remove.is_empty() {
+            let error_msg = if request.dp_aware {
+                format!("No workers found with prefix {}@", request.url)
+            } else {
+                format!("Worker {} not found", request.url)
+            };
+            return Err(WorkflowError::StepFailed {
+                step_id: StepId::new("find_workers_to_remove"),
+                message: error_msg,
+            });
+        }
+
+        debug!(
+            "Found {} worker(s) to remove for {}",
+            workers_to_remove.len(),
+            request.url
+        );
+
+        // Store workers and their model IDs for subsequent steps
+        let worker_urls: Vec<String> = workers_to_remove
+            .iter()
+            .map(|w| w.url().to_string())
+            .collect();
+
+        let affected_models: HashSet<String> = workers_to_remove
+            .iter()
+            .map(|w| w.model_id().to_string())
+            .collect();
+
+        context.set("workers_to_remove", workers_to_remove);
+        context.set("worker_urls", worker_urls);
+        context.set("affected_models", affected_models);
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false
+    }
+}

sgl-model-gateway/src/core/steps/worker/local/mod.rs

@@ -2,7 +2,13 @@ mod create_worker;
 mod detect_connection;
 mod discover_dp;
 mod discover_metadata;
-mod removal;
+mod find_worker_to_update;
+mod find_workers_to_remove;
+mod remove_from_policy_registry;
+mod remove_from_worker_registry;
+mod update_policies_for_worker;
+mod update_remaining_policies;
+mod update_worker_properties;
 
 use std::{sync::Arc, time::Duration};
 
@@ -10,17 +16,46 @@ pub use create_worker::CreateLocalWorkerStep;
 pub use detect_connection::DetectConnectionModeStep;
 pub use discover_dp::{get_dp_info, DiscoverDPInfoStep, DpInfo};
 pub use discover_metadata::DiscoverMetadataStep;
-pub use removal::{
+pub use find_worker_to_update::FindWorkerToUpdateStep;
-    FindWorkersToRemoveStep, RemoveFromPolicyRegistryStep, RemoveFromWorkerRegistryStep,
+pub use find_workers_to_remove::{FindWorkersToRemoveStep, WorkerRemovalRequest};
-    UpdateRemainingPoliciesStep, WorkerRemovalRequest,
+pub use remove_from_policy_registry::RemoveFromPolicyRegistryStep;
-};
+pub use remove_from_worker_registry::RemoveFromWorkerRegistryStep;
+pub use update_policies_for_worker::UpdatePoliciesForWorkerStep;
+pub use update_remaining_policies::UpdateRemainingPoliciesStep;
+pub use update_worker_properties::UpdateWorkerPropertiesStep;
 
 use super::shared::{ActivateWorkersStep, RegisterWorkersStep, UpdatePoliciesStep};
 use crate::{
     config::RouterConfig,
+    core::{Worker, WorkerRegistry},
     workflow::{BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, WorkflowDefinition},
 };
 
+/// Find workers by URL, supporting both DP-aware (prefix match) and regular (exact match) modes.
+///
+/// For DP-aware workers, finds all workers with URL prefix `{url}@`.
+/// For regular workers, finds the single worker with exact URL match.
+pub(crate) fn find_workers_by_url(
+    registry: &WorkerRegistry,
+    url: &str,
+    dp_aware: bool,
+) -> Vec<Arc<dyn Worker>> {
+    if dp_aware {
+        let worker_url_prefix = format!("{}@", url);
+        registry
+            .get_all()
+            .iter()
+            .filter(|worker| worker.url().starts_with(&worker_url_prefix))
+            .cloned()
+            .collect()
+    } else {
+        match registry.get_by_url(url) {
+            Some(worker) => vec![worker],
+            None => Vec::new(),
+        }
+    }
+}
+
 pub fn create_local_worker_workflow(router_config: &RouterConfig) -> WorkflowDefinition {
     let detect_timeout = Duration::from_secs(router_config.worker_startup_timeout_secs);
 
@@ -193,3 +228,55 @@ pub fn create_worker_removal_workflow() -> WorkflowDefinition {
             .depends_on(&["remove_from_worker_registry"]),
         )
 }
+
+/// Create a worker update workflow definition.
+///
+/// DAG structure:
+/// ```text
+///     find_worker_to_update
+///              │
+///     update_worker_properties
+///              │
+///     update_policies_for_worker
+/// ```
+pub fn create_worker_update_workflow() -> WorkflowDefinition {
+    WorkflowDefinition::new("worker_update", "Update worker properties")
+        .add_step(
+            StepDefinition::new(
+                "find_worker_to_update",
+                "Find worker to update",
+                Arc::new(FindWorkerToUpdateStep),
+            )
+            .with_timeout(Duration::from_secs(10))
+            .with_retry(RetryPolicy {
+                max_attempts: 1,
+                backoff: BackoffStrategy::Fixed(Duration::from_secs(0)),
+            }),
+        )
+        .add_step(
+            StepDefinition::new(
+                "update_worker_properties",
+                "Update worker properties",
+                Arc::new(UpdateWorkerPropertiesStep),
+            )
+            .with_timeout(Duration::from_secs(10))
+            .with_retry(RetryPolicy {
+                max_attempts: 1,
+                backoff: BackoffStrategy::Fixed(Duration::from_secs(0)),
+            })
+            .depends_on(&["find_worker_to_update"]),
+        )
+        .add_step(
+            StepDefinition::new(
+                "update_policies_for_worker",
+                "Update policies for updated worker",
+                Arc::new(UpdatePoliciesForWorkerStep),
+            )
+            .with_timeout(Duration::from_secs(10))
+            .with_retry(RetryPolicy {
+                max_attempts: 1,
+                backoff: BackoffStrategy::Fixed(Duration::from_secs(0)),
+            })
+            .depends_on(&["update_worker_properties"]),
+        )
+}

sgl-model-gateway/src/core/steps/worker/local/removal.rs

@@ -1,230 +0,0 @@
-//! Worker removal step implementations.
-
-use std::{collections::HashSet, sync::Arc};
-
-use async_trait::async_trait;
-use tracing::{debug, info};
-
-use crate::{
-    app_context::AppContext,
-    core::Worker,
-    workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
-};
-
-/// Request structure for worker removal.
-#[derive(Debug, Clone)]
-pub struct WorkerRemovalRequest {
-    pub url: String,
-    pub dp_aware: bool,
-}
-
-/// Step 1: Find workers to remove based on URL.
-pub struct FindWorkersToRemoveStep;
-
-#[async_trait]
-impl StepExecutor for FindWorkersToRemoveStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
-        let request: Arc<WorkerRemovalRequest> = context.get_or_err("removal_request")?;
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
-
-        debug!(
-            "Finding workers to remove for {} (dp_aware: {})",
-            request.url, request.dp_aware
-        );
-
-        let workers_to_remove: Vec<Arc<dyn Worker>> = if request.dp_aware {
-            // DP-aware: Find all workers with matching prefix
-            let worker_url_prefix = format!("{}@", request.url);
-            let all_workers = app_context.worker_registry.get_all();
-
-            all_workers
-                .iter()
-                .filter(|worker| worker.url().starts_with(&worker_url_prefix))
-                .cloned()
-                .collect()
-        } else {
-            // Non-DP-aware: Find single worker by exact URL
-            match app_context.worker_registry.get_by_url(&request.url) {
-                Some(worker) => vec![worker],
-                None => Vec::new(),
-            }
-        };
-
-        if workers_to_remove.is_empty() {
-            let error_msg = if request.dp_aware {
-                format!("No workers found with prefix {}@", request.url)
-            } else {
-                format!("Worker {} not found", request.url)
-            };
-            return Err(WorkflowError::StepFailed {
-                step_id: StepId::new("find_workers_to_remove"),
-                message: error_msg,
-            });
-        }
-
-        debug!(
-            "Found {} worker(s) to remove for {}",
-            workers_to_remove.len(),
-            request.url
-        );
-
-        // Store workers and their model IDs for subsequent steps
-        let worker_urls: Vec<String> = workers_to_remove
-            .iter()
-            .map(|w| w.url().to_string())
-            .collect();
-
-        let affected_models: HashSet<String> = workers_to_remove
-            .iter()
-            .map(|w| w.model_id().to_string())
-            .collect();
-
-        context.set("workers_to_remove", workers_to_remove);
-        context.set("worker_urls", worker_urls);
-        context.set("affected_models", affected_models);
-
-        Ok(StepResult::Success)
-    }
-
-    fn is_retryable(&self, _error: &WorkflowError) -> bool {
-        false
-    }
-}
-
-/// Step 2: Remove workers from policy registry.
-pub struct RemoveFromPolicyRegistryStep;
-
-#[async_trait]
-impl StepExecutor for RemoveFromPolicyRegistryStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
-        let workers_to_remove: Arc<Vec<Arc<dyn Worker>>> =
-            context.get_or_err("workers_to_remove")?;
-
-        debug!(
-            "Removing {} worker(s) from policy registry",
-            workers_to_remove.len()
-        );
-
-        for worker in workers_to_remove.iter() {
-            let model_id = worker.model_id().to_string();
-            let worker_url = worker.url();
-
-            // Remove from cache-aware policy
-            app_context
-                .policy_registry
-                .remove_worker_from_cache_aware(&model_id, worker_url);
-
-            // Notify policy registry
-            app_context.policy_registry.on_worker_removed(&model_id);
-
-            debug!(
-                "Removed worker {} from policy registry (model: {})",
-                worker_url, model_id
-            );
-        }
-
-        Ok(StepResult::Success)
-    }
-
-    fn is_retryable(&self, _error: &WorkflowError) -> bool {
-        false
-    }
-}
-
-/// Step 3: Remove workers from worker registry.
-pub struct RemoveFromWorkerRegistryStep;
-
-#[async_trait]
-impl StepExecutor for RemoveFromWorkerRegistryStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
-        let worker_urls: Arc<Vec<String>> = context.get_or_err("worker_urls")?;
-
-        debug!(
-            "Removing {} worker(s) from worker registry",
-            worker_urls.len()
-        );
-
-        let mut removed_count = 0;
-        for worker_url in worker_urls.iter() {
-            if app_context
-                .worker_registry
-                .remove_by_url(worker_url)
-                .is_some()
-            {
-                removed_count += 1;
-                debug!("Removed worker {} from registry", worker_url);
-            }
-        }
-
-        if removed_count != worker_urls.len() {
-            return Err(WorkflowError::StepFailed {
-                step_id: StepId::new("remove_from_worker_registry"),
-                message: format!(
-                    "Expected to remove {} workers but only removed {}",
-                    worker_urls.len(),
-                    removed_count
-                ),
-            });
-        }
-
-        Ok(StepResult::Success)
-    }
-
-    fn is_retryable(&self, _error: &WorkflowError) -> bool {
-        false
-    }
-}
-
-/// Step 4: Update cache-aware policies for remaining workers.
-pub struct UpdateRemainingPoliciesStep;
-
-#[async_trait]
-impl StepExecutor for UpdateRemainingPoliciesStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
-        let affected_models: Arc<HashSet<String>> = context.get_or_err("affected_models")?;
-        let worker_urls: Arc<Vec<String>> = context.get_or_err("worker_urls")?;
-
-        debug!(
-            "Updating cache-aware policies for {} affected model(s)",
-            affected_models.len()
-        );
-
-        for model_id in affected_models.iter() {
-            let remaining_workers = app_context.worker_registry.get_by_model_fast(model_id);
-
-            if let Some(policy) = app_context.policy_registry.get_policy(model_id) {
-                if policy.name() == "cache_aware" && !remaining_workers.is_empty() {
-                    app_context
-                        .policy_registry
-                        .init_cache_aware_policy(model_id, &remaining_workers);
-
-                    debug!(
-                        "Updated cache-aware policy for model {} ({} remaining workers)",
-                        model_id,
-                        remaining_workers.len()
-                    );
-                }
-            }
-        }
-
-        // Log final result at info level
-        if worker_urls.len() == 1 {
-            info!("Removed worker {}", worker_urls[0]);
-        } else {
-            info!(
-                "Removed {} DP-aware workers: {:?}",
-                worker_urls.len(),
-                worker_urls
-            );
-        }
-
-        Ok(StepResult::Success)
-    }
-
-    fn is_retryable(&self, _error: &WorkflowError) -> bool {
-        false
-    }
-}

sgl-model-gateway/src/core/steps/worker/local/remove_from_policy_registry.rs

@@ -0,0 +1,56 @@
+//! Step to remove workers from policy registry.
+
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use tracing::debug;
+
+use crate::{
+    app_context::AppContext,
+    core::Worker,
+    workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
+};
+
+/// Step to remove workers from the policy registry.
+///
+/// Removes each worker from cache-aware policies and notifies
+/// the policy registry of worker removal.
+pub struct RemoveFromPolicyRegistryStep;
+
+#[async_trait]
+impl StepExecutor for RemoveFromPolicyRegistryStep {
+    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+        let workers_to_remove: Arc<Vec<Arc<dyn Worker>>> =
+            context.get_or_err("workers_to_remove")?;
+
+        debug!(
+            "Removing {} worker(s) from policy registry",
+            workers_to_remove.len()
+        );
+
+        for worker in workers_to_remove.iter() {
+            let model_id = worker.model_id().to_string();
+            let worker_url = worker.url();
+
+            // Remove from cache-aware policy
+            app_context
+                .policy_registry
+                .remove_worker_from_cache_aware(&model_id, worker_url);
+
+            // Notify policy registry
+            app_context.policy_registry.on_worker_removed(&model_id);
+        }
+
+        debug!(
+            "Removed {} worker(s) from policy registry",
+            workers_to_remove.len()
+        );
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false
+    }
+}

sgl-model-gateway/src/core/steps/worker/local/remove_from_worker_registry.rs

@@ -0,0 +1,57 @@
+//! Step to remove workers from worker registry.
+
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use tracing::{debug, warn};
+
+use crate::{
+    app_context::AppContext,
+    workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
+};
+
+/// Step to remove workers from the worker registry.
+///
+/// Removes each worker by URL from the central worker registry.
+pub struct RemoveFromWorkerRegistryStep;
+
+#[async_trait]
+impl StepExecutor for RemoveFromWorkerRegistryStep {
+    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+        let worker_urls: Arc<Vec<String>> = context.get_or_err("worker_urls")?;
+
+        debug!(
+            "Removing {} worker(s) from worker registry",
+            worker_urls.len()
+        );
+
+        let mut removed_count = 0;
+        for worker_url in worker_urls.iter() {
+            if app_context
+                .worker_registry
+                .remove_by_url(worker_url)
+                .is_some()
+            {
+                removed_count += 1;
+            }
+        }
+
+        // Log if some workers were already removed (e.g., by another process)
+        if removed_count != worker_urls.len() {
+            warn!(
+                "Removed {} of {} workers (some may have been removed by another process)",
+                removed_count,
+                worker_urls.len()
+            );
+        } else {
+            debug!("Removed {} worker(s) from registry", removed_count);
+        }
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false
+    }
+}

sgl-model-gateway/src/core/steps/worker/local/update_policies_for_worker.rs

@@ -0,0 +1,65 @@
+//! Step to update policies for updated workers.
+
+use std::{collections::HashSet, sync::Arc};
+
+use async_trait::async_trait;
+use tracing::debug;
+
+use crate::{
+    app_context::AppContext,
+    core::Worker,
+    workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
+};
+
+/// Step to update policies for updated workers.
+///
+/// After workers are updated, this step re-initializes cache-aware policies
+/// for the affected models to reflect any priority or label changes.
+pub struct UpdatePoliciesForWorkerStep;
+
+#[async_trait]
+impl StepExecutor for UpdatePoliciesForWorkerStep {
+    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+        let updated_workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("updated_workers")?;
+
+        // Collect affected models
+        let affected_models: HashSet<String> = updated_workers
+            .iter()
+            .map(|w| w.model_id().to_string())
+            .collect();
+
+        debug!(
+            "Updating policies for {} affected model(s) after worker update",
+            affected_models.len()
+        );
+
+        for model_id in &affected_models {
+            let workers = app_context.worker_registry.get_by_model_fast(model_id);
+
+            if let Some(policy) = app_context.policy_registry.get_policy(model_id) {
+                if policy.name() == "cache_aware" && !workers.is_empty() {
+                    // Re-initialize cache-aware policy with updated workers
+                    app_context
+                        .policy_registry
+                        .init_cache_aware_policy(model_id, &workers);
+
+                    debug!(
+                        "Updated cache-aware policy for model {} ({} workers)",
+                        model_id,
+                        workers.len()
+                    );
+                }
+            }
+
+            // Notify policy registry of the update
+            app_context.policy_registry.on_worker_added(model_id, None);
+        }
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false
+    }
+}

sgl-model-gateway/src/core/steps/worker/local/update_remaining_policies.rs

@@ -0,0 +1,66 @@
+//! Step to update cache-aware policies for remaining workers after removal.
+
+use std::{collections::HashSet, sync::Arc};
+
+use async_trait::async_trait;
+use tracing::{debug, info};
+
+use crate::{
+    app_context::AppContext,
+    workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
+};
+
+/// Step to update cache-aware policies for remaining workers.
+///
+/// After workers are removed, this step re-initializes cache-aware policies
+/// for the affected models using the remaining workers.
+pub struct UpdateRemainingPoliciesStep;
+
+#[async_trait]
+impl StepExecutor for UpdateRemainingPoliciesStep {
+    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+        let affected_models: Arc<HashSet<String>> = context.get_or_err("affected_models")?;
+        let worker_urls: Arc<Vec<String>> = context.get_or_err("worker_urls")?;
+
+        debug!(
+            "Updating cache-aware policies for {} affected model(s)",
+            affected_models.len()
+        );
+
+        for model_id in affected_models.iter() {
+            let remaining_workers = app_context.worker_registry.get_by_model_fast(model_id);
+
+            if let Some(policy) = app_context.policy_registry.get_policy(model_id) {
+                if policy.name() == "cache_aware" && !remaining_workers.is_empty() {
+                    app_context
+                        .policy_registry
+                        .init_cache_aware_policy(model_id, &remaining_workers);
+
+                    debug!(
+                        "Updated cache-aware policy for model {} ({} remaining workers)",
+                        model_id,
+                        remaining_workers.len()
+                    );
+                }
+            }
+        }
+
+        // Log final result at info level
+        if worker_urls.len() == 1 {
+            info!("Removed worker {}", worker_urls[0]);
+        } else {
+            info!(
+                "Removed {} DP-aware workers: {:?}",
+                worker_urls.len(),
+                worker_urls
+            );
+        }
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false
+    }
+}

sgl-model-gateway/src/core/steps/worker/local/update_worker_properties.rs

@@ -0,0 +1,147 @@
+//! Step to update worker properties.
+
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use tracing::{debug, info};
+
+use crate::{
+    app_context::AppContext,
+    core::{BasicWorkerBuilder, HealthConfig, Worker},
+    protocols::worker_spec::WorkerUpdateRequest,
+    workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
+};
+
+/// Step to update worker properties.
+///
+/// This step creates new worker instances with updated properties and
+/// re-registers them to replace the old workers in the registry.
+///
+/// Expects the following context values:
+/// - "update_request": WorkerUpdateRequest (from protocols::worker_spec)
+/// - "app_context": Arc<AppContext>
+/// - "workers_to_update": Vec<Arc<dyn Worker>>
+///
+/// Sets the following context values:
+/// - "updated_workers": Vec<Arc<dyn Worker>>
+pub struct UpdateWorkerPropertiesStep;
+
+#[async_trait]
+impl StepExecutor for UpdateWorkerPropertiesStep {
+    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+        let request: Arc<WorkerUpdateRequest> = context.get_or_err("update_request")?;
+        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+        let workers_to_update: Arc<Vec<Arc<dyn Worker>>> =
+            context.get_or_err("workers_to_update")?;
+
+        debug!(
+            "Updating properties for {} worker(s)",
+            workers_to_update.len()
+        );
+
+        let mut updated_workers: Vec<Arc<dyn Worker>> = Vec::with_capacity(workers_to_update.len());
+
+        for worker in workers_to_update.iter() {
+            // Build updated labels - merge new labels into existing ones
+            let mut updated_labels = worker.metadata().labels.clone();
+            if let Some(ref new_labels) = request.labels {
+                for (key, value) in new_labels {
+                    updated_labels.insert(key.clone(), value.clone());
+                }
+            }
+
+            // Update priority if specified (stored in labels)
+            if let Some(priority) = request.priority {
+                updated_labels.insert("priority".to_string(), priority.to_string());
+            }
+
+            // Update cost if specified (stored in labels)
+            if let Some(cost) = request.cost {
+                updated_labels.insert("cost".to_string(), cost.to_string());
+            }
+
+            // Build updated health config
+            let existing_health = &worker.metadata().health_config;
+            let updated_health_config = HealthConfig {
+                timeout_secs: request
+                    .health_check_timeout_secs
+                    .unwrap_or(existing_health.timeout_secs),
+                check_interval_secs: request
+                    .health_check_interval_secs
+                    .unwrap_or(existing_health.check_interval_secs),
+                endpoint: existing_health.endpoint.clone(),
+                failure_threshold: request
+                    .health_failure_threshold
+                    .unwrap_or(existing_health.failure_threshold),
+                success_threshold: request
+                    .health_success_threshold
+                    .unwrap_or(existing_health.success_threshold),
+            };
+
+            // Determine API key: use new one if provided, otherwise keep existing
+            let updated_api_key = request
+                .api_key
+                .clone()
+                .or_else(|| worker.metadata().api_key.clone());
+
+            // Create a new worker with updated properties
+            let new_worker: Arc<dyn Worker> = if worker.is_dp_aware() {
+                // For DP-aware workers, extract DP info and rebuild
+                let dp_rank = worker.dp_rank().unwrap_or(0);
+                let dp_size = worker.dp_size().unwrap_or(1);
+                let base_url = worker.base_url().to_string();
+
+                let mut builder =
+                    crate::core::DPAwareWorkerBuilder::new(base_url, dp_rank, dp_size)
+                        .worker_type(worker.worker_type().clone())
+                        .connection_mode(worker.connection_mode().clone())
+                        .runtime_type(worker.metadata().runtime_type.clone())
+                        .labels(updated_labels)
+                        .health_config(updated_health_config.clone())
+                        .models(worker.metadata().models.clone());
+
+                if let Some(ref api_key) = updated_api_key {
+                    builder = builder.api_key(api_key.clone());
+                }
+
+                Arc::new(builder.build())
+            } else {
+                // For basic workers, rebuild with updated properties
+                let mut builder = BasicWorkerBuilder::new(worker.url())
+                    .worker_type(worker.worker_type().clone())
+                    .connection_mode(worker.connection_mode().clone())
+                    .runtime_type(worker.metadata().runtime_type.clone())
+                    .labels(updated_labels)
+                    .health_config(updated_health_config.clone())
+                    .models(worker.metadata().models.clone());
+
+                if let Some(ref api_key) = updated_api_key {
+                    builder = builder.api_key(api_key.clone());
+                }
+
+                Arc::new(builder.build())
+            };
+
+            // Re-register the worker (this replaces the old one)
+            app_context.worker_registry.register(new_worker.clone());
+
+            updated_workers.push(new_worker);
+        }
+
+        // Log result
+        if updated_workers.len() == 1 {
+            info!("Updated worker {}", updated_workers[0].url());
+        } else {
+            info!("Updated {} workers", updated_workers.len());
+        }
+
+        // Store updated workers for subsequent steps
+        context.set("updated_workers", updated_workers);
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false
+    }
+}

sgl-model-gateway/src/core/steps/worker/mod.rs

@@ -9,9 +9,10 @@ pub use external::{
 };
 pub use local::{
     create_local_worker_workflow as create_worker_registration_workflow,
-    create_local_worker_workflow, create_worker_removal_workflow, CreateLocalWorkerStep,
+    create_local_worker_workflow, create_worker_removal_workflow, create_worker_update_workflow,
-    DetectConnectionModeStep, DiscoverDPInfoStep, DiscoverMetadataStep, DpInfo,
+    CreateLocalWorkerStep, DetectConnectionModeStep, DiscoverDPInfoStep, DiscoverMetadataStep,
-    FindWorkersToRemoveStep, RemoveFromPolicyRegistryStep, RemoveFromWorkerRegistryStep,
+    DpInfo, FindWorkerToUpdateStep, FindWorkersToRemoveStep, RemoveFromPolicyRegistryStep,
-    UpdateRemainingPoliciesStep, WorkerRemovalRequest,
+    RemoveFromWorkerRegistryStep, UpdatePoliciesForWorkerStep, UpdateRemainingPoliciesStep,
+    UpdateWorkerPropertiesStep, WorkerRemovalRequest,
 };
 pub use shared::{ActivateWorkersStep, RegisterWorkersStep, UpdatePoliciesStep, WorkerList};

sgl-model-gateway/src/protocols/worker_spec.rs

@@ -210,7 +210,7 @@ pub struct WorkerTypeStats {
 }
 
 /// Worker update request
-#[derive(Debug, Clone, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct WorkerUpdateRequest {
     /// Update priority
     #[serde(skip_serializing_if = "Option::is_none")]
@@ -223,6 +223,26 @@ pub struct WorkerUpdateRequest {
     /// Update labels
     #[serde(skip_serializing_if = "Option::is_none")]
     pub labels: Option<HashMap<String, String>>,
+
+    /// Update API key (for key rotation)
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub api_key: Option<String>,
+
+    /// Update health check timeout in seconds
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub health_check_timeout_secs: Option<u64>,
+
+    /// Update health check interval in seconds
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub health_check_interval_secs: Option<u64>,
+
+    /// Update health success threshold
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub health_success_threshold: Option<u32>,
+
+    /// Update health failure threshold
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub health_failure_threshold: Option<u32>,
 }
 
 /// Generic API response

sgl-model-gateway/src/server.rs

@@ -27,6 +27,7 @@ use crate::{
             create_external_worker_registration_workflow, create_mcp_registration_workflow,
             create_wasm_module_registration_workflow, create_wasm_module_removal_workflow,
             create_worker_registration_workflow, create_worker_removal_workflow,
+            create_worker_update_workflow,
         },
         worker_to_info, Job, JobQueue, JobQueueConfig, WorkerManager, WorkerType,
     },
@@ -45,7 +46,7 @@ use crate::{
         rerank::{RerankRequest, V1RerankReqInput},
         responses::{ResponsesGetParams, ResponsesRequest},
         validated::ValidatedJson,
-        worker_spec::{WorkerConfigRequest, WorkerErrorResponse, WorkerInfo},
+        worker_spec::{WorkerConfigRequest, WorkerErrorResponse, WorkerInfo, WorkerUpdateRequest},
     },
     routers::{conversations, router_manager::RouterManager, RouterTrait},
     service_discovery::{start_service_discovery, ServiceDiscoveryConfig},
@@ -588,6 +589,41 @@ 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>,
+    Json(update): Json<WorkerUpdateRequest>,
+) -> Response {
+    let worker_id = url.clone();
+    let job = Job::UpdateWorker {
+        url,
+        update: Box::new(update),
+    };
+
+    let job_queue = state
+        .context
+        .worker_job_queue
+        .get()
+        .expect("JobQueue not initialized");
+    match job_queue.submit(job).await {
+        Ok(_) => {
+            let response = json!({
+                "status": "accepted",
+                "worker_id": worker_id,
+                "message": "Worker update queued for background processing"
+            });
+            (StatusCode::ACCEPTED, Json(response)).into_response()
+        }
+        Err(error) => {
+            let error_response = WorkerErrorResponse {
+                error,
+                code: "INTERNAL_SERVER_ERROR".to_string(),
+            };
+            (StatusCode::INTERNAL_SERVER_ERROR, Json(error_response)).into_response()
+        }
+    }
+}
+
 pub struct ServerConfig {
     pub host: String,
     pub port: u16,
@@ -679,8 +715,10 @@ pub fn build_app(
     let worker_routes = Router::new()
         .route("/workers", post(create_worker))
         .route("/workers", get(list_workers_rest))
-        .route("/workers/{url}", get(get_worker))
+        .route(
-        .route("/workers/{url}", delete(delete_worker))
+            "/workers/{url}",
+            get(get_worker).put(update_worker).delete(delete_worker),
+        )
         .route_layer(axum::middleware::from_fn_with_state(
             auth_config.clone(),
             middleware::auth_middleware,
@@ -779,6 +817,9 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box<dyn std::error::Err
     engine
         .register_workflow(create_worker_removal_workflow())
         .expect("worker_removal workflow should be valid");
+    engine
+        .register_workflow(create_worker_update_workflow())
+        .expect("worker_update workflow should be valid");
     engine
         .register_workflow(create_mcp_registration_workflow())
         .expect("mcp_registration workflow should be valid");