[Core feature] map_task to support ContainerTask by machichima · Pull Request #3249 · flyteorg/flytekit

machichima · 2025-05-17T15:30:05Z

Tracking issue

flyteorg/flyte#6277

Why are the changes needed?

Currently, we cannot use ContainerTask in map task. This PR is trying to make this integration works.

What changes were proposed in this pull request?

Make ContainerTask work with map task in local and remote
Add unit test for running ContainerTask with map task locally

How was this patch tested?

unit test & example code below:

Execute following code locally and remotly should pass:

local: pyflyte run containertask_maptask.py wf
remote: pyflyte run --remote containertask_maptask.py wf

# containertask_maptask.py
import functools

from flytekit import ContainerTask, kwtypes, map_task, workflow

calculate_ellipse_area_python_template_style = ContainerTask(
    name="calculate_ellipse_area_python_template_style",
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs=kwtypes(a=float, b=float),
    outputs=kwtypes(area=float),
    image="ghcr.io/flyteorg/rawcontainers-python:v2",
    command=[
        "python",
        "calculate-ellipse-area.py",
        "{{.inputs.a}}",
        "{{.inputs.b}}",
        "/var/outputs",
    ],
)


@workflow
def nomap_wf(a: float = 0.5, b: float = 0.4) -> float:
    res = calculate_ellipse_area_python_template_style(a, b)
    return res


@workflow
def wf(a: list[float] = [3.0, 4.0, 5.0], b: float = 4.0) -> list[float]:
    partial_task = functools.partial(calculate_ellipse_area_python_template_style, b=b)
    res = map_task(partial_task)(a=a)
    return res

Setup process

Screenshots

Result in remote

Result in local

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

Summary by Bito

This pull request enhances map task functionality by integrating ContainerTask support for both local and remote execution. It includes significant refactoring of the execute method to return native Python types, improving usability and maintainability. Comprehensive unit tests validate this integration and confirm expected behavior, ensuring expected functionality in both environments.

Signed-off-by: machichima <nary12321@gmail.com>

…ontainer -e Signed-off-by: machichima <nary12321@gmail.com>

-e Signed-off-by: machichima <nary12321@gmail.com>

codecov · 2025-05-17T15:54:06Z

Codecov Report

Attention: Patch coverage is 52.94118% with 8 lines in your changes missing coverage. Please review.

Project coverage is 76.36%. Comparing base (3c6b61d) to head (2b9387d).

Files with missing lines	Patch %	Lines
flytekit/core/array_node_map_task.py	44.44%	3 Missing and 2 partials ⚠️
flytekit/core/container_task.py	57.14%	1 Missing and 2 partials ⚠️

❗ There is a different number of reports uploaded between BASE (3c6b61d) and HEAD (2b9387d). Click for more details.

HEAD has 40 uploads less than BASE

Flag BASE (3c6b61d) HEAD (2b9387d)

41 1

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #3249      +/-   ##
==========================================
- Coverage   83.35%   76.36%   -7.00%     
==========================================
  Files         347      215     -132     
  Lines       28791    22528    -6263     
  Branches     2960     2966       +6     
==========================================
- Hits        23999    17203    -6796     
- Misses       3956     4507     +551     
+ Partials      836      818      -18

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

machichima · 2025-05-20T13:55:44Z

tests/flytekit/unit/core/test_array_node_map_task.py

            monkeypatch.undo()
+
+
+@pytest.mark.skipif(


Modification in this file before this line is simply caused by formatter

-e Signed-off-by: machichima <nary12321@gmail.com>

machichima · 2025-05-22T14:09:58Z

flytekit/tools/translator.py

            # tasks that rely on user code defined in the container. This should be encapsulated by the auto container
            # parent class
-            container._args = prefix_with_fast_execute(settings, container.args)
+            container._args = prefix_with_fast_execute(settings, container.args or [])


args for ContainerTask is None, adding this to prevent error

wild-endeavor · 2025-05-22T16:30:23Z

flytekit/core/array_node_map_task.py

            try:
                o = self._run_task.execute(**single_instance_inputs)
+                # For Container task, it will return the LiteralMap. We need to convert it to native
+                # type here.


@machichima can you add comment please that this is only for local execution for container tasks? this code doesn't run for backend cluster runs of container task.

I make ContainerTask return Python type instead of LiteralMap. So this part is not needed anymore.
Thanks!

…ontainer -e Signed-off-by: machichima <nary12321@gmail.com>

flytekit/core/base_task.py

-e Signed-off-by: machichima <nary12321@gmail.com>

machichima · 2025-05-25T14:57:43Z

I think let ConatinerTask return python type break tests for plugins. Will have a look

Update: I think it's not my problem. Merging the master branch and fixed

…ontainer -e Signed-off-by: machichima <nary12321@gmail.com>

flyte-bot · 2025-06-03T23:40:02Z

Bito Review Skipped - No Changes Detected

Bito didn't review this pull request because we did not detect any changes in the pull request to review.

flyte-bot · 2025-06-03T23:40:29Z

Bito Review Skipped - No Changes Detected

Bito didn't review this pull request because we did not detect any changes in the pull request to review.

pingsutw · 2025-06-08T06:31:37Z

flytekit/core/array_node_map_task.py

+        if isinstance(self._run_task, ContainerTask):
+            return self.python_function_task.get_container(settings)


nit:
We can modify self.prepare_target, do nothing for container task

Fixed! Thanks

…ontainer -e Signed-off-by: machichima <nary12321@gmail.com>

-e Signed-off-by: machichima <nary12321@gmail.com>

Future-Outlier

Can we add an intergration test to make sure it works well in remote execution?

Future-Outlier · 2025-06-09T10:47:23Z

flytekit/core/array_node_map_task.py

+        if isinstance(self._run_task, ContainerTask):
+            yield
+            return
+


~~why not just return directly?~~
update: we need to return a generator, please ignore my comment.

-e Signed-off-by: machichima <nary12321@gmail.com>

machichima · 2025-06-10T01:46:45Z

Can we add an intergration test to make sure it works well in remote execution?

Added! Thanks

machichima and others added 7 commits March 29, 2025 19:45

feat: update to support container task

281bd66

Signed-off-by: machichima <nary12321@gmail.com>

Merge branch 'flyteorg:master' into 6277-map-container

1d16454

Merge branch 'master' of github.com:flyteorg/flytekit into 6277-map-c…

137ab21

…ontainer -e Signed-off-by: machichima <nary12321@gmail.com>

feat: enable using container task in map task

70b8bef

-e Signed-off-by: machichima <nary12321@gmail.com>

feat: enable running map+container task in local

0d34130

-e Signed-off-by: machichima <nary12321@gmail.com>

test: for running map+container task locally

e96b480

-e Signed-off-by: machichima <nary12321@gmail.com>

style: lint

c53a64e

-e Signed-off-by: machichima <nary12321@gmail.com>

machichima requested review from Future-Outlier, cosmicBboy, davidmirror-ops, eapolinario, kumare3, pingsutw, samhita-alla and wild-endeavor as code owners May 17, 2025 15:30

machichima changed the title ~~6277 map container~~ [Core feature] map_task to support ContainerTask May 17, 2025

machichima commented May 20, 2025

View reviewed changes

fix: using Union intead of "|"

5c95774

-e Signed-off-by: machichima <nary12321@gmail.com>

machichima commented May 22, 2025

View reviewed changes

wild-endeavor reviewed May 22, 2025

View reviewed changes

Merge branch 'master' of github.com:flyteorg/flytekit into 6277-map-c…

6aae613

…ontainer -e Signed-off-by: machichima <nary12321@gmail.com>

machichima commented May 24, 2025

View reviewed changes

flytekit/core/base_task.py Outdated Show resolved Hide resolved

machichima commented May 24, 2025

View reviewed changes

flytekit/core/base_task.py Outdated Show resolved Hide resolved

docs: make comment more clear

0f6dd70

-e Signed-off-by: machichima <nary12321@gmail.com>

machichima force-pushed the 6277-map-container branch from 2c2bd66 to 0f6dd70 Compare May 24, 2025 13:17

feat: make container task return python type

4827e55

-e Signed-off-by: machichima <nary12321@gmail.com>

machichima requested a review from wild-endeavor May 25, 2025 09:35

Merge branch 'master' of github.com:flyteorg/flytekit into 6277-map-c…

006493a

…ontainer -e Signed-off-by: machichima <nary12321@gmail.com>

Merge branch 'master' of github.com:flyteorg/flytekit into 6277-map-c…

f15415b

…ontainer -e Signed-off-by: machichima <nary12321@gmail.com>

machichima force-pushed the 6277-map-container branch from a63386e to f15415b Compare June 3, 2025 21:49

pingsutw reviewed Jun 8, 2025

View reviewed changes

machichima added 2 commits June 8, 2025 20:41

Merge branch 'master' of github.com:flyteorg/flytekit into 6277-map-c…

0126bf3

…ontainer -e Signed-off-by: machichima <nary12321@gmail.com>

refactor: prepare_target do nothing for ContainerTask

d8b4362

-e Signed-off-by: machichima <nary12321@gmail.com>

Future-Outlier reviewed Jun 9, 2025

View reviewed changes

test: add container task map exec integration test

2b9387d

-e Signed-off-by: machichima <nary12321@gmail.com>

machichima requested a review from Future-Outlier June 10, 2025 01:46

		if isinstance(self._run_task, ContainerTask):
		return self.python_function_task.get_container(settings)

		monkeypatch.undo()


		@pytest.mark.skipif(

Conversation

machichima commented May 17, 2025 • edited by flyte-bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Tracking issue

Why are the changes needed?

What changes were proposed in this pull request?

How was this patch tested?

Setup process

Screenshots

Check all the applicable boxes

Related PRs

Docs link

Summary by Bito

Uh oh!

codecov bot commented May 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

machichima May 20, 2025

Choose a reason for hiding this comment

Uh oh!

machichima May 22, 2025

Choose a reason for hiding this comment

Uh oh!

wild-endeavor May 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

machichima May 24, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

machichima commented May 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

flyte-bot commented Jun 3, 2025

Uh oh!

flyte-bot commented Jun 3, 2025

Uh oh!

pingsutw Jun 8, 2025

Choose a reason for hiding this comment

Uh oh!

machichima Jun 8, 2025

Choose a reason for hiding this comment

Uh oh!

Future-Outlier left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Future-Outlier Jun 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

machichima commented Jun 10, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

machichima commented May 17, 2025 •

edited by flyte-bot

Loading

codecov bot commented May 17, 2025 •

edited

Loading

wild-endeavor May 22, 2025 •

edited

Loading

machichima commented May 25, 2025 •

edited

Loading

Future-Outlier left a comment •

edited

Loading

Future-Outlier Jun 9, 2025 •

edited

Loading