skypilot-nightly 1.0.0.dev20250509__py3-none-any.whl → 1.0.0.dev20251107__py3-none-any.whl

sky/client/cli/table_utils.py

@@ -0,0 +1,320 @@
+"""Utilities for formatting tables for CLI output."""
+import abc
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+import prettytable
+
+from sky import sky_logging
+from sky.jobs import utils as managed_jobs
+from sky.schemas.api import responses
+from sky.skylet import constants
+from sky.utils import common_utils
+from sky.utils import log_utils
+from sky.utils import volume
+
+logger = sky_logging.init_logger(__name__)
+
+
+def format_job_queue(jobs: List[responses.ClusterJobRecord]):
+    """Format the job queue for display.
+
+    Usage:
+        jobs = get_job_queue()
+        print(format_job_queue(jobs))
+    """
+    job_table = log_utils.create_table([
+        'ID', 'NAME', 'USER', 'SUBMITTED', 'STARTED', 'DURATION', 'RESOURCES',
+        'STATUS', 'LOG', 'GIT COMMIT'
+    ])
+    for job in jobs:
+        job_table.add_row([
+            job.job_id,
+            job.job_name,
+            job.username,
+            log_utils.readable_time_duration(job.submitted_at),
+            log_utils.readable_time_duration(job.start_at),
+            log_utils.readable_time_duration(job.start_at,
+                                             job.end_at,
+                                             absolute=True),
+            job.resources,
+            job.status.colored_str(),
+            job.log_path,
+            job.metadata.get('git_commit', '-'),
+        ])
+    return job_table
+
+
+def format_storage_table(storages: List[responses.StorageRecord],
+                         show_all: bool = False) -> str:
+    """Format the storage table for display.
+
+    Args:
+        storage_table (dict): The storage table.
+
+    Returns:
+        str: The formatted storage table.
+    """
+    storage_table = log_utils.create_table([
+        'NAME',
+        'UPDATED',
+        'STORE',
+        'COMMAND',
+        'STATUS',
+    ])
+
+    for row in storages:
+        launched_at = row.launched_at
+        if show_all:
+            command = row.last_use
+        else:
+            command = common_utils.truncate_long_string(
+                row.last_use, constants.LAST_USE_TRUNC_LENGTH)
+        storage_table.add_row([
+            # NAME
+            row.name,
+            # LAUNCHED
+            log_utils.readable_time_duration(launched_at),
+            # CLOUDS
+            ', '.join([s.value for s in row.store]),
+            # COMMAND,
+            command,
+            # STATUS
+            row.status.value,
+        ])
+    if storages:
+        return str(storage_table)
+    else:
+        return 'No existing storage.'
+
+
+def format_job_table(
+    jobs: List[responses.ManagedJobRecord],
+    show_all: bool,
+    show_user: bool,
+    pool_status: Optional[List[Dict[str, Any]]] = None,
+    max_jobs: Optional[int] = None,
+    status_counts: Optional[Dict[str, int]] = None,
+):
+    jobs = [job.model_dump() for job in jobs]
+    return managed_jobs.format_job_table(
+        jobs,
+        pool_status=pool_status,
+        show_all=show_all,
+        show_user=show_user,
+        max_jobs=max_jobs,
+        job_status_counts=status_counts,
+    )
+
+
+_BASIC_COLUMNS = [
+    'NAME',
+    'TYPE',
+    'INFRA',
+    'SIZE',
+    'USER',
+    'WORKSPACE',
+    'AGE',
+    'STATUS',
+    'LAST_USE',
+    'USED_BY',
+]
+
+
+def _get_infra_str(cloud: Optional[str], region: Optional[str],
+                   zone: Optional[str]) -> str:
+    """Get the infrastructure string for the volume."""
+    infra = ''
+    if cloud:
+        infra += cloud
+    if region:
+        infra += f'/{region}'
+    if zone:
+        infra += f'/{zone}'
+    return infra
+
+
+class VolumeTable(abc.ABC):
+    """The volume table."""
+
+    def __init__(self,
+                 volumes: List[responses.VolumeRecord],
+                 show_all: bool = False):
+        super().__init__()
+        self.table = self._create_table(show_all)
+        self._add_rows(volumes, show_all)
+
+    def _get_row_base_columns(self,
+                              row: responses.VolumeRecord,
+                              show_all: bool = False) -> List[str]:
+        """Get the base columns for a row."""
+        # Convert last_attached_at timestamp to human readable string
+        last_attached_at = row.get('last_attached_at')
+        if last_attached_at is not None:
+            last_attached_at_str = datetime.fromtimestamp(
+                last_attached_at).strftime('%Y-%m-%d %H:%M:%S')
+        else:
+            last_attached_at_str = '-'
+        size = row.get('size', '')
+        if size:
+            size = f'{size}Gi'
+        usedby_str = '-'
+        usedby_clusters = row.get('usedby_clusters')
+        usedby_pods = row.get('usedby_pods')
+        if usedby_clusters:
+            usedby_str = f'{", ".join(usedby_clusters)}'
+        elif usedby_pods:
+            usedby_str = f'{", ".join(usedby_pods)}'
+        if show_all:
+            usedby = usedby_str
+        else:
+            usedby = common_utils.truncate_long_string(
+                usedby_str, constants.USED_BY_TRUNC_LENGTH)
+        infra = _get_infra_str(row.get('cloud'), row.get('region'),
+                               row.get('zone'))
+        return [
+            row.get('name', ''),
+            row.get('type', ''),
+            infra,
+            size,
+            row.get('user_name', '-'),
+            row.get('workspace', '-'),
+            log_utils.human_duration(row.get('launched_at', 0)),
+            row.get('status', ''),
+            last_attached_at_str,
+            usedby,
+        ]
+
+    def _create_table(self, show_all: bool = False) -> prettytable.PrettyTable:
+        """Create the volume table."""
+        raise NotImplementedError
+
+    def _add_rows(self,
+                  volumes: List[responses.VolumeRecord],
+                  show_all: bool = False) -> None:
+        """Add rows to the volume table."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def format(self) -> str:
+        """Format the volume table for display."""
+        raise NotImplementedError
+
+
+class PVCVolumeTable(VolumeTable):
+    """The PVC volume table."""
+
+    def _create_table(self, show_all: bool = False) -> prettytable.PrettyTable:
+        """Create the PVC volume table."""
+        #  If show_all is False, show the table with the columns:
+        #   NAME, TYPE, INFRA, SIZE, USER, WORKSPACE,
+        #   AGE, STATUS, LAST_USE, USED_BY
+        #  If show_all is True, show the table with the columns:
+        #   NAME, TYPE, INFRA, SIZE, USER, WORKSPACE,
+        #   AGE, STATUS, LAST_USE, USED_BY, NAME_ON_CLOUD
+        #   STORAGE_CLASS, ACCESS_MODE
+
+        if show_all:
+            columns = _BASIC_COLUMNS + [
+                'NAME_ON_CLOUD',
+                'STORAGE_CLASS',
+                'ACCESS_MODE',
+            ]
+        else:
+            columns = _BASIC_COLUMNS
+
+        table = log_utils.create_table(columns)
+        return table
+
+    def _add_rows(self,
+                  volumes: List[responses.VolumeRecord],
+                  show_all: bool = False) -> None:
+        """Add rows to the PVC volume table."""
+        for row in volumes:
+            table_row = self._get_row_base_columns(row, show_all)
+            if show_all:
+                table_row.append(row.get('name_on_cloud', ''))
+                table_row.append(
+                    row.get('config', {}).get('storage_class_name', '-'))
+                table_row.append(row.get('config', {}).get('access_mode', ''))
+
+            self.table.add_row(table_row)
+
+    def format(self) -> str:
+        """Format the PVC volume table for display."""
+        return 'Kubernetes PVCs:\n' + str(self.table)
+
+
+class RunPodVolumeTable(VolumeTable):
+    """The RunPod volume table."""
+
+    def _create_table(self, show_all: bool = False) -> prettytable.PrettyTable:
+        """Create the RunPod volume table."""
+        #  If show_all is False, show the table with the columns:
+        #   NAME, TYPE, INFRA, SIZE, USER, WORKSPACE,
+        #   AGE, STATUS, LAST_USE, USED_BY
+        #  If show_all is True, show the table with the columns:
+        #   NAME, TYPE, INFRA, SIZE, USER, WORKSPACE,
+        #   AGE, STATUS, LAST_USE, USED_BY, NAME_ON_CLOUD
+
+        if show_all:
+            columns = _BASIC_COLUMNS + ['NAME_ON_CLOUD']
+        else:
+            columns = _BASIC_COLUMNS
+
+        table = log_utils.create_table(columns)
+        return table
+
+    def _add_rows(self,
+                  volumes: List[responses.VolumeRecord],
+                  show_all: bool = False) -> None:
+        """Add rows to the RunPod volume table."""
+        for row in volumes:
+            table_row = self._get_row_base_columns(row, show_all)
+            if show_all:
+                table_row.append(row.get('name_on_cloud', ''))
+
+            self.table.add_row(table_row)
+
+    def format(self) -> str:
+        """Format the RunPod volume table for display."""
+        return 'RunPod Network Volumes:\n' + str(self.table)
+
+
+def format_volume_table(volumes: List[responses.VolumeRecord],
+                        show_all: bool = False) -> str:
+    """Format the volume table for display.
+
+    Args:
+        volume_table (dict): The volume table.
+
+    Returns:
+        str: The formatted volume table.
+    """
+    volumes_per_type: Dict[str, List[responses.VolumeRecord]] = {}
+    supported_volume_types = [
+        volume_type.value for volume_type in volume.VolumeType
+    ]
+    for row in volumes:
+        volume_type = row.get('type', '')
+        if volume_type in supported_volume_types:
+            if volume_type not in volumes_per_type:
+                volumes_per_type[volume_type] = []
+            volumes_per_type[volume_type].append(row)
+        else:
+            logger.warning(f'Unknown volume type: {volume_type}')
+            continue
+    table_str = ''
+    for volume_type, volume_list in volumes_per_type.items():
+        if table_str:
+            table_str += '\n\n'
+        if volume_type == volume.VolumeType.PVC.value:
+            pvc_table = PVCVolumeTable(volume_list, show_all)
+            table_str += pvc_table.format()
+        elif volume_type == volume.VolumeType.RUNPOD_NETWORK_VOLUME.value:
+            runpod_table = RunPodVolumeTable(volume_list, show_all)
+            table_str += runpod_table.format()
+    if table_str:
+        return table_str
+    else:
+        return 'No existing volumes.'

sky/client/common.py

@@ -16,8 +16,10 @@ import zipfile
 
 from sky import sky_logging
 from sky.adaptors import common as adaptors_common
+from sky.client import service_account_auth
 from sky.data import data_utils
 from sky.data import storage_utils
+from sky.schemas.api import responses as api_responses
 from sky.server import common as server_common
 from sky.server.requests import payloads
 from sky.skylet import constants
@@ -31,7 +33,7 @@ if typing.TYPE_CHECKING:
     import requests
 
     import sky
-    import sky.dag as dag_lib
+    from sky import dag as dag_lib
 else:
     httpx = adaptors_common.LazyImport('httpx')
     requests = adaptors_common.LazyImport('requests')
@@ -42,8 +44,10 @@ logger = sky_logging.init_logger(__name__)
 _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
+# can be caused by NGINX's client_max_body_size or Cloudflare's upload limit.
+# As of 09/25/2025, the upload limit for Cloudflare's free plan is 100MiB:
+# https://developers.cloudflare.com/support/troubleshooting/http-status-codes/4xx-client-error/error-413/
+_UPLOAD_CHUNK_BYTES = 100 * 1024 * 1024
 
 FILE_UPLOAD_LOGS_DIR = os.path.join(constants.SKY_LOGS_DIRECTORY,
                                     'file_uploads')
@@ -80,10 +84,11 @@ def download_logs_from_api_server(
             local_machine_prefix) 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,
-                             cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST',
+        '/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
@@ -164,14 +169,19 @@ class UploadChunkParams:
     log_file: str
 
 
-def _upload_chunk_with_retry(params: UploadChunkParams) -> None:
-    """Uploads a chunk of a zip file to the API server."""
+def _upload_chunk_with_retry(params: UploadChunkParams) -> str:
+    """Uploads a chunk of a zip file to the API server.
+
+    Returns:
+        Status of the upload.
+    """
     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
+    sa_headers = service_account_auth.get_service_account_headers()
     with open(params.file_path, 'rb') as f:
         for attempt in range(max_attempts):
             response = params.client.post(
@@ -184,19 +194,23 @@ def _upload_chunk_with_retry(params: UploadChunkParams) -> None:
                 },
                 content=FileChunkIterator(f, _UPLOAD_CHUNK_BYTES,
                                           params.chunk_index),
-                headers={'Content-Type': 'application/octet-stream'},
+                headers={
+                    'Content-Type': 'application/octet-stream',
+                    **sa_headers,
+                },
                 cookies=server_common.get_api_cookie_jar())
             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':
+                       f'{params.chunk_index + 1} / {params.total_chunks} '
+                       f'(Status: {status})')
+                if status == api_responses.UploadStatus.UPLOADING.value:
                     missing_chunks = data.get('missing_chunks')
                     if missing_chunks:
                         msg += f' - Waiting for chunks: {missing_chunks}'
                 upload_logger.info(msg)
-                return
+                return status
             elif attempt < max_attempts - 1:
                 upload_logger.error(
                     f'Failed to upload chunk: '
@@ -204,17 +218,29 @@ def _upload_chunk_with_retry(params: UploadChunkParams) -> None:
                     f'{response.content.decode("utf-8")}')
                 upload_logger.info(
                     f'Retrying... ({attempt + 1} / {max_attempts})')
-                time.sleep(1)
+                if response.status_code == 503:
+                    # If the server is temporarily unavailable,
+                    # wait a little longer before retrying.
+                    time.sleep(10)
+                else:
+                    time.sleep(1)
             else:
+                try:
+                    response_details = response.json().get('detail')
+                except Exception:  # pylint: disable=broad-except
+                    response_details = response.content
                 error_msg = (
                     f'Failed to upload chunk: {params.chunk_index + 1} / '
-                    f'{params.total_chunks}: {response.json().get("detail")}')
+                    f'{params.total_chunks}: {response_details} '
+                    f'(Status code: {response.status_code})')
                 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))
+    # If we reach here, the upload failed.
+    return 'failed'
 
 
 @contextlib.contextmanager
@@ -267,7 +293,7 @@ def upload_mounts_to_api_server(dag: 'sky.Dag',
     upload_list = []
     for task_ in dag.tasks:
         task_.file_mounts_mapping = {}
-        if task_.workdir:
+        if task_.workdir and isinstance(task_.workdir, str):
             workdir = task_.workdir
             assert os.path.isabs(workdir)
             upload_list.append(workdir)
@@ -299,14 +325,12 @@ def upload_mounts_to_api_server(dag: 'sky.Dag',
                         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
+            keyfile = task_.service.tls_credential.keyfile
+            certfile = task_.service.tls_credential.certfile
+            upload_list.append(_full_path(keyfile))
+            upload_list.append(_full_path(certfile))
+            task_.file_mounts_mapping[keyfile] = _full_path(keyfile)
+            task_.file_mounts_mapping[certfile] = _full_path(certfile)
 
     if upload_list:
         os.makedirs(os.path.expanduser(FILE_UPLOAD_LOGS_DIR), exist_ok=True)
@@ -339,15 +363,29 @@ def upload_mounts_to_api_server(dag: 'sky.Dag',
                     log_file,
                     is_local=True))
 
+            upload_completed = False
             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)
+                total_retries = 3
+                for retry in range(total_retries):
+                    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)
+                    ]
+                    statuses = subprocess_utils.run_in_parallel(
+                        _upload_chunk_with_retry, chunk_params)
+                    if any(status == api_responses.UploadStatus.COMPLETED.value
+                           for status in statuses):
+                        upload_completed = True
+                        break
+                    else:
+                        upload_logger.info(
+                            f'No chunk upload returned completed status. '
+                            'Retrying entire upload... '
+                            f'({retry + 1} / {total_retries})')
+            if not upload_completed:
+                raise RuntimeError('Failed to upload files to API server.')
         os.unlink(temp_zip_file.name)
         upload_logger.info(f'Uploaded files: {upload_list}')
         logger.info(

sky/client/oauth.py

@@ -0,0 +1,82 @@
+"""Client-side OAuth module."""
+from http.server import BaseHTTPRequestHandler
+from http.server import HTTPServer
+import threading
+import time
+from typing import Dict, Optional
+
+AUTH_TIMEOUT = 300  # 5 minutes
+
+
+class _AuthCallbackHandler(BaseHTTPRequestHandler):
+    """HTTP request handler for OAuth callback."""
+
+    def __init__(self, token_container: Dict[str, Optional[str]],
+                 remote_endpoint: str, *args, **kwargs):
+        self.token_container = token_container
+        self.remote_endpoint = remote_endpoint
+        super().__init__(*args, **kwargs)
+
+    def do_POST(self):  # pylint: disable=invalid-name
+        """Handle POST request for OAuth callback."""
+        data = self.rfile.read(int(self.headers['Content-Length']))
+
+        if data:
+            token = data.decode('utf-8')
+            self.token_container['token'] = token
+
+            # Send success response
+            self.send_response(200)
+            self.send_header('Content-type', 'text/html')
+            self.send_header('Access-Control-Allow-Origin',
+                             self.remote_endpoint)
+            self.end_headers()
+        else:
+            # Send error response
+            self.send_response(400)
+            self.send_header('Content-type', 'text/html')
+            self.send_header('Access-Control-Allow-Origin',
+                             self.remote_endpoint)
+            self.end_headers()
+
+    def log_message(self, *args):  # pylint: disable=unused-argument
+        """Suppress default HTTP server logging."""
+        pass
+
+
+def start_local_auth_server(port: int,
+                            token_store: Dict[str, Optional[str]],
+                            remote_endpoint: str,
+                            timeout: int = AUTH_TIMEOUT) -> HTTPServer:
+    """Start a local HTTP server to handle OAuth callback.
+
+    Args:
+        port: Port to bind the server to.
+        token_container: Dict to store the received token.
+        remote_endpoint: The endpoint of the SkyPilot API server that will send
+          the token, needed for CORS.
+        timeout: Timeout in seconds to wait for the callback.
+
+    Returns:
+        The HTTP server instance.
+    """
+
+    def handler_factory(*args, **kwargs):
+        return _AuthCallbackHandler(token_store, remote_endpoint, *args,
+                                    **kwargs)
+
+    server = HTTPServer(('localhost', port), handler_factory)
+    server.timeout = timeout
+
+    def serve_until_token():
+        """Serve requests until token is received or timeout."""
+        start_time = time.time()
+        while (token_store['token'] is None and
+               time.time() - start_time < timeout):
+            server.handle_request()
+
+    # Start server in a separate thread
+    server_thread = threading.Thread(target=serve_until_token, daemon=True)
+    server_thread.start()
+
+    return server