Skip to content

feat: support heterogeneous prefill groups in disaggregated ArksApplication#81

Draft
LikiosSedo wants to merge 2 commits into
scitix:mainfrom
LikiosSedo:feat/prefill-groups
Draft

feat: support heterogeneous prefill groups in disaggregated ArksApplication#81
LikiosSedo wants to merge 2 commits into
scitix:mainfrom
LikiosSedo:feat/prefill-groups

Conversation

@LikiosSedo

Copy link
Copy Markdown
Contributor

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 with spec.prefill, mode=disaggregated only):

spec:
  mode: disaggregated
  prefillGroups:
    - name: fast
      model: { name: glm5-fp8-sharded-tp8dp8 }   # optional per-group model override
      size: 1
      runtimeCommonArgs: ["--tp", "8", "--dp", "8", "--enable-dp-attention"]
    - name: long
      size: 2
      runtimeCommonArgs: ["--tp", "16", "--context-length", "202752"]
  decode: { ... }
  router: { replicas: 1 }
  • Each group reuses the existing workload schema (inline) — tp/ep/etc. are expressed via the group's own runtimeCommonArgs, no new engine-specific fields.
  • name is a DNS-1123 label; the list is listType=map/listMapKey=name (admission-enforced uniqueness).
  • Optional group.model overrides spec.model, e.g. a pre-sharded weight copy living on an existing PVC.
  • CEL: exactly one of prefill / prefillGroups in disaggregated mode; prefillGroups forbidden in unified mode; coordinationPolicy rejected together with prefillGroups for now.

Controller

  • Each group renders as its own RBG role prefill-<name> + LeaderWorkerSet, with independent command, model, resources, and replicas.
  • Backward compatible: a single spec.prefill keeps the role name prefill (and thus its LWS name), so existing deployments are not recreated on upgrade. Pinned by a regression test.
  • Pod labels: all prefill 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 future group-aware routing. The arks.ai/model label always carries spec.model even when a group overrides its model.
  • Per-group model resolution → distinct --model-path and PVC volume; the controller resolves and waits for every referenced ArksModel.
  • Status: per-group counts in status.prefillGroups[], aggregate sums in status.prefill.
  • Duplicate group names are also rejected in 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):

  • CEL prefillGroups-XOR accepted by the API server.
  • ArksApplication → RBGS → RBG → 3 LWS rendered with roles router / prefill-fast / prefill-long / decode.
  • Per-group rendering verified on the generated objects: prefill-fast.../qwen-0.5b-070 with --mem-fraction-static 0.6; prefill-long.../qwen-0.5b-074 with --context-length 8192 --mem-fraction-static 0.5 (no arg cross-leakage); decode.../qwen-0.5b.
  • Live pods carry the correct labels (role=prefill + worker-group=fast|long; decode has no worker-group) and schedule to the per-group nodes; prefill/decode engines reach Running.

Follow-ups (not in this PR)

  • Router-side group-aware routing (route long vs short requests to specific prefill groups) is owned by the SGLang router; this PR only wires the per-group worker-group label the router can select on. First-stage e2e uses the default flat-pool selector.
  • coordinationPolicy + prefillGroups interaction left undefined until there is a concrete need.
  • Real-GLM-scale topology validation is a separate milestone.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant