diff --git a/flytekit/bin/entrypoint.py b/flytekit/bin/entrypoint.py index 49103319d0..d41c95d6da 100644 --- a/flytekit/bin/entrypoint.py +++ b/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") diff --git a/flytekit/core/base_task.py b/flytekit/core/base_task.py index 6430aa9eac..527f44a1df 100644 --- a/flytekit/core/base_task.py +++ b/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) diff --git a/flytekit/core/constants.py b/flytekit/core/constants.py index 903e5d5ced..4cad8a98cc 100644 --- a/flytekit/core/constants.py +++ b/flytekit/core/constants.py @@ -38,3 +38,5 @@ CACHE_KEY_METADATA = "cache-key-metadata" SERIALIZATION_FORMAT = "serialization-format" + +FLYTE_ENABLE_DECK = "FLYTE_ENABLE_DECK" diff --git a/flytekit/core/context_manager.py b/flytekit/core/context_manager.py index d804cbddc8..7aca647d84 100644 --- a/flytekit/core/context_manager.py +++ b/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 diff --git a/flytekit/deck/deck.py b/flytekit/deck/deck.py index 025306d47b..201ae61245 100644 --- a/flytekit/deck/deck.py +++ b/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) diff --git a/flytekit/tools/translator.py b/flytekit/tools/translator.py index 5c7a6d5eb4..7d12aec82e 100644 --- a/flytekit/tools/translator.py +++ b/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}