diff --git a/flytekit/core/base_task.py b/flytekit/core/base_task.py index 9c9607e0719..422be439665 100644 --- 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 index 31190ff1587..ce16e9634d5 100644 --- 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 index 1d4ae259849..1ca6c721b31 100644 --- 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")