Page MenuHomePhabricator

Support GPU Partition specification in deployments
Open, Needs TriagePublic

Description

As an ML engineer,

I want to be able to explicitly define the GPU partition in the deployment charts manifest so that:

  • scheduling aligns with the model’s resource requirements
  • larger GPU partitions are not allocated unnecessarily & GPU resources remain available for other workloads

Event Timeline

Current state.

There are two patterns in use right now:

  1. Default (untainted) GPU pool - used by e.g. edit-check, revertrisk (wikidata) and revise-tone-task-generator. These only set amd.com/gpu in the container resources and let the scheduler place the pod on any untainted node with a free GPU.
  2. MI300x pool - used by services in llm and experimental namespace. These set amd.com/gpu, a toleration for the dedicated=mi300x-experiments taint (NoSchedule + NoExecute), and a nodeAffinity to specific hostnames (e.g. kubernetes.io/hostname In [ml-serve1012]).

On the cluster side, we have amdgpu-node-labeller deployed in helmfile.d/admin_ng/values/ml-serve.yaml. It adds amd.com/gpu.* labels on the nodes, and for MI300X it also addsmemory partition mode (NPS1 or NPS4). The taints and base node labels live in operations/puppet.


Our Problem

GPU placement today works based on toleration + a hardcoded hostname allowlist, which has some downsides:

  • Using hostname might be flaky when we re-image, repurpose or add a node.
  • It doesn't fully express intent like "give me a small partition of MI300X node", you can only say "place me on this exact machine"
  • it doesn't prevent failure mode of a model requiring small amount of VRAM accidentally landing on full 192GB MI300X GPU (e.g. it tolerates dedicated=mi300x-experiments, but hostname affinity is missing/wrong/unschedulable)

We want an easy-to-use and reliable mechanism to specify the GPU type/partition we want our ISVC to run on. I'll share two proposals for discussion in a comment below.

First, let me give some context on options I'll share below - taints/tolerations and nodeAffinity do opposite jobs: a taint repels everything that does not explicitly tolerate it, while nodeAffinity only attracts a pod when it is present, but does not do any "blocking" during scheduling. The practical consequence is that if we want hard guarantees about which GPUs we land on, it has to come from taints as affinities cannot prevent accidental schedules.

Option 1 - taint every GPU pool

Every GPU node type would get its own taint e.g. dedicated=mi300x-spx_nps1, dedicated=mi300x-cpx_nps4, dedicated=mi210. To run on a given GPU type you tolerate exactly that taint, no nodeAffinity is involved.

  • Upside - strong guarantee. Pod can only ever land on the GPU type it explicitly tolerates - impossible to accidentally grab a full 192GB MI300x.
  • Upside - single dimension to reason about. We only deal with taints/tolerations, we don't care about affinities.
  • Downside - less flexible when we are short on GPUs. If the tolerated pool is full, the pod will stay pending (unless we add multiple taints)
  • Note - this needs SREs to taint all GPU pools, including the older nodes. We would need to add tolerations to existing non-MI300X services.

Example of how this would look like in deployment-charts code:

yaml
  container:
    resources:
      limits:
        amd.com/gpu: "1"
  config:
    tolerations:
      - { key: dedicated, operator: Equal, value: mi300x-cpx_nps4, effect: NoSchedule }
      - { key: dedicated, operator: Equal, value: mi300x-cpx_nps4, effect: NoExecute }

Option 2 - taint only the biggest GPUs

Only the largest/expensive nodes are tainted - 192GB MI300X today + MI350X in the future. All other GPU nodes stay untainted and we use nodeAffinity to express a preference for the GPU type we want.

  • Upside - flexibility when short on GPUs. If using preferred affinity, we can fall back to other GPU types under shortage.
  • Downside - no hard guarantee of GPU type. We only protect against accidentally taking the biggest nodes, but not against landing on a non-preferred one.

Example of how this would look like in deployment-charts code:

yaml
  container:
    resources:
      limits:
        amd.com/gpu: "1"
  config:
    affinity:
      nodeAffinity:
        requiredDuringSchedulingIgnoredDuringExecution:
          nodeSelectorTerms:
            - matchExpressions:
                - { key: amd.com/gpu.compute-memory-partition, operator: In, values: ["cpx_nps4"] }

In both cases the GPU count stays exactly where it is today and we keep the raw option to do one-off node pinning to specific hosts if required.
The question is - are we comfortable tainting all GPU pools (Option 1) to get a hard guarantee of GPU type we are running on, or do we prefer to have smaller GPUs untainted for flexibility (Option 2)? It's mostly a tradeoff between guarantee of knowing exactly which GPU we run on vs scheduling flexibility when GPUs are scarce.
I am leaning slightly towards option 1, because I like guarantees and single dimension of reasoning and I feel that we are currently not scarce on GPUs, however I see the upsides of 2nd option. Would love to hear more opinions!

A few GPU facts for the decision (gathered from deployment-charts, puppet, doc, gerrit — please correct me if anything's off):

Current schedulable GPU nodes: eqiad MI210 = ml-serve1009-11, MI300 = ml-serve1012–1015; codfw MI210 = ml-serve2009-11.

Usage: default/untainted pool (edit-check, revertrisk-wikidata, revise-tone) = 4/6 MI210 per DC; MI300 (llm + experimental) = ~11 of 200 partitions (8 × 192G SPX + 192 × 24G CPX, per the doc)

Partitioning (set up in T420507): puppet declaratively manages only the memory partition (NPS) — ml-serve1012 = NPS1, ml-serve1013-15 = NPS4. The compute mode (SPX/CPX) is not set via puppet's partition script. And there are no GPU taints in effect today, so the tolerations still in the llm/experimental manifests are no-ops and MI300 placement rests purely on the hostname allowlist.

Takeaway: I lean Option 1. We're not GPU-scarce (default pool ≈ 4/6 MI210 per DC, MI300 lightly used, MI350 in the future), so Option 2's fallback-when-short upside is marginal — while Option 1 gives a hard, hostname-free type guarantee. The main thing to work out with SRE is how best to apply the per-pool taints to the existing nodes — there was an earlier attempt that got reverted (3d7935d11c), so I'd value their input on the rollout.

I'm leaning towards option 1 as well. With the current status and the fact that we are starting to deploy services on these hosts and use partitioning for the first time I think we will be more on top of things by being explicit. It might introduce a maintenance burden in the future but we could solve that once it occurs instead of optimizing for it prematurely.
The risk I see with Option 2 is scheduling a new revision in a different partition (which has not been previously tested) which might result in unforeseen errors (with OOM being the most obvious one). This unpredictable behavior worries me a bit when talking about production services.

Thank you for your inputs @achou @isarantopoulos! It seems we are all agreeing on Option 1, which means tainting all current GPU nodes.

I'd propose the following taint scheme - every GPU pool gets dedicated=<vendor>-<model>-<vram>, with both NoSchedule and NoExecute. It should make the taints easy to understand (by hiding spx/cpx/nps details) and extends cleanly to future GPUs like MI350X.

Taint valueNodesNode Hardware
dedicated=amd-mi100-32gbml-staging2001AMD Instinct MI100
dedicated=amd-mi210-64gbml-serve1009-1011, ml-serve2009-2011, ml-staging2003AMD Instinct MI210
dedicated=amd-mi300x-192gbml-serve1012MI300X SPX/NPS1, 8 x 192GB
dedicated=amd-mi300x-24gbml-serve1013-1015MI300X CPX/NPS4, 64 x 24GB

I used this as basis for node/hardware information: https://wikitech.wikimedia.org/wiki/Machine_Learning/AMD_GPU

In the deployment-charts, our services would need to tolerate those taints, which would look like this:

config:
  tolerations:
    - key: dedicated
      operator: Equal
      value: amd-mi300x-24gb

Next steps would be:

  1. Make sure we agree on the taints proposed above
  2. Add tolerations to all GPU services in deployment-charts. Those won't change anything when nodes are untainted so those are safe to deploy.
  3. SRE applies taints pool by pool. We verify that pods keep running and schedule properly after tainting the nodes.
  4. Cleanup of nodeAffinity blocks and the old mi300x-experiments tolerations.