feat: support heterogeneous prefill groups in disaggregated ArksApplication#81
Draft
LikiosSedo wants to merge 2 commits into
Draft
feat: support heterogeneous prefill groups in disaggregated ArksApplication#81LikiosSedo wants to merge 2 commits into
LikiosSedo wants to merge 2 commits into
Conversation
Different request shapes want different engine configurations (e.g. one prefill instance tuned for fast TTFT with tp8+dp-attention, another with tp16 for long contexts) behind a single router. - Add spec.prefillGroups[]: named prefill variants, each rendered as its own RBG role (prefill-<name>) and LeaderWorkerSet with independent runtimeCommonArgs, resources, and replicas. Mutually exclusive with spec.prefill (CEL XOR + controller validation); the legacy single prefill keeps its role name so upgrades do not recreate workloads. - Per-group model override (group.model): a group can load its own ArksModel, e.g. a pre-sharded weight copy living on an existing PVC (ArksModel without spec.source). The controller now resolves and waits for every referenced model. - Pod labels: all groups share arks.ai/role=prefill so the router's service discovery selector keeps matching the whole prefill pool; each group additionally carries arks.ai/worker-group=<name> for group-aware routing policies. - Status: per-group counts in status.prefillGroups[], aggregate sums in status.prefill. - coordinationPolicy is rejected together with prefillGroups for now (cross-group coordination semantics are not defined yet). - Add unit tests, a 2P1D sample, and a usage-guide section.
The CRD already enforces uniqueness via listMapKey=name at admission; re-check in validate() so programmatic paths that bypass API validation cannot generate duplicate RBG roles. Also document the derived-name length budget for app + group names.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Support heterogeneous prefill groups in disaggregated
ArksApplication: a single PD-disaggregated deployment can now run several prefill variants with different engine configurations behind one router.Motivating use case (from the SGLang side): a 2P1D setup where requests of different lengths want different engine configs — e.g. one prefill instance tuned for fast TTFT (
tp8+ dp-attention), another for long contexts (tp16) — all sharing one decode pool and one router.API
New
spec.prefillGroups[](mutually exclusive withspec.prefill,mode=disaggregatedonly):tp/ep/etc. are expressed via the group's ownruntimeCommonArgs, no new engine-specific fields.nameis a DNS-1123 label; the list islistType=map/listMapKey=name(admission-enforced uniqueness).group.modeloverridesspec.model, e.g. a pre-sharded weight copy living on an existing PVC.prefill/prefillGroupsin disaggregated mode;prefillGroupsforbidden in unified mode;coordinationPolicyrejected together withprefillGroupsfor now.Controller
prefill-<name>+ LeaderWorkerSet, with independent command, model, resources, and replicas.spec.prefillkeeps the role nameprefill(and thus its LWS name), so existing deployments are not recreated on upgrade. Pinned by a regression test.arks.ai/role=prefillso the router's service-discovery selector keeps matching the whole prefill pool; each group additionally carriesarks.ai/worker-group=<name>for future group-aware routing. Thearks.ai/modellabel always carriesspec.modeleven when a group overrides its model.--model-pathand PVC volume; the controller resolves and waits for every referencedArksModel.status.prefillGroups[], aggregate sums instatus.prefill.validate()(defense-in-depth for programmatic paths that bypass API validation).Testing
Unit tests (
arksapplication_prefillgroups_test.go): 4-role render order, per-group LWS sizes, shared/per-group labels, per-group command + model path + PVC, legacy single-prefill role-name preservation, CEL XOR.End-to-end on a live cluster (qwen-0.5b, against RBG v0.6.0):
prefillGroups-XOR accepted by the API server.ArksApplication→ RBGS → RBG → 3 LWS rendered with rolesrouter / prefill-fast / prefill-long / decode.prefill-fast→.../qwen-0.5b-070with--mem-fraction-static 0.6;prefill-long→.../qwen-0.5b-074with--context-length 8192 --mem-fraction-static 0.5(no arg cross-leakage);decode→.../qwen-0.5b.role=prefill+worker-group=fast|long; decode has noworker-group) and schedule to the per-group nodes; prefill/decode engines reachRunning.Follow-ups (not in this PR)
worker-grouplabel the router can select on. First-stage e2e uses the default flat-pool selector.coordinationPolicy+prefillGroupsinteraction left undefined until there is a concrete need.