[Feature] Support multiple Ray containers per Pod

[ryanaoleary]

Search before asking

• I had searched in the issues and found no similar feature requirement.

Description

It may be beneficial to support creating multiple Ray workers on the same Pod created with KubeRay.

There is more context in this PR: ai-on-gke/kuberay-tpu-webhook#19, which supports multiple TPU containers with KubeRay, and specifically this comment: ai-on-gke/kuberay-tpu-webhook#19 (comment).

I was able to create a RayCluster with 2 workers, each with 2 Ray containers (so 4 Ray nodes total), and run a workload on it. However, in order to avoid port conflicts and pass the correct Ray resources it's necessary to manually construct a ray start command for the second container. The amount of manual intervention required by the user would be reduced if KubeRay updated the resource detection logic and port assignment to handle multiple containers automatically.

Use case

TPU v7x introduces a dual-chiplet architecture where a standard 4-chip VM spans two distinct NUMA nodes. To optimize memory bandwidth and avoid cross-NUMA latency, v7x workloads can now run as multiple NUMA-aligned containers within a single Pod. This would entail multiple Ray containers on the same Pod created by KubeRay.

More context on the new accelerator: https://docs.cloud.google.com/tpu/docs/tpu7x

Related issues

N/A

Are you willing to submit a PR?

• Yes I am willing to submit a PR!

[Future-Outlier]

this is related to this ray-project/ray#59551.
(single host, multi-node scenario)

[ryanaoleary]

cc: @andrewsykim @Future-Outlier @rueian We're thinking about trying to include this support in v1.6 since multiple containers per Pod for TPU workloads (2 containers each requesting 2 chips) will be the most common deployment pattern for Ironwood (tpu7x) TPU.

I can share a design document once I figure out the scope of the change, but at a high level I think it would involve a new API in the RayCluster or WorkerGroup indicating the number of Ray containers, and then logic in pod.go to assign unique Ports and detect the Ray resources for each Pod. Does that sound doable for v1.6?

[Future-Outlier]

Hi, @ryanaoleary

1. We are planning to release v1.6 on March 15, so I can’t guarantee this will be included.
2. I’d be happy to take a look at this in the next two weeks.
3. One concern I have is that Ray assumes single host / single node in many places, but this use case is single host with multiple nodes. I’m wondering whether there could be Ray-side bugs in this scenario.
4. Have you tried this setup and confirmed it works? Specifically, running multiple Ray containers per Pod with this TPU. Before moving forward, I’d like to make sure this actually runs and doesn’t hit any Ray edge cases.

[ryanaoleary]

Hi, @ryanaoleary

1. We are planning to release v1.6 on March 15, so I can’t guarantee this will be included.
2. I’d be happy to take a look at this in the next two weeks.
3. One concern I have is that Ray assumes single host / single node in many places, but this use case is single host with multiple nodes. I’m wondering whether there could be Ray-side bugs in this scenario.
4. Have you tried this setup and confirmed it works? Specifically, running multiple Ray containers per Pod with this TPU. Before moving forward, I’d like to make sure this actually runs and doesn’t hit any Ray edge cases.

@Future-Outlier Sounds good, and yeah I tested the setup in this comment: ai-on-gke/kuberay-tpu-webhook#19 (comment) running a simple test for JAX init, and one with Ray Train.

[ryanaoleary]

Hi, @ryanaoleary

1. We are planning to release v1.6 on March 15, so I can’t guarantee this will be included.
2. I’d be happy to take a look at this in the next two weeks.
3. One concern I have is that Ray assumes single host / single node in many places, but this use case is single host with multiple nodes. I’m wondering whether there could be Ray-side bugs in this scenario.
4. Have you tried this setup and confirmed it works? Specifically, running multiple Ray containers per Pod with this TPU. Before moving forward, I’d like to make sure this actually runs and doesn’t hit any Ray edge cases.

@Future-Outlier Sounds good, and yeah I tested the setup in this comment: ai-on-gke/kuberay-tpu-webhook#19 (comment) running a simple test for JAX init, and one with Ray Train. It worked for both examples but I agree it could run into bugs with Ray for more complex examples.

[Future-Outlier]

Hi, @ryanaoleary just an edge case comes to my mind.
this might not work properly with the history server project, is that ok in your usecase?
this is ray's problem, for more detail, you can see here.
ray-project/ray#58880 (comment)

[Future-Outlier]

Hi @ryanaoleary, if possible, please also let me know what the potential new API might look like in your mind.

[rueian]

Hi @ryanaoleary, do you know where the "NUMA-aligned" for containers happens? Can't that be done inside one container? Such as setting affinities to processes?

[ryanaoleary]

Hi @ryanaoleary, do you know where the "NUMA-aligned" for containers happens? Can't that be done inside one container? Such as setting affinities to processes?

NUMA alignment requires resource alignment across cpu, memory, TPUs and the NIC. The alignment can be done through creating containers that align on the NUMA boundaries and allowing the topology manager in the Kubelet to handle the scheduling. This is how frameworks like Jobset and XPK handle Ironwood. Alternatively, if we had the two worker processes run on the same pod (same Ray node, so sharing 1 Raylet) I think we could also try to ensure NUMA alignment through using numactl when the process starts but I haven't explored this option as much.

if possible, please also let me know what the potential new API might look like in your mind.

I think it could look something like:

workerGroupSpecs:
- replicas: 0
  rayContainers: 2
  ...
  template:
      spec:
        containers:
        <2+ containers>

I think the controller would then just have to validate that template.Spec.Containers contains at least rayContainers containers, and it can construct the ray start arguments/set unique ports for the first rayContainers out of the total defined for the Pod.

this might not work properly with the history server project, is that ok in your usecase?
this is ray's problem, for more detail, you can see here.

I think ideally it would work together, so maybe that's another reason we should try to get numactl to work within Ray when the worker process starts. Alternatively, it seems like your solution to write logs to /tmp/ray/session_latest/{node_id}/ would work for the multi-container setup since each container would have it's own node id.

[rueian]

I think we could also try to ensure NUMA alignment through using numactl when the process starts but I haven't explored this option as much.

Thanks for sharing! Happy to learn that. One more question: aren't multiple pods on the same host also NUMA-aligned?

maybe that's another reason we should try to get numactl to work within Ray when the worker process starts.

I think that will be ideal, in the sense of avoiding unnecessary Raylet, but that may be much more difficult to implement than multiple containers in KubeRay.