[BUG] support setting extended resources for array node map tasks

[pvditt]

Tracking issue

fixes: https://linear.app/unionai/issue/COR-1565/problem-with-array-node-map-tasks-and-gpu-accelerators-pod-toleration

Why are the changes needed?

gpu_accelerator is not set for tasks that are subtasks of ArrayNodes

(also impacts legacy map tasks - if we want to fix that too)

What changes were proposed in this pull request?

Return the underlying python function task's extended resources to be set when serializing an array node map task

How was this patch tested?

• added unit test

Setup process

Confirmed that the GPU Accelerator was set under the task's extended resources

from flytekit import Resources, task, workflow, map_task
from flytekit.extras.accelerators import T4, V100, L4, GPUAccelerator


@task(
    limits=Resources(gpu="1"),
    interruptible=False,
    accelerator=GPUAccelerator("nvidia-l4"),
)
def with_defined_gpu_accelerator_l4(input: str) -> str:
    return f"Hello {input}!"


@workflow
def wf() -> str:
    return with_defined_gpu_accelerator_l4(input="1")


@workflow
def wf1():
    map_task(with_defined_gpu_accelerator_l4)(input=["1", "2", "3"])

Generated flyteworkflow for the ArrayNode subtask

Tasks:
  resource_type:TASK  project:"flytesnacks"  domain:"development"  name:"tests.flytekit.integration.map_task_issue.map_with_defined_gpu_accelerator_l4_1_6b3bd0353da5de6e84d7982921ead2b3-arraynode"  version:"OLIadKy0uixtSn6C-3Y4FA":
    Container:
      Args:
        pyflyte-fast-execute
        --additional-distribution
        s3://my-s3-bucket/flytesnacks/development/TND7O55TRQU7DOG45ZVB2NN73I======/script_mode.tar.gz
        --dest-dir
        .
        --
        pyflyte-map-execute
        --inputs
        {{.input}}
        --output-prefix
        {{.outputPrefix}}
        --raw-output-data-prefix
        {{.rawOutputDataPrefix}}
        --checkpoint-path
        {{.checkpointOutputPrefix}}
        --prev-checkpoint
        {{.prevCheckpointPrefix}}
        --resolver
        flytekit.core.array_node_map_task.ArrayNodeMapTaskResolver
        --
        vars
        
        resolver
        flytekit.core.python_auto_container.default_task_resolver
        task-module
        tests.flytekit.integration.map_task_issue
        task-name
        with_defined_gpu_accelerator_l4_1
      Image:  cr.flyte.org/flyteorg/flytekit:py3.11-latest
      Resources:
        Limits:
          Name:   CPU
          Value:  500m
          Name:   MEMORY
          Value:  500Mi
          Name:   GPU
          Value:  1
        Requests:
          Name:   CPU
          Value:  500m
          Name:   MEMORY
          Value:  500Mi
          Name:   GPU
          Value:  1
    Custom:
      Min Success Ratio:  1
    Extended Resources:
      Gpu Accelerator:
        Device:  nvidia-l4
    Id:
      Domain:         development
      Name:           tests.flytekit.integration.map_task_issue.map_with_defined_gpu_accelerator_l4_1_6b3bd0353da5de6e84d7982921ead2b3-arraynode
      Project:        flytesnacks
      Resource Type:  TASK
      Version:        OLIadKy0uixtSn6C-3Y4FA
    Interface:
      Inputs:
        Variables:
          Input:
            Type:
              Collection Type:
                Simple:  STRING
      Outputs:
        Variables:
          o0:
            Type:
              Collection Type:
                Simple:  STRING
    Metadata:
      Interruptible:  false
      Retries:
      Runtime:
        Flavor:         python
        Type:           FLYTE_SDK
        Version:        1.12.0b7.dev4+g1c5e56358.d20240506
    Task Type Version:  1
    Type:               python-task

Screenshots

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Related PRs

Docs link

[pvditt]

Do we also want to fix this for legacy map tasks? Really simple fix, was just unsure what the procedure is with fixing bugs in legacy entities

[eapolinario]

Do we also want to fix this for legacy map tasks? Really simple fix, was just unsure what the procedure is with fixing bugs in legacy entities

No, let's leave legacy map tasks alone.