gpu.go 1.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. package gcp
  2. import (
  3. "regexp"
  4. "strings"
  5. )
  6. // ---- Original OpenCost regex fallback ----
  7. var (
  8. nvidiaTeslaGPURegex = regexp.MustCompile(`(?i)nvidia[\s-]*tesla[\s-]*([a-z0-9]+)`)
  9. nvidiaGPURegex = regexp.MustCompile(`(?i)nvidia[\s-]*([a-z0-9]+)`)
  10. )
  11. // Explicit substring → canonical GPU label
  12. var gpuSKUToGpuLabel = map[string]string{
  13. // A100
  14. "nvidia tesla a100 80gb": "nvidia-a100-80gb",
  15. "nvidia a100 80gb": "nvidia-a100-80gb",
  16. "nvidia tesla a100": "nvidia-tesla-a100",
  17. "nvidia a100": "nvidia-tesla-a100",
  18. // L4
  19. "nvidia l4": "nvidia-l4",
  20. // T4
  21. "tesla t4": "nvidia-tesla-t4",
  22. "nvidia t4": "nvidia-tesla-t4",
  23. // V100
  24. "tesla v100": "nvidia-tesla-v100",
  25. "nvidia v100": "nvidia-tesla-v100",
  26. // P100 (reviewer case)
  27. "tesla p100": "nvidia-tesla-p100",
  28. "nvidia p100": "nvidia-tesla-p100",
  29. }
  30. // ---- Main Normalizer ----
  31. func NormalizeGPULabel(desc string) string {
  32. d := strings.ToLower(desc)
  33. // --- Step 1: A100 detection first ---
  34. if strings.Contains(d, "a100") {
  35. has80 := strings.Contains(d, "80gb") || strings.Contains(d, "80 gb")
  36. has40 := strings.Contains(d, "40gb") || strings.Contains(d, "40 gb")
  37. if has80 {
  38. return "nvidia-a100-80gb"
  39. }
  40. if has40 {
  41. return "nvidia-tesla-a100"
  42. }
  43. return "nvidia-tesla-a100" // generic A100 → legacy
  44. }
  45. // --- Step 2: explicit substring mapping ---
  46. for key, model := range gpuSKUToGpuLabel {
  47. if strings.Contains(d, key) {
  48. return model
  49. }
  50. }
  51. // --- Step 3: regex fallback (original OpenCost behavior) ---
  52. if match := nvidiaTeslaGPURegex.FindStringSubmatch(desc); len(match) == 2 {
  53. return "nvidia-tesla-" + strings.ToLower(match[1])
  54. }
  55. if match := nvidiaGPURegex.FindStringSubmatch(desc); len(match) == 2 {
  56. return "nvidia-" + strings.ToLower(match[1])
  57. }
  58. return ""
  59. }