InferencePool's pod selector cannot select multiple ports within a pod needed for data parallel attention

[smarterclayton]

Data parallel attention in vLLM involves multiple processes on each host, typically one per GPU, communicating collectively with the other accelerators on the host and off the host (depending on the size of the wide expert parallelism config, it could be as many as 20 nodes). Each process is expected to be a serving replica and may experience non-uniform traffic, and it is desirable to load balance from IGW based on the discrete load of each.

The natural way to run data parallel attention in python would be a single pod per node using all GPUs, with standard python fork/spawn used to launch child processes. The processes are communicating over both NVLINK, CUDA IPC (and thus need shared memory), and at least today are expected to be started / restarted by the launching process.

While we have discussed a few options in llm-d involving pods per rank, the lowest low friction way to implement this in Kubernetes would be one pod containing multiple processes, each listening on a different port, which requires multiple container ports per pod, and then program the load balancers to see all ports equally (may require multiple services / other fixes).

The current design of the InferencePool pod selector cannot select multiple container ports per pod.

Going forward, I would expect to see

• more multi-process / multi-rank coordination within pods
• potentially distjoint port sets across variants within the pool
• a need to associate a port -> a container -> a GPU for failure detection and two way linking
• health checks per port (which is not possible on pods without multiple containers today)
• port metadata so the gateway can correlate port -> ranks for algorithms

This is likely a blocker to data-parallel attention support for inference gateway, although we could work around it inside the EPP at the cost of causing higher level systems to not have the right data (anyone building on top of the pool, metrics, etc)

[smarterclayton]

For context about why this would be important before v1, changing to support multiple options from singleton is a massive pain to get right and preserve forward compatibility. Long term, removing the singleton in favor of a ports array with be much cheaper at the cost of delaying v1.

[capri-xiyue]

@robscott for the the notice of potential additional work.

[robscott]

Thanks for raising this @smarterclayton! Thinking through what this could mean for upcoming v1.0 release. Here are a few options I can think of:

1. Continue on current path
This allows us to continue with existing release timeline at the cost of being stuck with at least one field that's likely to be replaced in the future. This would mean we'd have something like targetPort now, and in the future targetPorts with some CEL validation that allowed at most one to be set. From controller perspective this means extra work ~forever. We could potentially use a MutatingAdmissionPolicy to copy the contents of targetPort to targetPorts and allow controllers to only look at targetPorts going forward.

2. Replace targetPort with targetPorts in v1
This would be conceptually similar to #1330. We'd be changing the structure of the API with the goal of making it more future proof without meaningfully changing the behavior. We could optionally further limit the scope of this change by setting min + max length == 1 to start, though we may quickly regret that.

3. Make larger changes to API and leave and bump to v1alpha3 instead of v1
If we think there are some large changes lingering, we could stick in alpha for longer. Given the significant cost of each version bump, this is my least favorite option of the bunch. All the theoretical API changes I can think of could be accomplished with an approach similar to 1 or 2 described above.

What do others think here?

/cc @danehans @kfswain @nirrozenbaum

[nirrozenbaum]

I think the problem is bigger than what is described in this issue.
I'll start with the small problem, the port -
I think there is a more basic issue here that port is defined on the InferencePool while Port is not a property of the pool, it's a property of an endpoint that is associated with the pool.
To start simple, it was a nice way to define the port on the InferencePool and make an implicit assumption that all endpoints are completely identical. but port is not a property of the pool and with the fast pace that AI requirements are changing, I don't think we want to limit ourselves to the case where all pods are identical. the same can be said about metrics endpoint - we currently assume all pods expose metrics on the same relative path and same port + exposing the same set of metrics, which may be wrong in DP scenarios (e.g., we might have <rank-x> as prefix to the metrics, each may use different metrics port, etc).
I'd recommend removing the port completely from the InferencePool and putting it as an annotation on the pods, same as Prometheus is doing. e.g., prometheus.io/port: "8080".

we could have on pods:inference-gateway/port:"8080" annotation (or something along these lines). We could also have multiple ports, each with the <rank-x> prefix. As long as we don't put this in the InferencePool CRD, we have the flexibility to change this and improve it iteratively as we make progress.
Putting an annotation on the pod is easy and straight forward, and we can also keep the existing behavior by putting the annotation in the pod template inside a deployment (so all pods have the same port).

The bigger problem-
Thinking about the new use cases that we are currently starting to see and future wise, I think we should prefer making it right rather than getting stuck ~forever with the wrong APIs. My feeling is that we're making too much last minute changes to InferencePool CRD less than a week before we cut a GA release. that's not good.
One of the basic rules of thumb for going GA is to have a stable CRD for at least couple of months or couple of subsequent releases.
With the latest issues that have been opened recently, it seems like:

• we're changing the label selector;
• we need to change the port; and
• we change the EndpointPickerConfig not to be inline.

Looking at the InferencePool CRD, there are 3 fields in the spec - label selector, targetport and EndpointPickerConfig.
So it seems like the recent issues are suggesting to change all fields of the CRD, and although some of them are small changes, I admit I feel very uncomfortable with changing all fields of the CRD less than a week before we cut a GA release.
I agree with @robscott that this option is definitely not my favorite, but unfortunately given the last changes and requirements to change the CRD, I'd recommend stabling the CRD before rushing to GA and then regret it.

[kfswain]

My feeling is that we're making too much last minute changes to InferencePool CRD less than a week before we cut a GA release. that's not good.

I agree with this.

I don't think we want to limit ourselves to the case where all pods are identical.

I think sharing a port number doesnt necessarily mean all pods are identical, but the point is still understood. It depends on how Gateways have been implementing InfPool support. This would be a pretty substantial change to the conformance surface.

[kfswain]

Seperately: I have lots of thoughts and conflicting opinions at the moment. Still taking time to digest and consider what to do, also reasearching the specific topic more (I've not been plugged into this specific topic too much). Looking forward to discussing a bit more tomorrow

[elevran]

There have been multiple proposals on how to implement DP/rank in vLLM and how to expose to IGW (Pod per rank, Pod per node with internal LB, separate process and port per rank, etc.).
Each implementation could impact the CRD. Aligning the CRD fields with a specific implementation choice seems premature and it's unclear to me how stable that choice would be (with new optimization techniques being promoted constantly).

While I certainly see the merit of supporting multiple ports, prefer not making this part of the InferencePool to allow maximum forward flexibility.
Using annotations (suggested above, as an example) to define the relevant ports and removing that from the CRD, allows us to defer making a commitment. If this is too k8s centric, perhaps a /configuration endpoint on the serving pod would be an acceptable choice (I suspect there shall be additional configuration fields that we may care about in the future).
Either option would be aligned with the ideas of the Model Server Protocol, IMHO, where servers adhere to an expected implementation to play well with IGW.