[Core feature] Caching for non-flyte specific offloaded objects

[kumare3]

Motivation: Why do you think this is important?

Currently, if users use pandas.DataFrame or a pyspark.DataFrame or pander.DataFrameSchema, Flytekit simply extracts the data from the transport Literaltype.Schema. So consider the following function

@task(cache=True, cache_version="1.0")
def sub_foo(df: pandas.DataFrame) -> pandas.DataFrame:
    return df

The above function will be cached, because the input type has the file path that is cached and hence the function is not run.

But, now consider the following

@dynamic
def foo(df: pandas.DataFrame) -> pandas.DataFrame:
     return sub_foo(df=df)

This task will not be cached. This is because the dataframe is downloaded and then re-uploaded, as the underlying task transforms are not aware that the passed dataframe was not mutated. But, if FlyteSchema is used, this would work fine.

Goal: What should the final outcome look like, ideally?

Either this should work as the user expects, i.e., cache hit, else this abnormality should be documented.

Describe alternatives you've considered

NA

Propose: Link/Inline OR Additional context

No response

• Are you sure this issue hasn't been raised already?

• Have you read the Code of Conduct?

[kumare3]

cc @eapolinario While writing the issue I realized, it is not that big a problem?

[cosmicBboy]

I'm facing a similar dataframe caching issue in the weather forecasting project. I'm using a dynamic workflow to manage the training of a model and the tasks within it rely on training data in the form of a dataframe. Should I use FlyteSchema for this instead?

[cosmicBboy]

solution proposal: caching for complex data types, e.g. dataframes

For blob and schema types

1. expose a hash method in TypeTransformer, which would implement a cache-by-value system that we would implement and maintain. Users who define custom type transformers can implement a hash method too. This would enable default caching solutions for common types like pandas dataframes.
2. in cases where a hashing function is not provided for a particular type (e.g. new_pandas_like_library.DataFrame), introduce a cache_output_fn (naming TBD) to the @task decorator like so:

@task(
    cache=True,
    cache_output_fn={
        pd.DataFrame: lambda x: hash(x),
    },
    cache_version=CACHE_VERSION,
)
def func(x: int) -> pd.DataFrame:
    df = ...  # get a dataframe
    return df

[cosmicBboy]

Update:
Consider a mechanism that uses s3 etag (or whatever the blob storage equivalent is) as the hash metadata for a particular reference-based data artifact (dataframe, files, etc). This offloads the burden of computing that hash to blob storage instead of manually computing the hash in the container.