[Core feature] Convert List[Any] to a single pickle file

[Yicheng-Lu-llll]

TL;DR

Currently, flyte create N (size of the list) pickle files if the output type is List[Any]. This slows down serialization. it takes more than 15 mins to upload the pickles to s3 if the size of the list is 1000 (see Slack Thread)

This PR:

• converts the entire list into a single pickle file by default.
• takes batches of list items and uploads each batch as a separate pickle file if the user annotates the batch size(see below).

@task
def t1() -> Annotated[List[Any],  100]

Type

• Bug Fix
• Feature
• Plugin

Are all requirements met?

• Code completed
• Smoke tested
• Unit tests added
• Code documentation added
• Any pending items have an associated Issue

Complete description

Originally, in ListTransformer.to_literal, Items are converted one by one. It is time-consuming, especially for large lists.

flytekit/flytekit/core/type_engine.py

Line 977 in 8ccb5dd

lit_list = [TypeEngine.to_literal(ctx, x, t, expected.collection_type) for x in python_val] # type: ignore

This PR checks if python_type is list[flytePickle]. If so, batch the list of items and call FlytePickleTransformer.to_literal.

Measure the performance

Test file:

from typing import List, Dict, Any
from flytekit import task, Resources, workflow

@task(
    limits=Resources(mem="4Gi",cpu="1"),
    disable_deck=True,
)
def test() -> List[Any]:
    return  [{"a": {0: "foo"}}] * 10000

@workflow
def wf():
    test()

if __name__ == "__main__":
    wf()

# yichenglu/flytekit:mytest5 contains this PR
# pyflyte run  --image yichenglu/flytekit:mytest5 --remote ./test.py  wf
# cr.flyte.org/flyteorg/flytekit:py3.10-1.4.0b1 is the default image
# pyflyte run  --image cr.flyte.org/flyteorg/flytekit:py3.10-1.4.0b1 --remote ./test.py  wf

In an 4VCPU, 16G RAM EC2 instance, measuring the performance in flyte cluster using the above test file and the commands：

• Prior to these changes, , converting a list with type List[Any] with 10000 items needs 157 seconds.
• After this PR, convertinga list with type List[Dict[str, Any]] with 10000 items only needs 7 seconds.
  
  

backward compatibility:

• case1: old task's cached output -> new task

from typing import List, Dict, Any, Annotated
from flytekit import task, Resources, workflow
@task(
    limits=Resources(mem="4Gi",cpu="1"),
    disable_deck=True,
    container_image="{{.image.oldImage}}",
    cache=True, cache_version="1.0"   
)
def task0() -> List[List[Any]]:
    return   [ ["foo","foo","foo"] ] * 2
@task(
    limits=Resources(mem="4Gi",cpu="1"),
    disable_deck=True,
    container_image="{{.image.prImage}}",
)
def task1(data: List[List[Any]]) -> List[List[Any]]:
    print(data)
    return  data
@workflow
def wf():
    data = task0()
    task1(data=data)
# pyflyte run --image prImage="yichenglu/flytekit:mytest4" --image oldImage="cr.flyte.org/flyteorg/flytekit:py3.10-1.4.1" --remote ./test.py  wf

Backward compatibility is preserved in this case. Can see pod log ton double check.

• case2: new task(image include this PR) -> old task(image without this PR)

from typing import List, Dict, Any, Annotated
from flytekit import task, Resources, workflow
@task(
    limits=Resources(mem="4Gi",cpu="1"),
    disable_deck=True,
    container_image="{{.image.prImage}}",
)
def task0() -> List[List[Any]]:
    return   [ ["foo","foo","foo"] ] * 2
@task(
    limits=Resources(mem="4Gi",cpu="1"),
    disable_deck=True,
    container_image="{{.image.oldImage}}"
)
def task1(data: List[List[Any]]) -> List[List[Any]]:
    print(data)
    return  data
@workflow
def wf():
    data = task0()
    task1(data=data)
# pyflyte run --image prImage="yichenglu/flytekit:mytest4" --image oldImage="cr.flyte.org/flyteorg/flytekit:py3.10-1.4.1" --remote ./test.py  wf

[[['foo', 'foo', 'foo']], [['foo', 'foo', 'foo']]] will be printed if see the pod's log
This occurs because in Task0, the outputs are batched, while in Task 1, which does not include this PR, the results are not expanded.

While slightly impacting backward compatibility, it is justifiable to ask users to update their images for two primary reasons:

• Before the implementation of this PR, users faced lengthy wait times for tasks involving pickling to be completed.
• Only tasks with the type list[any] will be affected. In most scenarios, users can still utilize the old images or cached outputs.

Tracking Issue

flyteorg/flyte#3207

Follow-up issue

NA

[welcome]

Thank you for opening this pull request! 🙌

These tips will help get your PR across the finish line:

• Most of the repos have a PR template; if not, fill it out to the best of your knowledge.
• Sign off your commits (Reference: DCO Guide).

[codecov]

Codecov Report

Merging #1535 (ea7e1aa) into master (ee2714f) will increase coverage by 0.29%.
The diff coverage is 84.72%.

@@            Coverage Diff             @@
##           master    #1535      +/-   ##
==========================================
+ Coverage   69.65%   69.95%   +0.29%     
==========================================
  Files         315      319       +4     
  Lines       29093    29585     +492     
  Branches     2761     5335    +2574     
==========================================
+ Hits        20265    20696     +431     
- Misses       8314     8370      +56     
- Partials      514      519       +5     

Impacted Files	Coverage Δ
flytekit/types/pickle/__init__.py	0.00% <0.00%> (ø)
flytekit/types/pickle/pickle.py	45.16% <33.33%> (-1.27%)	⬇️
flytekit/core/type_engine.py	61.18% <80.00%> (+0.43%)	⬆️
flytekit/core/promise.py	52.70% <100.00%> (+0.17%)	⬆️
tests/flytekit/unit/core/test_flyte_pickle.py	91.80% <100.00%> (+0.42%)	⬆️
tests/flytekit/unit/core/test_promise.py	89.88% <100.00%> (+0.35%)	⬆️
tests/flytekit/unit/core/test_type_engine.py	98.62% <100.00%> (+0.03%)	⬆️

... and 31 files with indirect coverage changes

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[pingsutw]

LGTM

[pingsutw]

Prior to these changes, , converting a list with type List[Dict[str, Any]] with 1000 items needs 271.448063737 seconds.
After this PR, convertinga list with type List[Dict[str, Any]] with 1000 items only needs 2.990736610998283 seconds.

Amazing, thanks for testing it

[pingsutw]

Thanks, Have you tried to run it with cache enabled in the sandbox?

[wild-endeavor]

thanks! can we also do a quick end to end test please on sandbox if you haven't already

[Yicheng-Lu-llll]

thanks! can we also do a quick end to end test please on sandbox if you haven't already

Certainly! May I know which test you are referring to? I have tested the following two cases(shown in the PR description):

• old task's cached output -> new task
• new task(image include this PR) -> old task(image without this PR)

Do you mean to add cases that demonstrate successful execution?

[welcome]

Congrats on merging your first pull request! 🎉

[elibixby]

I think with the Annotation functionality I added in #1615 you can accomplish this much more clearly and cleanly without bringing the functionality into TypeEngine, where it adds a lot of complexity, and doesn't even really provide this feature in many cases (for example what if I want to batch types which are already registered in the TypeEngine).

BatchedList = Annotated[List[Any], FlytePickleTransformer()]

@task
def my_task(...) -> BatchedList:
   ...

Or even using a custom transformer with a fixed batch size, and explicit type;

class BatchedList:

    def __class_getitem__(self, type: Type, batch_size: int):
         return  Annotated[List[type],  MyBatchingTransformer(t=type, batch_size=int)]


@task
def my_batching_task(...) -> BatchedList[int, 5E10]:
   ...

@Yicheng-Lu-llll Does this meet your needs?

If so perhaps we can revert this as it adds significant complexity. @wild-endeavor
I'm sorry I didn't catch this before it went in, the Annotation override feature is newer and not well documented.

[Yicheng-Lu-llll]

Hi @elibixby, Thank you for commenting! We can further discuss about this.

I think with the Annotation functionality I added in #1615 you can accomplish this much more clearly and cleanly without bringing the functionality into TypeEngine, where it adds a lot of complexity, and doesn't even really provide this feature in many cases (for example what if I want to batch types which are already registered in the TypeEngine).

1. The primary goal of this PR is to address the performance issues. Some users find it is extremely slow if they have large list with type Any. This issue is not apparent with other types. Hence, this PR focuses on batching the List[Any] type to enhance performance, while leaving all other types unaffected.

2. In my initial implementation,I also change the behavior of get_transformer, However, as you pointed out, , this approach will lose the ability to change the batch size. From a user's perspective, creating a custom transformer to support batch size can be overly complex. Also, adding batch size to FlytePickleTransformer would introduce equivalent complexity.

3. If we do not care about the batch size, I do strongly agree your approach is much better. But even we choose to disregard batch size and proceed with your suggestions, more modifications will be required beyond add annotation option for serialization #1615. For example, We need to make sure translate_inputs_to_literals and TypeEngine are consistent.

[elibixby]

2. In my initial implementation,I also change the behavior of get_transformer, However, as you pointed out, , this approach will lose the ability to change the batch size. From a user's perspective, creating a custom transformer to support batch size can be overly complex. Also, adding batch size to FlytePickleTransformer would introduce equivalent complexity.

I agree if we were to use annotations, we should provide a convenience wrapper for batch size to users, so they don't need to implement their own type transformer. I disagree that it would introduce equivalent complexity to this PR. It would be a completely isolated subclass of FlytePickleTransformer. (as a side note I am working on a PR right now which alters the way FlytePickleTransformer is used to be a default in get_transformer rather than a special case only for function signatures as implmented in #1565 , I've pulled this into a separate change #1661 to make it clearer what I propose here)

3. If we do not care about the batch size, I do strongly agree your approach is much better. But even we choose to disregard batch size and proceed with your suggestions, more modifications will be required beyond add annotation option for serialization #1615. For example, We need to make sure translate_inputs_to_literals and TypeEngine are consistent.

This is something I've discussed pretty extensively with @wild-endeavor . I think the preferred solution here is to make it so that adding a TypeTransformer by itself (either through registration or annotation), is sufficient to fully the support the type in Flyte.

The main thrust of this is in supporting pyflyte run (as you point out), but it could fairly easily be supported by adding a from_json to the TypeTransformer base class which pyflyte uses to unmarshal yaml data.

Would you be open to a solution that addresses both of these issues through an Annotated type as I've just described?

[Yicheng-Lu-llll]

I disagree that it would introduce equivalent complexity to this PR. It would be a completely isolated subclass of FlytePickleTransformer

If no batch size is specified, when a list comes in, the FlytePickleTransformer can simply transform the entire list into a single pickle file. However, with a batch size, things may be different. For example, we can split the list into three, and then convert them into three separate pickle files. But then, how to store and handle it as a literal seems to be a problem. Because it appears we have to change the literal type (from blob to LiteralCollection), otherwise we can not simply deal with the three separate pickle files. This is why I said "may introduce equivalent complexity to the current FlytePickleTransformer". It would make things easier to handle the batching process within the ListTransformer, which naturally supports the LiteralCollection.

Would you be open to a solution that addresses both of these issues through an Annotated type as I've just described?

Overall make sense to me. But I am not sure how will you deal with the above concern.