Add caching to local execution #592

eapolinario · 2021-08-12T17:38:27Z

TL;DR

Enable the ability to cache the output of tasks in local executions

Type

Bug Fix
Feature
Plugin

Are all requirements met?

Complete description

joblib.Memory provides a persistent store to cache the result of function invocations. Similarly to how we do in the hosted case, we use the pair (task inputs, cache version) to memoize task outputs across different executions.

We're also adding another command to pyflyte to help users clear the cache. In the future we might extend this to allow for the inspection of the values in the local cache.

Tracking Issue

flyteorg/flyte#761

Follow-up issue

NA

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

welcome · 2021-08-12T17:38:29Z

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).

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

flytekit/core/base_task.py

codecov · 2021-08-12T18:10:59Z

Codecov Report

Merging #592 (9cbc70a) into master (fc65e97) will increase coverage by 0.06%.
The diff coverage is 97.05%.

@@            Coverage Diff             @@
##           master     #592      +/-   ##
==========================================
+ Coverage   85.51%   85.57%   +0.06%     
==========================================
  Files         376      379       +3     
  Lines       29311    29682     +371     
  Branches     2357     2376      +19     
==========================================
+ Hits        25064    25400     +336     
- Misses       3611     3640      +29     
- Partials      636      642       +6

Impacted Files	Coverage Δ
flytekit/clis/sdk_in_container/local_cache.py	`77.77% <77.77%> (ø)`
flytekit/core/local_cache.py	`80.95% <80.95%> (ø)`
flytekit/clis/sdk_in_container/pyflyte.py	`82.45% <100.00%> (+0.63%)`	⬆️
flytekit/core/base_task.py	`88.26% <100.00%> (+0.26%)`	⬆️
tests/flytekit/unit/core/test_local_cache.py	`100.00% <100.00%> (ø)`
tests/flytekit/integration/remote/test_remote.py	`82.26% <0.00%> (-5.43%)`	⬇️
tests/flytekit/unit/core/test_node_creation.py	`95.58% <0.00%> (-2.09%)`	⬇️
flytekit/remote/remote.py	`72.06% <0.00%> (-2.06%)`	⬇️
tests/flytekit/unit/core/test_launch_plan.py	`93.96% <0.00%> (-0.72%)`	⬇️
... and 16 more

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update fc65e97...9cbc70a. Read the comment docs.

Makefile

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

katrogan

looks great, can we add some more tests!

katrogan · 2021-08-13T22:44:58Z

flytekit/core/base_task.py

-        outputs_literal_map = self.dispatch_execute(ctx, input_literal_map)
+        # if metadata.cache is set, check memoized version
+        if self._metadata.cache:
+            # The cache key is composed only of '(input_literal_map, cache_version)', i.e. all other parameters


should we also include the task name as part of the key?

similar to what we do with the identifier on hosted flyte https://github.com/flyteorg/flyteplugins/blob/master/go/tasks/pluginmachinery/catalog/client.go#L26

The key is already implicitly part of the cache key (in other words, joblib.Memory only caches the result of each function call). I'm going to update the comment to reflect that.

Actually, this is a really good call. I'm going to add a test case to demonstrate the problem and subsequent fix.

Fixed by 130b432.

katrogan · 2021-08-13T22:53:45Z

tests/flytekit/unit/core/test_local_cache.py

@@ -0,0 +1,78 @@
+from pytest import fixture


can we test some more complex input and output types, like dataclasses, schemas, files?

katrogan · 2021-08-13T22:55:21Z

flytekit/core/local_cache.py

+
+    @staticmethod
+    def initialize():
+        LocalCache._memory = Memory(CACHE_LOCATION, verbose=5)


should we make the verbosity configurable too?

I don't think we gain a lot from making this configurable, especially given that configuration files are not used at all in local executions. I'm going to leave a TODO as well.

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

katrogan · 2021-08-17T16:29:32Z

tests/flytekit/unit/core/test_local_cache.py

+
+def test_wf_custom_types():
+    @dataclass_json
+    @dataclass


awesome, this works? ran into issues with dataclass serialization last time i tried this but maybe joblib has since worked around this. woohoo!

question around this, are we storing LiteralMap or are we storing the native python values. Also what happens to file refrences, do we store a shallow copy or a deep copy?

eduardo apolinario and others added 5 commits August 12, 2021 00:03

Add initial structure of local cache using joblib

6b58da3

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

Add bogus unit tests

e8012e8

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

Add clear-cache make target

0529f7c

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

Run make requirements

1cc9bdb

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

Remove a few TODO

fa951a5

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

eapolinario requested review from katrogan, kumare3 and wild-endeavor as code owners August 12, 2021 17:38

eapolinario added 2 commits August 12, 2021 10:45

Linting

5b137a8

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

More linting

ce947a2

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

kumare3 reviewed Aug 12, 2021

View reviewed changes

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

kumare3 reviewed Aug 12, 2021

View reviewed changes

Makefile Outdated Show resolved Hide resolved

eapolinario added 10 commits August 12, 2021 15:41

Define LocalCache and replace uses of it in base_task.py

ec9fdfe

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

LocalCache type hints

8dfdd1e

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

Comment use of LocalCache in base_task.py

790284e

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

Remove use of Optional from declaration of _memory

ffc4f52

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

Move comment closer to invocation of 'dispatch_execute_func'

9a486e1

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

Use global counter to validate cache hits

44b4858

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

Use ~/.flyte/local-cache as the default location

35d4a2b

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

Force initialization in LocalCache.clear()

01ca49c

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

Add pyflyte local-cache command

ffb66a0

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

Remove clear-cache make target

823dffc

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

eapolinario changed the title ~~[wip] Add caching to local execution~~ Add caching to local execution Aug 13, 2021

eapolinario added 2 commits August 13, 2021 14:23

Linting.

9deff96

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

More linting

f0c8aed

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

kumare3 previously approved these changes Aug 13, 2021

View reviewed changes

katrogan reviewed Aug 13, 2021

View reviewed changes

Add more tests

f9fd960

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

eapolinario added 4 commits August 16, 2021 21:06

Add constant for cache verbosity

a00ead1

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

Add task name to cache key definition

130b432

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

Add tests containing complex inputs and outputs

5c991e3

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

Linting

e3170b0

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

eapolinario dismissed kumare3’s stale review via e3170b0 August 17, 2021 04:10

eapolinario added 3 commits August 16, 2021 21:14

Comment the test used to confirm the use of task names in cache keys

5b83e76

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

More linting

045cc3b

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

Fix linting

9cbc70a

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

katrogan approved these changes Aug 17, 2021

View reviewed changes

eapolinario merged commit 86d3368 into master Aug 17, 2021

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add caching to local execution #592

Add caching to local execution #592

eapolinario commented Aug 12, 2021 •

edited

Loading

welcome bot commented Aug 12, 2021

codecov bot commented Aug 12, 2021 •

edited

Loading

katrogan left a comment

katrogan Aug 13, 2021

eapolinario Aug 17, 2021

eapolinario Aug 17, 2021

eapolinario Aug 17, 2021

katrogan Aug 13, 2021

katrogan Aug 13, 2021

eapolinario Aug 17, 2021

katrogan Aug 17, 2021

kumare3 Aug 17, 2021

Add caching to local execution #592

Add caching to local execution #592

Conversation

eapolinario commented Aug 12, 2021 • edited Loading

TL;DR

Type

Are all requirements met?

Complete description

Tracking Issue

Follow-up issue

welcome bot commented Aug 12, 2021

codecov bot commented Aug 12, 2021 • edited Loading

Codecov Report

katrogan left a 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

eapolinario commented Aug 12, 2021 •

edited

Loading

codecov bot commented Aug 12, 2021 •

edited

Loading