fix(pd): Fix PD cache-aware policy lifecycle

[aurickq]

Description

Problem

PD cache-aware routing uses dedicated prefill and decode policy instances, separate from the per-model policy used by the regular routing path. The regular path initializes and refreshes its cache-aware policy during worker registration, update, and removal, but the PD cache-aware policy instances were not refreshed through the same lifecycle hooks.

This shows up for PD configurations that use cache_aware routing for prefill and/or decode workers, especially when there is more than one worker in a PD role and requests have overlapping prefixes. The request path can reach a cache-aware policy whose string tree was never initialized, producing logs like:

Warning: No string tree found for model '<model-id>', using random worker selection

At that point, prefix-aware routing is bypassed for the affected PD policy and worker selection falls back to random selection. The symptom is that repeated or prefix-overlapping requests do not consistently preserve locality even though cache_aware routing is configured.

Solution

Mirror the existing regular cache-aware lifecycle behavior for PD cache-aware policy instances:

• initialize PD cache-aware policies when workers are registered
• refresh PD cache-aware policies when workers are updated
• remove stale worker URLs from PD cache-aware policies when workers are removed
• refresh PD cache-aware policies after removal using the remaining worker set

This keeps the PD path aligned with the existing non-PD cache-aware path.

Changes

• Added remove_worker_from_pd_cache_aware(...) next to the existing regular cache-aware removal helper.
• Call init_pd_cache_aware_policies(...) from worker registration, update, and post-removal policy refresh paths.
• Call the new PD cache-aware removal helper from worker removal.
• Added a focused unit test covering PD cache-aware policy initialization and repeated-prefix selection.

Test Plan

Repro before this change:

1. Configure PD routing with cache_aware selected for prefill and/or decode routing and more than one worker in at least one PD role.
2. Register workers through the normal worker lifecycle.
3. Submit repeated requests with overlapping prefixes.
4. Observe that the affected PD cache-aware policy can log a missing string tree and fall back to random worker selection:

Warning: No string tree found for model '<model-id>', using random worker selection

Expected behavior after this change:

1. Worker registration initializes the PD cache-aware prefill/decode policy trees.
2. Worker update/removal refreshes those PD policy trees using the current worker registry state.
3. Repeated requests with overlapping prefixes can use the initialized PD cache-aware policy instead of falling back due to missing prefix trees.

Commands run:

cargo fmt --all
cargo test -p smg test_pd_cache_aware_policy_initialization
cargo check -p smg

cargo fmt --all completed successfully, while printing existing rustfmt warnings about unstable import-formatting options in the repo config.

Checklist

• cargo +nightly fmt passes
• cargo clippy --all-targets --all-features -- -D warnings passes
• (Optional) Documentation updated
• (Optional) Please join us on Slack #sig-smg to discuss, review, and merge PRs

Summary by CodeRabbit

• Improvements

  • Enhanced worker management for PD-mode cache-aware policies, including deterministic worker selection and safer removal of workers.
  • Ensures policies reinitialize and stay synchronized after worker additions or removals.

• Tests

  • Added tests for cache-aware policy initialization and worker lifecycle operations, including deterministic selection and removal scenarios.

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds PolicyRegistry::remove_worker_from_pd_cache_aware to remove a worker from PD-mode cache-aware prefill/decode policies, and updates workflow steps to reinitialize PD cache-aware policies using current prefill/decode worker lists after worker lifecycle events.

Changes

PD Cache-Aware Policy Management

Layer / File(s)	Summary
PD cache-aware policy removal and tests model_gateway/src/policies/registry.rs	PolicyRegistry::remove_worker_from_pd_cache_aware removes a worker URL from both PD prefill and decode cache-aware policy instances. Tests add helpers to build workers and a CacheAwarePolicy, plus test_pd_cache_aware_policy_initialization validating init, deterministic select_worker, and worker removal.
Workflow re-initialization hooks model_gateway/src/workflow/steps/local/remove_from_policy_registry.rs, model_gateway/src/workflow/steps/local/update_policies_for_worker.rs, model_gateway/src/workflow/steps/local/update_remaining_policies.rs, model_gateway/src/workflow/steps/shared/update_policies.rs	Workflow steps now call remove_worker_from_pd_cache_aware when removing a worker and call init_pd_cache_aware_policies(&prefill_workers, &decode_workers) to refresh PD cache-aware policies after worker additions/removals or final updates.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• lightseekorg/smg#1113: Both PRs touch retrieval and use of prefill/decode worker lists from worker_registry.

Suggested labels

tests

Suggested reviewers

• CatherineSue
• slin1237

Poem

🐰 I nibble code and hop with glee,
PD caches learn which worker to be,
Init, select, then remove with care,
Prefill and decode now balanced fair,
Hooray — the registry hums with me!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 61.54% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The pull request title accurately summarizes the main change: fixing the PD cache-aware policy lifecycle by adding initialization, refresh, and removal logic across worker registration, update, and removal workflows.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[mergify]

Hi @aurickq, the DCO sign-off check has failed. All commits must include a Signed-off-by line.

To fix existing commits:

# Sign off the last N commits (replace N with the number of unsigned commits)
git rebase HEAD~N --signoff
git push --force-with-lease

To sign off future commits automatically:

• Use git commit -s every time, or
• VSCode: enable Git: Always Sign Off in Settings
• PyCharm: enable Sign-off commit in the Commit tool window

[gemini-code-assist]

Code Review

This pull request introduces functionality to manage worker removal and initialization for Prefill-Decode (PD) cache-aware policies. Specifically, it adds a remove_worker_from_pd_cache_aware method to the PolicyRegistry and ensures that init_pd_cache_aware_policies is called during worker registration and removal workflows to maintain policy consistency. Unit tests have been included to verify the initialization and removal logic. The reviewer noted that integrating these calls into shared workflow steps correctly handles the global nature of PD policies and prevents code duplication.

[gemini-code-assist]

Calling init_pd_cache_aware_policies here ensures that PD policies are refreshed whenever a new worker is registered. Since this step is shared across different registration paths, it correctly handles the global nature of PD policies and avoids code duplication, ensuring consistent behavior for all worker types.

References

1. If an optimization or logic is applicable to multiple code paths, extract it into a shared helper function to ensure consistency and avoid code duplication.

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@model_gateway/src/policies/registry.rs`:
- Around line 585-587: The test currently calls
registry.remove_worker_from_pd_cache_aware("http://prefill-1:8000") and
registry.remove_worker_from_pd_cache_aware("http://decode-1:8000") but never
asserts the observable routing/resulting state; update the PD cache-aware test
to validate post-removal behavior by invoking the routing/resolution method used
elsewhere (e.g., the registry's route selection or resolve_worker calls) and
assert that lookups no longer return the removed endpoints and that requests
route to expected remaining workers (or return an error/none when none remain);
use the same registry instance and the unique helper methods
remove_worker_from_pd_cache_aware and the registry's resolution method names to
locate where to add these assertions.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: 0cfb4078-047c-4803-958b-7e047ecc26ec

📥 Commits

Reviewing files that changed from the base of the PR and between 12413dc and 4134c6d.

📒 Files selected for processing (5)

• model_gateway/src/policies/registry.rs
• model_gateway/src/workflow/steps/local/remove_from_policy_registry.rs
• model_gateway/src/workflow/steps/local/update_policies_for_worker.rs
• model_gateway/src/workflow/steps/local/update_remaining_policies.rs
• model_gateway/src/workflow/steps/shared/update_policies.rs

[slin1237]

thanks!
lgtm, please fix DCO