Preserve DRA resource claims when overriding task resources (#1324)

Signed-off-by: Alex Wu <c.alexwu@gmail.com>
Signed-off-by: Kevin Su <pingsutw@apache.org>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Kevin Su <pingsutw@apache.org>

Author: 3 people

src/flyte/_internal/runtime/task_serde.py

@@ -418,7 +418,11 @@ def _get_k8s_pod(primary_container: tasks_pb2.Container, pod_template: PodTempla
                 existing = container.resources or V1ResourceRequirements()
                 merged_limits = {**(existing.limits or {}), **limits}
                 merged_requests = {**(existing.requests or {}), **requests}
-                container.resources = V1ResourceRequirements(limits=merged_limits, requests=merged_requests)
+                container.resources = V1ResourceRequirements(
+                    limits=merged_limits,
+                    requests=merged_requests,
+                    claims=existing.claims,
+                )
 
             if primary_container.env is not None:
                 container.env = [V1EnvVar(name=e.key, value=e.value) for e in primary_container.env] + (

src/flyte/remote/_task.py

@@ -73,11 +73,14 @@ def _apply_overrides_to_primary_container(
         if proto_res is not None:
             requests = {_sanitize_resource_name(e): e.value for e in proto_res.requests}
             limits = {_sanitize_resource_name(e): e.value for e in proto_res.limits}
+            existing = target.get("resources") or {}
             rr: Dict[str, Any] = {}
             if requests:
                 rr["requests"] = requests
             if limits:
                 rr["limits"] = limits
+            if existing.get("claims"):
+                rr["claims"] = existing["claims"]
             target["resources"] = rr
 
     if env_vars:

tests/flyte/internal/runtime/test_task_serde.py

@@ -9,10 +9,13 @@
 from flyteidl2.core.security_pb2 import SecurityContext
 from flyteidl2.task import common_pb2, environment_pb2
 from kubernetes.client import (
+    CoreV1ResourceClaim,
     V1Container,
     V1EnvVar,
     V1LocalObjectReference,
+    V1PodResourceClaim,
     V1PodSpec,
+    V1ResourceRequirements,
 )
 
 import flyte
@@ -279,6 +282,59 @@ async def t2(a: int, b: str) -> str:
         get_proto_task(t2, context)
 
 
+def test_get_k8s_pod_preserves_dra_claims_with_task_resources():
+    """A DRA resource claim on the pod template's primary container must survive
+    even when the task also declares cpu/memory via Resources(...)."""
+    pod_template = PodTemplate(
+        pod_spec=V1PodSpec(
+            resource_claims=[V1PodResourceClaim(name="gpu", resource_claim_template_name="flyte-task-gpu-x2")],
+            containers=[
+                V1Container(
+                    name="primary",
+                    resources=V1ResourceRequirements(claims=[CoreV1ResourceClaim(name="gpu")]),
+                )
+            ],
+        ),
+    )
+
+    env = flyte.TaskEnvironment(
+        name="test_env",
+        image="python:3.10",
+        # Declaring cpu/memory is exactly what used to strip the .claims link.
+        resources=flyte.Resources(cpu="24", memory="490Gi"),
+        pod_template=pod_template,
+    )
+
+    @env.task(short_name="dra_task")
+    async def t1(a: int) -> int:
+        return a
+
+    context = SerializationContext(
+        project="test-project",
+        domain="test-domain",
+        version="test-version",
+        org="test-org",
+        input_path="/tmp/inputs",
+        output_path="/tmp/outputs",
+        image_cache=None,
+        code_bundle=None,
+        root_dir=pathlib.Path.cwd(),
+    )
+
+    from google.protobuf import json_format
+
+    k8s_pod = _get_k8s_pod(_get_urun_container(context, t1), pod_template)
+    pod_spec = json_format.MessageToDict(k8s_pod.pod_spec)
+    primary = next(c for c in pod_spec["containers"] if c["name"] == "primary")
+    resources = primary["resources"]
+
+    # The DRA claim link is preserved...
+    assert resources["claims"] == [{"name": "gpu"}]
+    # ...alongside the task-declared cpu/memory.
+    assert resources["requests"]["cpu"] == "24"
+    assert resources["requests"]["memory"] == "490Gi"
+
+
 @pytest.fixture(scope="module")
 def env_task_ctx():
     # Create a real task environment

tests/flyte/remote/test_task.py

@@ -289,6 +289,45 @@ def test_gpu_override_on_pod_template_sets_container_request(self):
         side = next(c for c in spec["containers"] if c["name"] == "sidecar")
         assert not side.get("resources")
 
+    def test_resource_override_on_pod_template_preserves_dra_claims(self):
+        # Regression: a resource override on a pod-template task must not drop a
+        # DRA claim (e.g. a GPU claimed through a ResourceClaimTemplate). The
+        # claim can only be expressed via the pod template; without it the pod
+        # schedules with no device.
+        from flyteidl2.task import task_definition_pb2
+        from google.protobuf import json_format
+
+        from flyte._pod import _PRIMARY_CONTAINER_NAME_FIELD
+
+        pb2 = task_definition_pb2.TaskDetails()
+        tmpl = pb2.spec.task_template
+        tmpl.config[_PRIMARY_CONTAINER_NAME_FIELD] = "primary"
+        spec = {
+            "resourceClaims": [{"name": "gpu", "resourceClaimTemplateName": "flyte-task-gpu-x2"}],
+            "containers": [
+                {
+                    "name": "primary",
+                    "image": "example:latest",
+                    "resources": {
+                        "requests": {"cpu": "1", "memory": "1Gi"},
+                        "claims": [{"name": "gpu"}],
+                    },
+                },
+            ],
+        }
+        json_format.ParseDict(spec, tmpl.k8s_pod.pod_spec)
+        td = TaskDetails(pb2)
+
+        out = td.override(resources=flyte.Resources(cpu="24", memory="490Gi"))
+        tmpl = out.pb2.spec.task_template
+
+        spec = json_format.MessageToDict(tmpl.k8s_pod.pod_spec)
+        primary = next(c for c in spec["containers"] if c["name"] == "primary")
+        resources = primary["resources"]
+        assert resources["claims"] == [{"name": "gpu"}], f"DRA claim dropped: {resources}"
+        assert resources["requests"]["cpu"] == "24"
+        assert resources["requests"]["memory"] == "490Gi"
+
     def test_env_override_on_pod_template_sets_container_env(self):
         # Regression: env_vars overrides on a pod-template task were silently
         # dropped — they were only applied to template.container, which a