Offload literals #2872

eapolinario · 2024-10-28T19:30:41Z

Tracking issue

Why are the changes needed?

We need to be able to offload literals from flytekit.

What changes were proposed in this pull request?

This PR adds support for offloading of literals gated by the environment variables _F_L_MIN_SIZE_MB and _F_L_MAX_SIZE_MB, the minimal literal size for offloading and the max size, respectively.

Since we need to use the output name (i.e. "o0", "o1", etc) in the name of the offloaded literal, we use the entrypoint (since it's there that we have natural access to the metadata bucket) to write the offloaded literals to files named "o0_offloaded_metadata.pb", "o1_offloaded_metadata.pb", etc.

How hashing is handled

We need a stable representation of offloaded literals for caching purposes. In golang we use jsonpb.Marshaler + objecthash to generate a stable hash for a proto message, but unfortunately there's no readily available equivalent in python. Instead, we rely on the stable json representation of the proto Literal and use SHA256 to produce the hash.

In case a task makes use of HashMethod, then we use that hash in the offloaded literal.

How was this patch tested?

Unit tests and testing locally in sandbox.

Setup process

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

Signed-off-by: Eduardo Apolinario <[email protected]>

eapolinario · 2024-10-28T19:31:11Z

flytekit/bin/entrypoint.py

+                        offloaded_metadata=_literal_models.LiteralOffloadedMetadata(
+                            uri=f"{ctx.user_space_params.output_metadata_prefix}/{offloaded_filename}",
+                            size_bytes=lit.ByteSize(),
+                            # TODO: do I have to set the inferred literal type?


@pmahindrakar-oss , in what conditions we have to set this?

This is needed when we have launcplan execution created by propeller
eg : here
https://demo.hosted.unionai.cloud/console/projects/flytesnacks/domains/development/executions/akzthr8gxbdvpnq6f5lr/nodes

[UserError] failed to launch workflow, caused by: rpc error: code = InvalidArgument desc = invalid input input wrong type. Expected collection_type:{simple:STRING}, but got collection_type:{}

Code link in my comments below https://github.com/flyteorg/flytekit/pull/2872/files#r1819666906

I commented in about separating those discussions and how we use the task typed interface to set inferred_type with the appropriate literal type.

flytekit/bin/entrypoint.py

Signed-off-by: Eduardo Apolinario <[email protected]>

pmahindrakar-oss · 2024-10-28T19:51:22Z

flytekit/bin/entrypoint.py

+                        offloaded_metadata=_literal_models.LiteralOffloadedMetadata(
+                            uri=f"{ctx.user_space_params.output_metadata_prefix}/{offloaded_filename}",
+                            size_bytes=lit.ByteSize(),
+                            # TODO: do I have to set the inferred literal type?


We do need to add inferred type, since during consumption we validate the input type

https://github.com/unionai/flyte/blob/5e9ec71a3b55003955642aeeaaf6dbeb43b3a279/flytepropeller/pkg/compiler/validators/utils.go#L271-L272

https://github.com/unionai/flyte/blob/2c070c301ccb37c5a336dd7022a1d9846d9bb173/flyteadmin/pkg/manager/impl/validation/execution_validator.go#L100

We're trying to get rid of LiteralTypeForLiteral in flyteorg/flyte#5909. Still pending a review.

Wont we still need this in the new isInstance form https://github.com/flyteorg/flyte/pull/5909/files#diff-af9c7a5cd03e0e860c155861edcf0b96863e8781526f51d1b13c7f9278ffa72fR415

if we get rid of the literaltypeforliteral function, and we don't need inferredtype, we should deprecate the field in the idl. don't want to keep it around and have future us wonder if it's necessary.

Let's separate those two discussions. We might be able to remove LiteralTypeForLiteral, but that shouldn't block us from getting the actual literal type from the typed interface in flytekit and use it here.

Also, if/when we end up removing LiteralTypeForLiteral we should use a different field to store the offloaded literal literal type and not use inferred in the name since it'll be the exact type.

pmahindrakar-oss · 2024-10-28T19:55:15Z

flytekit/bin/entrypoint.py

+            min_offloaded_size = int(os.environ.get("FK_L_MIN_SIZE_MB", "10")) * 1024 * 1024
+            max_offloaded_size = int(os.environ.get("FK_L_MAX_SIZE_MB", "1000")) * 1024 * 1024


Do we use the same names here that we use in propeller
https://github.com/unionai/flyte/blob/c10b55373e9dea2bcce493f978499648c3f89d8f/flytepropeller/pkg/controller/config/config.go#L204-L207

we'll need to update the injection of environment variables to match those names.

The need for short env vars is described flyteorg/flyte#5665.

Ok with keeping the current ones .

pmahindrakar-oss · 2024-10-28T19:55:31Z

flytekit/bin/entrypoint.py

+                if lit.ByteSize() >= min_offloaded_size:
+                    logger.debug(f"Literal {k} is too large to be inlined, offloading to metadata bucket")
+
+                    # TODO: hash calculation


pmahindrakar-oss · 2024-10-28T19:56:19Z

flytekit/interaction/string_literals.py

+    if lit.offloaded_metadata:
+        # TODO: load literal from offloaded literal?
+        return f"Offloaded literal metadata: {lit.offloaded_metadata}"


Is there an example we can try to see if we need to load the literal here

this is relevant for the pyflyte fetch command:

❯ pyflyte --config ~/.flyte/config-sandbox.yaml fetch flyte://v1/flytesnacks/development/asvskwn766f5v492pzgt/n0-0-n1/o Fetching data from flyte://v1/flytesnacks/development/asvskwn766f5v492pzgt/n0-0-n1/o... ╭───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮ │ { │ │ 'o0': [ │ │ 'Offloaded literal metadata: Flyte Serialized object (LiteralOffloadedMetadata):\n uri: s3://my-s3-bucket/metadata/propeller/flytesnacks-development- [...]\n size_bytes: 39936015', │ │ 'Offloaded literal metadata: Flyte Serialized object (LiteralOffloadedMetadata):\n uri: s3://my-s3-bucket/metadata/propeller/flytesnacks-development- [...]\n size_bytes: 39936015', │ ... │ 'Offloaded literal metadata: Flyte Serialized object (LiteralOffloadedMetadata):\n uri: s3://my-s3-bucket/metadata/propeller/flytesnacks-development- [...]\n size_bytes: 39936015' │ │ ] │ │ } │ ╰───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯

Signed-off-by: Eduardo Apolinario <[email protected]>

wild-endeavor · 2024-11-07T01:10:48Z

flytekit/bin/entrypoint.py

+                if lit.ByteSize() >= min_offloaded_size:
+                    logger.debug(f"Literal {k} is too large to be inlined, offloading to metadata bucket")
+
+                    # TODO: hash calculation


how will this hash interoperate with hashmethod as specified by the user? how does caching work again?

Since the hash calculation in entrypoint only happens after dispatch_execute (in other words, if a user defined a HashMethod then we can reuse it here). I just pushed an update that implements this idea.

wild-endeavor · 2024-11-07T01:12:06Z

flytekit/bin/entrypoint.py

+                        offloaded_metadata=_literal_models.LiteralOffloadedMetadata(
+                            uri=f"{ctx.user_space_params.output_metadata_prefix}/{offloaded_filename}",
+                            size_bytes=lit.ByteSize(),
+                            # TODO: do I have to set the inferred literal type?


if we get rid of the literaltypeforliteral function, and we don't need inferredtype, we should deprecate the field in the idl. don't want to keep it around and have future us wonder if it's necessary.

flytekit/bin/entrypoint.py

…aded literals Signed-off-by: Eduardo Apolinario <[email protected]>

Signed-off-by: Eduardo Apolinario <[email protected]>

codecov · 2024-11-07T12:39:59Z

Codecov Report

Attention: Patch coverage is 12.50000% with 21 lines in your changes missing coverage. Please review.

Project coverage is 51.45%. Comparing base (faee3da) to head (b3a1b0d).
Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
flytekit/utils/pbhash.py	0.00%	19 Missing ⚠️
flytekit/interaction/string_literals.py	0.00%	2 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (faee3da) and HEAD (b3a1b0d). Click for more details.

HEAD has 6 uploads less than BASE

Flag BASE (faee3da) HEAD (b3a1b0d)

8 2

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #2872       +/-   ##
===========================================
- Coverage   79.32%   51.45%   -27.88%     
===========================================
  Files         199      200        +1     
  Lines       20870    20894       +24     
  Branches     2684     2687        +3     
===========================================
- Hits        16555    10750     -5805     
- Misses       3566     9544     +5978     
+ Partials      749      600      -149

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

🚨 Try these New Features:

Flaky Tests Detection - Detect and resolve failed and flaky tests

Signed-off-by: Eduardo Apolinario <[email protected]>

…to the metadata bucket Signed-off-by: Eduardo Apolinario <[email protected]>

Signed-off-by: Eduardo Apolinario <[email protected]>

eapolinario · 2024-11-19T19:30:43Z

hold on, I had a bad merge.

Signed-off-by: Eduardo Apolinario <[email protected]>

wild-endeavor

+1 with dm

Signed-off-by: Eduardo Apolinario <[email protected]>

…e flakiness) Signed-off-by: Eduardo Apolinario <[email protected]>

Signed-off-by: Eduardo Apolinario <[email protected]>

* wip - Implement offloading of literals Signed-off-by: Eduardo Apolinario <[email protected]> * Fix use of metadata bucket prefix Signed-off-by: Eduardo Apolinario <[email protected]> * Fix repeated use of uri Signed-off-by: Eduardo Apolinario <[email protected]> * Add temporary representation for offloaded literal Signed-off-by: Eduardo Apolinario <[email protected]> * Add one unit test Signed-off-by: Eduardo Apolinario <[email protected]> * Add another test Signed-off-by: Eduardo Apolinario <[email protected]> * Stylistic changes to the two tests Signed-off-by: Eduardo Apolinario <[email protected]> * Add test for min offloading threshold set to 1MB Signed-off-by: Eduardo Apolinario <[email protected]> * Pick a unique engine-dir for tests Signed-off-by: Eduardo Apolinario <[email protected]> * s/new_outputs/literal_map_copy/ Signed-off-by: Eduardo Apolinario <[email protected]> * Remove unused constant Signed-off-by: Eduardo Apolinario <[email protected]> * Use output_prefix in definition of offloaded literals Signed-off-by: Eduardo Apolinario <[email protected]> * Add initial version of pbhash.py Signed-off-by: Eduardo Apolinario <[email protected]> * Add tests to verify that overriding the hash is carried over to offloaded literals Signed-off-by: Eduardo Apolinario <[email protected]> * Add a few more tests Signed-off-by: Eduardo Apolinario <[email protected]> * Always import ParamSpec from `typing_extensions` Signed-off-by: Eduardo Apolinario <[email protected]> * Fix lint warnings Signed-off-by: Eduardo Apolinario <[email protected]> * Set inferred_type using the task type interface Signed-off-by: Eduardo Apolinario <[email protected]> * Add comment about offloaded literals files and how they are uploaded to the metadata bucket Signed-off-by: Eduardo Apolinario <[email protected]> * Add offloading_enabled Signed-off-by: Eduardo Apolinario <[email protected]> * Add more unit tests including a negative test Signed-off-by: Eduardo Apolinario <[email protected]> * Fix bad merge Signed-off-by: Eduardo Apolinario <[email protected]> * Incorporate feedback. Signed-off-by: Eduardo Apolinario <[email protected]> * Fix image name (unrelated to this PR - just a nice-to-have to decrease flakiness) Signed-off-by: Eduardo Apolinario <[email protected]> * Add `is_map_task` to `_dispatch_execute` Signed-off-by: Eduardo Apolinario <[email protected]> --------- Signed-off-by: Eduardo Apolinario <[email protected]> Co-authored-by: Eduardo Apolinario <[email protected]> Signed-off-by: Eduardo Apolinario <[email protected]>

* wip - Implement offloading of literals * Fix use of metadata bucket prefix * Fix repeated use of uri * Add temporary representation for offloaded literal * Add one unit test * Add another test * Stylistic changes to the two tests * Add test for min offloading threshold set to 1MB * Pick a unique engine-dir for tests * s/new_outputs/literal_map_copy/ * Remove unused constant * Use output_prefix in definition of offloaded literals * Add initial version of pbhash.py * Add tests to verify that overriding the hash is carried over to offloaded literals * Add a few more tests * Always import ParamSpec from `typing_extensions` * Fix lint warnings * Set inferred_type using the task type interface * Add comment about offloaded literals files and how they are uploaded to the metadata bucket * Add offloading_enabled * Add more unit tests including a negative test * Fix bad merge * Incorporate feedback. * Fix image name (unrelated to this PR - just a nice-to-have to decrease flakiness) * Add `is_map_task` to `_dispatch_execute` --------- Signed-off-by: Eduardo Apolinario <[email protected]> Co-authored-by: Eduardo Apolinario <[email protected]>

eapolinario added 11 commits October 18, 2024 17:54

wip - Implement offloading of literals

12f9edf

Signed-off-by: Eduardo Apolinario <[email protected]>

Fix use of metadata bucket prefix

5fefa85

Signed-off-by: Eduardo Apolinario <[email protected]>

Fix repeated use of uri

1668fda

Signed-off-by: Eduardo Apolinario <[email protected]>

Add temporary representation for offloaded literal

853df62

Signed-off-by: Eduardo Apolinario <[email protected]>

Add one unit test

5e53a1b

Signed-off-by: Eduardo Apolinario <[email protected]>

Add another test

177368d

Signed-off-by: Eduardo Apolinario <[email protected]>

Stylistic changes to the two tests

5fc2e84

Signed-off-by: Eduardo Apolinario <[email protected]>

Add test for min offloading threshold set to 1MB

db48d18

Signed-off-by: Eduardo Apolinario <[email protected]>

Pick a unique engine-dir for tests

6884ee0

Signed-off-by: Eduardo Apolinario <[email protected]>

s/new_outputs/literal_map_copy/

5a6423c

Signed-off-by: Eduardo Apolinario <[email protected]>

Remove unused constant

dbfea93

Signed-off-by: Eduardo Apolinario <[email protected]>

eapolinario requested review from wild-endeavor, kumare3, pingsutw, cosmicBboy, samhita-alla, thomasjpfan and Future-Outlier as code owners October 28, 2024 19:30

eapolinario commented Oct 28, 2024

View reviewed changes

Merge remote-tracking branch 'origin' into offload-literals

adeed34

Signed-off-by: Eduardo Apolinario <[email protected]>

pmahindrakar-oss reviewed Oct 28, 2024

View reviewed changes

eapolinario added 2 commits October 28, 2024 22:30

Use output_prefix in definition of offloaded literals

25908d9

Signed-off-by: Eduardo Apolinario <[email protected]>

Add initial version of pbhash.py

e827693

Signed-off-by: Eduardo Apolinario <[email protected]>

wild-endeavor reviewed Nov 7, 2024

View reviewed changes

eapolinario added 3 commits November 6, 2024 21:44

Add tests to verify that overriding the hash is carried over to offlo…

e0e2016

…aded literals Signed-off-by: Eduardo Apolinario <[email protected]>

Add a few more tests

b284492

Signed-off-by: Eduardo Apolinario <[email protected]>

Always import ParamSpec from typing_extensions

b579b83

Signed-off-by: Eduardo Apolinario <[email protected]>

eapolinario added 2 commits November 7, 2024 07:41

Fix lint warnings

8c2336e

Signed-off-by: Eduardo Apolinario <[email protected]>

Merge remote-tracking branch 'origin' into offload-literals

c28d537

Signed-off-by: Eduardo Apolinario <[email protected]>

eapolinario added 2 commits November 7, 2024 20:52

Set inferred_type using the task type interface

37b2bb4

Signed-off-by: Eduardo Apolinario <[email protected]>

Add comment about offloaded literals files and how they are uploaded …

e25496d

…to the metadata bucket Signed-off-by: Eduardo Apolinario <[email protected]>

pmahindrakar-oss previously approved these changes Nov 11, 2024

View reviewed changes

eapolinario added 2 commits November 18, 2024 18:45

Add offloading_enabled

9276e96

Signed-off-by: Eduardo Apolinario <[email protected]>

Add more unit tests including a negative test

a8bdbca

Signed-off-by: Eduardo Apolinario <[email protected]>

eapolinario dismissed pmahindrakar-oss’s stale review via a8bdbca November 19, 2024 19:26

Merge remote-tracking branch 'origin' into offload-literals

fe822b9

Signed-off-by: Eduardo Apolinario <[email protected]>

Fix bad merge

a4fcfab

Signed-off-by: Eduardo Apolinario <[email protected]>

wild-endeavor previously approved these changes Nov 19, 2024

View reviewed changes

Incorporate feedback.

b3a1b0d

Signed-off-by: Eduardo Apolinario <[email protected]>

eapolinario dismissed wild-endeavor’s stale review via b3a1b0d November 19, 2024 20:53

wild-endeavor previously approved these changes Nov 19, 2024

View reviewed changes

eapolinario added 2 commits November 21, 2024 17:17

Fix image name (unrelated to this PR - just a nice-to-have to decreas…

32c5896

…e flakiness) Signed-off-by: Eduardo Apolinario <[email protected]>

Add is_map_task to _dispatch_execute

12e194a

Signed-off-by: Eduardo Apolinario <[email protected]>

eapolinario dismissed wild-endeavor’s stale review via 12e194a November 21, 2024 22:30

eapolinario mentioned this pull request Nov 21, 2024

Revert "[COR-2297/] Fix nested offloaded type validation (#552) (#5996)" flyteorg/flyte#6045

Merged

3 tasks

wild-endeavor approved these changes Nov 22, 2024

View reviewed changes

eapolinario merged commit 40c9540 into master Nov 22, 2024
103 of 104 checks passed

eapolinario mentioned this pull request Nov 22, 2024

Offload literals (#2872) #2950

Merged

3 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Offload literals #2872

Offload literals #2872

eapolinario commented Oct 28, 2024 •

edited

Loading

eapolinario Oct 28, 2024

pmahindrakar-oss Oct 28, 2024

eapolinario Nov 8, 2024

pmahindrakar-oss Oct 28, 2024

eapolinario Oct 29, 2024

pmahindrakar-oss Oct 30, 2024

wild-endeavor Nov 7, 2024

eapolinario Nov 8, 2024

pmahindrakar-oss Oct 28, 2024

eapolinario Oct 29, 2024

pmahindrakar-oss Oct 30, 2024

pmahindrakar-oss Oct 28, 2024

pmahindrakar-oss Oct 28, 2024

eapolinario Oct 29, 2024

wild-endeavor Nov 7, 2024

eapolinario Nov 7, 2024

wild-endeavor Nov 7, 2024

codecov bot commented Nov 7, 2024 •

edited

Loading

eapolinario commented Nov 19, 2024

wild-endeavor left a comment

		min_offloaded_size = int(os.environ.get("FK_L_MIN_SIZE_MB", "10")) * 1024 * 1024
		max_offloaded_size = int(os.environ.get("FK_L_MAX_SIZE_MB", "1000")) * 1024 * 1024

Offload literals #2872

Offload literals #2872

Conversation

eapolinario commented Oct 28, 2024 • edited Loading

Tracking issue

Why are the changes needed?

What changes were proposed in this pull request?

How hashing is handled

How was this patch tested?

Setup process

Screenshots

Check all the applicable boxes

Related PRs

Docs link

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented Nov 7, 2024 • edited Loading

Codecov Report

eapolinario commented Nov 19, 2024

wild-endeavor left a comment

Choose a reason for hiding this comment

eapolinario commented Oct 28, 2024 •

edited

Loading

codecov bot commented Nov 7, 2024 •

edited

Loading