[flyte deck] Streaming Decks

flytekit/bin/entrypoint.py

@@ -324,7 +324,7 @@ def _dispatch_execute(
     logger.info(f"Engine folder written successfully to the output prefix {output_prefix}")
 
     if task_def is not None and not getattr(task_def, "disable_deck", True):
-        _output_deck(task_def.name.split(".")[-1], ctx.user_space_params)
+        _output_deck(task_name=task_def.name.split(".")[-1], new_user_params=ctx.user_space_params)
 
     logger.debug("Finished _dispatch_execute")
 

flytekit/core/base_task.py

@@ -720,7 +720,11 @@ def dispatch_execute(
           may be none
         * ``DynamicJobSpec`` is returned when a dynamic workflow is executed
         """
-        if DeckField.TIMELINE.value in self.deck_fields and ctx.user_space_params is not None:
+        # if not self.disable_deck:
+        #     ctx.user_space_params.builder().add_attr("FLYTE_ENABLE_DECK", True)
+        # else:
+        #     ctx.user_space_params.builder().add_attr("FLYTE_ENABLE_DECK", False)
+        if DeckField.TIMELINE.value in self.deck_fields and ctx.user_space_params is not None and not self.disable_deck:
             ctx.user_space_params.decks.append(ctx.user_space_params.timeline_deck)
         # Invoked before the task is executed
         new_user_params = self.pre_execute(ctx.user_space_params)

flytekit/core/constants.py

@@ -38,3 +38,5 @@
 CACHE_KEY_METADATA = "cache-key-metadata"
 
 SERIALIZATION_FORMAT = "serialization-format"
+
+FLYTE_ENABLE_DECK = "FLYTE_ENABLE_DECK"

flytekit/core/context_manager.py

@@ -93,6 +93,7 @@ class Builder(object):
         logging: Optional[_logging.Logger] = None
         task_id: typing.Optional[_identifier.Identifier] = None
         output_metadata_prefix: Optional[str] = None
+        # enable_deck: bool = False
 
         def __init__(self, current: typing.Optional[ExecutionParameters] = None):
             self.stats = current.stats if current else None

flytekit/deck/deck.py

@@ -86,6 +86,12 @@ def name(self) -> str:
     def html(self) -> str:
         return self._html
 
+    @classmethod
+    def publish(cls):
+        params = FlyteContextManager.current_context().user_space_params
+        task_name = params.task_id.name
+        _output_deck(task_name=task_name, new_user_params=params)
+
 
 class TimeLineDeck(Deck):
     """
@@ -148,7 +154,8 @@ def generate_time_table(data: dict) -> str:
 
 
 def _get_deck(
-    new_user_params: ExecutionParameters, ignore_jupyter: bool = False
+    new_user_params: ExecutionParameters,
+    ignore_jupyter: bool = False,
 ) -> typing.Union[str, "IPython.core.display.HTML"]:  # type:ignore
     """
     Get flyte deck html string
@@ -180,7 +187,7 @@ def _output_deck(task_name: str, new_user_params: ExecutionParameters):
     local_path = f"{local_dir}{os.sep}{DECK_FILE_NAME}"
     try:
         with open(local_path, "w", encoding="utf-8") as f:
-            f.write(_get_deck(new_user_params, ignore_jupyter=True))
+            f.write(_get_deck(new_user_params=new_user_params, ignore_jupyter=True))
         logger.info(f"{task_name} task creates flyte deck html to file://{local_path}")
         if ctx.execution_state.mode == ExecutionState.Mode.TASK_EXECUTION:
             fs = ctx.file_access.get_filesystem_for_path(new_user_params.output_metadata_prefix)
@@ -197,6 +204,7 @@ def _output_deck(task_name: str, new_user_params: ExecutionParameters):
 def get_deck_template() -> Template:
     root = os.path.dirname(os.path.abspath(__file__))
     templates_dir = os.path.join(root, "html", "template.html")
+
     with open(templates_dir, "r") as f:
         template_content = f.read()
     return Template(template_content)

flytekit/tools/translator.py

@@ -13,6 +13,7 @@
 from flytekit.core.array_node_map_task import ArrayNodeMapTask
 from flytekit.core.base_task import PythonTask
 from flytekit.core.condition import BranchNode
+from flytekit.core.constants import FLYTE_ENABLE_DECK
 from flytekit.core.container_task import ContainerTask
 from flytekit.core.gate import Gate
 from flytekit.core.launch_plan import LaunchPlan, ReferenceLaunchPlan
@@ -188,8 +189,13 @@ def get_serializable_task(
 
     extra_config = {}
 
-    if hasattr(entity, "task_function") and isinstance(entity.task_function, ClassDecorator):
-        extra_config = entity.task_function.get_extra_config()
+    if hasattr(entity, "task_function"):
+        if isinstance(entity.task_function, ClassDecorator):
+            extra_config = entity.task_function.get_extra_config()
+        if not entity.disable_deck:
+            extra_config[FLYTE_ENABLE_DECK] = "true"
+        else:
+            extra_config[FLYTE_ENABLE_DECK] = "false"
 
     merged_config = {**entity_config, **extra_config}