pod_spec_from_resources

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

src/flyte/_resources.py

k8s_gpu_resource_key: str = "<http://nvidia.com/gpu|nvidia.com/gpu>"

k8s_gpu_resource_key: Optional[str] = None

FLYTE_K8S_GPU_RESOURCE_KEY

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

tests/flyte/test_pod_spec_gpu_key.py

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

Summary

Make the Kubernetes GPU resource key used by pod_spec_from_resources configurable via an environment variable and still allow explicit overrides. Add unit tests that verify default behavior, environment-variable override, and explicit override precedence.

Why

Currently pod_spec_from_resources hardcodes the GPU resource key in its signature as "nvidia.com/gpu". Some users run on clusters that expose GPU resources under a different resource name (for example, custom device plugins). We want to allow changing that key globally without requiring callers to construct PodTemplate objects everywhere.

What to change

1. src/flyte/_resources.py

• Add an import for os at the top of the file.

• Change the pod_spec_from_resources signature to accept k8s_gpu_resource_key: Optional[str] = None instead of the hardcoded default.

• Inside pod_spec_from_resources resolve the effective GPU resource key in this order:

1. explicit k8s_gpu_resource_key argument if not None

2. environment variable FLYTE_K8S_GPU_RESOURCE_KEY if set

3. fallback to string "nvidia.com/gpu"

• Add a short docstring to pod_spec_from_resources documenting the resolution order.

• Keep the rest of the function behavior unchanged.

Exact code changes (illustrative):

• At the top of the file, add:

import os

• Change the function signature from:

def pod_spec_from_resources(..., k8s_gpu_resource_key: str = "nvidia.com/gpu") -> "V1PodSpec":

to:

def pod_spec_from_resources(..., k8s_gpu_resource_key: Optional[str] = None) -> "V1PodSpec":

• Immediately inside the function, add resolution logic:

if k8s_gpu_resource_key is None:

k8s_gpu_resource_key = os.environ.get("FLYTE_K8S_GPU_RESOURCE_KEY", "nvidia.com/gpu")

• Add a short docstring describing the resolution order.

Notes:

• Only change the signature and add the resolution step; do not change the resources_map usage or other logic.

• This keeps backward compatibility: callers that pass an explicit value continue to work; callers that relied on the default behavior still get the same default unless they set the env var.

1. Tests: add new test file

Create tests/flyte/test_pod_spec_gpu_key.py with three pytest tests:

• test_pod_spec_from_resources_default_key

• Ensure FLYTE_K8S_GPU_RESOURCE_KEY is not set (monkeypatch.delenv(...)).

• Call pod_spec_from_resources(requests=Resources(cpu="1", memory="1Gi", gpu=1)).

• Assert the returned V1PodSpec's primary container resources.requests includes key "nvidia.com/gpu" with value 1.

• test_pod_spec_from_resources_env_override

• Set monkeypatch.setenv("FLYTE_K8S_GPU_RESOURCE_KEY", "custom.com/gpu").

• Call pod_spec_from_resources with same Resources.

• Assert requests includes "custom.com/gpu" with value 1.

• test_pod_spec_from_resources_explicit_overrides_env

• Set env var to "custom.com/gpu" via monkeypatch.setenv.

• Call pod_spec_from_resources(..., k8s_gpu_resource_key="explicit.com/gpu").

• Assert the requests map contains "explicit.com/gpu" and not the env value.

Test implementation notes

• Use monkeypatch fixture to set/del environment variables.

• Inspect the returned pod spec via pod.containers[0].resources.requests.

• Be defensive: if requests is a dict-like object ensure assertions coerce keys/values (e.g., use dict(requests) when necessary).

Run & validation

• Run pytest (or the repo's full test suite) to ensure tests pass.

Deliverable

Open a draft pull request against flyteorg/flyte-sdk with:

• src/flyte/_resources.py updated as described.

• New test file tests/flyte/test_pod_spec_gpu_key.py containing the three tests.

• A concise PR description explaining purpose, backwards compatibility, and how to use the env var (FLYTE_K8S_GPU_RESOURCE_KEY).

If you want I can also:

• Target a specific branch (provide branch name) instead of letting the agent pick.

• Use a different env var name or a configuration module instead of an env var.

References

• Original hardcoded usage: src/flyte/_resources.py currently has k8s_gpu_resource_key: str = "nvidia.com/gpu" in pod_spec_from_resources. Repository: flyteorg/flyte-sdk

Summary

Make the Kubernetes GPU resource key used by pod_spec_from_resources configurable via an environment variable and still allow explicit overrides. Add unit tests that verify default behavior, environment-variable override, and explicit override precedence.

Why

Currently pod_spec_from_resources hardcodes the GPU resource key in its signature as "nvidia.com/gpu". Some users run on clusters that expose GPU resources under a different resource name (for example, custom device plugins). We want to allow changing that key globally without requiring callers to construct PodTemplate objects everywhere.

What to change

1. src/flyte/_resources.py

• Add an import for os at the top of the file.

• Change the pod_spec_from_resources signature to accept k8s_gpu_resource_key: Optional[str] = None instead of the hardcoded default.

• Inside pod_spec_from_resources resolve the effective GPU resource key in this order:

1. explicit k8s_gpu_resource_key argument if not None

2. environment variable FLYTE_K8S_GPU_RESOURCE_KEY if set

3. fallback to string "nvidia.com/gpu"

• Add a short docstring to pod_spec_from_resources documenting the resolution order.

• Keep the rest of the function behavior unchanged.

Exact code changes (illustrative):

• At the top of the file, add:

import os

• Change the function signature from:

def pod_spec_from_resources(..., k8s_gpu_resource_key: str = "nvidia.com/gpu") -> "V1PodSpec":

to:

def pod_spec_from_resources(..., k8s_gpu_resource_key: Optional[str] = None) -> "V1PodSpec":

• Immediately inside the function, add resolution logic:

if k8s_gpu_resource_key is None:

k8s_gpu_resource_key = os.environ.get("FLYTE_K8S_GPU_RESOURCE_KEY", "nvidia.com/gpu")

• A…