[KED-2744] Memory Leakage - Unexpected Caching behaviour with CacheDataset

[noklam]

Description

When I have a pipeline like this, I expected once the execution of the node is finished, it should release the input to reduce memory footprint. I found this is not the case and causing Memory Error for my pipeline.
A -> B # Step1
B -> C # Step2
C -> D # Step3

For example, as soon as step1 is finished, there is no reason why we should keep A in memory anymore.

Context

It keeps unnecessary variables in memory. Originally I was trying to write some notes to dive deep into Kedro dataset management and found that the behavior is not what I would expect.

Steps to Reproduce

https://noklam.github.io/blog/posts/2021-07-02-kedro-datacatalog.html

Expected Result

For a non-cached dataset, it should exist in memory only in its node. For CachedDataSet, it should be deleted as soon as it is not needed anymore.

Actual Result

-- If you received an error, place it here.

-- Separate them if you have more than one.

Your Environment

Include as many relevant details about the environment in which you experienced the bug:

• Kedro version used (pip show kedro or kedro -V):
• Python version used (python -V):
• Operating system and version:

[datajoely]

See full discussion on Discord

[merelcht]

Hi @noklam, thanks for flagging this and writing up such a thorough explanation in your blog post 👏 😄
I will get to the bottom of why it's been implemented in this way and keep you up to date about what I find.

[merelcht]

Hi @noklam, I've found that this is not actually a bug but explicitly designed this way.

The main reason being that pipeline inputs and outputs are special datasets. They can be used for debugging after the run of a pipeline. So this means that the release functionality has explicitly been developed for the intermediate datasets.

On top of that, input parameters are handled as MemoryDataSets and so if the check for data_set not in pipeline.inputs() didn't exist, parameters would be released from memory as well. The code was written a couple of years ago, and we haven't further investigated whether it would be a problem for the parameters to be released, but regardless pipeline inputs and outputs should be treated differently according to our philosophy.

[noklam]

@MerelTheisenQB

Thanks for the explanation. If it is only for debugging purposes, there should be a flag to enable this behavior rather than setting it as default? Maybe I am missing something here, but I don't see why parameters need to be treated specially. It should be okay to release it as long as it is not needed in the remaining nodes. A function should only depend on its inputs.

In addition, CacheDataSet is just MemoryDataSet, which means CacheDataSet will stay in memory forever. The purpose of using Cache is to reduce the I/O instead of writing to disk in 1 node then reading the same data from disk immediately. This is a big problem for large pipelines.

This is a common pattern for my training pipeline.

1. preprocess raw data -> processed data
2. processed data ->Train/test split
3. processed data concat with some metadata or prediction label for analysis

I'll use CacheDataSet for these functions for sure because I don't want to write 10GB of data and reading it off from Disk (it could take minutes).

Thoughts?

[idanov]

@noklam Thanks for flagging this. This functionality was implemented long ago as @MerelTheisenQB pointed out and the main purpose is to enable interactive workflows more generally (and as extension to that debugging). E.g. the idea is that you could run a pipeline multiple times with the same catalog instance in the process of experimenting with different pipelines and if you start off with a MemoryDataSet inputs which are released after the run, that will make for a very painful experience.

Most of the time this shouldn't cause problems, since one would normally use MemoryDataSet as pipeline inputs only in interactive workflows. However you have correctly identified that CachedDataSet can be used in non-interactive workflows for inputs quite often and the desired behaviour would be to release the data once no longer needed.

For the case of pipeline outputs unconsumed by anyone after the run, there will be no reason for someone to define them as MemoryDataSet unless they work in interactive sessions (notebooks). Therefore the behaviour there is as expected.

CachedDataSet was not developed as a core Kedro component and it was done as internal contribution, so we might not have considered the issue you point out. That means that it's probably a good idea for us to revisit this and change the behaviour as follows:

• Both MemoryDataSet and CachedDataSet should be released when used as intermediate datasets (Kedro already behaves like this to my knowledge) ✅
• MemoryDataSet's used as pipeline inputs or outputs should not be released after runs to support interactive workflows (Kedro already behaves like this as claimed in this GitHub issue) ✅
• CachedDataSet should be released after being used as pipeline input, but not as pipeline output (this is where Kedro fails user expectation) ❌

Would such behaviour suit your needs @noklam ?

[noklam]

@idanov Thanks, I rarely use Kedro in interactive mode and I can now see why this was implemented this way.

Yes, I think this match my expectation. Any datasets that are not in remanining nodes inputs should be released.

[datajoely]

Created ticket on backlog

[acnazarejr]

I am facing the same problem, however, my dataset is a pandas.ParquetDataSet. Even after the node is terminated and the dataframe is saved to disk, the data is still in memory. The memory is only released after all pipeline execution.

Maybe I missed something, but should I indicate something to kedro so it doesn't keep the data in memory?

[datajoely]

@acnazarejr - to confirm is this also using the CachedDataSet? Persisted writes to Parquet should discard the memory and be garbage collected by Python.

[acnazarejr]

@datajoely, sorry for the delayed response. I'm not using the CachedDataSet.

[datajoely]

Hi @acnazarejr I don't think it should be behaving this way - how do you know the memory is not being released?

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[merelcht]

This ticket hasn't had any recent activity, so I'm closing it.

[noklam]

Reopen as this issue still exists, though usage of CacheDataset is not known. Partly due to this is an internal contribution and no docs has ever been written about it other than the API docs #3616.

E.g. the idea is that you could run a pipeline multiple times with the same catalog instance in the process of experimenting with different pipelines and if you start off with a MemoryDataSet inputs which are released after the run, that will make for a very painful experience.

This is a very interesting point and I would love to re-visit this at some points. It seems that years ago Kedro focus more on the interactive flow witht he idea of re-run pipeline with the same DataCatalog, which pretty much is now blocked by the KedroSession unless users use Runner ,DataCatalog directly.

• Both MemoryDataSet and CachedDataSet should be released when used as intermediate datasets (Kedro already behaves like this to my knowledge) ✅
• MemoryDataSet's used as pipeline inputs or outputs should not be released after runs to support interactive workflows (Kedro already behaves like this as claimed in this GitHub issue) ✅
• CachedDataSet should be released after being used as pipeline input, but not as pipeline output (this is where Kedro fails user expectation) ❌

The last point is the real problem here. We may not need to fix this issue immediately, but I think it's important if we think about re-running pipeline in a notebook.

[datajoely]

Could I also ask - why not make MemoryDatasets cache by default?

[noklam]

@datajoely Can you clarify a bit? MemoryDataset do "cache" in the sense that it only get load once only when passing around nodes. CacheDataset is slightly different because it is used as a wrapper dataset, under the hood it use MemoryDataset as its cache.

[datajoely]

I guess my question is - if we can, why not make regular datasets cache by default?

[noklam]

@datajoely I think it's because there are 1% edge case that is not solved. I very much welcome this to be changed.

df.to_csv("my_csv.csv")
df2 = pd.read_csv("my_csv.csv")

In most case they will be identical, in badly typed cases they are different, read/load separately is more closed to an orchestrator mode I guess. This is very bad because depends where you start running your pipeline, you get different outputs.

[datajoely]

Got it! They say cache invalidation is the 2nd hardest problem for a reason!