Repository: PrefectHQ/prefect

Issue:
Opt-out of checkpointing on a single Task
## Current behavior
With 0.9.0 Prefect Flows will begin aggressively checkpointing, taking a hint to enable or disable checkpointing based on the existence of a result handler. It may be desired to have select tasks to not checkpoint by specifying `@prefect.task(checkpointing=False)`, however, this option is now ignored.

## Proposed behavior
Allow for select tasks to disable checkpointing with `@prefect.task(checkpointing=False)`. The Flow should continue to checkpoint all other tasks.



Discussion:

Resolving patch:
diff --git a/src/prefect/core/task.py b/src/prefect/core/task.py
--- a/src/prefect/core/task.py
+++ b/src/prefect/core/task.py
@@ -125,11 +125,9 @@ def run(self, x, y):
         - cache_key (str, optional): if provided, a `cache_key` serves as a unique identifier for this Task's cache, and can
             be shared across both Tasks _and_ Flows; if not provided, the Task's _name_ will be used if running locally, or the
             Task's database ID if running in Cloud
-        - checkpoint (bool, optional, DEPRECATED): if this Task is successful, whether to
-            store its result using the `result_handler` available during the run; defaults to the value of
-            `tasks.defaults.checkpoint` in your user config.  Note that in the future, all tasks with result handlers
-            will be checkpointed.  Also note that checkpointing will only occur locally if `prefect.config.flows.checkpointing` is
-            set to `True`
+        - checkpoint (bool, optional): if this Task is successful, whether to
+            store its result using the `result_handler` available during the run; Also note that 
+            checkpointing will only occur locally if `prefect.config.flows.checkpointing` is set to `True`
         - result_handler (ResultHandler, optional): the handler to use for
             retrieving and storing state results during execution; if not provided, will default to the
             one attached to the Flow
@@ -232,16 +230,7 @@ def __init__(
             else prefect.engine.cache_validators.duration_only
         )
         self.cache_validator = cache_validator or default_validator
-        if checkpoint is not None:
-            warnings.warn(
-                "The `checkpoint` keyword will be deprecated in an upcoming release - all tasks will be checkpointed in the future.",
-                UserWarning,
-            )
-        self.checkpoint = (
-            checkpoint
-            if checkpoint is not None
-            else prefect.config.tasks.defaults.checkpoint
-        )
+        self.checkpoint = checkpoint
         self.result_handler = result_handler
 
         if state_handlers and not isinstance(state_handlers, collections.Sequence):
@@ -1028,7 +1017,11 @@ def __init__(
         from prefect.engine.result_handlers import JSONResultHandler
 
         super().__init__(
-            name=name, slug=name, tags=tags, result_handler=JSONResultHandler(),
+            name=name,
+            slug=name,
+            tags=tags,
+            result_handler=JSONResultHandler(),
+            checkpoint=True,
         )
 
     def __repr__(self) -> str:
diff --git a/src/prefect/engine/task_runner.py b/src/prefect/engine/task_runner.py
--- a/src/prefect/engine/task_runner.py
+++ b/src/prefect/engine/task_runner.py
@@ -904,10 +904,11 @@ def get_task_run_state(
             result=result, message="Task run succeeded.", cached_inputs=inputs
         )
 
-        ## checkpoint tasks if a result_handler is present
+        ## checkpoint tasks if a result_handler is present, except for when the user has opted out by disabling checkpointing
         if (
             state.is_successful()
             and prefect.context.get("checkpointing") is True
+            and self.task.checkpoint is not False
             and self.result_handler is not None
         ):
             state._result.store_safe_value()


Test patch:
diff --git a/tests/core/test_task.py b/tests/core/test_task.py
--- a/tests/core/test_task.py
+++ b/tests/core/test_task.py
@@ -253,11 +253,10 @@ def test_create_parameter_uses_json_result_handler(self):
 
     def test_create_task_with_and_without_checkpoint(self):
         t = Task()
-        assert t.checkpoint is False
+        assert t.checkpoint is None
 
-        with pytest.warns(UserWarning, match="deprecated"):
-            s = Task(checkpoint=True)
-            assert s.checkpoint is True
+        s = Task(checkpoint=True)
+        assert s.checkpoint is True
 
 
 def test_task_has_logger():
diff --git a/tests/engine/cloud/test_cloud_task_runner.py b/tests/engine/cloud/test_cloud_task_runner.py
--- a/tests/engine/cloud/test_cloud_task_runner.py
+++ b/tests/engine/cloud/test_cloud_task_runner.py
@@ -655,10 +655,13 @@ def add(x, y):
         assert states[1].is_failed()
         assert "unsupported operand" in states[1].message
 
-    def test_task_runner_sends_checkpointed_success_states_to_cloud(self, client):
+    @pytest.mark.parametrize("checkpoint", [True, None])
+    def test_task_runner_sends_checkpointed_success_states_to_cloud(
+        self, client, checkpoint
+    ):
         handler = JSONResultHandler()
 
-        @prefect.task(checkpoint=True, result_handler=handler)
+        @prefect.task(checkpoint=checkpoint, result_handler=handler)
         def add(x, y):
             return x + y
 
diff --git a/tests/engine/test_task_runner.py b/tests/engine/test_task_runner.py
--- a/tests/engine/test_task_runner.py
+++ b/tests/engine/test_task_runner.py
@@ -1281,10 +1281,11 @@ def fn(x):
         assert new_state.is_successful()
         assert new_state._result.safe_value is NoResult
 
-    def test_success_state_with_checkpointing_in_config(self):
+    @pytest.mark.parametrize("checkpoint", [True, None])
+    def test_success_state_with_checkpointing_in_config(self, checkpoint):
         handler = JSONResultHandler()
 
-        @prefect.task(checkpoint=True, result_handler=handler)
+        @prefect.task(checkpoint=checkpoint, result_handler=handler)
         def fn(x):
             return x + 1
 
@@ -1296,10 +1297,11 @@ def fn(x):
         assert new_state.is_successful()
         assert new_state._result.safe_value == SafeResult("3", result_handler=handler)
 
-    def test_success_state_with_checkpointing_in_context(self):
+    @pytest.mark.parametrize("checkpoint", [True, None])
+    def test_success_state_with_checkpointing_in_context(self, checkpoint):
         handler = JSONResultHandler()
 
-        @prefect.task(checkpoint=True, result_handler=handler)
+        @prefect.task(checkpoint=checkpoint, result_handler=handler)
         def fn(x):
             return x + 1
 
@@ -1310,10 +1312,11 @@ def fn(x):
         assert new_state.is_successful()
         assert new_state._result.safe_value == SafeResult("3", result_handler=handler)
 
-    def test_success_state_is_checkpointed_if_result_handler_present(self):
+    @pytest.mark.parametrize("checkpoint", [True, None])
+    def test_success_state_is_checkpointed_if_result_handler_present(self, checkpoint):
         handler = JSONResultHandler()
 
-        @prefect.task(checkpoint=False, result_handler=handler)
+        @prefect.task(checkpoint=checkpoint, result_handler=handler)
         def fn():
             return 1
 
@@ -1332,6 +1335,31 @@ def fn():
         assert new_state.is_successful()
         assert new_state._result.safe_value == SafeResult("1", result_handler=handler)
 
+    def test_success_state_is_not_checkpointed_if_result_handler_present_and_opted_out(
+        self,
+    ):
+        handler = JSONResultHandler()
+
+        @prefect.task(checkpoint=False, result_handler=handler)
+        def fn():
+            return 1
+
+        ## checkpointing allows users to toggle behavior for local testing
+        with prefect.context(checkpointing=False):
+            new_state = TaskRunner(task=fn, result_handler=handler).get_task_run_state(
+                state=Running(), inputs={}, timeout_handler=None
+            )
+        assert new_state.is_successful()
+        assert new_state._result.safe_value == NoResult
+
+        # We should expect that globally enabling checkpointing does not override the task checkout option (if set)
+        with prefect.context(checkpointing=True):
+            new_state = TaskRunner(task=fn, result_handler=handler).get_task_run_state(
+                state=Running(), inputs={}, timeout_handler=None
+            )
+        assert new_state.is_successful()
+        assert new_state._result.safe_value == NoResult
+
     def test_success_state_for_parameter(self):
         handler = JSONResultHandler()
         p = prefect.Parameter("p", default=2)
@@ -1342,7 +1370,8 @@ def test_success_state_for_parameter(self):
         assert new_state.is_successful()
         assert new_state._result.safe_value == SafeResult("2", result_handler=handler)
 
-    def test_success_state_with_bad_handler_results_in_failed_state(self):
+    @pytest.mark.parametrize("checkpoint", [True, None])
+    def test_success_state_with_bad_handler_results_in_failed_state(self, checkpoint):
         class BadHandler(ResultHandler):
             def read(self, val):
                 pass
@@ -1350,7 +1379,7 @@ def read(self, val):
             def write(self, val):
                 raise SyntaxError("Oh boy")
 
-        @prefect.task(checkpoint=True, result_handler=BadHandler())
+        @prefect.task(checkpoint=checkpoint, result_handler=BadHandler())
         def fn(x):
             return x + 1
 
@@ -1361,6 +1390,24 @@ def fn(x):
         assert new_state.is_failed()
         assert "SyntaxError" in new_state.message
 
+    def test_success_state_with_bad_handler_and_checkpointing_disabled(self):
+        class BadHandler(ResultHandler):
+            def read(self, val):
+                pass
+
+            def write(self, val):
+                raise SyntaxError("Oh boy")
+
+        @prefect.task(checkpoint=False, result_handler=BadHandler())
+        def fn(x):
+            return x + 1
+
+        with prefect.context(checkpointing=True):
+            new_state = TaskRunner(task=fn).get_task_run_state(
+                state=Running(), inputs={"x": Result(1)}, timeout_handler=None
+            )
+        assert new_state.is_successful()
+
 
 class TestCheckRetryStep:
     @pytest.mark.parametrize(
@@ -1921,10 +1968,11 @@ def run(self, x):
     assert state.cached_inputs == dict(x=Result(15))
 
 
-def test_failures_arent_checkpointed():
+@pytest.mark.parametrize("checkpoint", [True, None])
+def test_failures_arent_checkpointed(checkpoint):
     handler = MagicMock(store_safe_value=MagicMock(side_effect=SyntaxError))
 
-    @prefect.task(checkpoint=True, result_handler=handler)
+    @prefect.task(checkpoint=checkpoint, result_handler=handler)
     def fn():
         raise TypeError("Bad types")
 
@@ -1934,10 +1982,11 @@ def fn():
     assert isinstance(new_state.result, TypeError)
 
 
-def test_skips_arent_checkpointed():
+@pytest.mark.parametrize("checkpoint", [True, None])
+def test_skips_arent_checkpointed(checkpoint):
     handler = MagicMock(store_safe_value=MagicMock(side_effect=SyntaxError))
 
-    @prefect.task(checkpoint=True, result_handler=handler)
+    @prefect.task(checkpoint=checkpoint, result_handler=handler)
     def fn():
         return 2
 
@@ -2040,7 +2089,8 @@ def my_task():
         assert state.is_successful()
         assert state.result == 3
 
-    def test_looping_only_checkpoints_the_final_result(self):
+    @pytest.mark.parametrize("checkpoint", [True, None])
+    def test_looping_only_checkpoints_the_final_result(self, checkpoint):
         class Handler(ResultHandler):
             data = []
 
@@ -2053,7 +2103,7 @@ def read(self, idx):
 
         handler = Handler()
 
-        @prefect.task(checkpoint=True, result_handler=handler)
+        @prefect.task(checkpoint=checkpoint, result_handler=handler)
         def my_task():
             curr = prefect.context.get("task_loop_result", 0)
             if prefect.context.get("task_loop_count", 1) < 3: