More time info for time line deck (#1680)

* more visualization

Signed-off-by: Yicheng-Lu-llll <luyc58576@gmail.com>

* more visualization

Signed-off-by: Yicheng-Lu-llll <luyc58576@gmail.com>

---------

Signed-off-by: Yicheng-Lu-llll <luyc58576@gmail.com>

flytekit/bin/entrypoint.py

@@ -24,6 +24,7 @@
 from flytekit.core.data_persistence import FileAccessProvider
 from flytekit.core.map_task import MapTaskResolver
 from flytekit.core.promise import VoidPromise
+from flytekit.deck.deck import _output_deck
 from flytekit.exceptions import scopes as _scoped_exceptions
 from flytekit.exceptions import scopes as _scopes
 from flytekit.interfaces.stats.taggable import get_stats as _get_stats
@@ -157,6 +158,7 @@ def _dispatch_execute(
         utils.write_proto_to_file(v.to_flyte_idl(), os.path.join(ctx.execution_state.engine_dir, k))
 
     ctx.file_access.put_data(ctx.execution_state.engine_dir, output_prefix, is_multipart=True)
+    _output_deck(task_def.name.split(".")[-1], ctx.user_space_params)
     logger.info(f"Engine folder written successfully to the output prefix {output_prefix}")
     logger.debug("Finished _dispatch_execute")
 

flytekit/core/base_task.py

@@ -499,7 +499,6 @@ def dispatch_execute(
 
         # Invoked before the task is executed
         new_user_params = self.pre_execute(ctx.user_space_params)
-        from flytekit.deck.deck import _output_deck
 
         # Create another execution context with the new user params, but let's keep the same working dir
         with FlyteContextManager.with_context(
@@ -590,8 +589,6 @@ def dispatch_execute(
                 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)))
 
-                _output_deck(self.name.split(".")[-1], new_user_params)
-
             outputs_literal_map = _literal_models.LiteralMap(literals=literals)
             # After the execute has been successfully completed
             return outputs_literal_map

flytekit/core/data_persistence.py

@@ -280,7 +280,6 @@ def upload_directory(self, local_path: str, remote_path: str):
         """
         return self.put_data(local_path, remote_path, is_multipart=True)
 
-    @timeit("Download data to local from remote")
     def get_data(self, remote_path: str, local_path: str, is_multipart: bool = False):
         """
         :param remote_path:
@@ -289,14 +288,14 @@ def get_data(self, remote_path: str, local_path: str, is_multipart: bool = False
         """
         try:
             pathlib.Path(local_path).parent.mkdir(parents=True, exist_ok=True)
-            self.get(remote_path, to_path=local_path, recursive=is_multipart)
+            with timeit(f"Download data to local from {remote_path}"):
+                self.get(remote_path, to_path=local_path, recursive=is_multipart)
         except Exception as ex:
             raise FlyteAssertion(
                 f"Failed to get data from {remote_path} to {local_path} (recursive={is_multipart}).\n\n"
                 f"Original exception: {str(ex)}"
             )
 
-    @timeit("Upload data to remote")
     def put_data(self, local_path: Union[str, os.PathLike], remote_path: str, is_multipart: bool = False, **kwargs):
         """
         The implication here is that we're always going to put data to the remote location, so we .remote to ensure
@@ -308,8 +307,8 @@ def put_data(self, local_path: Union[str, os.PathLike], remote_path: str, is_mul
         """
         try:
             local_path = str(local_path)
-
-            self.put(cast(str, local_path), remote_path, recursive=is_multipart, **kwargs)
+            with timeit(f"Upload data to {remote_path}"):
+                self.put(cast(str, local_path), remote_path, recursive=is_multipart, **kwargs)
         except Exception as ex:
             raise FlyteAssertion(
                 f"Failed to put data from {local_path} to {remote_path} (recursive={is_multipart}).\n\n"