improve tracing reporting and coverage (#4331)

* [tracing] report stack backtrace as part of the tracing timeline

* [tracing] add tracing to more expensive functions

sky/backends/backend_utils.py

@@ -1418,6 +1418,7 @@ def check_network_connection():
                                               'Network seems down.') from e
 
 
+@timeline.event
 def check_owner_identity(cluster_name: str) -> None:
     """Check if current user is the same as the user who created the cluster.
 

sky/clouds/cloud.py

@@ -18,6 +18,7 @@
 from sky.clouds import service_catalog
 from sky.utils import log_utils
 from sky.utils import resources_utils
+from sky.utils import timeline
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
@@ -366,6 +367,7 @@ def is_label_valid(cls, label_key: str,
         del label_key, label_value
         return True, None
 
+    @timeline.event
     def get_feasible_launchable_resources(
             self,
             resources: 'resources_lib.Resources',

sky/jobs/core.py

@@ -26,9 +26,11 @@
 from sky.utils import dag_utils
 from sky.utils import rich_utils
 from sky.utils import subprocess_utils
+from sky.utils import timeline
 from sky.utils import ux_utils
 
 
+@timeline.event
 @usage_lib.entrypoint
 def launch(
     task: Union['sky.Task', 'sky.Dag'],

sky/optimizer.py

@@ -22,6 +22,7 @@
 from sky.utils import resources_utils
 from sky.utils import rich_utils
 from sky.utils import subprocess_utils
+from sky.utils import timeline
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
@@ -105,6 +106,7 @@ def _egress_time(src_cloud: clouds.Cloud, dst_cloud: clouds.Cloud,
         return egress_time
 
     @staticmethod
+    @timeline.event
     def optimize(dag: 'dag_lib.Dag',
                  minimize: OptimizeTarget = OptimizeTarget.COST,
                  blocked_resources: Optional[Iterable[

sky/provision/__init__.py

@@ -24,6 +24,7 @@
 from sky.provision import runpod
 from sky.provision import vsphere
 from sky.utils import command_runner
+from sky.utils import timeline
 
 logger = sky_logging.init_logger(__name__)
 
@@ -59,6 +60,7 @@ def _wrapper(*args, **kwargs):
 # pylint: disable=unused-argument
 
 
+@timeline.event
 @_route_to_cloud_impl
 def query_instances(
     provider_name: str,

sky/utils/timeline.py

@@ -9,6 +9,7 @@
 import os
 import threading
 import time
+import traceback
 from typing import Callable, Optional, Union
 
 import filelock
@@ -48,8 +49,9 @@ def begin(self):
             'ph': 'B',
             'ts': f'{time.time() * 10 ** 6: .3f}',
         })
+        event_begin['args'] = {'stack': '\n'.join(traceback.format_stack())}
         if self._message is not None:
-            event_begin['args'] = {'message': self._message}
+            event_begin['args']['message'] = self._message
         _events.append(event_begin)
 
     def end(self):