#827 Add optional device_name to GPU() for MIG hybrid partition support Pull request opened by davidmirror-ops Summary • Adds an optional

device_name

parameter to

GPU()

and the

Device

dataclass, allowing users to specify the Kubernetes resource key used in the pod spec • Validated to start with

<http://nvidia.com/|nvidia.com/>

— the required prefix for NVIDIA device plugin resources •

pod_spec_from_resources()

now reads

device.device_name

and uses it as the k8s GPU resource key, overriding the hardcoded default

"<http://nvidia.com/gpu|nvidia.com/gpu>"

•

resources_serde.py

is unchanged — the device_name is a pod spec concern, not a backend protobuf concern Motivation When using NVIDIA MIG with mixed partition strategy, the Kubernetes resource name is partition-specific (e.g. `nvidia.com/mig-1g.5gb`) (ref) rather than the generic

<http://nvidia.com/gpu|nvidia.com/gpu>

. The key was previously hardcoded in

pod_spec_from_resources()

, making it impossible for users to request the correct MIG slice unless using Pod templates. This change intends to make the device name customizable. One other option was to automatically build the device name from the partition profile the user request in

Resources

but that's only useful if the

mixed

strategy is configured in the GPU nodes, for the

single

strategy the operator labels all the devices with the default

<http://nvidia.com/gpu|nvidia.com/gpu>

. Introducing an additional arg to declare the partitioning strategy exceeded the scope of this change. Usage env = flyte.TaskEnvironment( name="mig_training", resources=flyte.Resources( gpu=flyte.GPU( device="A100", quantity=1, partition="1g.5gb", device_name="nvidia.com/mig-1g.5gb", # k8s resource key for this MIG profile ) ), ) Without

device_name

, behaviour is unchanged — defaults to

"<http://nvidia.com/gpu|nvidia.com/gpu>"

. Test plan •

test_gpu_device_name_valid

— accepts valid

<http://nvidia.com/*|nvidia.com/*>

device_name •

test_gpu_device_name_invalid_prefix

/

test_gpu_device_name_invalid_no_prefix

— rejects bad values •

test_gpu_device_name_mig_profiles

— all A100 MIG profiles accepted (parametrized) •

test_device_dataclass_device_name_validation

— validation in

Device.__post_init__

•

test_pod_spec_uses_custom_device_name_from_requests/limits

— pod spec uses the override key •

test_pod_spec_default_gpu_key_when_no_device_name

— backward-compatible fallback •

test_pod_spec_requests_device_name_takes_precedence_over_limits

•

test_pod_spec_custom_k8s_key_overridden_by_device_name

• All 108 existing tests continue to pass 🤖 Generated with Claude Code flyteorg/flyte-sdk