Skip to content

Commit e1741bc

Browse files
Jenever-feitunJenever-feitun
authored
feat: add tpu hardware (#1235)
1 parent 86fe94e commit e1741bc

10 files changed

Lines changed: 130 additions & 3 deletions

File tree

‎builder/deploy/cluster/cluster_manager.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -643,6 +643,10 @@ func getXPULabel(labels map[string]string, config *config.Config) (string, strin
643643
//for amd gpu
644644
return "amd.com/gpu", "amd.com/gpu.product-name", []string{"amd.com/gpu.vram"}
645645
}
646+
if _, found := labels["chipltech.com/tpu.product"]; found {
647+
//for chipltech tpu
648+
return "chipltech.com/dlc-lyp", "chipltech.com/tpu.product", []string{"chipltech.com/tpu.mem"}
649+
}
646650
//check custom gpu model label
647651
if config.Space.GPUModelLabel != "" {
648652
var gpuLabels []types.GPUModel

‎builder/deploy/cluster/cluster_manager_test.go‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -243,6 +243,15 @@ func TestGetXPULabel(t *testing.T) {
243243
wantCapacity: "nvidia.com/gpu",
244244
wantTypeLabel: "nvidia.com/nvidia_name",
245245
},
246+
{
247+
name: "chipltech TPU label",
248+
labels: map[string]string{
249+
"chipltech.com/tpu.product": "TPU-X1",
250+
},
251+
config: &config.Config{},
252+
wantCapacity: "chipltech.com/dlc-lyp",
253+
wantTypeLabel: "chipltech.com/tpu.product",
254+
},
246255
{
247256
name: "no matching label",
248257
labels: map[string]string{},

‎builder/deploy/common/utils.go‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,9 @@ func UpdateEvaluationEnvHardware(env map[string]string, hardware types.HardWare)
6464
} else if hardware.GPGpu.Num != "" {
6565
env["GPGPU_NUM"] = hardware.GPGpu.Num
6666
xpuNum = hardware.GPGpu.Num
67+
} else if hardware.Tpu.Num != "" {
68+
env["TPU_NUM"] = hardware.Tpu.Num
69+
xpuNum = hardware.Tpu.Num
6770
}
6871
if xpuNum != "0" {
6972
env["GPU_NUM"] = xpuNum
@@ -84,6 +87,8 @@ func ResourceType(hardware types.HardWare) types.ResourceType {
8487
resourceType = types.ResourceTypeDCU
8588
} else if hardware.GPGpu.Num != "" {
8689
resourceType = types.ResourceTypeGPGPU
90+
} else if hardware.Tpu.Num != "" {
91+
resourceType = types.ResourceTypeTPU
8792
}
8893
return resourceType
8994
}
@@ -106,6 +111,9 @@ func GetResourceAndType(hardware types.HardWare) (string, string) {
106111
} else if hardware.Dcu.Num != "" {
107112
resourceType = hardware.Dcu.Type
108113
resource += "·" + hardware.Dcu.Num + "Dcu"
114+
} else if hardware.Tpu.Num != "" {
115+
resourceType = hardware.Tpu.Type
116+
resource += "·" + hardware.Tpu.Num + "Tpu"
109117
} else {
110118
resourceType = hardware.Cpu.Type
111119
}
@@ -115,7 +123,8 @@ func GetResourceAndType(hardware types.HardWare) (string, string) {
115123
func ContainsGraphicResource(hardware types.HardWare) bool {
116124
if hardware.Gpu.Num != "" || hardware.Npu.Num != "" ||
117125
hardware.Gcu.Num != "" || hardware.Mlu.Num != "" ||
118-
hardware.Dcu.Num != "" || hardware.GPGpu.Num != "" {
126+
hardware.Dcu.Num != "" || hardware.GPGpu.Num != "" ||
127+
hardware.Tpu.Num != "" {
119128
return true
120129
}
121130
return false
@@ -135,6 +144,8 @@ func GetXPUNumber(hardware types.HardWare) (int, error) {
135144
xpuNumStr = hardware.Dcu.Num
136145
} else if hardware.GPGpu.Num != "" {
137146
xpuNumStr = hardware.GPGpu.Num
147+
} else if hardware.Tpu.Num != "" {
148+
xpuNumStr = hardware.Tpu.Num
138149
}
139150
xpuNum, err := strconv.Atoi(xpuNumStr)
140151
return xpuNum, err

‎builder/deploy/common/utils_test.go‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,22 @@ func TestUpdateEvaluationEnvHardware(t *testing.T) {
100100
"DCU_NUM": "1",
101101
},
102102
},
103+
{
104+
name: "TPU set when others are empty",
105+
env: make(map[string]string),
106+
hardware: types.HardWare{
107+
Gpu: types.Processor{Num: ""},
108+
Npu: types.Processor{Num: ""},
109+
Gcu: types.Processor{Num: ""},
110+
Mlu: types.Processor{Num: ""},
111+
Dcu: types.Processor{Num: ""},
112+
Tpu: types.Processor{Num: "1"},
113+
},
114+
expected: map[string]string{
115+
"TPU_NUM": "1",
116+
"GPU_NUM": "1",
117+
},
118+
},
103119
{
104120
name: "No hardware set",
105121
env: make(map[string]string),
@@ -210,6 +226,14 @@ func TestGetXPUNumber(t *testing.T) {
210226
want: 1,
211227
wantErr: false,
212228
},
229+
{
230+
name: "Only Tpu specified",
231+
hardware: types.HardWare{
232+
Tpu: types.Processor{Num: "3"},
233+
},
234+
want: 3,
235+
wantErr: false,
236+
},
213237
{
214238
name: "GPU takes priority over others",
215239
hardware: types.HardWare{
@@ -434,6 +458,18 @@ func TestResourceType(t *testing.T) {
434458
},
435459
expected: types.ResourceTypeGPU,
436460
},
461+
{
462+
name: "TPU when TPU Num is not empty and others are empty",
463+
hardware: types.HardWare{
464+
Gpu: types.Processor{Num: ""},
465+
Npu: types.Processor{Num: ""},
466+
Mlu: types.Processor{Num: ""},
467+
Gcu: types.Processor{Num: ""},
468+
Dcu: types.Processor{Num: ""},
469+
Tpu: types.Processor{Num: "1"},
470+
},
471+
expected: types.ResourceTypeTPU,
472+
},
437473
}
438474

439475
for _, tt := range tests {
@@ -445,3 +481,30 @@ func TestResourceType(t *testing.T) {
445481
})
446482
}
447483
}
484+
485+
func TestContainsGraphicResource(t *testing.T) {
486+
tests := []struct {
487+
name string
488+
hardware types.HardWare
489+
want bool
490+
}{
491+
{
492+
name: "cpu only",
493+
hardware: types.HardWare{Cpu: types.CPU{Num: "2"}, Memory: "4Gi"},
494+
want: false,
495+
},
496+
{
497+
name: "tpu is graphic resource",
498+
hardware: types.HardWare{
499+
Tpu: types.Processor{Num: "1"},
500+
},
501+
want: true,
502+
},
503+
}
504+
505+
for _, tt := range tests {
506+
t.Run(tt.name, func(t *testing.T) {
507+
require.Equal(t, tt.want, ContainsGraphicResource(tt.hardware))
508+
})
509+
}
510+
}

‎builder/deploy/deployer_resource.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,8 @@ func isCPUOnlyWorkload(hardware *types.HardWare) bool {
174174
hardware.Gcu.Num == "" &&
175175
hardware.Mlu.Num == "" &&
176176
hardware.Dcu.Num == "" &&
177-
hardware.GPGpu.Num == ""
177+
hardware.GPGpu.Num == "" &&
178+
hardware.Tpu.Num == ""
178179
}
179180

180181
func isXPUNode(node types.NodeResourceInfo) bool {

‎common/types/hardware.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ type (
5050
Mlu Processor `json:"mlu,omitempty"` // cambricon
5151
Dcu Processor `json:"dcu,omitempty"` // hygon
5252
GPGpu Processor `json:"gpgpu,omitempty"` // iluvatar,metax
53+
Tpu Processor `json:"tpu,omitempty"` // chipltech
5354
Cpu CPU `json:"cpu,omitempty"`
5455
Memory string `json:"memory,omitempty"`
5556
EphemeralStorage string `json:"ephemeral_storage,omitempty"`
@@ -79,6 +80,8 @@ func (h *HardWare) GetResXPUMode() string {
7980
xpuProc = h.Dcu
8081
case strings.TrimSpace(h.GPGpu.Num) != "":
8182
xpuProc = h.GPGpu
83+
case strings.TrimSpace(h.Tpu.Num) != "":
84+
xpuProc = h.Tpu
8285
}
8386

8487
xpuModel := strings.TrimSpace(xpuProc.Type)

‎common/types/space_resource.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ const (
1515
ResourceTypeGPGPU ResourceType = "gpgpu"
1616
ResourceTypeMLU ResourceType = "mlu"
1717
ResourceTypeDCU ResourceType = "dcu"
18+
ResourceTypeTPU ResourceType = "tpu"
1819
PayModeFree PayMode = "free"
1920
PayModeMinute PayMode = "minute"
2021
PayModeMonth PayMode = "month"
@@ -28,7 +29,8 @@ func ResourceTypeValid(resourceType ResourceType) bool {
2829
resourceType == ResourceTypeGCU ||
2930
resourceType == ResourceTypeGPGPU ||
3031
resourceType == ResourceTypeMLU ||
31-
resourceType == ResourceTypeDCU
32+
resourceType == ResourceTypeDCU ||
33+
resourceType == ResourceTypeTPU
3234
}
3335

3436
var ResourceTypeValidator validator.Func = func(fl validator.FieldLevel) bool {

‎runner/common/resource.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ func GenerateResources(params rtypes.ResourceGeneratorParams) *rtypes.GeneratedR
3838
{hardware.Dcu.Labels},
3939
{hardware.GPGpu.Labels},
4040
{hardware.Cpu.Labels},
41+
{hardware.Tpu.Labels},
4142
}
4243

4344
for _, hw := range hardwareTypes {

‎runner/common/resource_ce.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ func handleAccelerator(hardware types.HardWare, resReq map[corev1.ResourceName]r
2121
{hardware.Mlu.ResourceName, hardware.Mlu.Num},
2222
{hardware.Dcu.ResourceName, hardware.Dcu.Num},
2323
{hardware.GPGpu.ResourceName, hardware.GPGpu.Num},
24+
{hardware.Tpu.ResourceName, hardware.Tpu.Num},
2425
}
2526

2627
for _, acc := range accelerators {

‎runner/common/resource_ce_test.go‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,18 @@ func Test_handleAccelerator_CE(t *testing.T) {
4343
"huawei.com/npu": resource.MustParse("2"),
4444
},
4545
},
46+
{
47+
name: "tpu with resource name",
48+
hardware: types.HardWare{
49+
Tpu: types.Processor{
50+
Num: "1",
51+
ResourceName: "chipltech.com/dlc-lyp",
52+
},
53+
},
54+
wantReq: map[corev1.ResourceName]resource.Quantity{
55+
"chipltech.com/dlc-lyp": resource.MustParse("1"),
56+
},
57+
},
4658
{
4759
name: "no resource name",
4860
hardware: types.HardWare{
@@ -93,6 +105,26 @@ func Test_GenerateResources_CE(t *testing.T) {
93105
Tolerations: nil,
94106
},
95107
},
108+
{
109+
name: "tpu resources with resource name (CE)",
110+
hardware: types.HardWare{
111+
Tpu: types.Processor{
112+
Num: "2",
113+
ResourceName: "chipltech.com/dlc-lyp",
114+
},
115+
},
116+
nodes: []types.Node{},
117+
deployExt: types.DeployExtend{},
118+
config: &config.Config{},
119+
want: &rtypes.GeneratedResources{
120+
ResourceRequirements: map[corev1.ResourceName]resource.Quantity{
121+
"chipltech.com/dlc-lyp": resource.MustParse("2"),
122+
},
123+
NodeSelector: map[string]string{},
124+
NodeAffinity: nil,
125+
Tolerations: nil,
126+
},
127+
},
96128
}
97129

98130
for _, tt := range tests {

0 commit comments

Comments
 (0)