Skip to content

Commit 950aad2

Browse files
committed
Add an environment variable and CLI option to disable default caching
Signed-off-by: ddalvi <ddalvi@redhat.com>
1 parent 36cf066 commit 950aad2

File tree

3 files changed

+22
-3
lines changed

3 files changed

+22
-3
lines changed

sdk/python/kfp/cli/compile_.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,14 +133,32 @@ def parse_parameters(parameters: Optional[str]) -> Dict:
133133
is_flag=True,
134134
default=False,
135135
help='Whether to disable type checking.')
136+
@click.option(
137+
'--execution-caching-enabled-by-default',
138+
type=click.Choice(['enabled', 'disabled'], case_sensitive=False),
139+
default=None,
140+
help="Enable task-level caching. Enabled by default, set it to disabled to disable caching."
141+
)
136142
def compile_(
137143
py: str,
138144
output: str,
139145
function_name: Optional[str] = None,
140146
pipeline_parameters: Optional[str] = None,
141147
disable_type_check: bool = False,
148+
execution_caching_enabled_by_default: Optional[bool] = None,
142149
) -> None:
143150
"""Compiles a pipeline or component written in a .py file."""
151+
152+
env_enable_caching = os.getenv('KFP_EXECUTION_CACHING_ENABLED_BY_DEFAULT', 'enabled').lower() == 'enabled'
153+
if execution_caching_enabled_by_default is None:
154+
execution_caching_enabled_by_default = env_enable_caching
155+
else:
156+
execution_caching_enabled_by_default = execution_caching_enabled_by_default.lower() == 'enabled'
157+
if execution_caching_enabled_by_default:
158+
os.environ['KFP_EXECUTION_CACHING_ENABLED_BY_DEFAULT'] = 'enabled'
159+
else:
160+
os.environ['KFP_EXECUTION_CACHING_ENABLED_BY_DEFAULT'] = 'disabled'
161+
144162
pipeline_func = collect_pipeline_or_component_func(
145163
python_file=py, function_name=function_name)
146164
parsed_parameters = parse_parameters(parameters=pipeline_parameters)

sdk/python/kfp/dsl/pipeline_task.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import functools
1919
import inspect
2020
import itertools
21+
import os
2122
import re
2223
from typing import Any, Dict, List, Mapping, Optional, Union
2324
import warnings
@@ -130,7 +131,7 @@ def __init__(
130131
inputs=dict(args.items()),
131132
dependent_tasks=[],
132133
component_ref=component_spec.name,
133-
enable_caching=True)
134+
enable_caching=os.getenv('KFP_EXECUTION_CACHING_ENABLED_BY_DEFAULT', 'enabled').lower() == 'enabled')
134135
self._run_after: List[str] = []
135136

136137
self.importer_spec = None

sdk/python/kfp/dsl/structures.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import collections
1818
import dataclasses
1919
import itertools
20+
import os
2021
import re
2122
from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
2223
import uuid
@@ -420,11 +421,10 @@ class TaskSpec:
420421
trigger_strategy: Optional[str] = None
421422
iterator_items: Optional[Any] = None
422423
iterator_item_input: Optional[str] = None
423-
enable_caching: bool = True
424+
enable_caching: bool = os.getenv('KFP_EXECUTION_CACHING_ENABLED_BY_DEFAULT', 'enabled').lower() == 'enabled'
424425
display_name: Optional[str] = None
425426
retry_policy: Optional[RetryPolicy] = None
426427

427-
428428
@dataclasses.dataclass
429429
class ImporterSpec:
430430
"""ImporterSpec definition.

0 commit comments

Comments
 (0)