feat: Add accelerator type for Slurm [FOUNDENG-483] (#711)

If the Slrum hpcPartitionDetails contains accelerator type set it for the resource pool.
determined-ai · Jun 28, 2023 · 6cc3157 · 6cc3157
1 parent 5167f60
commit 6cc3157
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/master/internal/rm/dispatcherrm/dispatcher_resource_manager.go b/master/internal/rm/dispatcherrm/dispatcher_resource_manager.go
@@ -97,6 +97,7 @@ type hpcPartitionDetails struct {
 	TotalGpuSlots          int    `json:"totalGpuSlots"`
 	TotalAvailableCPUSlots int    `json:"totalAvailableCpuSlots"`
 	TotalCPUSlots          int    `json:"totalCpuSlots"`
+	Accelerator            string `json:"accelerator"`
 }
 
 // hpcNodeDetails holds HPC Slurm node details.
@@ -1379,6 +1380,7 @@ func (m *dispatcherResourceManager) summarizeResourcePool(
 			ImageId:                      "",
 			InstanceType:                 wlmName,
 			Details:                      &resourcepoolv1.ResourcePoolDetail{},
+			Accelerator:                  v.Accelerator,
 		}
 		poolNameMap[pool.Name] = &pool
 		result = append(result, &pool)

diff --git a/master/internal/rm/dispatcherrm/dispatcher_resource_manager_test.go b/master/internal/rm/dispatcherrm/dispatcher_resource_manager_test.go
@@ -243,6 +243,7 @@ func Test_summarizeResourcePool(t *testing.T) {
 		TotalAvailableGpuSlots: 2,
 		TotalNodes:             10,
 		TotalGpuSlots:          5,
+		Accelerator:            "tesla",
 	}
 	p2 := hpcPartitionDetails{
 		TotalAvailableNodes:    0,
@@ -339,6 +340,7 @@ func Test_summarizeResourcePool(t *testing.T) {
 						SlotsAvailable: 5,
 						SlotsUsed:      3,
 						NumAgents:      10,
+						Accelerator:    "tesla",
 					},
 					{
 						Name:           "partition 2",