Repository navigation
feat(runtime): add hierarchical Model/WorkerSet architecture for multi-namespace support - #6054
Conversation
github-actions
Bot
added
backend::vllm
WalkthroughThis PR centralizes worker namespace resolution (get_worker_namespace), introduces NamespaceFilter, adds Model and WorkerSet abstractions, refactors ModelManager for per-model/per-namespace orchestration, plumbs namespace_prefix through CLI and bindings, and updates discovery, entrypoints, and prefill routing to be namespace-aware. Changes
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Tip Issue Planner is now in beta. Read the docs and try it out! Share your feedback on Discord. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
lib/llm/src/discovery/model_manager.rs (1)
579-621:⚠️ Potential issue | 🔴 CriticalConcurrent
register_prefill_routercalls can lose a channel—including unguarded paths in C bindings.The
remove→ inspect →insertpattern onDashMapis not atomic. Two concurrent threads callingregister_prefill_router("llama", "ns1")can both seeNonefromremove, both create channels, and the secondinsertoverwrites the first sender, droppingtx1and causingrx1to resolve with aRecvError.The Rust watcher code is protected:
registering_worker_sets(aDashSet<String>keyed onmodel_namespace_key(model_name, namespace)) in watcher.rs guardsdo_worker_set_registration, which contains the only production call toregister_prefill_routerat line 424. This serializes concurrent registrations for the same model+namespace.However, lib/bindings/c/src/lib.rs:1421 calls
register_prefill_routerdirectly without the DashSet guard, creating an unprotected concurrent path. If the C bindings are used to register prefill routers concurrently for the same model+namespace, the race condition manifests.Either move the guard into
register_prefill_routeritself, or document that this function must only be called from within aregistering_worker_setsguard and audit all call sites (including C bindings).
🤖 Fix all issues with AI agents
In `@components/src/dynamo/common/utils/namespace.py`:
- Around line 7-18: get_namespace may return values containing "." (from
DYN_NAMESPACE or DYN_NAMESPACE_WORKER_SUFFIX) which breaks downstream parsing of
dyn://{namespace}.component.endpoint; update get_namespace to validate and
reject or sanitize dots: check os.environ.get("DYN_NAMESPACE") and
os.environ.get("DYN_NAMESPACE_WORKER_SUFFIX") and either raise a clear
ValueError/log an error if either contains a "." or replace "." with "-" (choose
one policy consistently), and ensure the function get_namespace documents this
behavior so callers (and CI) won't produce malformed dyn endpoints.
In `@lib/llm/src/discovery/model_manager.rs`:
- Around line 572-573: The current model_namespace_key(model_name: &str,
namespace: &str) builds keys with ":" which can collide when model names include
":"; change the key format to a collision-safe scheme (e.g., use a fixed
multi-char delimiter like "::ns::" or a non-printable separator like '\0', or
better: encode each component with percent-encoding/base64/hex) and update any
code that parses or compares these keys to use the new scheme; specifically
modify model_namespace_key to produce the safe-encoded/delimiter-separated
string and ensure any consumers that split or match keys use the corresponding
decode/split logic so existing lookups remain correct.
In `@lib/llm/src/entrypoint/input/common.rs`:
- Around line 83-86: The code spawns a watcher task that currently hardcodes
NamespaceFilter::Global (inner_watch_obj/watcher_task/watch/discovery_stream),
which ignores model flags; instead build the NamespaceFilter from the model
configuration (use the same logic as grpc.rs/http.rs that derives from
model.namespace and model.namespace_prefix), create the filter variable before
spawning the tokio::spawn, and then pass or move that filter into the async
closure so the watch call uses the constructed filter rather than
NamespaceFilter::Global; ensure the filter is clonable or moved into the closure
to avoid lifetime issues.
In `@lib/llm/src/namespace.rs`:
- Line 4: The change of GLOBAL_NAMESPACE to "default" breaks tests and runtime
expectations that use "dynamo" as the global namespace; restore GLOBAL_NAMESPACE
to "dynamo" and update the global-check logic so is_global_namespace(...) treats
both "dynamo" and the empty string (and optionally "default" if backward
compatibility desired) as global; specifically modify the constant
GLOBAL_NAMESPACE back to "dynamo" and adjust the is_global_namespace function to
return true for namespace == GLOBAL_NAMESPACE || namespace.is_empty() (and add
"default" to the allowed set only if you want to accept both).
🧹 Nitpick comments (10)
components/src/dynamo/common/utils/__init__.py (1)
10-25: Docstring doesn't list the newnamespacesubmodule.The module docstring (lines 10-15) enumerates submodules but omits the newly added
namespacemodule, while__all__and imports include it.Suggested fix
Submodules: - endpoint_types: Endpoint type parsing utilities + - namespace: Namespace resolution utilities - otel_tracing: OpenTelemetry tracing header utilities - paths: Workspace directory detection and path utilities - prometheus: Prometheus metrics collection and logging utilitiescomponents/src/dynamo/vllm/args.py (1)
366-368: Inconsistent import location compared to other backends.
get_namespaceis imported at module level inmocker/args.py,sglang/args.py, andtrtllm_utils.py, but here it's imported insideparse_args(). Consider moving it to the top-level imports for consistency.Suggested refactor
Add to the top-level imports (e.g., near line 19):
from dynamo.common.utils.namespace import get_namespaceThen at line 368, simply:
- from dynamo.common.utils.namespace import get_namespace - - config.namespace = get_namespace() + config.namespace = get_namespace()lib/llm/tests/http_metrics.rs (1)
506-564: Commented-out cleanup block — consider removing or converting to a TODO.This entire block is wrapped in
/* ... */and has no effect. If the intent is to implement this later, a brief// TODO:comment would be clearer than leaving ~60 lines of dead code.lib/llm/src/namespace.rs (1)
48-53:Prefixmatching may be broader than intended.
starts_with("ns")matches"ns-abc123"but also"nsomething". The doc comment on line 18 implies delimiter-separated variants ("ns","ns-abc123"). If the intent is to match only the exact prefix or prefix followed by a delimiter (e.g.,-), consider tightening the match:NamespaceFilter::Prefix(prefix) => namespace == prefix || namespace.starts_with(&format!("{}-", prefix)),If the broader
starts_withbehavior is intentional, this is fine as-is — just noting the subtlety.lib/llm/src/http/service/openai.rs (1)
1010-1011: Unnecessary.clone()onparsing_options.
parsing_optionsis not used after being passed tofrom_annotated_streamin eitherchat_completionsorresponses, so the.clone()is redundant — you can pass it by value directly.🔧 Remove unnecessary clones
- NvCreateChatCompletionResponse::from_annotated_stream(stream, parsing_options.clone()) + NvCreateChatCompletionResponse::from_annotated_stream(stream, parsing_options)(Same for line 1272.)
Also applies to: 1271-1272
lib/llm/src/discovery/model.rs (1)
185-201:load_threshold_configupdates all monitors but returns from the first — document this intentional asymmetry.When
configisSome, all WorkerSet monitors are updated (correct for consistent thresholds during rollout), but only the first monitor's config is returned. This is fine, but a brief inline comment would clarify the intent for future readers.📝 Suggested comment
pub fn load_threshold_config( &self, config: Option<&LoadThresholdConfig>, ) -> Option<LoadThresholdConfig> { + // When setting: update ALL monitors so all WorkerSets use consistent thresholds. + // When getting: return from any monitor (they should all agree). let mut result = None;lib/llm/src/discovery/model_manager.rs (4)
349-368: TOCTOU betweenhas_chat_engine()check andadd_worker_set()in alladd_*_modelmethods.The check-then-act pattern (
if model_entry.has_chat_engine() → return Errthenadd_worker_set) is not atomic. Two concurrent callers could both pass the check and both add WorkerSets. For in-process models this is likely single-threaded initialization so the practical risk is low, but it's worth noting. The same pattern repeats inadd_completions_model,add_embeddings_model,add_tensor_model,add_images_model, andadd_prefill_model.
117-138: Subtle race window inadd_worker_set: newly created empty model can be removed before WorkerSet is added.
get_or_create_model(line 123) may create an emptyModel. Between that call andmodel.add_worker_set(line 124), a concurrentremove_worker_set→remove_model_if_emptycould find the model empty and remove it from the DashMap. Theadd_worker_setwould then mutate a detachedArc<Model>invisible to lookups.This is a narrow window, and the watcher-level
DashSetguard likely prevents it, but worth documenting or adding an assertion that the model is still in the map after the add.
345-347: Acknowledged: TODO for passing realModelDeploymentCardto in-process WorkerSets.The TODO is well-documented. Just noting that until this is addressed,
ParsingOptions(tool_call_parser, reasoning_parser) will always be defaults for in-process models, which could cause subtle behavioral differences compared to dynamically-discovered models.Would you like me to open an issue to track passing the real
ModelDeploymentCardfrom callers?
484-517:remove_*_modelmethods rely on matching the synthetic namespace format fromadd_*_model.The namespace strings (e.g.,
__local_chat_{}) are duplicated between add and remove methods. If one side changes format without the other, removal silently fails. Consider extracting a shared helper.Proposed refactor
+ fn local_namespace(kind: &str, model: &str) -> String { + format!("__local_{}_{}", kind, model) + } + pub fn add_chat_completions_model( &self, model: &str, card_checksum: &str, engine: OpenAIChatCompletionsStreamingEngine, ) -> Result<(), ModelManagerError> { let model_entry = self.get_or_create_model(model); if model_entry.has_chat_engine() { return Err(ModelManagerError::ModelAlreadyExists(model.to_string())); } - let namespace = format!("__local_chat_{}", model); + let namespace = Self::local_namespace("chat", model); // ... } pub fn remove_chat_completions_model(&self, model: &str) -> Result<(), ModelManagerError> { - let namespace = format!("__local_chat_{}", model); + let namespace = Self::local_namespace("chat", model); self.remove_worker_set(model, &namespace) .map(|_| ()) .ok_or_else(|| ModelManagerError::ModelNotFound(model.to_string())) }
79725ce to
5ed7ea7
Compare
tmonty12
changed the title
|
@coderabbitai full review |
✅ Actions performedFull review triggered. |
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
lib/bindings/c/src/lib.rs (2)
1267-1284:⚠️ Potential issue | 🔴 CriticalBuild failure:
activate_prefill_routerrequires 3 arguments, but only 2 are provided.The method requires
model_name,namespace, andendpointparameters. The namespace parameter must be added—&endpoint_id.namespaceis available in scope.🐛 Proposed fix
- if let Err(e) = model_manager.activate_prefill_router(card.name(), endpoint) + if let Err(e) = model_manager.activate_prefill_router(card.name(), &endpoint_id.namespace, endpoint)
1383-1388:⚠️ Potential issue | 🔴 CriticalFix signature mismatch in
cards_for_modelcall.The call
cards_for_model(model_name, Some(namespace), false)is incorrect. The function signature expects only two parameters:model_name: &strandnamespace_filter: &NamespaceFilter. The current call passes three arguments with mismatched types—the second parameter should be a reference to aNamespaceFilter(notOption<&str>), and the third boolean argument doesn't exist.Proposed fix
+ use dynamo_llm::namespace::NamespaceFilter; + + let namespace_filter = NamespaceFilter::Exact(namespace.to_string()); let cards = watcher - .cards_for_model(model_name, Some(namespace), false) + .cards_for_model(model_name, &namespace_filter) .await .with_context(|| format!("Failed to discover model: {}", model_name))?;lib/llm/tests/http_metrics.rs (1)
356-359:⚠️ Potential issue | 🟠 MajorFix type error:
watchexpectsNamespaceFilter, notOption.The
watchmethod inModelWatcher(lib/llm/src/discovery/watcher.rs:117) takes aNamespaceFilterparameter, notOption. PassingNonehere is a type mismatch. UseNamespaceFilter::Globalinstead, which matches all namespaces (semantically equivalent to no filtering).Fix
+ use dynamo_llm::namespace::NamespaceFilter; + let _watcher_task = tokio::spawn(async move { - model_watcher.watch(discovery_stream, None).await; + model_watcher.watch(discovery_stream, NamespaceFilter::Global).await; });
🤖 Fix all issues with AI agents
In `@lib/llm/src/discovery/model_manager.rs`:
- Around line 343-347: The WorkerSet creation for in-process models currently
uses ModelDeploymentCard::default(), causing parsing_options() to be wrong and
losing tool/reasoning parsers; update the add_*_model functions that build
WorkerSet (the in-process convenience methods) to accept a ModelDeploymentCard
parameter instead of using ModelDeploymentCard::default(), pass that card into
WorkerSet construction so parsing_options() reflects the real config, and update
call-sites in http.rs and grpc.rs to supply the actual ModelDeploymentCard (and
adjust any engine parameter as needed) so parsing behavior matches the deployed
model.
In `@lib/llm/src/discovery/watcher.rs`:
- Around line 692-718: Call site is passing three args to cards_for_model but
the function signature expects (&self, model_name: &str, namespace_filter:
&NamespaceFilter); replace the erroneous call like .cards_for_model(model_name,
Some(namespace), false) with .cards_for_model(model_name,
&NamespaceFilter::Exact(namespace.to_string())), and ensure NamespaceFilter is
in scope (import or fully qualify) so the reference type matches the function
signature.
In `@lib/llm/src/namespace.rs`:
- Around line 25-83: The tests expect "default" to be treated as a global
namespace but is_global_namespace currently only checks GLOBAL_NAMESPACE and
empty string, causing from_namespace_and_prefix to produce Exact("default");
update is_global_namespace to treat "default" as a global alias (e.g., return
true when namespace == GLOBAL_NAMESPACE || namespace.is_empty() || namespace ==
"default") so from_namespace_and_prefix and the test
(test_from_namespace_and_prefix_global) will align and produce
NamespaceFilter::Global for Some("default"); ensure references to
GLOBAL_NAMESPACE and is_global_namespace are used so behavior remains consistent
across from_namespace_and_prefix and other callers.
🧹 Nitpick comments (1)
components/src/dynamo/common/utils/__init__.py (1)
10-26: Docstring is missing entries fornamespaceandruntimesubmodules.The module docstring (lines 10-15) lists
endpoint_types,otel_tracing,paths, andprometheusbut omits the newly addednamespaceand the existingruntimesubmodule, both of which are now exported.Proposed fix
Submodules: - endpoint_types: Endpoint type parsing utilities + - namespace: Worker namespace resolution utilities - otel_tracing: OpenTelemetry tracing header utilities - paths: Workspace directory detection and path utilities - prometheus: Prometheus metrics collection and logging utilities + - runtime: Runtime endpoint parsing utilities """
071545c to
9df3b9d
Compare
Signed-off-by: tmontfort <tmontfort@nvidia.com>
Signed-off-by: tmontfort <tmontfort@nvidia.com>
Signed-off-by: tmontfort <tmontfort@nvidia.com>
Signed-off-by: tmontfort <tmontfort@nvidia.com>
Signed-off-by: tmontfort <tmontfort@nvidia.com>
Signed-off-by: tmontfort <tmontfort@nvidia.com>
Signed-off-by: tmontfort <tmontfort@nvidia.com>
Signed-off-by: tmontfort <tmontfort@nvidia.com>
Signed-off-by: tmontfort <tmontfort@nvidia.com>
Signed-off-by: tmontfort <tmontfort@nvidia.com>
ff4723d to
ddaff5d
Compare
Summary
Model → WorkerSetarchitecture in the frontend'sModelManager, replacing the previous flat engine map designModelholds aDashMap<namespace, Arc<WorkerSet>>where eachWorkerSetowns its own complete pipeline (engines, KV router, prefill router, worker monitor)DYN_NAMESPACE+ optionalDYN_NAMESPACE_WORKER_SUFFIXenvironment variableMotivation
This change enables rolling upgrades with full isolation between old and new workers. During a rolling update:
dynamo-v2) while old workers remain under the original namespace (dynamo)WorkerSetwith a dedicated pipeline — KV router, prefill router, and worker monitor are all isolatedWorkerSetis cleaned up automaticallyWithout this isolation, old and new workers would share a single KV cache router and prefill router, leading to compatibility issues during transitions (e.g., KV cache blocks being routed between incompatible model versions).
Key Changes
Modelstruct (discovery/model.rs): Manages multipleWorkerSets per model, weighted selection, aggregated engine viewsWorkerSetstruct (discovery/worker_set.rs): Owns a complete pipeline with atomic worker countingModelManager(discovery/model_manager.rs): Delegates through theModellayer; handles per-namespace checksum validationModelWatcher(discovery/watcher.rs): Builds per-WorkerSet pipelines, cleans up empty namespaces/models on worker departureDYN_NAMESPACE_WORKER_SUFFIXsupport for worker-side namespace configurationhttp.rs,grpc.rs): Convenience methods that create synthetic single-namespace WorkerSetsTest plan
Modelweighted selection, worker counting, and engine aggregationWorkerSetlifecyclehttp_metrics.rsupdated for new API)🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Improvements