Add enable_deck to task decorator (flyteorg#1898)

* Add enable_deck to task decorator Signed-off-by: Eduardo Apolinario <[email protected]> * Actually handle the case where none are set Signed-off-by: Eduardo Apolinario <[email protected]> --------- Signed-off-by: Eduardo Apolinario <[email protected]> Co-authored-by: Eduardo Apolinario <[email protected]> Signed-off-by: Future Outlier <[email protected]>
Future-Outlier · Oct 20, 2023 · e84f001 · e84f001
1 parent 0e31010
commit e84f001
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 19 deletions.
diff --git a/flytekit/core/base_task.py b/flytekit/core/base_task.py
@@ -413,7 +413,8 @@ def __init__(
         task_config: Optional[T],
         interface: Optional[Interface] = None,
         environment: Optional[Dict[str, str]] = None,
-        disable_deck: bool = True,
+        disable_deck: Optional[bool] = None,
+        enable_deck: Optional[bool] = None,
         use_sync_plugin: bool = False,
         **kwargs,
     ):
@@ -428,7 +429,8 @@ def __init__(
                 signature of the task
             environment (Optional[Dict[str, str]]): Any environment variables that should be supplied during the
                 execution of the task. Supplied as a dictionary of key/value pairs
-            disable_deck (bool): If true, this task will not output deck html file
+            disable_deck (bool): (deprecated) If true, this task will not output deck html file
+            enable_deck (bool): If true, this task will output deck html file
             use_sync_plugin (bool): If true, this task will invoke sync plugin in flytepropeller and flyteplugin
         """
         super().__init__(
@@ -441,7 +443,17 @@ def __init__(
         self._python_interface = interface if interface else Interface()
         self._environment = environment if environment else {}
         self._task_config = task_config
-        self._disable_deck = disable_deck
+
+        # Confirm that disable_deck and enable_deck do not contradict each other
+        if disable_deck is not None and enable_deck is not None:
+            raise ValueError("disable_deck and enable_deck cannot both be set at the same time")
+
+        if enable_deck is not None:
+            self._disable_deck = not enable_deck
+        elif disable_deck is not None:
+            self._disable_deck = disable_deck
+        else:
+            self._disable_deck = True
         if self._python_interface.docstring:
             if self.docs is None:
                 self._docs = Documentation(

diff --git a/flytekit/core/task.py b/flytekit/core/task.py
@@ -98,7 +98,8 @@ def task(
     execution_mode: PythonFunctionTask.ExecutionBehavior = ...,
     task_resolver: Optional[TaskResolverMixin] = ...,
     docs: Optional[Documentation] = ...,
-    disable_deck: bool = ...,
+    disable_deck: Optional[bool] = ...,
+    enable_deck: Optional[bool] = ...,
     pod_template: Optional["PodTemplate"] = ...,
     pod_template_name: Optional[str] = ...,
 ) -> Callable[[Callable[..., FuncOut]], PythonFunctionTask[T]]:
@@ -124,7 +125,8 @@ def task(
     execution_mode: PythonFunctionTask.ExecutionBehavior = ...,
     task_resolver: Optional[TaskResolverMixin] = ...,
     docs: Optional[Documentation] = ...,
-    disable_deck: bool = ...,
+    disable_deck: Optional[bool] = ...,
+    enable_deck: Optional[bool] = ...,
     pod_template: Optional["PodTemplate"] = ...,
     pod_template_name: Optional[str] = ...,
 ) -> Union[PythonFunctionTask[T], Callable[..., FuncOut]]:
@@ -149,7 +151,8 @@ def task(
     execution_mode: PythonFunctionTask.ExecutionBehavior = PythonFunctionTask.ExecutionBehavior.DEFAULT,
     task_resolver: Optional[TaskResolverMixin] = None,
     docs: Optional[Documentation] = None,
-    disable_deck: bool = True,
+    disable_deck: Optional[bool] = None,
+    enable_deck: Optional[bool] = None,
     pod_template: Optional["PodTemplate"] = None,
     pod_template_name: Optional[str] = None,
 ) -> Union[Callable[[Callable[..., FuncOut]], PythonFunctionTask[T]], PythonFunctionTask[T], Callable[..., FuncOut]]:
@@ -240,7 +243,8 @@ def foo2():
                      may change based on the backend provider.
     :param execution_mode: This is mainly for internal use. Please ignore. It is filled in automatically.
     :param task_resolver: Provide a custom task resolver.
-    :param disable_deck: If true, this task will not output deck html file
+    :param disable_deck: (deprecated) If true, this task will not output deck html file
+    :param enable_deck: If true, this task will output deck html file
     :param docs: Documentation about this task
     :param pod_template: Custom PodTemplate for this task.
     :param pod_template_name: The name of the existing PodTemplate resource which will be used in this task.
@@ -269,6 +273,7 @@ def wrapper(fn: Callable[..., Any]) -> PythonFunctionTask[T]:
             execution_mode=execution_mode,
             task_resolver=task_resolver,
             disable_deck=disable_deck,
+            enable_deck=enable_deck,
             docs=docs,
             pod_template=pod_template,
             pod_template_name=pod_template_name,

diff --git a/tests/flytekit/unit/deck/test_deck.py b/tests/flytekit/unit/deck/test_deck.py
@@ -67,28 +67,46 @@ def t1(a: int) -> str:
 
 
 @pytest.mark.parametrize(
-    "disable_deck, expected_decks",
+    "enable_deck,disable_deck, expected_decks, expect_error",
     [
-        (None, 2),  # default deck and time line deck
-        (False, 4),  # default deck and time line deck + input and output decks
-        (True, 2),  # default deck and time line deck
+        (None, None, 2, False),  # default deck and time line deck
+        (None, False, 4, False),  # default deck and time line deck + input and output decks
+        (None, True, 2, False),  # default deck and time line deck
+        (True, None, 4, False),  # default deck and time line deck + input and output decks
+        (False, None, 2, False),  # default deck and time line deck
+        (True, True, -1, True),  # Set both disable_deck and enable_deck to True and confirm that it fails
+        (False, False, -1, True),  # Set both disable_deck and enable_deck to False and confirm that it fails
     ],
 )
-def test_deck_pandas_dataframe(disable_deck, expected_decks):
+def test_deck_pandas_dataframe(enable_deck, disable_deck, expected_decks, expect_error):
     ctx = FlyteContextManager.current_context()
 
     kwargs = {}
     if disable_deck is not None:
         kwargs["disable_deck"] = disable_deck
 
-    @task(**kwargs)
-    def t_df(a: str) -> int:
-        df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
-        flytekit.current_context().default_deck.append(TopFrameRenderer().to_html(df))
-        return int(a)
+    if enable_deck is not None:
+        kwargs["enable_deck"] = enable_deck
 
-    t_df(a="42")
-    assert len(ctx.user_space_params.decks) == expected_decks
+    if expect_error:
+        with pytest.raises(ValueError):
+
+            @task(**kwargs)
+            def t_df(a: str) -> int:
+                df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+                flytekit.current_context().default_deck.append(TopFrameRenderer().to_html(df))
+                return int(a)
+
+    else:
+
+        @task(**kwargs)
+        def t_df(a: str) -> int:
+            df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+            flytekit.current_context().default_deck.append(TopFrameRenderer().to_html(df))
+            return int(a)
+
+        t_df(a="42")
+        assert len(ctx.user_space_params.decks) == expected_decks
 
 
 @mock.patch("flytekit.deck.deck.ipython_check")