From 38967f6886df4f2a2c4b86a25c0409a5ecb0e133 Mon Sep 17 00:00:00 2001 From: Aylei Date: Tue, 14 Jan 2025 04:02:43 +0800 Subject: [PATCH] [UX] minor optimizations for launch and introduce py-spy (#4495) * [UX] minor optimizations for launch and introduce py-spy Signed-off-by: Aylei * fix Signed-off-by: Aylei --------- Signed-off-by: Aylei --- CONTRIBUTING.md | 7 ++++--- sky/adaptors/common.py | 24 +++++++++++++++--------- sky/usage/usage_lib.py | 12 ++++++++++-- sky/utils/log_utils.py | 2 ++ 4 files changed, 31 insertions(+), 14 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 85ca90b2c4a..25c6421c347 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -50,9 +50,10 @@ pytest tests/test_smoke.py --generic-cloud azure For profiling code, use: ``` -pip install tuna # Tuna is used for visualization of profiling data. -python3 -m cProfile -o sky.prof -m sky.cli status # Or some other command -tuna sky.prof +pip install py-spy # py-spy is a sampling profiler for Python programs +py-spy record -t -o sky.svg -- python -m sky.cli status # Or some other command +py-spy top -- python -m sky.cli status # Get a live top view +py-spy -h # For more options ``` #### Testing in a container diff --git a/sky/adaptors/common.py b/sky/adaptors/common.py index 0cfb91cb587..d039813cd32 100644 --- a/sky/adaptors/common.py +++ b/sky/adaptors/common.py @@ -1,6 +1,7 @@ """Lazy import for modules to avoid import error when not used.""" import functools import importlib +import threading from typing import Any, Callable, Optional, Tuple @@ -24,17 +25,22 @@ def __init__(self, self._module = None self._import_error_message = import_error_message self._set_loggers = set_loggers + self._lock = threading.RLock() def load_module(self): - if self._module is None: - try: - self._module = importlib.import_module(self._module_name) - if self._set_loggers is not None: - self._set_loggers() - except ImportError as e: - if self._import_error_message is not None: - raise ImportError(self._import_error_message) from e - raise + # Avoid extra imports when multiple threads try to import the same + # module. The overhead is minor since import can only run in serial + # due to GIL even in multi-threaded environments. + with self._lock: + if self._module is None: + try: + self._module = importlib.import_module(self._module_name) + if self._set_loggers is not None: + self._set_loggers() + except ImportError as e: + if self._import_error_message is not None: + raise ImportError(self._import_error_message) from e + raise return self._module def __getattr__(self, name: str) -> Any: diff --git a/sky/usage/usage_lib.py b/sky/usage/usage_lib.py index 07867939ee5..3cc630b3a98 100644 --- a/sky/usage/usage_lib.py +++ b/sky/usage/usage_lib.py @@ -3,7 +3,6 @@ import contextlib import datetime import enum -import inspect import json import os import time @@ -12,19 +11,28 @@ from typing import Any, Callable, Dict, List, Optional, Union import click -import requests import sky from sky import sky_logging +from sky.adaptors import common as adaptors_common from sky.usage import constants from sky.utils import common_utils from sky.utils import env_options from sky.utils import ux_utils if typing.TYPE_CHECKING: + import inspect + + import requests + from sky import resources as resources_lib from sky import status_lib from sky import task as task_lib +else: + # requests and inspect cost ~100ms to load, which can be postponed to + # collection phase or skipped if user specifies no collection + requests = adaptors_common.LazyImport('requests') + inspect = adaptors_common.LazyImport('inspect') logger = sky_logging.init_logger(__name__) diff --git a/sky/utils/log_utils.py b/sky/utils/log_utils.py index a5884333609..5a7d8cfd5f7 100644 --- a/sky/utils/log_utils.py +++ b/sky/utils/log_utils.py @@ -5,6 +5,8 @@ from typing import Callable, Iterator, List, Optional, TextIO, Type import colorama +# slow due to https://github.com/python-pendulum/pendulum/issues/808 +# FIXME(aylei): bump pendulum if it get fixed import pendulum import prettytable