[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

@@ -129,6 +129,7 @@ class TaskMetadata(object):
         timeout (Optional[Union[datetime.timedelta, int]]): the max amount of time for which one execution of this task
             should be executed for. The execution will be terminated if the runtime exceeds the given timeout
             (approximately)
+        :param bool generates_deck: Whether the task will generate a Deck URI.
         pod_template_name (Optional[str]): the name of existing PodTemplate resource in the cluster which will be used in this task.
     """
 
@@ -141,6 +142,7 @@ class TaskMetadata(object):
     retries: int = 0
     timeout: Optional[Union[datetime.timedelta, int]] = None
     pod_template_name: Optional[str] = None
+    generates_deck: bool = False
     is_eager: bool = False
 
     def __post_init__(self):
@@ -179,6 +181,7 @@ def to_taskmetadata_model(self) -> _task_model.TaskMetadata:
             discovery_version=self.cache_version,
             deprecated_error_message=self.deprecated,
             cache_serializable=self.cache_serialize,
+            generates_deck=self.generates_deck,
             pod_template_name=self.pod_template_name,
             cache_ignore_input_vars=self.cache_ignore_input_vars,
             is_eager=self.is_eager,
@@ -720,7 +723,9 @@ 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._enable_deck = True  # type: ignore
+        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
@@ -125,6 +126,7 @@ def build(self) -> ExecutionParameters:
                 raw_output_prefix=self.raw_output_prefix,
                 task_id=self.task_id,
                 output_metadata_prefix=self.output_metadata_prefix,
+                enable_deck=self.enable_deck,
                 **self.attrs,
             )
 
@@ -161,6 +163,7 @@ def __init__(
         checkpoint=None,
         decks=None,
         task_id: typing.Optional[_identifier.Identifier] = None,
+        enable_deck: bool = False,
         **kwargs,
     ):
         """
@@ -188,8 +191,13 @@ def __init__(
         self._checkpoint = checkpoint
         self._decks = decks
         self._task_id = task_id
+        self._enable_deck = enable_deck
         self._timeline_deck = None
 
+    @property
+    def enable_deck(self) -> bool:
+        return self._enable_deck
+
     @property
     def stats(self) -> taggable.TaggableStats:
         """

flytekit/deck/deck.py

@@ -86,6 +86,16 @@ def name(self) -> str:
     def html(self) -> str:
         return self._html
 
+    @classmethod
+    def publish(cls):
+        params = FlyteContextManager.current_context().user_space_params
+        if params.enable_deck:
+            task_name = params.task_id.name
+            _output_deck(task_name=task_name, new_user_params=params)
+        else:
+            # todo: change to a more proper error
+            raise ValueError("Deck is disabled for this task, please don't call Deck.publish()")
+
 
 class TimeLineDeck(Deck):
     """
@@ -148,7 +158,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 +191,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 +208,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/models/task.py

@@ -180,6 +180,7 @@ def __init__(
         pod_template_name,
         cache_ignore_input_vars,
         is_eager: bool = False,
+        generates_deck: bool = False,
     ):
         """
         Information needed at runtime to determine behavior such as whether or not outputs are discoverable, timeouts,
@@ -199,6 +200,7 @@ def __init__(
             receive deprecation warnings.
         :param bool cache_serializable: Whether or not caching operations are executed in serial. This means only a
             single instance over identical inputs is executed, other concurrent executions wait for the cached results.
+        :param bool generates_deck: Whether the task will generate a Deck URI.
         :param pod_template_name: The name of the existing PodTemplate resource which will be used in this task.
         :param cache_ignore_input_vars: Input variables that should not be included when calculating hash for cache.
         :param is_eager:
@@ -214,6 +216,7 @@ def __init__(
         self._pod_template_name = pod_template_name
         self._cache_ignore_input_vars = cache_ignore_input_vars
         self._is_eager = is_eager
+        self._generates_deck = generates_deck
 
     @property
     def is_eager(self):
@@ -295,6 +298,14 @@ def pod_template_name(self):
         """
         return self._pod_template_name
 
+    @property
+    def generates_deck(self) -> bool:
+        """
+        Whether the task will generate a Deck URI.
+        :rtype: bool
+        """
+        return self._generates_deck
+
     @property
     def cache_ignore_input_vars(self):
         """
@@ -315,6 +326,7 @@ def to_flyte_idl(self):
             discovery_version=self.discovery_version,
             deprecated_error_message=self.deprecated_error_message,
             cache_serializable=self.cache_serializable,
+            generates_deck=self.generates_deck,
             pod_template_name=self.pod_template_name,
             cache_ignore_input_vars=self.cache_ignore_input_vars,
             is_eager=self.is_eager,
@@ -338,6 +350,7 @@ def from_flyte_idl(cls, pb2_object: _core_task.TaskMetadata):
             discovery_version=pb2_object.discovery_version,
             deprecated_error_message=pb2_object.deprecated_error_message,
             cache_serializable=pb2_object.cache_serializable,
+            generates_deck=pb2_object.generates_deck,
             pod_template_name=pb2_object.pod_template_name,
             cache_ignore_input_vars=pb2_object.cache_ignore_input_vars,
             is_eager=pb2_object.is_eager,

flytekit/tools/translator.py

@@ -185,11 +185,13 @@ def get_serializable_task(
                 entity.reset_command_fn()
 
     entity_config = entity.get_config(settings) or {}
-
     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:
+            entity.metadata.generates_deck = True
 
     merged_config = {**entity_config, **extra_config}