From ae2e91e8ee4ecf20167621d90b842bc14bb1b82d Mon Sep 17 00:00:00 2001
From: ddalvi <ddalvi@redhat.com>
Date: Mon, 26 Aug 2024 08:41:10 -0400
Subject: [PATCH] Add an environment variable and CLI option to disable default
 caching

Signed-off-by: ddalvi <ddalvi@redhat.com>
---
 sdk/python/kfp/cli/compile_.py      | 16 ++++++++++++++++
 sdk/python/kfp/dsl/pipeline_task.py |  3 ++-
 sdk/python/kfp/dsl/structures.py    |  4 ++--
 3 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/sdk/python/kfp/cli/compile_.py b/sdk/python/kfp/cli/compile_.py
index 2bd3bab18c23..1218c20f2bd8 100644
--- a/sdk/python/kfp/cli/compile_.py
+++ b/sdk/python/kfp/cli/compile_.py
@@ -133,14 +133,30 @@ def parse_parameters(parameters: Optional[str]) -> Dict:
     is_flag=True,
     default=False,
     help='Whether to disable type checking.')
+@click.option(
+    '--disable-caching',
+    is_flag=True,
+    default=None,
+    help="Disable task-level caching by default."
+)
 def compile_(
     py: str,
     output: str,
     function_name: Optional[str] = None,
     pipeline_parameters: Optional[str] = None,
     disable_type_check: bool = False,
+    disable_caching: Optional[bool] = None,
 ) -> None:
     """Compiles a pipeline or component written in a .py file."""
+
+    env_disable_caching = os.getenv('KFP_DISABLE_CACHING', 'false').lower() == 'true'
+    if disable_caching is None:
+        disable_caching = env_disable_caching
+    elif disable_caching:
+        os.environ['KFP_DISABLE_CACHING'] = 'true'
+    else:
+        os.environ['KFP_DISABLE_CACHING'] = 'false'
+
     pipeline_func = collect_pipeline_or_component_func(
         python_file=py, function_name=function_name)
     parsed_parameters = parse_parameters(parameters=pipeline_parameters)
diff --git a/sdk/python/kfp/dsl/pipeline_task.py b/sdk/python/kfp/dsl/pipeline_task.py
index 2e82d23378aa..0c7519143f58 100644
--- a/sdk/python/kfp/dsl/pipeline_task.py
+++ b/sdk/python/kfp/dsl/pipeline_task.py
@@ -18,6 +18,7 @@
 import functools
 import inspect
 import itertools
+import os
 import re
 from typing import Any, Dict, List, Mapping, Optional, Union
 import warnings
@@ -130,7 +131,7 @@ def __init__(
             inputs=dict(args.items()),
             dependent_tasks=[],
             component_ref=component_spec.name,
-            enable_caching=True)
+            enable_caching=not os.getenv('KFP_DISABLE_CACHING', 'false').lower() == 'true')
         self._run_after: List[str] = []
 
         self.importer_spec = None
diff --git a/sdk/python/kfp/dsl/structures.py b/sdk/python/kfp/dsl/structures.py
index 440f9a3940af..6485cc7c7975 100644
--- a/sdk/python/kfp/dsl/structures.py
+++ b/sdk/python/kfp/dsl/structures.py
@@ -17,6 +17,7 @@
 import collections
 import dataclasses
 import itertools
+import os
 import re
 from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
 import uuid
@@ -420,11 +421,10 @@ class TaskSpec:
     trigger_strategy: Optional[str] = None
     iterator_items: Optional[Any] = None
     iterator_item_input: Optional[str] = None
-    enable_caching: bool = True
+    enable_caching: bool = not os.getenv('KFP_DISABLE_CACHING', 'false').lower() == 'true'
     display_name: Optional[str] = None
     retry_policy: Optional[RetryPolicy] = None
 
-
 @dataclasses.dataclass
 class ImporterSpec:
     """ImporterSpec definition.