skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/server/stream_utils.py

@@ -0,0 +1,141 @@
+"""Utilities for streaming logs from response."""
+
+import asyncio
+import collections
+import pathlib
+from typing import AsyncGenerator, Deque, Optional
+
+import aiofiles
+import fastapi
+
+from sky import sky_logging
+from sky.server.requests import requests as requests_lib
+from sky.utils import message_utils
+from sky.utils import rich_utils
+
+logger = sky_logging.init_logger(__name__)
+
+
+async def _yield_log_file_with_payloads_skipped(
+        log_file) -> AsyncGenerator[str, None]:
+    async for line in log_file:
+        if not line:
+            return
+        is_payload, line_str = message_utils.decode_payload(
+            line.decode('utf-8'), raise_for_mismatch=False)
+        if is_payload:
+            continue
+
+        yield line_str
+
+
+async def log_streamer(request_id: Optional[str],
+                       log_path: pathlib.Path,
+                       plain_logs: bool = False,
+                       tail: Optional[int] = None,
+                       follow: bool = True) -> AsyncGenerator[str, None]:
+    """Streams the logs of a request."""
+
+    if request_id is not None:
+        status_msg = rich_utils.EncodedStatusMessage(
+            f'[dim]Checking request: {request_id}[/dim]')
+        request_task = requests_lib.get_request(request_id)
+
+        if request_task is None:
+            raise fastapi.HTTPException(
+                status_code=404, detail=f'Request {request_id} not found')
+        request_id = request_task.request_id
+
+        # Do not show the waiting spinner if the request is a fast, non-blocking
+        # request.
+        show_request_waiting_spinner = (not plain_logs and
+                                        request_task.schedule_type
+                                        == requests_lib.ScheduleType.LONG)
+
+        if show_request_waiting_spinner:
+            yield status_msg.init()
+            yield status_msg.start()
+        is_waiting_msg_logged = False
+        waiting_msg = (f'Waiting for {request_task.name!r} request to be '
+                       f'scheduled: {request_id}')
+        while request_task.status < requests_lib.RequestStatus.RUNNING:
+            if show_request_waiting_spinner:
+                yield status_msg.update(f'[dim]{waiting_msg}[/dim]')
+            elif plain_logs and not is_waiting_msg_logged:
+                is_waiting_msg_logged = True
+                # Use smaller padding (1024 bytes) to force browser rendering
+                yield f'{waiting_msg}' + ' ' * 4096 + '\n'
+            # Sleep 0 to yield, so other coroutines can run. This busy waiting
+            # loop is performance critical for short-running requests, so we do
+            # not want to yield too long.
+            await asyncio.sleep(0.1)
+            request_task = requests_lib.get_request(request_id)
+            if not follow:
+                break
+        if show_request_waiting_spinner:
+            yield status_msg.stop()
+
+    # Find last n lines of the log file. Do not read the whole file into memory.
+    async with aiofiles.open(log_path, 'rb') as f:
+        if tail is not None:
+            # TODO(zhwu): this will include the control lines for rich status,
+            # which may not lead to exact tail lines when showing on the client
+            # side.
+            lines: Deque[str] = collections.deque(maxlen=tail)
+            async for line_str in _yield_log_file_with_payloads_skipped(f):
+                lines.append(line_str)
+            for line_str in lines:
+                yield line_str
+
+        while True:
+            # Sleep 0 to yield control to allow other coroutines to run,
+            # while keeps the loop tight to make log stream responsive.
+            await asyncio.sleep(0)
+            line: Optional[bytes] = await f.readline()
+            if not line:
+                if request_id is not None:
+                    request_task = requests_lib.get_request(request_id)
+                    if request_task.status > requests_lib.RequestStatus.RUNNING:
+                        if (request_task.status ==
+                                requests_lib.RequestStatus.CANCELLED):
+                            yield (f'{request_task.name!r} request {request_id}'
+                                   ' cancelled\n')
+                        break
+                if not follow:
+                    break
+                # Sleep shortly to avoid storming the DB and CPU, this has
+                # little impact on the responsivness here since we are waiting
+                # for a new line to come in.
+                await asyncio.sleep(0.1)
+                continue
+            line_str = line.decode('utf-8')
+            if plain_logs:
+                is_payload, line_str = message_utils.decode_payload(
+                    line_str, raise_for_mismatch=False)
+                if is_payload:
+                    continue
+            yield line_str
+
+
+def stream_response(
+    request_id: str, logs_path: pathlib.Path,
+    background_tasks: fastapi.BackgroundTasks
+) -> fastapi.responses.StreamingResponse:
+
+    async def on_disconnect():
+        logger.info(f'User terminated the connection for request '
+                    f'{request_id}')
+        requests_lib.kill_requests([request_id])
+
+    # The background task will be run after returning a response.
+    # https://fastapi.tiangolo.com/tutorial/background-tasks/
+    background_tasks.add_task(on_disconnect)
+
+    return fastapi.responses.StreamingResponse(
+        log_streamer(request_id, logs_path),
+        media_type='text/plain',
+        headers={
+            'Cache-Control': 'no-cache, no-transform',
+            'X-Accel-Buffering': 'no',
+            'Transfer-Encoding': 'chunked'
+        })

sky/setup_files/MANIFEST.in

@@ -1,15 +1,11 @@
 include sky/backends/monkey_patches/*.py
 exclude sky/clouds/service_catalog/data_fetchers/analyze.py
 include sky/provision/kubernetes/manifests/*
+include sky/provision/azure/*
 include sky/setup_files/*
 include sky/skylet/*.sh
 include sky/skylet/LICENSE
-include sky/skylet/providers/azure/*
-include sky/skylet/providers/gcp/*
 include sky/skylet/providers/ibm/*
-include sky/skylet/providers/kubernetes/*
-include sky/skylet/providers/lambda_cloud/*
-include sky/skylet/providers/oci/*
 include sky/skylet/providers/scp/*
 include sky/skylet/providers/*.py
 include sky/skylet/ray_patches/*.patch
@@ -18,3 +14,4 @@ include sky/jobs/dashboard/templates/*
 include sky/jobs/dashboard/static/*
 include sky/templates/*
 include sky/utils/kubernetes/*
+include sky/server/html/*

sky/setup_files/dependencies.py

@@ -0,0 +1,159 @@
+"""Dependencies for SkyPilot.
+
+This file is imported by setup.py, so:
+- It may not be able to import other skypilot modules, since sys.path may not be
+  correct.
+- It should not import any dependencies, as they may not be installed yet.
+"""
+import sys
+from typing import Dict, List
+
+install_requires = [
+    'wheel',
+    'cachetools',
+    # NOTE: ray requires click>=7.0.
+    'click >= 7.0',
+    'colorama',
+    'cryptography',
+    # Jinja has a bug in older versions because of the lack of pinning
+    # the version of the underlying markupsafe package. See:
+    # https://github.com/pallets/jinja/issues/1585
+    'jinja2 >= 3.0',
+    'jsonschema',
+    'networkx',
+    'pandas>=1.3.0',
+    'pendulum',
+    # PrettyTable with version >=2.0.0 is required for the support of
+    # `add_rows` method.
+    'PrettyTable >= 2.0.0',
+    'python-dotenv',
+    'rich',
+    'tabulate',
+    # Light weight requirement, can be replaced with "typing" once
+    # we deprecate Python 3.7 (this will take a while).
+    'typing_extensions',
+    'filelock >= 3.6.0',
+    'packaging',
+    'psutil',
+    'pulp',
+    # Cython 3.0 release breaks PyYAML 5.4.*
+    # (https://github.com/yaml/pyyaml/issues/601)
+    # <= 3.13 may encounter https://github.com/ultralytics/yolov5/issues/414
+    'pyyaml > 3.13, != 5.4.*',
+    'requests',
+    'fastapi',
+    'uvicorn[standard]',
+    # Some pydantic versions are not compatible with ray. Adopted from ray's
+    # setup.py:
+    # https://github.com/ray-project/ray/blob/ray-2.9.3/python/setup.py#L254
+    # We need pydantic>=2.0.0 for API server and client.
+    'pydantic!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,<3,>2',
+    # Required for Form data by pydantic
+    'python-multipart',
+    'aiofiles',
+    'httpx',
+    'setproctitle',
+]
+
+local_ray = [
+    # Lower version of ray will cause dependency conflict for
+    # click/grpcio/protobuf.
+    # Excluded 2.6.0 as it has a bug in the cluster launcher:
+    # https://github.com/ray-project/ray/releases/tag/ray-2.6.1
+    'ray[default] >= 2.2.0, != 2.6.0',
+]
+
+remote = [
+    # Adopted from ray's setup.py:
+    # https://github.com/ray-project/ray/blob/ray-2.9.3/python/setup.py#L251-L252
+    # SkyPilot: != 1.48.0 is required to avoid the error where ray dashboard
+    # fails to start when ray start is called (#2054).
+    # Tracking issue: https://github.com/ray-project/ray/issues/30984
+    'grpcio >= 1.32.0, != 1.48.0; python_version < \'3.10\'',
+    'grpcio >= 1.42.0, != 1.48.0; python_version >= \'3.10\'',
+    # Adopted from ray's setup.py:
+    # https://github.com/ray-project/ray/blob/ray-2.9.3/python/setup.py#L343
+    'protobuf >= 3.15.3, != 3.19.5',
+]
+
+# NOTE: Change the templates/jobs-controller.yaml.j2 file if any of the
+# following packages dependencies are changed.
+aws_dependencies = [
+    # botocore does not work with urllib3>=2.0.0, according to
+    # https://github.com/boto/botocore/issues/2926
+    # We have to explicitly pin the version to optimize the time for
+    # poetry install. See https://github.com/orgs/python-poetry/discussions/7937
+    'urllib3<2',
+    # NOTE: this installs CLI V1. To use AWS SSO (e.g., `aws sso login`), users
+    # should instead use CLI V2 which is not pip-installable. See
+    # https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html.
+    'awscli>=1.27.10',
+    'botocore>=1.29.10',
+    'boto3>=1.26.1',
+    # NOTE: required by awscli. To avoid ray automatically installing
+    # the latest version.
+    'colorama < 0.4.5',
+]
+
+# azure-cli cannot be installed normally by uv, so we need to work around it in
+# a few places.
+AZURE_CLI = 'azure-cli>=2.65.0'
+
+extras_require: Dict[str, List[str]] = {
+    'aws': aws_dependencies,
+    # TODO(zongheng): azure-cli is huge and takes a long time to install.
+    # Tracked in: https://github.com/Azure/azure-cli/issues/7387
+    # azure-identity is needed in node_provider.
+    # We need azure-identity>=1.13.0 to enable the customization of the
+    # timeout of AzureCliCredential.
+    'azure': [
+        AZURE_CLI,
+        'azure-core>=1.31.0',
+        'azure-identity>=1.19.0',
+        'azure-mgmt-network>=27.0.0',
+        'azure-mgmt-compute>=33.0.0',
+        'azure-storage-blob>=12.23.1',
+        'msgraph-sdk',
+    ] + local_ray,
+    # We need google-api-python-client>=2.69.0 to enable 'discardLocalSsd'
+    # parameter for stopping instances. Reference:
+    # https://github.com/googleapis/google-api-python-client/commit/f6e9d3869ed605b06f7cbf2e8cf2db25108506e6
+    'gcp': ['google-api-python-client>=2.69.0', 'google-cloud-storage'],
+    'ibm': [
+        'ibm-cloud-sdk-core', 'ibm-vpc', 'ibm-platform-services', 'ibm-cos-sdk'
+    ] + local_ray,
+    'docker': ['docker'] + local_ray,
+    'lambda': local_ray,
+    'cloudflare': aws_dependencies,
+    'scp': local_ray,
+    'oci': ['oci'] + local_ray,
+    # Kubernetes 32.0.0 has an authentication bug: https://github.com/kubernetes-client/python/issues/2333 # pylint: disable=line-too-long
+    'kubernetes': ['kubernetes>=20.0.0,!=32.0.0'],
+    'remote': remote,
+    # For the container registry auth api. Reference:
+    # https://github.com/runpod/runpod-python/releases/tag/1.6.1
+    'runpod': ['runpod>=1.6.1'],
+    'fluidstack': [],  # No dependencies needed for fluidstack
+    'cudo': ['cudo-compute>=0.1.10'],
+    'paperspace': [],  # No dependencies needed for paperspace
+    'do': ['pydo>=0.3.0', 'azure-core>=1.24.0', 'azure-common'],
+    'vast': ['vastai-sdk>=0.1.12'],
+    'vsphere': [
+        'pyvmomi==8.0.1.0.2',
+        # vsphere-automation-sdk is also required, but it does not have
+        # pypi release, which cause failure of our pypi release.
+        # https://peps.python.org/pep-0440/#direct-references
+        # We have the instruction for its installation in our
+        # docs instead.
+        # 'vsphere-automation-sdk @ git+https://github.com/vmware/vsphere-automation-sdk-python.git@v8.0.1.0' pylint: disable=line-too-long
+    ],
+    'nebius': ['nebius>=0.2.0',]
+}
+
+# Nebius needs python3.10. If python 3.9 [all] will not install nebius
+if sys.version_info < (3, 10):
+    filtered_keys = [k for k in extras_require if k != 'nebius']
+    extras_require['all'] = sum(
+        [v for k, v in extras_require.items() if k != 'nebius'], [])
+else:
+    extras_require['all'] = sum(extras_require.values(), [])

sky/setup_files/setup.py

@@ -1,3 +1,4 @@
+# pylint: skip-file
 """SkyPilot.
 
 SkyPilot is a framework for easily running machine learning* workloads on any
@@ -12,25 +13,33 @@ and would love to hear more about how we can better support your requirements -
 please join us in [this
 discussion](https://github.com/skypilot-org/skypilot/discussions/1016)
 """
-
 import atexit
 import io
 import os
 import platform
 import re
+import runpy
 import subprocess
 import sys
-from typing import Dict, List
 
 import setuptools
 
+# __file__ is setup.py at the root of the repo. We shouldn't assume it's a
+# symlink - e.g. in the sdist it's resolved to a normal file.
 ROOT_DIR = os.path.dirname(__file__)
+DEPENDENCIES_FILE_PATH = os.path.join(ROOT_DIR, 'sky', 'setup_files',
+                                      'dependencies.py')
 INIT_FILE_PATH = os.path.join(ROOT_DIR, 'sky', '__init__.py')
 _COMMIT_FAILURE_MESSAGE = (
     'WARNING: SkyPilot fail to {verb} the commit hash in '
     f'{INIT_FILE_PATH!r} (SkyPilot can still be normally used): '
     '{error}')
 
+# setuptools does not include the script dir on the search path, so we can't
+# just do `import dependencies`. Instead, use runpy to manually load it. Note:
+# dependencies here is a dict, not a module, so we access it by subscripting.
+dependencies = runpy.run_path(DEPENDENCIES_FILE_PATH)
+
 original_init_content = None
 
 system = platform.system()
@@ -130,126 +139,6 @@ def parse_readme(readme: str) -> str:
     return readme
 
 
-install_requires = [
-    'wheel',
-    'cachetools',
-    # NOTE: ray requires click>=7.0.
-    'click >= 7.0',
-    'colorama',
-    'cryptography',
-    # Jinja has a bug in older versions because of the lack of pinning
-    # the version of the underlying markupsafe package. See:
-    # https://github.com/pallets/jinja/issues/1585
-    'jinja2 >= 3.0',
-    'jsonschema',
-    'networkx',
-    'pandas>=1.3.0',
-    'pendulum',
-    # PrettyTable with version >=2.0.0 is required for the support of
-    # `add_rows` method.
-    'PrettyTable >= 2.0.0',
-    'python-dotenv',
-    'rich',
-    'tabulate',
-    # Light weight requirement, can be replaced with "typing" once
-    # we deprecate Python 3.7 (this will take a while).
-    "typing_extensions",
-    'filelock >= 3.6.0',
-    'packaging',
-    'psutil',
-    'pulp',
-    # Cython 3.0 release breaks PyYAML 5.4.* (https://github.com/yaml/pyyaml/issues/601)
-    # <= 3.13 may encounter https://github.com/ultralytics/yolov5/issues/414
-    'pyyaml > 3.13, != 5.4.*',
-    'requests',
-]
-
-local_ray = [
-    # Lower version of ray will cause dependency conflict for
-    # click/grpcio/protobuf.
-    # Excluded 2.6.0 as it has a bug in the cluster launcher:
-    # https://github.com/ray-project/ray/releases/tag/ray-2.6.1
-    'ray[default] >= 2.2.0, != 2.6.0',
-]
-
-remote = [
-    # Adopted from ray's setup.py: https://github.com/ray-project/ray/blob/ray-2.4.0/python/setup.py
-    # SkyPilot: != 1.48.0 is required to avoid the error where ray dashboard fails to start when
-    # ray start is called (#2054).
-    # Tracking issue: https://github.com/ray-project/ray/issues/30984
-    "grpcio >= 1.32.0, <= 1.49.1, != 1.48.0; python_version < '3.10' and sys_platform == 'darwin'",  # noqa:E501
-    "grpcio >= 1.42.0, <= 1.49.1, != 1.48.0; python_version >= '3.10' and sys_platform == 'darwin'",  # noqa:E501
-    # Original issue: https://github.com/ray-project/ray/issues/33833
-    "grpcio >= 1.32.0, <= 1.51.3, != 1.48.0; python_version < '3.10' and sys_platform != 'darwin'",  # noqa:E501
-    "grpcio >= 1.42.0, <= 1.51.3, != 1.48.0; python_version >= '3.10' and sys_platform != 'darwin'",  # noqa:E501
-    # Adopted from ray's setup.py:
-    # https://github.com/ray-project/ray/blob/ray-2.9.3/python/setup.py#L343
-    'protobuf >= 3.15.3, != 3.19.5',
-    # Some pydantic versions are not compatible with ray. Adopted from ray's
-    # setup.py: https://github.com/ray-project/ray/blob/ray-2.9.3/python/setup.py#L254
-    'pydantic!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,<3',
-]
-
-# NOTE: Change the templates/jobs-controller.yaml.j2 file if any of the
-# following packages dependencies are changed.
-aws_dependencies = [
-    # botocore does not work with urllib3>=2.0.0, according to https://github.com/boto/botocore/issues/2926
-    # We have to explicitly pin the version to optimize the time for
-    # poetry install. See https://github.com/orgs/python-poetry/discussions/7937
-    'urllib3<2',
-    # NOTE: this installs CLI V1. To use AWS SSO (e.g., `aws sso login`), users
-    # should instead use CLI V2 which is not pip-installable. See
-    # https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html.
-    'awscli>=1.27.10',
-    'botocore>=1.29.10',
-    'boto3>=1.26.1',
-    # NOTE: required by awscli. To avoid ray automatically installing
-    # the latest version.
-    'colorama < 0.4.5',
-]
-
-extras_require: Dict[str, List[str]] = {
-    'aws': aws_dependencies,
-    # TODO(zongheng): azure-cli is huge and takes a long time to install.
-    # Tracked in: https://github.com/Azure/azure-cli/issues/7387
-    # azure-identity is needed in node_provider.
-    # We need azure-identity>=1.13.0 to enable the customization of the
-    # timeout of AzureCliCredential.
-    'azure': [
-        'azure-cli>=2.31.0', 'azure-core', 'azure-identity>=1.13.0',
-        'azure-mgmt-network'
-    ] + local_ray,
-    # We need google-api-python-client>=2.69.0 to enable 'discardLocalSsd'
-    # parameter for stopping instances.
-    # Reference: https://github.com/googleapis/google-api-python-client/commit/f6e9d3869ed605b06f7cbf2e8cf2db25108506e6
-    'gcp': ['google-api-python-client>=2.69.0', 'google-cloud-storage'],
-    'ibm': [
-        'ibm-cloud-sdk-core', 'ibm-vpc', 'ibm-platform-services', 'ibm-cos-sdk'
-    ] + local_ray,
-    'docker': ['docker'] + local_ray,
-    'lambda': local_ray,
-    'cloudflare': aws_dependencies,
-    'scp': local_ray,
-    'oci': ['oci'] + local_ray,
-    'kubernetes': ['kubernetes>=20.0.0'],
-    'remote': remote,
-    'runpod': ['runpod>=1.5.1'],
-    'fluidstack': [],  # No dependencies needed for fluidstack
-    'cudo': ['cudo-compute>=0.1.10'],
-    'paperspace': [],  # No dependencies needed for paperspace
-    'vsphere': [
-        'pyvmomi==8.0.1.0.2',
-        # vsphere-automation-sdk is also required, but it does not have
-        # pypi release, which cause failure of our pypi release.
-        # https://peps.python.org/pep-0440/#direct-references
-        # We have the instruction for its installation in our
-        # docs instead.
-        # 'vsphere-automation-sdk @ git+https://github.com/vmware/vsphere-automation-sdk-python.git@v8.0.1.0'
-    ],
-}
-
-extras_require['all'] = sum(extras_require.values(), [])
-
 long_description = ''
 readme_filepath = 'README.md'
 # When sky/backends/wheel_utils.py builds wheels, it will not contain the
@@ -276,8 +165,8 @@ setuptools.setup(
     long_description_content_type='text/markdown',
     setup_requires=['wheel'],
     requires_python='>=3.7',
-    install_requires=install_requires,
-    extras_require=extras_require,
+    install_requires=dependencies['install_requires'],
+    extras_require=dependencies['extras_require'],
     entry_points={
         'console_scripts': ['sky = sky.cli:cli'],
     },
@@ -297,6 +186,6 @@ setuptools.setup(
         'Homepage': 'https://github.com/skypilot-org/skypilot',
         'Issues': 'https://github.com/skypilot-org/skypilot/issues',
         'Discussion': 'https://github.com/skypilot-org/skypilot/discussions',
-        'Documentation': 'https://skypilot.readthedocs.io/en/latest/',
+        'Documentation': 'https://docs.skypilot.co/',
     },
 )

sky/sky_logging.py

@@ -1,20 +1,27 @@
 """Logging utilities."""
 import builtins
 import contextlib
+from datetime import datetime
 import logging
+import os
 import sys
 import threading
 
 import colorama
 
+from sky.skylet import constants
 from sky.utils import env_options
 from sky.utils import rich_utils
 
-# If the SKYPILOT_MINIMIZE_LOGGING environment variable is set to True,
-# remove logging prefixes and unnecessary information in optimizer
-_FORMAT = (None if env_options.Options.MINIMIZE_LOGGING.get() else
-           '%(levelname).1s %(asctime)s %(filename)s:%(lineno)d] %(message)s')
+# UX: Should we show logging prefixes and some extra information in optimizer?
+_FORMAT = '%(levelname).1s %(asctime)s %(filename)s:%(lineno)d] %(message)s'
 _DATE_FORMAT = '%m-%d %H:%M:%S'
+_SENSITIVE_LOGGER = ['sky.provisioner', 'sky.optimizer']
+
+
+def _show_logging_prefix():
+    return env_options.Options.SHOW_DEBUG_INFO.get(
+    ) or not env_options.Options.MINIMIZE_LOGGING.get()
 
 
 class NewLineFormatter(logging.Formatter):
@@ -34,17 +41,11 @@ class NewLineFormatter(logging.Formatter):
         return msg
 
 
-class RichSafeStreamHandler(logging.StreamHandler):
-
-    def emit(self, record: logging.LogRecord) -> None:
-        with rich_utils.safe_logger():
-            return super().emit(record)
-
-
 _root_logger = logging.getLogger('sky')
 _default_handler = None
 _logging_config = threading.local()
 
+NO_PREFIX_FORMATTER = NewLineFormatter(None, datefmt=_DATE_FORMAT)
 FORMATTER = NewLineFormatter(_FORMAT, datefmt=_DATE_FORMAT)
 DIM_FORMATTER = NewLineFormatter(_FORMAT, datefmt=_DATE_FORMAT, dim=True)
 
@@ -60,17 +61,37 @@ def _setup_logger():
     _root_logger.setLevel(logging.DEBUG)
     global _default_handler
     if _default_handler is None:
-        _default_handler = RichSafeStreamHandler(sys.stdout)
+        _default_handler = rich_utils.RichSafeStreamHandler(sys.stdout)
         _default_handler.flush = sys.stdout.flush  # type: ignore
         if env_options.Options.SHOW_DEBUG_INFO.get():
             _default_handler.setLevel(logging.DEBUG)
         else:
             _default_handler.setLevel(logging.INFO)
         _root_logger.addHandler(_default_handler)
-    _default_handler.setFormatter(FORMATTER)
+    if _show_logging_prefix():
+        _default_handler.setFormatter(FORMATTER)
+    else:
+        _default_handler.setFormatter(NO_PREFIX_FORMATTER)
     # Setting this will avoid the message
     # being propagated to the parent logger.
     _root_logger.propagate = False
+    if env_options.Options.SUPPRESS_SENSITIVE_LOG.get():
+        # If the sensitive log is enabled, we reinitialize a new handler
+        # and force set the level to INFO to suppress the debug logs
+        # for certain loggers.
+        for logger_name in _SENSITIVE_LOGGER:
+            logger = logging.getLogger(logger_name)
+            handler_to_logger = rich_utils.RichSafeStreamHandler(sys.stdout)
+            handler_to_logger.flush = sys.stdout.flush  # type: ignore
+            logger.addHandler(handler_to_logger)
+            logger.setLevel(logging.INFO)
+            if _show_logging_prefix():
+                handler_to_logger.setFormatter(FORMATTER)
+            else:
+                handler_to_logger.setFormatter(NO_PREFIX_FORMATTER)
+            # Do not propagate to the parent logger to avoid parent
+            # logger printing the logs.
+            logger.propagate = False
 
 
 def reload_logger():
@@ -91,10 +112,21 @@ def reload_logger():
 _setup_logger()
 
 
-def init_logger(name: str):
+def init_logger(name: str) -> logging.Logger:
     return logging.getLogger(name)
 
 
+@contextlib.contextmanager
+def set_logging_level(logger: str, level: int):
+    logger = logging.getLogger(logger)
+    original_level = logger.level
+    logger.setLevel(level)
+    try:
+        yield
+    finally:
+        logger.setLevel(original_level)
+
+
 @contextlib.contextmanager
 def silent():
     """Make all sky_logging.print() and logger.{info, warning...} silent.
@@ -128,3 +160,16 @@ def is_silent():
         # threads.
         _logging_config.is_silent = False
     return _logging_config.is_silent
+
+
+def get_run_timestamp() -> str:
+    return 'sky-' + datetime.now().strftime('%Y-%m-%d-%H-%M-%S-%f')
+
+
+def generate_tmp_logging_file_path(file_name: str) -> str:
+    """Generate an absolute path of a tmp file for logging."""
+    run_timestamp = get_run_timestamp()
+    log_dir = os.path.join(constants.SKY_LOGS_DIRECTORY, run_timestamp)
+    log_path = os.path.expanduser(os.path.join(log_dir, file_name))
+
+    return log_path

sky/skylet/autostop_lib.py

@@ -9,7 +9,7 @@ import psutil
 from sky import sky_logging
 from sky.skylet import configs
 from sky.skylet import constants
-from sky.utils import common_utils
+from sky.utils import message_utils
 
 logger = sky_logging.init_logger(__name__)
 
@@ -85,7 +85,7 @@ def get_is_autostopping() -> bool:
 def get_is_autostopping_payload() -> str:
     """Payload for whether the cluster is in the process of autostopping."""
     is_autostopping = get_is_autostopping()
-    return common_utils.encode_payload(is_autostopping)
+    return message_utils.encode_payload(is_autostopping)
 
 
 def get_last_active_time() -> float: