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

sky/utils/kubernetes/rsync_helper.sh

@@ -0,0 +1,24 @@
+# When using pod@namespace+context, rsync passes args as: {us} -l pod namespace+context
+# We need to split the pod@namespace+context into pod, namespace and context
+# For backward compatibility, we use + as the separator between namespace and context and add handling when context is not provided
+shift
+pod=$1
+shift
+echo "pod: $pod" >&2
+encoded_namespace_context=$1
+# Revert the encoded namespace+context to the original string.
+namespace_context=$(echo "$encoded_namespace_context" | sed 's|%40|@|g' | sed 's|%3A|:|g' | sed 's|%2B|+|g' | sed 's|%2F|/|g')
+echo "namespace_context: $namespace_context" >&2
+namespace=$(echo $namespace_context | cut -d+ -f1)
+echo "namespace: $namespace" >&2
+context=$(echo $namespace_context | grep '+' >/dev/null && echo $namespace_context | cut -d+ -f2- || echo "")
+echo "context: $context" >&2
+context_lower=$(echo "$context" | tr '[:upper:]' '[:lower:]')
+shift
+if [ -z "$context" ] || [ "$context_lower" = "none" ]; then
+    # If context is none, it means we are using incluster auth. In this case,
+    # use need to set KUBECONFIG to /dev/null to avoid using kubeconfig file.
+    kubectl exec -i $pod -n $namespace --kubeconfig=/dev/null -- "$@"
+else
+    kubectl exec -i $pod -n $namespace --context=$context -- "$@"
+fi

sky/utils/kubernetes/ssh_jump_lifecycle_manager.py

@@ -126,7 +126,7 @@ def manage_lifecycle():
                          f'error: {e}\n')
         raise
 
-    if len(ret.items) == 0:
+    if not ret.items:
         sys.stdout.write(
             f'[Lifecycle] Did not find pods with label '
             f'"{label_selector}" in namespace {current_namespace}\n')

sky/utils/log_utils.py

@@ -1,13 +1,18 @@
 """Logging utils."""
 import enum
-from typing import List, Optional
+import time
+import types
+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
 
 from sky import sky_logging
 from sky.utils import rich_utils
+from sky.utils import ux_utils
 
 logger = sky_logging.init_logger(__name__)
 
@@ -15,13 +20,15 @@ logger = sky_logging.init_logger(__name__)
 class LineProcessor(object):
     """A processor for log lines."""
 
-    def __enter__(self):
+    def __enter__(self) -> None:
         pass
 
-    def process_line(self, log_line):
+    def process_line(self, log_line: str) -> None:
         pass
 
-    def __exit__(self, except_type, except_value, traceback):
+    def __exit__(self, except_type: Optional[Type[BaseException]],
+                 except_value: Optional[BaseException],
+                 traceback: Optional[types.TracebackType]) -> None:
         del except_type, except_value, traceback  # unused
         pass
 
@@ -34,33 +41,39 @@ class RayUpLineProcessor(LineProcessor):
         RUNTIME_SETUP = 1
         PULLING_DOCKER_IMAGES = 2
 
-    def __enter__(self):
+    def __init__(self, log_path: str):
+        self.log_path = log_path
+
+    def __enter__(self) -> None:
         self.state = self.ProvisionStatus.LAUNCH
-        self.status_display = rich_utils.safe_status('[bold cyan]Launching')
+        self.status_display = rich_utils.safe_status(
+            ux_utils.spinner_message('Launching', self.log_path))
         self.status_display.start()
 
-    def process_line(self, log_line):
+    def process_line(self, log_line: str) -> None:
         if ('Success.' in log_line and
                 self.state == self.ProvisionStatus.LAUNCH):
-            logger.info(f'{colorama.Fore.GREEN}Head node is up.'
-                        f'{colorama.Style.RESET_ALL}')
+            logger.info('  Head VM is up.')
             self.status_display.update(
-                '[bold cyan]Launching - Preparing SkyPilot runtime')
+                ux_utils.spinner_message(
+                    'Launching - Preparing SkyPilot runtime', self.log_path))
             self.state = self.ProvisionStatus.RUNTIME_SETUP
         if ('Pulling from' in log_line and
                 self.state == self.ProvisionStatus.RUNTIME_SETUP):
             self.status_display.update(
-                '[bold cyan]Launching - Pulling docker images')
+                ux_utils.spinner_message(
+                    'Launching - Initializing docker container', self.log_path))
             self.state = self.ProvisionStatus.PULLING_DOCKER_IMAGES
         if ('Status: Downloaded newer image' in log_line and
                 self.state == self.ProvisionStatus.PULLING_DOCKER_IMAGES):
-            logger.info(f'{colorama.Fore.GREEN}Docker image is downloaded.'
-                        f'{colorama.Style.RESET_ALL}')
             self.status_display.update(
-                '[bold cyan]Launching - Preparing SkyPilot runtime')
+                ux_utils.spinner_message(
+                    'Launching - Preparing SkyPilot runtime', self.log_path))
             self.state = self.ProvisionStatus.RUNTIME_SETUP
 
-    def __exit__(self, except_type, except_value, traceback):
+    def __exit__(self, except_type: Optional[Type[BaseException]],
+                 except_value: Optional[BaseException],
+                 traceback: Optional[types.TracebackType]) -> None:
         del except_type, except_value, traceback  # unused
         self.status_display.stop()
 
@@ -68,42 +81,69 @@ class RayUpLineProcessor(LineProcessor):
 class SkyLocalUpLineProcessor(LineProcessor):
     """A processor for `sky local up` log lines."""
 
+    def __init__(self, log_path: str, is_local: bool):
+        self.log_path = log_path
+        self.is_local = is_local
+
     def __enter__(self):
-        status = rich_utils.safe_status('[bold cyan]Creating local cluster - '
-                                        'initializing Kubernetes')
+        # TODO(romilb): Use ux_utils.INDENT_SYMBOL to be consistent with other
+        #  messages.
+        msg = 'Creating local cluster - initializing Kubernetes'
+        status = rich_utils.safe_status(
+            ux_utils.spinner_message(msg,
+                                     log_path=self.log_path,
+                                     is_local=self.is_local))
         self.status_display = status
         self.status_display.start()
 
-    def process_line(self, log_line):
+    def process_line(self, log_line: str) -> None:
         if 'Kind cluster created.' in log_line:
             logger.info(f'{colorama.Fore.GREEN}Kubernetes is running.'
                         f'{colorama.Style.RESET_ALL}')
         if 'Installing NVIDIA GPU operator...' in log_line:
-            self.status_display.update('[bold cyan]Creating local cluster - '
-                                       'Installing NVIDIA GPU operator')
+            self.status_display.update(
+                ux_utils.spinner_message(
+                    'Creating local cluster - '
+                    'Installing NVIDIA GPU operator',
+                    log_path=self.log_path,
+                    is_local=self.is_local))
         if 'Starting wait for GPU operator installation...' in log_line:
             self.status_display.update(
-                '[bold cyan]Creating local cluster - '
-                'waiting for NVIDIA GPU operator installation to complete')
+                ux_utils.spinner_message(
+                    'Creating local cluster - '
+                    'waiting for NVIDIA GPU operator installation to complete',
+                    log_path=self.log_path,
+                    is_local=self.is_local))
             logger.info('To check NVIDIA GPU operator status, '
                         'see pods: kubectl get pods -n gpu-operator')
         if 'GPU operator installed' in log_line:
             logger.info(f'{colorama.Fore.GREEN}NVIDIA GPU Operator installed.'
                         f'{colorama.Style.RESET_ALL}')
         if 'Pulling SkyPilot GPU image...' in log_line:
-            self.status_display.update('[bold cyan]Creating local cluster - '
-                                       'pulling and loading SkyPilot GPU image')
+            self.status_display.update(
+                ux_utils.spinner_message(
+                    'Creating local cluster - '
+                    'pulling and loading SkyPilot GPU image',
+                    log_path=self.log_path,
+                    is_local=self.is_local))
         if 'SkyPilot GPU image loaded into kind cluster' in log_line:
             logger.info(f'{colorama.Fore.GREEN}SkyPilot GPU image pulled.'
                         f'{colorama.Style.RESET_ALL}')
         if 'Labelling nodes with GPUs...' in log_line:
-            self.status_display.update('[bold cyan]Creating local cluster - '
-                                       'launching GPU labelling jobs')
+            self.status_display.update(
+                ux_utils.spinner_message(
+                    'Creating local cluster - '
+                    'launching GPU labelling jobs',
+                    log_path=self.log_path,
+                    is_local=self.is_local))
         if ('Starting wait for SkyPilot GPU labeling jobs to complete'
                 in log_line):
             self.status_display.update(
-                '[bold cyan]Creating local cluster - '
-                'waiting for GPU labelling jobs to complete')
+                ux_utils.spinner_message(
+                    'Creating local cluster - '
+                    'waiting for GPU labelling jobs to complete',
+                    log_path=self.log_path,
+                    is_local=self.is_local))
             logger.info(
                 'To check GPU labelling status, see jobs: '
                 'kubectl get jobs -n kube-system -l job=sky-gpu-labeler')
@@ -111,20 +151,136 @@ class SkyLocalUpLineProcessor(LineProcessor):
             logger.info(f'{colorama.Fore.GREEN}GPU labelling complete.'
                         f'{colorama.Style.RESET_ALL}')
         if 'Pulling SkyPilot CPU image...' in log_line:
-            self.status_display.update('[bold cyan]Creating local cluster - '
-                                       'pulling and loading SkyPilot CPU image')
+            self.status_display.update(
+                ux_utils.spinner_message(
+                    'Creating local cluster - '
+                    'pulling and loading SkyPilot CPU image',
+                    log_path=self.log_path,
+                    is_local=self.is_local))
         if 'SkyPilot CPU image loaded into kind cluster' in log_line:
             logger.info(f'{colorama.Fore.GREEN}SkyPilot CPU image pulled.'
                         f'{colorama.Style.RESET_ALL}')
         if 'Starting installation of Nginx Ingress Controller...' in log_line:
             self.status_display.update(
-                '[bold cyan]Creating Nginx Ingress Controller')
+                ux_utils.spinner_message(
+                    'Creating local cluster - '
+                    'creating Nginx Ingress Controller',
+                    log_path=self.log_path,
+                    is_local=self.is_local))
         if 'Nginx Ingress Controller installed' in log_line:
             logger.info(
                 f'{colorama.Fore.GREEN}Nginx Ingress Controller installed.'
                 f'{colorama.Style.RESET_ALL}')
+            self.status_display.update(
+                ux_utils.spinner_message('Wrapping up local cluster setup',
+                                         log_path=self.log_path,
+                                         is_local=self.is_local))
 
-    def __exit__(self, except_type, except_value, traceback):
+    def __exit__(self, except_type: Optional[Type[BaseException]],
+                 except_value: Optional[BaseException],
+                 traceback: Optional[types.TracebackType]) -> None:
+        del except_type, except_value, traceback  # unused
+        self.status_display.stop()
+
+
+class SkyRemoteUpLineProcessor(LineProcessor):
+    """A processor for deploy_remote_cluster.sh log lines."""
+
+    def __init__(self, log_path: str, is_local: bool):
+        self.log_path = log_path
+        self.is_local = is_local
+
+    def __enter__(self) -> None:
+        # TODO(romilb): Use ux_utils.INDENT_SYMBOL to be consistent with other
+        #  messages.
+        status = rich_utils.safe_status(
+            ux_utils.spinner_message('Creating remote cluster',
+                                     log_path=self.log_path,
+                                     is_local=self.is_local))
+        self.status_display = status
+        self.status_display.start()
+
+    def process_line(self, log_line: str) -> None:
+        # Pre-flight checks
+        if 'SSH connection successful' in log_line:
+            logger.info(f'{colorama.Fore.GREEN}SSH connection established.'
+                        f'{colorama.Style.RESET_ALL}')
+
+        # Kubernetes installation steps
+        if 'Deploying Kubernetes on head node' in log_line:
+            self.status_display.update(
+                ux_utils.spinner_message(
+                    'Creating remote cluster - '
+                    'deploying Kubernetes on head node',
+                    log_path=self.log_path,
+                    is_local=self.is_local))
+        if 'K3s deployed on head node.' in log_line:
+            logger.info(f'{colorama.Fore.GREEN}'
+                        '✔ K3s successfully deployed on head node.'
+                        f'{colorama.Style.RESET_ALL}')
+
+        # Worker nodes
+        if 'Deploying Kubernetes on worker node' in log_line:
+            self.status_display.update(
+                ux_utils.spinner_message(
+                    'Creating remote cluster - '
+                    'deploying Kubernetes on worker nodes',
+                    log_path=self.log_path,
+                    is_local=self.is_local))
+        if 'Kubernetes deployed on worker node' in log_line:
+            logger.info(f'{colorama.Fore.GREEN}'
+                        '✔ K3s successfully deployed on worker node.'
+                        f'{colorama.Style.RESET_ALL}')
+
+        # Cluster configuration
+        if 'Configuring local kubectl to connect to the cluster...' in log_line:
+            self.status_display.update(
+                ux_utils.spinner_message(
+                    'Creating remote cluster - '
+                    'configuring local kubectl',
+                    log_path=self.log_path,
+                    is_local=self.is_local))
+        if 'kubectl configured to connect to the cluster.' in log_line:
+            logger.info(f'{colorama.Fore.GREEN}'
+                        '✔ kubectl configured for the remote cluster.'
+                        f'{colorama.Style.RESET_ALL}')
+
+        # GPU operator installation
+        if 'Installing Nvidia GPU Operator...' in log_line:
+            self.status_display.update(
+                ux_utils.spinner_message(
+                    'Creating remote cluster - '
+                    'installing Nvidia GPU Operator',
+                    log_path=self.log_path,
+                    is_local=self.is_local))
+        if 'GPU Operator installed.' in log_line:
+            logger.info(f'{colorama.Fore.GREEN}'
+                        '✔ Nvidia GPU Operator installed successfully.'
+                        f'{colorama.Style.RESET_ALL}')
+
+        # Cleanup steps
+        if 'Cleaning up head node' in log_line:
+            self.status_display.update(
+                ux_utils.spinner_message('Cleaning up head node',
+                                         log_path=self.log_path,
+                                         is_local=self.is_local))
+        if 'Cleaning up node' in log_line:
+            self.status_display.update(
+                ux_utils.spinner_message('Cleaning up worker node',
+                                         log_path=self.log_path,
+                                         is_local=self.is_local))
+        if 'cleaned up successfully' in log_line:
+            logger.info(f'{colorama.Fore.GREEN}'
+                        f'{log_line.strip()}{colorama.Style.RESET_ALL}')
+
+        # Final status
+        if 'Cluster deployment completed.' in log_line:
+            logger.info(f'{colorama.Fore.GREEN}✔ Remote k3s is running.'
+                        f'{colorama.Style.RESET_ALL}')
+
+    def __exit__(self, except_type: Optional[Type[BaseException]],
+                 except_value: Optional[BaseException],
+                 traceback: Optional[types.TracebackType]) -> None:
         del except_type, except_value, traceback  # unused
         self.status_display.stop()
 
@@ -157,7 +313,8 @@ def readable_time_duration(start: Optional[float],
           e.g. "1h 2m 23s"
     """
     # start < 0 means that the starting time is not specified yet.
-    # It is only used in spot_utils.show_jobs() for job duration calculation.
+    # It is only used in jobs_utils.format_job_table() for job duration
+    # calculation.
     if start is None or start < 0:
         return '-'
     if end == start == 0:
@@ -191,3 +348,53 @@ def readable_time_duration(start: Optional[float],
         diff = diff.replace('hour', 'hr')
 
     return diff
+
+
+def follow_logs(
+    file: TextIO,
+    *,
+    should_stop: Callable[[], bool],
+    stop_on_eof: bool = False,
+    process_line: Optional[Callable[[str], Iterator[str]]] = None,
+    idle_timeout_seconds: Optional[int] = None,
+) -> Iterator[str]:
+    """Streams and processes logs line by line from a file.
+
+    Args:
+        file: File object to read logs from.
+        should_stop: Callback that returns True when streaming should stop.
+        stop_on_eof: If True, stop when reaching end of file.
+        process_line: Optional callback to transform/filter each line.
+        idle_timeout_seconds: If set, stop after these many seconds without
+            new content.
+
+    Yields:
+        Log lines, possibly transformed by process_line if provided.
+    """
+    current_line: str = ''
+    seconds_without_content: int = 0
+
+    while True:
+        content = file.readline()
+
+        if not content:
+            if stop_on_eof or should_stop():
+                break
+
+            if idle_timeout_seconds is not None:
+                if seconds_without_content >= idle_timeout_seconds:
+                    break
+                seconds_without_content += 1
+
+            time.sleep(1)
+            continue
+
+        seconds_without_content = 0
+        current_line += content
+
+        if '\n' in current_line or '\r' in current_line:
+            if process_line is not None:
+                yield from process_line(current_line)
+            else:
+                yield current_line
+            current_line = ''

sky/utils/message_utils.py

@@ -0,0 +1,81 @@
+"""Utilities for encoding and decoding messages."""
+import json
+import re
+import typing
+from typing import Any, Literal, Optional, Tuple, Union
+
+_PAYLOAD_PATTERN = re.compile(r'<sky-payload(.*?)>(.*?)</sky-payload>')
+_PAYLOAD_STR = '<sky-payload{type}>{content}</sky-payload>\n'
+
+
+def encode_payload(payload: Any, payload_type: Optional[str] = None) -> str:
+    """Encode a payload to make it more robust for parsing.
+
+    This makes message transfer more robust to any additional strings added to
+    the message during transfer.
+
+    An example message that is polluted by the system warning:
+    "LC_ALL: cannot change locale (en_US.UTF-8)\n<sky-payload>hello, world</sky-payload>" # pylint: disable=line-too-long
+
+    Args:
+        payload: A str, dict or list to be encoded.
+
+    Returns:
+        A string that is encoded from the payload.
+    """
+    payload_str = json.dumps(payload)
+    if payload_type is None:
+        payload_type = ''
+    payload_str = _PAYLOAD_STR.format(type=payload_type, content=payload_str)
+    return payload_str
+
+
+@typing.overload
+def decode_payload(payload_str: str,
+                   payload_type: Optional[str] = None,
+                   raise_for_mismatch: Literal[True] = True) -> Any:
+    ...
+
+
+@typing.overload
+def decode_payload(
+        payload_str: str,
+        payload_type: Optional[str] = None,
+        raise_for_mismatch: Literal[False] = False) -> Tuple[bool, Any]:
+    ...
+
+
+def decode_payload(
+        payload_str: str,
+        payload_type: Optional[str] = None,
+        raise_for_mismatch: bool = True) -> Union[Tuple[bool, Any], Any]:
+    """Decode a payload string.
+
+    Args:
+        payload_str: A string that is encoded from a payload.
+        payload_type: The type of the payload.
+        raise_for_mismatch: Whether to raise an error if the payload string is
+            not valid.
+
+    Returns:
+        A tuple of (bool, Any). The bool indicates whether it is a payload
+        string. The Any is the decoded payload, which is a str, dict or list.
+    """
+    matched = _PAYLOAD_PATTERN.findall(payload_str)
+    if not matched:
+        if raise_for_mismatch:
+            raise ValueError(f'Invalid payload string: \n{payload_str}')
+        else:
+            return False, payload_str
+
+    for payload_type_str, payload_str in matched:
+        if payload_type is None or payload_type == payload_type_str:
+            if raise_for_mismatch:
+                return json.loads(payload_str)
+            else:
+                return True, json.loads(payload_str)
+
+    if raise_for_mismatch:
+        raise ValueError(f'Invalid payload string: \n{payload_str}')
+    else:
+        return False, payload_str

sky/utils/registry.py

@@ -0,0 +1,127 @@
+"""Registry for classes to be discovered"""
+
+import typing
+from typing import Callable, Dict, List, Optional, Set, Type, Union
+
+from sky.utils import ux_utils
+
+if typing.TYPE_CHECKING:
+    from sky.backends import backend
+    from sky.clouds import cloud
+    from sky.jobs import recovery_strategy
+
+T = typing.TypeVar('T')
+
+
+class _Registry(dict, typing.Generic[T]):
+    """Registry."""
+
+    def __init__(self,
+                 registry_name: str,
+                 exclude: Optional[Set[str]],
+                 type_register: bool = False):
+        super().__init__()
+        self._registry_name = registry_name
+        self._exclude = exclude or set()
+        self._default: Optional[str] = None
+        self._type_register: bool = type_register
+        self._aliases: Dict[str, str] = {}
+
+    def from_str(self, name: Optional[str]) -> Optional[T]:
+        """Returns the cloud instance from the canonical name or alias."""
+        if name is None:
+            return None
+
+        search_name = name.lower()
+        if search_name in self._exclude:
+            return None
+
+        if search_name in self:
+            return self[search_name]
+
+        if search_name in self._aliases:
+            return self[self._aliases[search_name]]
+
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(
+                f'{self._registry_name.capitalize()} {name!r} is not a '
+                f'valid {self._registry_name} among '
+                f'{[*self.keys(), *self._aliases.keys()]}')
+
+    def type_register(self,
+                      name: str,
+                      default: bool = False) -> Callable[[Type[T]], Type[T]]:
+
+        name = name.lower()
+
+        def decorator(cls: Type[T]) -> Type[T]:
+            assert self._type_register, ('type_register can only be used '
+                                         'when type_register is True')
+            assert name not in self, f'{name} already registered'
+            self[name] = cls
+            if default:
+                self._default = name
+            return cls
+
+        return decorator
+
+    @typing.overload
+    def register(self, cls: Type[T]) -> Type[T]:
+        ...
+
+    @typing.overload
+    def register(
+            self,
+            cls: None = None,
+            aliases: Optional[List[str]] = None
+    ) -> Callable[[Type[T]], Type[T]]:
+        ...
+
+    def register(
+        self,
+        cls: Optional[Type[T]] = None,
+        aliases: Optional[List[str]] = None
+    ) -> Union[Type[T], Callable[[Type[T]], Type[T]]]:
+        assert not self._type_register, ('register can only be used when '
+                                         'type_register is False')
+
+        def _register(cls: Type[T]) -> Type[T]:
+            name = cls.__name__.lower()
+            assert name not in self, f'{name} already registered'
+            self[name] = cls()
+
+            for alias in aliases or []:
+                alias = alias.lower()
+                assert alias not in self._aliases, f'{alias} already registered'
+                self._aliases[alias] = name
+            return cls
+
+        if cls is not None:
+            # Invocation without parentheses (e.g. @register)
+            return _register(cls)
+
+        # Invocation with parentheses (e.g. @register(aliases=['alias']))
+        return _register
+
+    @property
+    def default(self) -> str:
+        assert self._default is not None, ('default is not set', self)
+        return self._default
+
+
+# Backward compatibility. global_user_state's DB may have recorded
+# Local cloud, and we've just removed it from the registry, and
+# global_user_state.get_enabled_clouds() would call into this func
+# and fail.
+
+CLOUD_REGISTRY: _Registry = _Registry['cloud.Cloud'](registry_name='cloud',
+                                                     exclude={'local'})
+
+BACKEND_REGISTRY: _Registry = _Registry['backend.Backend'](
+    registry_name='backend', type_register=True, exclude=None)
+
+JOBS_RECOVERY_STRATEGY_REGISTRY: _Registry = (
+    _Registry['recovery_strategy.StrategyExecutor'](
+        registry_name='jobs recovery strategy',
+        exclude=None,
+        type_register=True))