Auto Cache Plugin #2971
Auto Cache Plugin #2971
Conversation
Signed-off-by: Daniel Sola <[email protected]>
Signed-off-by: Daniel Sola <[email protected]>
Signed-off-by: Daniel Sola <[email protected]>
Signed-off-by: Daniel Sola <[email protected]>
Signed-off-by: Daniel Sola <[email protected]>
Signed-off-by: Daniel Sola <[email protected]>
Signed-off-by: Daniel Sola <[email protected]>
Signed-off-by: Daniel Sola <[email protected]>
Signed-off-by: Daniel Sola <[email protected]>
Signed-off-by: Daniel Sola <[email protected]>
Signed-off-by: Daniel Sola <[email protected]>
Signed-off-by: Daniel Sola <[email protected]>
dansola
requested review from
wild-endeavor,
kumare3,
eapolinario,
pingsutw,
cosmicBboy,
samhita-alla,
thomasjpfan and
Future-Outlier
as code owners
flytekit/core/auto_cache.py
Outdated
| self.cache_serialize = cache_serialize | ||
| self.cache_version = cache_version | ||
| self.cache_ignore_input_vars = cache_ignore_input_vars |
There was a problem hiding this comment.
what's the purpose of saving this state here? aren't these just forwarded to the underlying TaskMetadata?
There was a problem hiding this comment.
The idea with this is the user could use the CachePolicy to define all the arguments relating to caching. This simplifies the UX a bit as opposed to having a CachePolicy and a cache_ignore_input_vars, cache_serialize, etc.
There was a problem hiding this comment.
I think this is a little confusing:
cache_versionshould not be exposed, since theAutoCacheprotocol is meant to produce this value automatically, andsaltis meant to fulfill the need of manually bumping the cache.- I think it makes sense to keep
cache_serializeandcache_ignore_input_varsas options to specify in the@taskdecorator as opposed to introducing this redundancy here.
There was a problem hiding this comment.
Sounds good. Sounds like there is a separate effort aimed at collecting all of the caching arguments here: flyteorg/flyte#6143
Happy to use that instead and simplify the arguments here!
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #2971 +/- ##
==========================================
+ Coverage 76.49% 77.96% +1.46%
==========================================
Files 200 202 +2
Lines 20901 21324 +423
Branches 2689 2739 +50
==========================================
+ Hits 15989 16625 +636
+ Misses 4195 3904 -291
- Partials 717 795 +78 ☔ View full report in Codecov by Sentry. |
Code Review Agent Run #bc105bActionable Suggestions - 9
Additional Suggestions - 7
Review Details
AI Code Review powered by |
Changelist by BitoThis pull request implements the following key changes.
|
flytekit/core/auto_cache.py
Outdated
| hash_obj = hashlib.sha256(task_hash.encode()) | ||
| return hash_obj.hexdigest() | ||
|
|
||
| return None |
There was a problem hiding this comment.
Consider returning an empty string instead of None for consistency in return types. The method signature indicates it returns str but can return None.
Code suggestion
Check the AI-generated fix before applying
| return None | |
| return "" |
Code Review Run #bc105b
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
| numpy==1.24.3 | ||
| pandas==2.0.3 | ||
| requests==2.31.0 | ||
| matplotlib==3.7.2 | ||
| pillow==10.0.0 | ||
| scipy==1.11.2 | ||
| pytest==7.4.0 | ||
| urllib3==2.0.4 | ||
| cryptography==41.0.3 | ||
| setuptools==68.0.0 | ||
| flask==2.3.2 | ||
| django==4.2.4 | ||
| scikit-learn==1.3.0 | ||
| beautifulsoup4==4.12.2 | ||
| pyyaml==6.0 | ||
| fastapi==0.100.0 | ||
| sqlalchemy==2.0.36 | ||
| tqdm==4.65.0 | ||
| pytest-mock==3.11.0 | ||
| jinja2==3.1.2 |
There was a problem hiding this comment.
Consider pinning dependencies to compatible versions using ~= or >= instead of == to allow for minor version updates that include security patches while maintaining compatibility. This helps keep dependencies up-to-date with security fixes.
Code suggestion
Check the AI-generated fix before applying
| numpy==1.24.3 | |
| pandas==2.0.3 | |
| requests==2.31.0 | |
| matplotlib==3.7.2 | |
| pillow==10.0.0 | |
| scipy==1.11.2 | |
| pytest==7.4.0 | |
| urllib3==2.0.4 | |
| cryptography==41.0.3 | |
| setuptools==68.0.0 | |
| flask==2.3.2 | |
| django==4.2.4 | |
| scikit-learn==1.3.0 | |
| beautifulsoup4==4.12.2 | |
| pyyaml==6.0 | |
| fastapi==0.100.0 | |
| sqlalchemy==2.0.36 | |
| tqdm==4.65.0 | |
| pytest-mock==3.11.0 | |
| jinja2==3.1.2 | |
| numpy~=1.24.3 | |
| pandas~=2.0.3 | |
| requests~=2.31.0 | |
| matplotlib~=3.7.2 | |
| pillow~=10.0.0 | |
| scipy~=1.11.2 | |
| pytest~=7.4.0 | |
| urllib3~=2.0.4 | |
| cryptography~=41.0.3 | |
| setuptools~=68.0.0 | |
| flask~=2.3.2 | |
| django~=4.2.4 | |
| scikit-learn~=1.3.0 | |
| beautifulsoup4~=4.12.2 | |
| pyyaml~=6.0 | |
| fastapi~=0.100.0 | |
| sqlalchemy~=2.0.36 | |
| tqdm~=4.65.0 | |
| pytest-mock~=3.11.0 | |
| jinja2~=3.1.2 |
Code Review Run #bc105b
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
| packages = cache.get_version_dict().keys() | ||
|
|
||
| expected_packages = {'PIL', 'bs4', 'numpy', 'pandas', 'scipy', 'sklearn'} | ||
| set(packages) == expected_packages, f"Expected keys {expected_packages}, but got {set(packages)}" |
There was a problem hiding this comment.
The assertion statement appears to be missing the assert keyword, which means this comparison won't actually validate the test condition. Consider adding the assert keyword.
Code suggestion
Check the AI-generated fix before applying
| set(packages) == expected_packages, f"Expected keys {expected_packages}, but got {set(packages)}" | |
| assert set(packages) == expected_packages, f"Expected keys {expected_packages}, but got {set(packages)}" |
Code Review Run #bc105b
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
plugins/flytekit-auto-cache/flytekitplugins/auto_cache/cache_private_modules.py
Show resolved
Hide resolved
plugins/flytekit-auto-cache/flytekitplugins/auto_cache/cache_function_body.py
Show resolved
Hide resolved
| def get_version(self, params: VersionParameters) -> str: | ||
| if params.func is None: | ||
| raise ValueError("Function-based cache requires a function parameter") | ||
| return self._get_version(func=params.func) |
There was a problem hiding this comment.
The get_version method could benefit from type checking params.func before accessing it to provide a more descriptive error message.
Code suggestion
Check the AI-generated fix before applying
| def get_version(self, params: VersionParameters) -> str: | |
| if params.func is None: | |
| raise ValueError("Function-based cache requires a function parameter") | |
| return self._get_version(func=params.func) | |
| def get_version(self, params: VersionParameters) -> str: | |
| if params.func is None: | |
| raise ValueError("Function-based cache requires a function parameter") | |
| if not callable(params.func): | |
| raise TypeError("params.func must be a callable function") | |
| return self._get_version(func=params.func) |
Code Review Run #bc105b
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
| module_b.another_helper() | ||
| result = norm([1, 2, 3]) | ||
| print(result) | ||
| sum([SOME_CONSTANT, utils.THIRD_CONSTANT]) |
There was a problem hiding this comment.
The sum() function call result is not being stored or used, which may indicate meaningless executed code. Consider either storing the result or removing if not needed.
Code suggestion
Check the AI-generated fix before applying
| sum([SOME_CONSTANT, utils.THIRD_CONSTANT]) | |
| result = sum([SOME_CONSTANT, utils.THIRD_CONSTANT]) |
Code Review Run #bc105b
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
| except Exception as e: | ||
| click.secho(f"Could not get version for {package_name} using importlib.metadata: {str(e)}", fg="yellow") |
There was a problem hiding this comment.
Catching a broad 'Exception' may hide bugs. Consider catching specific exceptions instead.
Code suggestion
Check the AI-generated fix before applying
| except Exception as e: | |
| click.secho(f"Could not get version for {package_name} using importlib.metadata: {str(e)}", fg="yellow") | |
| except (ImportError, AttributeError) as e: | |
| click.secho(f"Could not get version for {package_name} using importlib.metadata: {str(e)}", fg="yellow") |
Code Review Run #bc105b
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
flytekit/core/auto_cache.py
Outdated
| ... | ||
|
|
||
|
|
||
| class CachePolicy: |
There was a problem hiding this comment.
can CachePolicy live in the plugin? It makes sense for the abstract AutoCache protocol to be defined in flytekit core, but any implementation of it should be in the plugin.
There was a problem hiding this comment.
good point, thanks. i refactored this!
Signed-off-by: Daniel Sola <[email protected]>
Signed-off-by: Daniel Sola <[email protected]>
Code Review Agent Run #fa2b0bActionable Suggestions - 1
Review Details
AI Code Review powered by |
Why are the changes needed?
Make caching easier to use in flytekit by reducing cognitive burden of specifying cache versions
What changes were proposed in this pull request?
To use the caching mechanism in a Flyte task, you can define a
CachePolicythat combines multiple caching strategies. Here’s an example of how to set it up:Salt Parameter
The
saltparameter in theCachePolicyadds uniqueness to the generated hash. It can be used to differentiate between different versions of the same task. This ensures that even if the underlying code remains unchanged, the hash will vary if a different salt is provided. This feature is particularly useful for invalidating the cache for specific versions of a task.Cache Implementations
Users can add any number of cache policies that implement the
AutoCacheprotocol defined in@auto_cache.py. Below are the implementations available so far:1. CacheFunctionBody
This implementation hashes the contents of the function of interest, ignoring any formatting or comment changes. It ensures that the core logic of the function is considered for versioning.
2. CacheImage
This implementation includes the hash of the
container_imageobject passed. If the image is specified as a name, that string is hashed. If it is anImageSpec, the parametrization of theImageSpecis hashed, allowing for precise versioning of the container image used in the task.3. CachePrivateModules
This implementation recursively searches the task of interest for all callables and constants used. The contents of any callable (function or class) utilized by the task are hashed, ignoring formatting or comments. The values of the literal constants used are also included in the hash.
It accounts for both
importandfrom-importstatements at the global and local levels within a module or function. Any callables that are within site-packages (i.e., external libraries) are ignored.4. CacheExternalDependencies
This implementation recursively searches through all the callables like
CachePrivateModules, but when an external package is found, it records the version of the package, which is included in the hash. This ensures that changes in external dependencies are reflected in the task's versioning.How was this patch tested?
Unit tests for the following:
Setup process
Screenshots
Check all the applicable boxes
Related PRs
Docs link
Summary by Bito
This PR refactors Flytekit's caching mechanism by introducing a comprehensive auto-cache plugin that implements multiple strategies including function body hashing, container image versioning, and dependency tracking. The implementation migrates CachePolicy to a dedicated plugin, simplifying the core auto_cache module and streamlining cache parameter types while maintaining seamless integration with existing task and workflow decorators.Unit tests added: True
Estimated effort to review (1-5, lower is better): 5