Skip to content

TAS: nothing enforces the pod-index-offset contract on MPIJob #14475

Description

@thc1006

What happened:

kueue.x-k8s.io/pod-index-offset carries a documented contract. From site/data/labels_and_annotations.yaml:

This annotation is set by the Kueue webhook and used by the TAS topology ungater to correctly map Pod indexes to topology assignments.

note: This annotation is not added when kueue.x-k8s.io/podset-group-name is specified, as offset management is delegated to the PodSet Group mechanism in that case.

Nothing enforces either half of that on MPIJob.

The webhook only writes it, and only on CREATE. The mutating webhook is registered verbs=create (pkg/controller/jobs/mpijob/mpijob_webhook.go), while the validating webhook takes verbs=create;update. validateTopologyRequest never looks at the annotation, so no path checks the value or the conditions under which it may appear.

The ungater adds it to whatever the group already computed. In pkg/controller/tas/topology_ungater.go, rankOffsets[psa.Name] is filled from the PodSet group layout, and the annotation is then applied on top:

rankOffsets[psa.Name] += int32(offset)
maxRank[psa.Name] += int32(offset)

A non-integer value is a hard stop rather than a fallback, which #8677 made deliberate:

return reconcile.Result{}, errors.Join(err, errParseOffsetAnnotation)

Three shapes follow. The first two are admission-reachable and I checked them against the webhook directly; a podset-group-name needs a topology annotation beside it to pass validation, which a real TAS job carries anyway.

  1. A grouped MPIJob with pod-index-offset: "invalid" is accepted. The defaulting returns without touching it because the job is grouped, and validation does not look at it. The ungater then fails to parse it on every reconcile and the Pods stay gated until someone edits the object.
  2. A grouped MPIJob whose Worker carries a numeric offset is accepted, and the ungater applies both the group-derived offset and the annotation.
  3. Because the mutating webhook does not run on UPDATE, the "1" it wrote on CREATE can be removed or changed afterwards. Validation does not object, and Worker Pods created after that are read as starting at index 0 again, which is the rank-ordering failure TAS: TopologyUngater can not recognize rank-based ordering for MPIJob with runLauncherAsWorker #8471 describes and TAS: Fix a bug that MPIJob with runLauncherAsWorker Pod indexes are not correctly evaluated. #8618 added the offset for.

What you expected to happen:

  • An ungrouped MPIJob with runLauncherAsWorker: true and both replica specs present carries pod-index-offset: "1" on the Worker.
  • The annotation is absent when podset-group-name is set, as the note says.
  • Both hold on UPDATE as well as CREATE, either by validating them or by letting an update-capable mutator reconcile them.

How to reproduce it (as minimally and precisely as possible):

apiVersion: kubeflow.org/v2beta1
kind: MPIJob
metadata:
  name: grouped
  labels:
    kueue.x-k8s.io/queue-name: queue
spec:
  runLauncherAsWorker: true
  mpiReplicaSpecs:
    Launcher:
      template:
        metadata:
          annotations:
            kueue.x-k8s.io/podset-preferred-topology: rack
            kueue.x-k8s.io/podset-group-name: group
            kueue.x-k8s.io/pod-index-offset: invalid
    Worker:
      template:
        metadata:
          annotations:
            kueue.x-k8s.io/podset-preferred-topology: rack
            kueue.x-k8s.io/podset-group-name: group

Passing this through MpiJobWebhook.ValidateCreate with TopologyAwareScheduling enabled returns no error. Replacing invalid with "5" on the Worker is likewise accepted.

Anything else we need to know?:

The same annotation is written by the LeaderWorkerSet webhook (pkg/controller/jobs/leaderworkerset/leaderworkerset_webhook.go), and LWS has its own history with rank ordering in #8661. Whatever shape the enforcement takes is probably worth deciding for both rather than as an MPIJob-local check.

Not blocking #14421, which fixes a different defect in the same branch: that one decided on the size of mpiReplicaSpecs and dereferenced an absent Launcher. This is about the annotation's contract once the branch has run, or once someone edits the object afterwards.

Environment:

  • Kueue version (use git describe --tags --dirty --always): observed on main

Metadata

Metadata

Assignees

No one assigned

    Labels

    area/tasTopology-Aware Schedulingkind/bugCategorizes issue or PR as related to a bug.

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions