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

sky/client/common.py

@@ -0,0 +1,345 @@
+"""Common utilities for the client."""
+
+import contextlib
+import dataclasses
+import json
+import logging
+import math
+import os
+import pathlib
+import tempfile
+import time
+import typing
+from typing import Dict, Generator, Iterable
+import uuid
+import zipfile
+
+import httpx
+import requests
+
+from sky import sky_logging
+from sky.data import data_utils
+from sky.data import storage_utils
+from sky.server import common as server_common
+from sky.server.requests import payloads
+from sky.skylet import constants
+from sky.utils import common_utils
+from sky.utils import rich_utils
+from sky.utils import subprocess_utils
+from sky.utils import ux_utils
+
+if typing.TYPE_CHECKING:
+    import sky
+    import sky.dag as dag_lib
+
+logger = sky_logging.init_logger(__name__)
+
+# The chunk size for downloading the logs from the API server.
+_DOWNLOAD_CHUNK_BYTES = 8192
+# The chunk size for the zip file to be uploaded to the API server. We split
+# the zip file into chunks to avoid network issues for large request body that
+# can be caused by NGINX's client_max_body_size.
+_UPLOAD_CHUNK_BYTES = 512 * 1024 * 1024
+
+FILE_UPLOAD_LOGS_DIR = os.path.join(constants.SKY_LOGS_DIRECTORY,
+                                    'file_uploads')
+
+# Connection timeout when sending requests to the API server.
+API_SERVER_REQUEST_CONNECTION_TIMEOUT_SECONDS = 5
+
+
+def download_logs_from_api_server(
+        paths_on_api_server: Iterable[str]) -> Dict[str, str]:
+    """Downloads the logs from the API server.
+
+    Args:
+        paths_on_api_server: The paths on the API server to download.
+
+    Returns:
+        A dictionary mapping the remote path on API server to the local path.
+    """
+    remote2local_path_dict = {
+        remote_path: remote_path.replace(
+            # TODO(zhwu): handling the replacement locally is not stable, and
+            # may cause issues when we change the pattern of the remote path.
+            # This should be moved to remote API server. A proper way might be
+            # set the returned path to be started with a special prefix, instead
+            # of using the `api_server_user_logs_dir_prefix()`.
+            str(server_common.api_server_user_logs_dir_prefix()),
+            constants.SKY_LOGS_DIRECTORY) for remote_path in paths_on_api_server
+    }
+    body = payloads.DownloadBody(folder_paths=list(paths_on_api_server),)
+    response = requests.post(f'{server_common.get_server_url()}/download',
+                             json=json.loads(body.model_dump_json()),
+                             stream=True)
+    if response.status_code == 200:
+        remote_home_path = response.headers.get('X-Home-Path')
+        assert remote_home_path is not None, response.headers
+        with tempfile.NamedTemporaryFile(prefix='skypilot-logs-download-',
+                                         delete=True) as temp_file:
+            # Download the zip file from the API server to the local machine.
+            for chunk in response.iter_content(
+                    chunk_size=_DOWNLOAD_CHUNK_BYTES):
+                temp_file.write(chunk)
+            temp_file.flush()
+
+            # Unzip the downloaded file and save the logs to the correct local
+            # directory.
+            with zipfile.ZipFile(temp_file, 'r') as zipf:
+                for member in zipf.namelist():
+                    # Determine the new path
+                    zipped_filename = os.path.basename(member)
+                    zipped_dir = os.path.dirname('/' + member)
+                    local_dir = zipped_dir.replace(remote_home_path, '~')
+                    for remote_path, local_path in remote2local_path_dict.items(
+                    ):
+                        if local_dir.startswith(remote_path):
+                            local_dir = local_dir.replace(
+                                remote_path, local_path)
+                            break
+                    else:
+                        raise ValueError(f'Invalid folder path: {zipped_dir}')
+                    new_path = pathlib.Path(
+                        local_dir).expanduser().resolve() / zipped_filename
+                    new_path.parent.mkdir(parents=True, exist_ok=True)
+                    if member.endswith('/'):
+                        # If it is a directory, we need to create it.
+                        new_path.mkdir(parents=True, exist_ok=True)
+                    else:
+                        with zipf.open(member) as member_file:
+                            new_path.write_bytes(member_file.read())
+
+        return remote2local_path_dict
+    else:
+        raise Exception(
+            f'Failed to download logs: {response.status_code} {response.text}')
+
+
+# === Upload files to API server ===
+
+
+class FileChunkIterator:
+    """A file-like object that reads from a file in chunks."""
+
+    def __init__(self, file_obj, chunk_size: int, chunk_index: int):
+        self.file_obj = file_obj
+        self.chunk_size = chunk_size
+        self.chunk_index = chunk_index
+        self.bytes_read = 0
+
+    def __iter__(self):
+        # Seek to the correct position for this chunk
+        self.file_obj.seek(self.chunk_index * self.chunk_size)
+        while self.bytes_read < self.chunk_size:
+            # Read a smaller buffer size to keep memory usage low
+            buffer_size = min(64 * 1024,
+                              self.chunk_size - self.bytes_read)  # 64KB buffer
+            data = self.file_obj.read(buffer_size)
+            if not data:
+                break
+            self.bytes_read += len(data)
+            yield data
+
+
+@dataclasses.dataclass
+class UploadChunkParams:
+    client: httpx.Client
+    upload_id: str
+    chunk_index: int
+    total_chunks: int
+    file_path: str
+    upload_logger: logging.Logger
+    log_file: str
+
+
+def _upload_chunk_with_retry(params: UploadChunkParams) -> None:
+    """Uploads a chunk of a zip file to the API server."""
+    upload_logger = params.upload_logger
+    upload_logger.info(
+        f'Uploading chunk: {params.chunk_index + 1} / {params.total_chunks}')
+
+    server_url = server_common.get_server_url()
+    max_attempts = 3
+    with open(params.file_path, 'rb') as f:
+        for attempt in range(max_attempts):
+            response = params.client.post(
+                f'{server_url}/upload',
+                params={
+                    'user_hash': common_utils.get_user_hash(),
+                    'upload_id': params.upload_id,
+                    'chunk_index': str(params.chunk_index),
+                    'total_chunks': str(params.total_chunks),
+                },
+                content=FileChunkIterator(f, _UPLOAD_CHUNK_BYTES,
+                                          params.chunk_index),
+                headers={'Content-Type': 'application/octet-stream'})
+            if response.status_code == 200:
+                data = response.json()
+                status = data.get('status')
+                msg = ('Uploaded chunk: '
+                       f'{params.chunk_index + 1} / {params.total_chunks}')
+                if status == 'uploading':
+                    missing_chunks = data.get('missing_chunks')
+                    if missing_chunks:
+                        msg += f' - Waiting for chunks: {missing_chunks}'
+                upload_logger.info(msg)
+                return
+            elif attempt < max_attempts - 1:
+                upload_logger.error(
+                    f'Failed to upload chunk: '
+                    f'{params.chunk_index + 1} / {params.total_chunks}: '
+                    f'{response.content.decode("utf-8")}')
+                upload_logger.info(
+                    f'Retrying... ({attempt + 1} / {max_attempts})')
+                time.sleep(1)
+            else:
+                error_msg = (
+                    f'Failed to upload chunk: {params.chunk_index + 1} / '
+                    f'{params.total_chunks}: {response.json().get("detail")}')
+                upload_logger.error(error_msg)
+                with ux_utils.print_exception_no_traceback():
+                    raise RuntimeError(
+                        ux_utils.error_message(error_msg + '\n',
+                                               params.log_file,
+                                               is_local=True))
+
+
+@contextlib.contextmanager
+def _setup_upload_logger(
+        log_file: str) -> Generator[logging.Logger, None, None]:
+    try:
+        upload_logger = logging.getLogger('sky.upload')
+        upload_logger.propagate = False
+        handler = logging.FileHandler(os.path.expanduser(log_file),
+                                      encoding='utf-8')
+        handler.setFormatter(sky_logging.FORMATTER)
+        upload_logger.addHandler(handler)
+        upload_logger.setLevel(logging.DEBUG)
+        yield upload_logger
+    finally:
+        upload_logger.removeHandler(handler)
+        handler.close()
+
+
+def upload_mounts_to_api_server(dag: 'sky.Dag',
+                                workdir_only: bool = False) -> 'dag_lib.Dag':
+    """Upload user files to remote API server.
+
+    This function needs to be called after sdk.validate(),
+    as the file paths need to be expanded to keep file_mounts_mapping
+    aligned with the actual task uploaded to SkyPilot API server.
+
+    We don't use FastAPI's built-in multipart upload, as nginx's
+    client_max_body_size can block the request due to large request body, i.e.,
+    even though the multipart upload streams the file to the server, there is
+    only one HTTP request, and a large request body will be blocked by nginx.
+
+    Args:
+        dag: The dag where the file mounts are defined.
+        workdir_only: Whether to only upload the workdir, which is used for
+            `exec`, as it does not need other files/folders in file_mounts.
+
+    Returns:
+        The dag with the file_mounts_mapping updated, which maps the original
+        file paths to the full path, so that on API server, the file paths can
+        be retrieved by adding prefix to the full path.
+    """
+
+    if server_common.is_api_server_local():
+        return dag
+
+    def _full_path(src: str) -> str:
+        return os.path.abspath(os.path.expanduser(src))
+
+    upload_list = []
+    for task_ in dag.tasks:
+        task_.file_mounts_mapping = {}
+        if task_.workdir:
+            workdir = task_.workdir
+            assert os.path.isabs(workdir)
+            upload_list.append(workdir)
+            task_.file_mounts_mapping[workdir] = workdir
+        if workdir_only:
+            continue
+        if task_.file_mounts is not None:
+            for src in task_.file_mounts.values():
+                if not data_utils.is_cloud_store_url(src):
+                    assert os.path.isabs(src)
+                    upload_list.append(src)
+                    task_.file_mounts_mapping[src] = src
+                if src == constants.LOCAL_SKYPILOT_CONFIG_PATH_PLACEHOLDER:
+                    # The placeholder for the local skypilot config path is in
+                    # file mounts for controllers. It will be replaced with the
+                    # real path for config file on API server.
+                    pass
+        if task_.storage_mounts is not None:
+            for storage in task_.storage_mounts.values():
+                storage_source = storage.source
+                is_cloud_store_url = (
+                    isinstance(storage_source, str) and
+                    data_utils.is_cloud_store_url(storage_source))
+                if (storage_source is not None and not is_cloud_store_url):
+                    if isinstance(storage_source, str):
+                        storage_source = [storage_source]
+                    for src in storage_source:
+                        upload_list.append(_full_path(src))
+                        task_.file_mounts_mapping[src] = _full_path(src)
+        if (task_.service is not None and
+                task_.service.tls_credential is not None):
+            upload_list.append(task_.service.tls_credential.keyfile)
+            upload_list.append(task_.service.tls_credential.certfile)
+            task_.file_mounts_mapping[
+                task_.service.tls_credential.
+                keyfile] = task_.service.tls_credential.keyfile
+            task_.file_mounts_mapping[
+                task_.service.tls_credential.
+                certfile] = task_.service.tls_credential.certfile
+
+    if upload_list:
+        os.makedirs(os.path.expanduser(FILE_UPLOAD_LOGS_DIR), exist_ok=True)
+        upload_id = sky_logging.get_run_timestamp()
+        upload_id = f'{upload_id}-{uuid.uuid4().hex[:8]}'
+        log_file = os.path.join(FILE_UPLOAD_LOGS_DIR, f'{upload_id}.log')
+
+        logger.info(ux_utils.starting_message('Uploading files to API server'))
+        with rich_utils.client_status(
+                ux_utils.spinner_message(
+                    'Uploading files to API server (1/2 - Zipping)',
+                    log_file,
+                    is_local=True)) as status, _setup_upload_logger(
+                        log_file) as upload_logger:
+            with tempfile.NamedTemporaryFile(suffix='.zip',
+                                             delete=False) as temp_zip_file:
+                upload_logger.info(
+                    f'Zipping files to be uploaded: {upload_list}')
+                storage_utils.zip_files_and_folders(upload_list,
+                                                    temp_zip_file.name)
+                upload_logger.info(f'Zipped files to: {temp_zip_file.name}')
+
+            zip_file_size = os.path.getsize(temp_zip_file.name)
+            # Per chunk size 512 MB
+            total_chunks = int(math.ceil(zip_file_size / _UPLOAD_CHUNK_BYTES))
+            timeout = httpx.Timeout(None, read=180.0)
+            status.update(
+                ux_utils.spinner_message(
+                    'Uploading files to API server (2/2 - Uploading)',
+                    log_file,
+                    is_local=True))
+
+            with httpx.Client(timeout=timeout) as client:
+                chunk_params = [
+                    UploadChunkParams(client, upload_id, chunk_index,
+                                      total_chunks, temp_zip_file.name,
+                                      upload_logger, log_file)
+                    for chunk_index in range(total_chunks)
+                ]
+                subprocess_utils.run_in_parallel(_upload_chunk_with_retry,
+                                                 chunk_params)
+        os.unlink(temp_zip_file.name)
+        upload_logger.info(f'Uploaded files: {upload_list}')
+        logger.info(
+            ux_utils.finishing_message('Files uploaded',
+                                       log_file,
+                                       is_local=True))
+
+    return dag