Skip to content

[Feature] Add resource mapping for Huawei Ascend NPUs #5071

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
bakhovaddinov wants to merge 1 commit into
base: master
Choose a base branch
from
+264 −1
Open

[Feature] Add resource mapping for Huawei Ascend NPUs #5071

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 34 additions & 1 deletion ray-operator/controllers/ray/common/pod.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,13 +41,40 @@ const (
NeuronCoreRayResourceName = "neuron_cores"
TPUContainerResourceName = "google.com/tpu"
TPURayResourceName = "TPU"
AscendRayResourceName = "NPU"
)

var customAcceleratorToRayResourceMap = map[string]string{
NeuronCoreContainerResourceName: NeuronCoreRayResourceName,
TPUContainerResourceName: TPURayResourceName,
}

func isNPUResourceKey(key string) bool {
lowerKey := strings.ToLower(key)
if strings.HasPrefix(lowerKey, "huawei.com/ascend") {
// Skip metadata resource keys like huawei.com/Ascend910B-memory and huawei.com/Ascend910B-core,
// which describe accelerator capacity rather than being actual Ray compute resources.
if strings.HasSuffix(lowerKey, "-memory") || strings.HasSuffix(lowerKey, "-core") {
return false
}
Comment thread
bakhovaddinov marked this conversation as resolved.
return true
}
if lowerKey == "huawei.com/npu" {
return true
}
return false
}
Comment thread
win5923 marked this conversation as resolved.

func getCustomAcceleratorRayResourceName(resourceKeyString string) string {
if rayResourceName, ok := customAcceleratorToRayResourceMap[resourceKeyString]; ok {
return rayResourceName
}
if isNPUResourceKey(resourceKeyString) {
return AscendRayResourceName
Comment thread
cursor[bot] marked this conversation as resolved.
}
return ""
}

// Get the port required to connect to the Ray cluster by worker nodes and drivers
// started within the cluster.
// For Ray >= 1.11.0 this is the GCS server port. For Ray < 1.11.0 it is the Redis port.
Expand Down Expand Up @@ -1446,7 +1473,7 @@ func addWellKnownAcceleratorResources(rayStartParams map[string]string, resource

// Add the first encountered custom accelerator resource from the resource limits to the rayStartParams if not already present
if !isCustomAcceleratorResourceAdded {
if rayResourceName, ok := customAcceleratorToRayResourceMap[resourceKeyString]; ok && !resourceValue.IsZero() {
if rayResourceName := getCustomAcceleratorRayResourceName(resourceKeyString); rayResourceName != "" && !resourceValue.IsZero() {
if _, exists := resourcesMap[rayResourceName]; !exists {
resourcesMap[rayResourceName] = resourceValue.AsApproximateFloat64()

Expand Down Expand Up @@ -1474,6 +1501,10 @@ func isCustomAcceleratorPresentInResources(resourcesMap map[string]float64) bool
return true
}
}

if _, ok := resourcesMap[AscendRayResourceName]; ok {
return true
}
}

return false
Expand Down Expand Up @@ -1657,6 +1688,8 @@ func updateRayStartParamsResources(ctx context.Context, rayStartParams map[strin
rayStartParams["memory"] = strconv.FormatInt(q.Value(), 10)
} else if utils.IsGPUResourceKey(normalizedName) {
rayStartParams["num-gpus"] = strconv.FormatInt(q.Value(), 10)
} else if rayResourceName := getCustomAcceleratorRayResourceName(name); rayResourceName != "" {
rayResourcesJson[rayResourceName] = q.AsApproximateFloat64()
Comment thread
bakhovaddinov marked this conversation as resolved.
} else {
rayResourcesJson[name] = q.AsApproximateFloat64()
}
Expand Down
220 changes: 220 additions & 0 deletions ray-operator/controllers/ray/common/pod_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2364,6 +2364,80 @@ func TestGenerateRayStartCommand(t *testing.T) {
resource: corev1.ResourceRequirements{},
expected: "ray start --head --include-log-monitor=true ",
},
{
name: "WorkerNode with Ascend910B NPU",
nodeType: rayv1.WorkerNode,
rayStartParams: map[string]string{},
resource: corev1.ResourceRequirements{
Limits: corev1.ResourceList{
"huawei.com/Ascend910B": resource.MustParse("2"),
},
},
expected: `ray start --resources='{"NPU":2}' `,
},
{
name: "WorkerNode with Ascend NPU (huawei.com/npu)",
nodeType: rayv1.WorkerNode,
rayStartParams: map[string]string{},
resource: corev1.ResourceRequirements{
Limits: corev1.ResourceList{
"huawei.com/npu": resource.MustParse("2"),
},
},
expected: `ray start --resources='{"NPU":2}' `,
},
{
name: "HeadNode with Ascend910B and GPU",
nodeType: rayv1.HeadNode,
rayStartParams: map[string]string{},
resource: corev1.ResourceRequirements{
Limits: corev1.ResourceList{
"huawei.com/Ascend910B": resource.MustParse("2"),
"nvidia.com/gpu": resource.MustParse("1"),
},
},
expected: `ray start --head --num-gpus=1 --resources='{"NPU":2}' `,
},
{
name: "HeadNode with existing resources and Ascend910B",
nodeType: rayv1.HeadNode,
rayStartParams: map[string]string{
"resources": `'{"custom_resource":2}'`,
},
resource: corev1.ResourceRequirements{
Limits: corev1.ResourceList{
"huawei.com/Ascend910B": resource.MustParse("2"),
},
},
expected: `ray start --head --resources='{"NPU":2,"custom_resource":2}' `,
},
{
name: "HeadNode with existing NPU resources",
nodeType: rayv1.HeadNode,
rayStartParams: map[string]string{
"resources": `'{"NPU":3,"custom_resource":2}'`,
},
resource: corev1.ResourceRequirements{
Limits: corev1.ResourceList{
"huawei.com/Ascend910B": resource.MustParse("2"),
},
},
expected: `ray start --head --resources='{"NPU":3,"custom_resource":2}' `,
},
{
name: "HeadNode with multiple accelerators including Ascend910B",
nodeType: rayv1.HeadNode,
rayStartParams: map[string]string{},
resource: corev1.ResourceRequirements{
Limits: corev1.ResourceList{
"huawei.com/Ascend910B": resource.MustParse("2"),
"google.com/tpu": resource.MustParse("8"),
"aws.amazon.com/neuroncore": resource.MustParse("4"),
"nvidia.com/gpu": resource.MustParse("1"),
},
},
expected: `ray start --head --num-gpus=1 --resources='{"neuron_cores":4}' `,
},
}

for _, tt := range tests {
Expand Down Expand Up @@ -2646,6 +2720,35 @@ func TestUpdateRayStartParamsResources(t *testing.T) {
"resources": "'{\"Custom-Resource\":5}'",
},
},
"Ascend910B NPU resource set in `Resources`": {
initialRayStartParams: map[string]string{},
groupResources: map[string]string{
"huawei.com/Ascend910B": "2",
},
expectedRayStartParams: map[string]string{
"resources": "'{\"NPU\":2}'",
},
},
"Ascend910c NPU resource set in `Resources`": {
initialRayStartParams: map[string]string{},
groupResources: map[string]string{
"huawei.com/Ascend910c": "4",
},
expectedRayStartParams: map[string]string{
"resources": "'{\"NPU\":4}'",
},
},
"GPU and Ascend910B NPU resource set in `Resources`": {
initialRayStartParams: map[string]string{},
groupResources: map[string]string{
"nvidia.com/gpu": "1",
"huawei.com/Ascend910B": "2",
},
expectedRayStartParams: map[string]string{
"num-gpus": "1",
"resources": "'{\"NPU\":2}'",
},
Comment thread
cursor[bot] marked this conversation as resolved.
},
}

for name, tc := range tests {
Expand Down Expand Up @@ -2984,3 +3087,120 @@ func TestBuildCollectorContainerAndPodInjection(t *testing.T) {
assert.True(t, ok)
assert.Equal(t, utils.DEFAULT_RAY_EXPOSABLE_EVENT_TYPES, workerEventTypesV2Env.Value)
}

func TestIsNPUResourceKey(t *testing.T) {
tests := []struct {
name string
resourceKey string
expected bool
}{
{
name: "huawei.com/Ascend910B",
resourceKey: "huawei.com/Ascend910B",
expected: true,
},
{
name: "huawei.com/Ascend910B-power",
resourceKey: "huawei.com/Ascend910B-power",
expected: true,
},
{
name: "huawei.com/npu",
resourceKey: "huawei.com/npu",
expected: true,
},
{
name: "huawei.com/NPU",
resourceKey: "huawei.com/NPU",
expected: true,
},
{
name: "huawei.com/ascend-memory excluded",
resourceKey: "huawei.com/ascend-memory",
expected: false,
},
{
name: "huawei.com/ascend-core excluded",
resourceKey: "huawei.com/ascend-core",
expected: false,
},
{
name: "huawei.com/Ascend910B-memory excluded",
resourceKey: "huawei.com/Ascend910B-memory",
expected: false,
},
{
name: "huawei.com/Ascend910B-core excluded",
resourceKey: "huawei.com/Ascend910B-core",
expected: false,
},
{
name: "nvidia gpu not NPU",
resourceKey: "nvidia.com/gpu",
expected: false,
},
{
name: "cpu not NPU",
resourceKey: "cpu",
expected: false,
},
{
name: "memory not NPU",
resourceKey: "memory",
expected: false,
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := isNPUResourceKey(tt.resourceKey)
assert.Equal(t, tt.expected, result)
})
}
}

func TestGetCustomAcceleratorRayResourceName(t *testing.T) {
tests := []struct {
name string
key string
expected string
}{
{
name: "nvidia neuron core direct match",
key: "aws.amazon.com/neuroncore",
expected: "neuron_cores",
},
{
name: "google TPU direct match",
key: "google.com/tpu",
expected: "TPU",
},
{
name: "huawei Ascend910B NPU",
key: "huawei.com/Ascend910B",
expected: "NPU",
},
{
name: "huawei.com/npu NPU",
key: "huawei.com/npu",
expected: "NPU",
},
{
name: "huawei ascend-memory returns empty",
key: "huawei.com/ascend-memory",
expected: "",
},
{
name: "unknown resource returns empty",
key: "unknown.com/accelerator",
expected: "",
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := getCustomAcceleratorRayResourceName(tt.key)
assert.Equal(t, tt.expected, result)
})
}
}
10 changes: 10 additions & 0 deletions ray-operator/controllers/ray/utils/resources_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,16 @@ func TestIsGPUResourceKey(t *testing.T) {
resourceKey: "memory",
expected: false,
},
{
name: "huawei Ascend910B",
resourceKey: "huawei.com/Ascend910B",
expected: false,
},
{
name: "huawei npu",
resourceKey: "huawei.com/npu",
expected: false,
},
}

for _, tt := range tests {
Expand Down
Loading