Add support for parsing neuron core resource limit and pass it as ray… #2409
Conversation
| rayStartParams["num-gpus"] = strconv.FormatInt(resource.Value(), 10) | ||
| // For now, only support one GPU type. Break on first match. | ||
| break | ||
| } else if resourceKeyString == "aws.amazon.com/neuroncore" && !resource.IsZero() { |
There was a problem hiding this comment.
Can you put aws.amazon.com/neuroncore into a constant?
| rayStartParams["num-gpus"] = strconv.FormatInt(resource.Value(), 10) | ||
| // For now, only support one GPU type. Break on first match. | ||
| break | ||
| } else if resourceKeyString == "aws.amazon.com/neuroncore" && !resource.IsZero() { | ||
| if err := addNeuronCoresToResourcesIfNotExists(rayStartParams, resource.Value()); err != nil { |
There was a problem hiding this comment.
This would be the first hardware accelerator (outside GPUs) that we auto detect in container resource and pass into the ray start command as custom resource. I can see the appeal of doing this but it could also mean we end up supporting an ever growing list of custom resources in kuberay.
My question here is whether having to specify the custom resource in startParams is a big enough pain point for KubeRay to auto detect, commonly used custom resources.
There was a problem hiding this comment.
Thinking about it more, it's probably worth doing for the well-known accelerators at least. Can we generalize the implementation so it's easy to add new accelerators here? cc @ryanaoleary for TPUs
There was a problem hiding this comment.
Native support for neuron and other accelerators already exists for Ray's VM cluster-launchers.
See discussion in ray-project/ray#44361
This PR would move KubeRay closer to parity with the VM-based solution.
There was a problem hiding this comment.
I agree that there needs to be some styling here that could generalize to other hardware,
again along the same lines as what exists for Ray's VM support.
Why are these changes needed?
The PR is needed to parse the neuroncore resource limits fro the pod spec and pass them to
ray startcommand as it expects.Expected format: https://docs.ray.io/en/latest/ray-core/scheduling/accelerators.html
Related issue number
ray-project/ray#44361
Checks