func tpuTopologyToNodeCount()

in tpu-provisioner/internal/cloud/gke.go [571:597]


func tpuTopologyToNodeCount(accelerator, topo string) (int, error) {
	var expectedDims int
	switch accelerator {
	case V4PodSliceAccelerator, V5pPodSliceAccelerator:
		expectedDims = 3
	case V5ePodSliceAccelerator, V6eSliceAccelerator:
		expectedDims = 2
	default:
		return 0, fmt.Errorf("invalid accelerator: %v", accelerator)
	}

	split := strings.Split(topo, "x")
	if len(split) != expectedDims {
		return 0, fmt.Errorf("invalid topology: %v, expected %v dimensions", topo, expectedDims)
	}

	product := 1
	for _, s := range split {
		x, err := strconv.Atoi(s)
		if err != nil {
			return 0, fmt.Errorf("invalid topology: %v, could not convert %q to int: %w", topo, s, err)
		}
		product *= x
	}

	return int(math.Ceil(float64(product) / 4)), nil
}