kubeflow / trainer / 33192220415

70%

master: 69%

Pull #3997

github

fix(framework): make GPU count detection deterministic and unit-aware

numGPU scanned corev1.ResourceList for the first key whose lowercased name
contained "gpu" and returned its quantity. ResourceList is a map, so the
result varied between calls whenever a container requested more than one
resource with "gpu" in its name, silently and without any log or event.

GPU sharing stacks make that the normal case rather than an edge case. HAMi
requests nvidia.com/gpu alongside gpumem and gpucores, Volcano vGPU requests
vgpu-number alongside vgpu-memory and vgpu-cores, and every Intel resource
sits under the gpu.intel.com domain. Most of the colliding keys are not
device counts at all, so the helper could return megabytes or a core
percentage where a device count was expected.

The odds are not even. Go stores a map of this size in a single group and
places entries in insertion order, while iteration begins at a random slot,
so for n keys the first-inserted wins (8-n+1)/8 of the time and each of the
others wins 1/8. Measured on the previous implementation, nvidia.com/gpu
alongside nvidia.com/gpumem returned the memory value in 12% of calls, and
in 87% once the two were inserted the other way round. A cluster can
therefore look healthy for a long time and then misbehave for whichever
users happen to order their resources differently.

The five callers all derive process counts from this value: torch and MPI
size numProcPerNode, XGBoost sizes DMLC_NUM_WORKER, Flux sizes gpus-per-task,
and torchtune selects the recipe. TorchTune reads it once during webhook
validation and again during JobSet generation, so the two could disagree
about the same TrainJob.

Skip the names that denote memory, cores, tiles, or percentages, and visit
the remaining candidates in sorted order rather than ranging over the map.
Matching the unit rather than enumerating known device-count resources keeps
in-house device plugins working: unit names are few and shared across
vendors, whereas... (continued)

14 of 14 new or added lines in 1 file covered. (100.0%)

2724 of 3893 relevant lines covered (69.97%)

0.82 hits per line