#395 Add k8s_resource_key parameter to GPU() for explicit Kubernetes resource naming Pull request opened by Copilot on <!date^1765304973^{date_short}|2025-12-09T18:29:33Z> Allows users to specify exact Kubernetes GPU resource names without requiring PodTemplates. Supports MIG partition auto-synthesis. ## Changes • Device dataclass: Added optional

k8s_resource_key

field • GPU() function: Accepts

k8s_resource_key

parameter for explicit resource naming • Resource key resolution: Priority order - explicit key > synthesized MIG key from partition > default • pod_spec_from_resources: Uses device's k8s_resource_key when determining K8s resource name ## Usage # Explicit resource key gpu = flyte.GPU(device="A100", quantity=1, k8s_resource_key="nvidia.com/mig-1g.10gb") # Auto-synthesized from partition gpu = flyte.GPU(device="A100", quantity=1, partition="1g.5gb") # → Uses "nvidia.com/mig-1g.5gb" # Default behavior (unchanged) gpu = flyte.GPU(device="T4", quantity=1) # → Uses "nvidia.com/gpu" ## Implementation Helper function

_get_gpu_resource_key()

implements priority resolution: 1. Explicit

k8s_resource_key

parameter 2. Synthesized

<http://nvidia.com/mig-{partition}|nvidia.com/mig-{partition}>

when partition present 3. Default

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

or custom default from

pod_spec_from_resources(k8s_gpu_resource_key=...)

Backward compatible - all existing behavior preserved when new parameter not used. Warning Firewall rules blocked me from connecting to one or more addresses (expand for details) #### I tried to connect to the following addresses, but was blocked by firewall rules: •

astral.sh

• Triggering command:

/usr/bin/curl curl -LsSf REDACTED

(dns block) If you need me to access, download, or install something from one of these locations, you can either: • Configure Actions setup steps to set up my environment, which run before the firewall is enabled • Add the appropriate URLs or hosts to the custom allowlist in this repository's Copilot coding agent settings (admins only) Original prompt Summary Allow users to specify the exact Kubernetes resource name for GPUs when they call flyte.GPU(...). The Device produced by GPU(...) should carry an optional k8s_resource_key and serialization code should prefer that key (or synthesize a MIG key from partition) when producing PodSpecs for both pod_spec_from_resources helper and TaskEnvironment/App serialization path. This avoids requiring users to use PodTemplate just to change the GPU device resource name. Behavioral goals • Backwards compatible: default behavior remains the same unless users pass the new argument or have a partition present. • Per-GPU control: users can call flyte.GPU(..., k8s_resource_key="nvidia.com/mig-1g.10gb") to request a specific k8s resource key. • Partition-aware: when a Device has a partition (e.g., "1g.10gb") and no explicit k8s_resource_key, synthesize a MIG-style resource name: "nvidia.com/mig-" (e.g., "nvidia.com/mig-1g.10gb"). • Support both codepaths that call pod_spec_from_resources and the TaskEnvironment -> protobuf -> serialized pod path. Files to change 1. src/flyte/_resources.py • Add an import for os near the top (if not already present). • Modify the Device dataclass to include k8s_resource_key: str | None = None. • Change the GPU(...) factory signature to accept k8s_resource_key: Optional[str] = None and pass it into the returned Device. • Update pod_spec_from_resources to: • accept k8s_gpu_resource_key: Optional[str] = None (existing default allowed) and resolve it from env var FLYTE_K8S_GPU_RESOURCE_KEY if not passed. • when constructing the resources_map, for the gpu entry, instead of always using the single key, detect if the provided Resources.gpu is a Device with Device.k8s_resource_key or Device.partition and choose the appropriate key: • If Device.k8s_resource_key is set, use it. • Elif Device.partition is set, use f"nvidia.com/mig-{device.partition}". • Else use the resolved k8s_gpu_resource_key as before. Exact changes inside src/flyte/_resources.py (high-level): • Device dataclass quantity: int device_class: DeviceClass device: str | None = None partition: str | None = None • k8s_resource_key: str | None = None • GPU signature def GPU(device: GPUType, quantity: GPUQuantity, partition: A100Parts | A100_80GBParts | H200Parts | None = None, k8s_resource_key: str | None = None) -> Device: ... return Device(device=device, quantity=quantity, partition=partition, device_class="GPU", k8s_resource_key=k8s_resource_key) • pod_spec_from_resources: resolve effective k8s key and when building k8s_pod_resources, if resource.name == "gpu" and resources.gpu is a Device use device.k8s_resource_key / partition-derived key. 1. src/flyte/app/_runtime/app_serde.py • Modify _sanitize_resource_name to accept an optional 'app_resources: flyte.Resources | None = None' parameter (default None) and use it when resolving ResourceEntry.GPU: • If resource.name == tasks_pb2.Resources.ResourceName.GPU and app_resources is provided: • device = app_resources.get_device() • if device and getattr(device, "k8s_resource_key", None): return device.k8s_resource_key • elif device and device.partition: return f"nvidia.com/mig-{device.partition}" • else: fall back to env var FLYTE_K8S_GPU_RESOURCE_KEY or "nvidia.com/gpu" • Otherwise the existing behavior: return tasks_pb2.Resources.ResourceName.Name(resource.name).lower().replace("_", "-") • Update the two call sites within _serialized_pod_spec where _sanitize_resource_name(resource) is used to pass app_env.resources: limits[_sanitize_resource_name(resource, app_env.resources)] = resource.value requests[_sanitize_resource_name(resource, app_env.resources)] = resource.value Note: _serialized_pod_spec is the central serialization method that builds the pod spec from app_env and pod_template and is the right place to apply this mapping so TaskEnvironment cases work without PodTemplate. 1. Tests: add tests to validate the new behavior Create new file: tests/flyte/test_gpu_k8s_key_from_gpu_function.py Add tests: • test_device_k8s_key_on_sanitize_resource_name(monkeypatch): device = flyte.GPU(device="A100 80G", quantity=1, partition="1g.10gb", k8s_resource_key="nvidia.com/mig-1g.10gb") resources = flyte.Resources(gpu=device) entry = tasks_pb2.Resources.ResourceEntry(name=tasks_pb2.Resources.ResourceName.GPU, value="1") assert app_serde._sanitize_resource_name(entry, resources) == "nvidia.com/mig-1g.10gb" • test_partition_based_mig_key_when_no_explicit_k8s_key(monkeypatch): device = flyte.GPU(device="A100 80G", quantity=1, partition="1g.10gb") resources = flyte.Resources(gpu=device) entry = tasks_pb2.Resources.ResourceEntry(name=tasks_pb2.Resources.ResourceName.GPU, value="1") assert app_serde._sanitize_resource_name(entry, resources) == "nvidia.com/mig-1g.10gb" • test_task_environment_serialization_uses_device_k8s_key(monkeypatch): # Create... This pull request was created as a result of the following prompt from Copilot chat.

Summary

Allow users to specify the exact Kubernetes resource name for GPUs when they call flyte.GPU(...). The Device produced by GPU(...) should carry an optional k8s_resource_key and serialization code should prefer that key (or synthesize a MIG key from partition) when producing PodSpecs for both pod_spec_from_resources helper and TaskEnvironment/App serialization path. This avoids requiring users to use PodTemplate just to change the GPU device resource name.

>… flyteorg/flyte-sdk