{{ message }}
Repository navigation
fix(k8s): merge accelerator-devices over the defaults and add shared GPU entries - #8163
Open
samhita-alla wants to merge 1 commit into
Open
samhita-alla wants to merge 1 commit into
samhita-alla wants to merge 1 commit into
Conversation
…GPU entries Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Samhita Alla <aallasamhita@gmail.com>
This was referenced Oct 12, 2026
samhita-alla
marked this pull request as ready for review
October 14, 2026 05:52
EngHabu
added a commit
that referenced
this pull request
Oct 14, 2026
cosmicbboy
added a commit
to flyteorg/flyte-sdk
that referenced
this pull request
Oct 18, 2026
flyte.GPU("RTX_PRO_6000", n, partition=...) now accepts and validates
the RTX PRO 6000 Blackwell Server Edition MIG profiles: 1g.24gb, 2g.48gb
and 4g.96gb. They are exposed as RTXPro6000Parts next to the existing
A100, H100 and H200 partition types. The +gfx and +me profiles are left
out because the backend turns the partition into a node label value, and
label values cannot contain "+". The -me profiles (1g.24gb-me,
2g.48gb-me) are left out until the backend can parse them.
The GPU device type also gains shared devices: "T4 shared", "L4 shared",
"A10G shared", "A10 shared" and "L40S shared", grouped as SharedGPUType.
A "<model> shared" device requests time-sliced (or MPS) capacity, a
share of one physical GPU on a node pool that advertises the shared
accelerator for that model, so several tasks can run on the same GPU at
once. A plain "<model>" request never lands on such a pool. A shared
device takes no partition, and it works in the Resources string form too
(gpu="T4 shared:1"). The SDK sends the device name as written. The
shared names keep the space, unlike the underscore device names, because
the backend accelerator mapping uses that spelling. These devices need a
backend whose accelerator mapping includes the shared names, which comes
with flyteorg/flyte#8163 (flyteorg/flyte#8163).
Three short examples are added under examples/accelerators. Each one
only sets the GPU request, and the backend's accelerator and partition
node selectors do the rest.
- mig_slice.py runs a task on one 1g.24gb slice of an RTX PRO 6000.
- mig_mixed.py runs two 1g.24gb tasks and one 2g.48gb task, which fit on
one GPU under the MIG strategy mixed with the all-balanced layout.
- gpu_shared.py runs three tasks that share one T4 through gpu="T4
shared:1".
Test plan
- uv run pytest tests/user_api/test_resources.py
tests/internal/runtime/test_resources_serde.py (188 passed), covering
valid and invalid RTX PRO 6000 partitions, the partition_size sent to
the backend, the shared devices, and the error for a shared device with
a partition
- ruff check and ruff format --check pass on the changed files, and all
three examples import
- Ran RTX PRO 6000 MIG slices under the single and mixed strategies, and
shared T4 tasks, on a cluster with those node pools
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Signed-off-by: Samhita Alla <aallasamhita@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Niels Bantilan <niels.bantilan@gmail.com>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why are the changes needed?
Setting
accelerator-devicesin the k8s plugin config is close to unusable today. The config loader replaces the whole default map rather than merging into it, so adding one entry drops every built-in device. Viper also lowercases plain map keys, and since the lookup inGetNormalizedAcceleratorDeviceupper-cases the device name, a configured key likeT4is stored ast4and never matches.There's also no default entry for time-sliced or MPS-shared GPUs. A cluster that runs shared GPU pools can label those nodes with a distinct accelerator value, so a task asking for "T4 shared" lands only on the shared pool and a task asking for "T4" keeps getting a whole card.
What changes were proposed in this pull request?
GetK8sPluginConfignow merges the configuredaccelerator-devicesentries over the built-in defaults, with every key upper-cased so it matches the upper-cased lookup. The defaults are upper-cased once at package init. A configured entry wins over a default with the same name, and anyone who doesn't set the map gets the same defaults as before.-shared:T4 SHARED,L4 SHARED,A10G SHARED,A10 SHAREDandL40S SHARED.Because default keys are now upper-cased too, the existing
L40sentry resolves for the first time. Before this change it could never match the upper-cased lookup, so a task asking forL40sgot the raw string as its node label instead ofnvidia-l40s.How was this patch tested?
New unit tests in
flytek8s/utils_test.gocover a partial override keeping the defaults, lowercase configured keys resolving, the mixed-caseL40sdefault resolving, and each shared name resolving to its shared label.TestLoadConfignow loads anaccelerator-devicesentry through the real viper loader and checks that the override applies and the defaults survive. That assertion fails without the fix.go test ./flyteplugins/...passes, and golangci-lint reports no new issues on the changed packages.Labels
Check all the applicable boxes
🤖 Generated with Claude Code
main