fix UI

Signed-off-by: novahow <[email protected]>

flytekit/core/base_task.py

@@ -491,15 +491,6 @@ def __init__(
         self._python_interface = interface if interface else Interface()
         self._environment = environment if environment else {}
         self._task_config = task_config
-        self._decks = list(decks) if (decks is not None and (enable_deck is True or disable_deck is False)) else []
-
-        deck_members = set([_field.value for _field in DeckFields])
-        # enumerate additional decks, check if any of them are invalid
-        for deck in self._decks:
-            if deck not in deck_members:
-                raise ValueError(
-                    f"Element {deck} from decks param is not a valid deck field. Please use one of {deck_members}"
-                )
 
         # first we resolve the conflict between params regarding decks, if any two of [disable_deck, enable_deck]
         # are set, we raise an error
@@ -520,6 +511,16 @@ def __init__(
         else:
             self._disable_deck = True
 
+        self._decks = list(decks) if (decks is not None and self.disable_deck is False) else []
+
+        deck_members = set([_field.value for _field in DeckFields])
+        # enumerate additional decks, check if any of them are invalid
+        for deck in self._decks:
+            if deck not in deck_members:
+                raise ValueError(
+                    f"Element {deck} from decks param is not a valid deck field. Please use one of {deck_members}"
+                )
+
         if self._python_interface.docstring:
             if self.docs is None:
                 self._docs = Documentation(
@@ -664,13 +665,15 @@ def _write_decks(self, native_inputs, native_outputs_as_map, ctx, new_user_param
             INPUT = DeckFields.INPUT
             OUTPUT = DeckFields.OUTPUT
 
-            input_deck = Deck(INPUT.value, auto_add_to_deck=DeckFields.INPUT in self.decks)
-            for k, v in native_inputs.items():
-                input_deck.append(TypeEngine.to_html(ctx, v, self.get_type_for_input_var(k, v)))
+            if DeckFields.INPUT in self.decks:
+                input_deck = Deck(INPUT.value)
+                for k, v in native_inputs.items():
+                    input_deck.append(TypeEngine.to_html(ctx, v, self.get_type_for_input_var(k, v)))
 
-            output_deck = Deck(OUTPUT.value, auto_add_to_deck=DeckFields.OUTPUT in self.decks)
-            for k, v in native_outputs_as_map.items():
-                output_deck.append(TypeEngine.to_html(ctx, v, self.get_type_for_output_var(k, v)))
+            if DeckFields.OUTPUT in self.decks:
+                output_deck = Deck(OUTPUT.value)
+                for k, v in native_outputs_as_map.items():
+                    output_deck.append(TypeEngine.to_html(ctx, v, self.get_type_for_output_var(k, v)))
 
             if ctx.execution_state and ctx.execution_state.is_local_execution():
                 # When we run the workflow remotely, flytekit outputs decks at the end of _dispatch_execute

flytekit/core/context_manager.py

@@ -289,7 +289,7 @@ def timeline_deck(self) -> "TimeLineDeck":  # type: ignore
             if self._timeline_deck is not None:
                 time_line_deck = self._timeline_deck
             else:
-                time_line_deck = TimeLineDeck(DeckFields.TIMELINE.value)
+                time_line_deck = TimeLineDeck(DeckFields.TIMELINE.value, auto_add_to_deck=False)
 
         self._timeline_deck = time_line_deck
         return time_line_deck

flytekit/core/python_function_task.py

@@ -359,16 +359,14 @@ def _write_decks(self, native_inputs, native_outputs_as_map, ctx, new_user_param
                 source_code = inspect.getsource(self._task_function)
                 from flytekit.deck.renderer import SourceCodeRenderer
 
-                source_code_deck = Deck(
-                    DeckFields.SOURCE_CODE.value, auto_add_to_deck=DeckFields.SOURCE_CODE in self.decks
-                )
-                renderer = SourceCodeRenderer()
-                source_code_deck.append(renderer.to_html(source_code))
-
-            python_dependencies_deck = Deck(
-                DeckFields.DEPENDENCIES.value, auto_add_to_deck=DeckFields.DEPENDENCIES in self.decks
-            )
-            renderer = PythonDependencyRenderer()
-            python_dependencies_deck.append(renderer.to_html())
+                if DeckFields.SOURCE_CODE in self.decks:
+                    source_code_deck = Deck(DeckFields.SOURCE_CODE.value)
+                    renderer = SourceCodeRenderer()
+                    source_code_deck.append(renderer.to_html(source_code))
+
+            if DeckFields.DEPENDENCIES in self.decks:
+                python_dependencies_deck = Deck(DeckFields.DEPENDENCIES.value)
+                renderer = PythonDependencyRenderer()
+                python_dependencies_deck.append(renderer.to_html())
 
         return super()._write_decks(native_inputs, native_outputs_as_map, ctx, new_user_params)

flytekit/core/task.py

@@ -191,7 +191,7 @@ def task(
     docs: Optional[Documentation] = None,
     disable_deck: Optional[bool] = None,
     enable_deck: Optional[bool] = None,
-    decks: Optional[Tuple[str, ...]] = ("source_code", "dependencies"),
+    decks: Optional[Tuple[str, ...]] = ("Source Code", "Dependencies"),
     pod_template: Optional["PodTemplate"] = None,
     pod_template_name: Optional[str] = None,
     accelerator: Optional[BaseAccelerator] = None,

flytekit/deck/deck.py

@@ -16,11 +16,11 @@ class DeckFields(str, enum.Enum):
     DeckFields is used to specify the fields that will be rendered in the deck.
     """
 
-    INPUT = "input"
-    OUTPUT = "output"
-    SOURCE_CODE = "source_code"
+    INPUT = "Input"
+    OUTPUT = "Output"
+    SOURCE_CODE = "Source Code"
     TIMELINE = "Timeline"
-    DEPENDENCIES = "dependencies"
+    DEPENDENCIES = "Dependencies"
 
 
 class Deck:
@@ -93,7 +93,7 @@ class TimeLineDeck(Deck):
     Instead, the complete data set is used to create a comprehensive visualization of the execution time of each part of the task.
     """
 
-    def __init__(self, name: str, html: Optional[str] = "", auto_add_to_deck: bool = False):
+    def __init__(self, name: str, html: Optional[str] = "", auto_add_to_deck: bool = True):
         super().__init__(name, html, auto_add_to_deck)
         self.time_info = []