skypilot-nightly 1.0.0.dev20250806__py3-none-any.whl → 1.0.0.dev20250808__py3-none-any.whl

sky/serve/serve_utils.py

@@ -20,6 +20,7 @@ import uuid
 
 import colorama
 import filelock
+import yaml
 
 from sky import backends
 from sky import exceptions
@@ -65,13 +66,12 @@ def get_num_service_threshold():
 
 _CONTROLLER_URL = 'http://localhost:{CONTROLLER_PORT}'
 
-# NOTE(dev): We assume log paths are either in ~/sky_logs/... or ~/.sky/...
-# and always appear after a space. Be careful when changing UX as this
-# assumption is used to expand some log files while ignoring others.
-_SKYPILOT_LOG_DIRS = r'~/(sky_logs|\.sky)'
-_SKYPILOT_PROVISION_LOG_PATTERN = (
-    fr'.* ({_SKYPILOT_LOG_DIRS}/.*provision\.log)')
-_SKYPILOT_LOG_PATTERN = fr'.* ({_SKYPILOT_LOG_DIRS}/.*\.log)'
+# NOTE(dev): We assume log are print with the hint 'sky api logs -l'. Be careful
+# when changing UX as this assumption is used to expand some log files while
+# ignoring others.
+_SKYPILOT_LOG_HINT = r'.*sky api logs -l'
+_SKYPILOT_PROVISION_LOG_PATTERN = (fr'{_SKYPILOT_LOG_HINT} (.*/provision\.log)')
+_SKYPILOT_LOG_PATTERN = fr'{_SKYPILOT_LOG_HINT} (.*\.log)'
 
 # TODO(tian): Find all existing replica id and print here.
 _FAILED_TO_FIND_REPLICA_MSG = (
@@ -668,12 +668,18 @@ def _get_service_status(
     if record['pool']:
         latest_yaml_path = generate_task_yaml_file_name(service_name,
                                                         record['version'])
-        original_config = common_utils.read_yaml(latest_yaml_path)
-        original_config.pop('run', None)
-        svc: Dict[str, Any] = original_config.pop('service')
-        if svc is not None:
-            svc.pop('pool', None)
-            original_config['pool'] = svc
+        raw_yaml_config = common_utils.read_yaml(latest_yaml_path)
+        original_config = raw_yaml_config.get('_user_specified_yaml')
+        if original_config is None:
+            # Fall back to old display format.
+            original_config = raw_yaml_config
+            original_config.pop('run', None)
+            svc: Dict[str, Any] = original_config.pop('service')
+            if svc is not None:
+                svc.pop('pool', None)  # Remove pool from service config
+                original_config['pool'] = svc  # Add pool to root config
+        else:
+            original_config = yaml.safe_load(original_config)
         record['pool_yaml'] = common_utils.dump_yaml_str(original_config)
 
     record['target_num_replicas'] = 0
@@ -959,8 +965,10 @@ def wait_service_registration(service_name: str, job_id: int,
     """
     start_time = time.time()
     setup_completed = False
+    noun = 'pool' if pool else 'service'
     while True:
-        # TODO(tian): PID-based tracking.
+        # Only do this check for non-consolidation mode as consolidation mode
+        # has no setup process.
         if not is_consolidation_mode(pool):
             job_status = job_lib.get_status(job_id)
             if job_status is None or job_status < job_lib.JobStatus.RUNNING:
@@ -971,7 +979,7 @@ def wait_service_registration(service_name: str, job_id: int,
                     with ux_utils.print_exception_no_traceback():
                         raise RuntimeError(
                             f'Failed to start the controller process for '
-                            f'the service {service_name!r} within '
+                            f'the {noun} {service_name!r} within '
                             f'{constants.CONTROLLER_SETUP_TIMEOUT_SECONDS}'
                             f' seconds.')
                 # No need to check the service status as the controller process
@@ -979,22 +987,26 @@ def wait_service_registration(service_name: str, job_id: int,
                 time.sleep(1)
                 continue
 
-            if not setup_completed:
-                setup_completed = True
-                # Reset the start time to wait for the service to be registered.
-                start_time = time.time()
+        if not setup_completed:
+            setup_completed = True
+            # Reset the start time to wait for the service to be registered.
+            start_time = time.time()
 
-        record = serve_state.get_service_from_name(service_name)
+        record = _get_service_status(service_name,
+                                     pool=pool,
+                                     with_replica_info=False)
         if record is not None:
-            # TODO(tian): PID-based tracking.
-            if (not is_consolidation_mode(pool) and
-                    job_id != record['controller_job_id']):
+            if job_id != record['controller_job_id']:
+                if pool:
+                    command_to_run = 'sky jobs pool apply --pool'
+                else:
+                    command_to_run = 'sky serve update'
                 with ux_utils.print_exception_no_traceback():
                     raise ValueError(
-                        f'The service {service_name!r} is already running. '
-                        'Please specify a different name for your service. '
-                        'To update an existing service, run: sky serve update '
-                        f'{service_name} <new-service-yaml>')
+                        f'The {noun} {service_name!r} is already running. '
+                        f'Please specify a different name for your {noun}. '
+                        f'To update an existing {noun}, run: {command_to_run}'
+                        f' {service_name} <new-{noun}-yaml>')
             lb_port = record['load_balancer_port']
             if lb_port is not None:
                 return message_utils.encode_payload(lb_port)
@@ -1023,12 +1035,16 @@ def load_service_initialization_result(payload: str) -> int:
     return message_utils.decode_payload(payload)
 
 
-def check_service_status_healthy(service_name: str) -> Optional[str]:
-    service_record = serve_state.get_service_from_name(service_name)
+def _check_service_status_healthy(service_name: str,
+                                  pool: bool) -> Optional[str]:
+    service_record = _get_service_status(service_name,
+                                         pool,
+                                         with_replica_info=False)
+    capnoun = 'Service' if not pool else 'Pool'
     if service_record is None:
-        return f'Service {service_name!r} does not exist.'
+        return f'{capnoun} {service_name!r} does not exist.'
     if service_record['status'] == serve_state.ServiceStatus.CONTROLLER_INIT:
-        return (f'Service {service_name!r} is still initializing its '
+        return (f'{capnoun} {service_name!r} is still initializing its '
                 'controller. Please try again later.')
     return None
 
@@ -1067,7 +1083,10 @@ def _process_line(line: str,
     log_prompt = re.match(_SKYPILOT_LOG_PATTERN, line)
 
     if provision_log_prompt is not None:
-        nested_log_path = os.path.expanduser(provision_log_prompt.group(1))
+        log_path = provision_log_prompt.group(1)
+        nested_log_path = pathlib.Path(
+            skylet_constants.SKY_LOGS_DIRECTORY).expanduser().joinpath(
+                log_path).resolve()
 
         try:
             with open(nested_log_path, 'r', newline='', encoding='utf-8') as f:
@@ -1159,12 +1178,14 @@ def _capped_follow_logs_with_provision_expanding(
 
 
 def stream_replica_logs(service_name: str, replica_id: int, follow: bool,
-                        tail: Optional[int]) -> str:
-    msg = check_service_status_healthy(service_name)
+                        tail: Optional[int], pool: bool) -> str:
+    msg = _check_service_status_healthy(service_name, pool=pool)
     if msg is not None:
         return msg
+    repnoun = 'worker' if pool else 'replica'
+    caprepnoun = repnoun.capitalize()
     print(f'{colorama.Fore.YELLOW}Start streaming logs for launching process '
-          f'of replica {replica_id}.{colorama.Style.RESET_ALL}')
+          f'of {repnoun} {replica_id}.{colorama.Style.RESET_ALL}')
     log_file_name = generate_replica_log_file_name(service_name, replica_id)
     if os.path.exists(log_file_name):
         if tail is not None:
@@ -1181,7 +1202,7 @@ def stream_replica_logs(service_name: str, replica_id: int, follow: bool,
     launch_log_file_name = generate_replica_launch_log_file_name(
         service_name, replica_id)
     if not os.path.exists(launch_log_file_name):
-        return (f'{colorama.Fore.RED}Replica {replica_id} doesn\'t exist.'
+        return (f'{colorama.Fore.RED}{caprepnoun} {replica_id} doesn\'t exist.'
                 f'{colorama.Style.RESET_ALL}')
 
     replica_cluster_name = generate_replica_cluster_name(
@@ -1231,6 +1252,10 @@ def stream_replica_logs(service_name: str, replica_id: int, follow: bool,
                 print(line, end='', flush=True)
         return ''
 
+    # For pools, we don't stream the job logs as the run section is ignored.
+    if pool:
+        return ''
+
     backend = backends.CloudVmRayBackend()
     handle = global_user_state.get_handle_from_cluster_name(
         replica_cluster_name)
@@ -1245,13 +1270,13 @@ def stream_replica_logs(service_name: str, replica_id: int, follow: bool,
 
     # Notify user here to make sure user won't think the log is finished.
     print(f'{colorama.Fore.YELLOW}Start streaming logs for task job '
-          f'of replica {replica_id}...{colorama.Style.RESET_ALL}')
+          f'of {repnoun} {replica_id}...{colorama.Style.RESET_ALL}')
 
     # Always tail the latest logs, which represent user setup & run.
     if tail is None:
         returncode = backend.tail_logs(handle, job_id=None, follow=follow)
         if returncode != 0:
-            return (f'{colorama.Fore.RED}Failed to stream logs for replica '
+            return (f'{colorama.Fore.RED}Failed to stream logs for {repnoun} '
                     f'{replica_id}.{colorama.Style.RESET_ALL}')
     elif not follow and tail > 0:
         final = backend.tail_logs(handle,
@@ -1278,8 +1303,9 @@ def stream_replica_logs(service_name: str, replica_id: int, follow: bool,
 
 
 def stream_serve_process_logs(service_name: str, stream_controller: bool,
-                              follow: bool, tail: Optional[int]) -> str:
-    msg = check_service_status_healthy(service_name)
+                              follow: bool, tail: Optional[int],
+                              pool: bool) -> str:
+    msg = _check_service_status_healthy(service_name, pool)
     if msg is not None:
         return msg
     if stream_controller:
@@ -1288,7 +1314,9 @@ def stream_serve_process_logs(service_name: str, stream_controller: bool,
         log_file = generate_remote_load_balancer_log_file_name(service_name)
 
     def _service_is_terminal() -> bool:
-        record = serve_state.get_service_from_name(service_name)
+        record = _get_service_status(service_name,
+                                     pool,
+                                     with_replica_info=False)
         if record is None:
             return True
         return record['status'] in serve_state.ServiceStatus.failed_statuses()
@@ -1531,21 +1559,24 @@ class ServeCodeGen:
 
     @classmethod
     def stream_replica_logs(cls, service_name: str, replica_id: int,
-                            follow: bool, tail: Optional[int]) -> str:
+                            follow: bool, tail: Optional[int],
+                            pool: bool) -> str:
         code = [
+            f'kwargs={{}} if serve_version < 5 else {{"pool": {pool}}}',
             'msg = serve_utils.stream_replica_logs('
-            f'{service_name!r}, {replica_id!r}, follow={follow}, tail={tail})',
-            'print(msg, flush=True)'
+            f'{service_name!r}, {replica_id!r}, follow={follow}, tail={tail}, '
+            '**kwargs)', 'print(msg, flush=True)'
         ]
         return cls._build(code)
 
     @classmethod
     def stream_serve_process_logs(cls, service_name: str,
                                   stream_controller: bool, follow: bool,
-                                  tail: Optional[int]) -> str:
+                                  tail: Optional[int], pool: bool) -> str:
         code = [
+            f'kwargs={{}} if serve_version < 5 else {{"pool": {pool}}}',
             f'msg = serve_utils.stream_serve_process_logs({service_name!r}, '
-            f'{stream_controller}, follow={follow}, tail={tail})',
+            f'{stream_controller}, follow={follow}, tail={tail}, **kwargs)',
             'print(msg, flush=True)'
         ]
         return cls._build(code)

sky/serve/server/core.py

@@ -1,9 +1,6 @@
 """SkyServe core APIs."""
-import pathlib
-import signal
-import threading
 import typing
-from typing import Any, Dict, List, Optional, Set, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from sky import backends
 from sky import exceptions
@@ -12,11 +9,8 @@ from sky.backends import backend_utils
 from sky.serve import serve_utils
 from sky.serve.server import impl
 from sky.usage import usage_lib
-from sky.utils import command_runner
 from sky.utils import controller_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
@@ -24,42 +18,6 @@ if typing.TYPE_CHECKING:
 logger = sky_logging.init_logger(__name__)
 
 
-def _get_all_replica_targets(
-    service_name: str, backend: backends.CloudVmRayBackend,
-    handle: backends.CloudVmRayResourceHandle
-) -> Set[serve_utils.ServiceComponentTarget]:
-    """Helper function to get targets for all live replicas."""
-    code = serve_utils.ServeCodeGen.get_service_status([service_name],
-                                                       pool=False)
-    returncode, serve_status_payload, stderr = backend.run_on_head(
-        handle,
-        code,
-        require_outputs=True,
-        stream_logs=False,
-        separate_stderr=True)
-
-    try:
-        subprocess_utils.handle_returncode(returncode,
-                                           code,
-                                           'Failed to fetch services',
-                                           stderr,
-                                           stream_logs=True)
-    except exceptions.CommandError as e:
-        raise RuntimeError(e.error_msg) from e
-
-    service_records = serve_utils.load_service_status(serve_status_payload)
-    if not service_records:
-        raise ValueError(f'Service {service_name!r} not found.')
-    assert len(service_records) == 1
-    service_record = service_records[0]
-
-    return {
-        serve_utils.ServiceComponentTarget(serve_utils.ServiceComponent.REPLICA,
-                                           replica_info['replica_id'])
-        for replica_info in service_record['replica_info']
-    }
-
-
 @usage_lib.entrypoint
 def up(
     task: 'sky.Task',
@@ -277,59 +235,12 @@ def tail_logs(
         sky.exceptions.ClusterNotUpError: the sky serve controller is not up.
         ValueError: arguments not valid, or failed to tail the logs.
     """
-    if isinstance(target, str):
-        target = serve_utils.ServiceComponent(target)
-    if not isinstance(target, serve_utils.ServiceComponent):
-        with ux_utils.print_exception_no_traceback():
-            raise ValueError(f'`target` must be a string or '
-                             f'sky.serve.ServiceComponent, got {type(target)}.')
-
-    if target == serve_utils.ServiceComponent.REPLICA:
-        if replica_id is None:
-            with ux_utils.print_exception_no_traceback():
-                raise ValueError(
-                    '`replica_id` must be specified when using target=REPLICA.')
-    else:
-        if replica_id is not None:
-            with ux_utils.print_exception_no_traceback():
-                raise ValueError('`replica_id` must be None when using '
-                                 'target=CONTROLLER/LOAD_BALANCER.')
-
-    controller_type = controller_utils.Controllers.SKY_SERVE_CONTROLLER
-    handle = backend_utils.is_controller_accessible(
-        controller=controller_type,
-        stopped_message=controller_type.value.default_hint_if_non_existent)
-
-    backend = backend_utils.get_backend_from_handle(handle)
-    assert isinstance(backend, backends.CloudVmRayBackend), backend
-
-    if target != serve_utils.ServiceComponent.REPLICA:
-        code = serve_utils.ServeCodeGen.stream_serve_process_logs(
-            service_name,
-            stream_controller=(
-                target == serve_utils.ServiceComponent.CONTROLLER),
-            follow=follow,
-            tail=tail)
-    else:
-        assert replica_id is not None, service_name
-        code = serve_utils.ServeCodeGen.stream_replica_logs(service_name,
-                                                            replica_id,
-                                                            follow,
-                                                            tail=tail)
-
-    # With the stdin=subprocess.DEVNULL, the ctrl-c will not directly
-    # kill the process, so we need to handle it manually here.
-    if threading.current_thread() is threading.main_thread():
-        signal.signal(signal.SIGINT, backend_utils.interrupt_handler)
-        signal.signal(signal.SIGTSTP, backend_utils.stop_handler)
-
-    # Refer to the notes in
-    # sky/backends/cloud_vm_ray_backend.py::CloudVmRayBackend::tail_logs.
-    backend.run_on_head(handle,
-                        code,
-                        stream_logs=True,
-                        process_stream=False,
-                        ssh_mode=command_runner.SshMode.INTERACTIVE)
+    return impl.tail_logs(service_name,
+                          target=target,
+                          replica_id=replica_id,
+                          follow=follow,
+                          tail=tail,
+                          pool=False)
 
 
 @usage_lib.entrypoint
@@ -374,104 +285,9 @@ def sync_down_logs(
         sky.exceptions.ClusterNotUpError: If the controller is not up.
         ValueError: Arguments not valid.
     """
-    # Step 0) get the controller handle
-    with rich_utils.safe_status(
-            ux_utils.spinner_message('Checking service status...')):
-        controller_type = controller_utils.Controllers.SKY_SERVE_CONTROLLER
-        handle = backend_utils.is_controller_accessible(
-            controller=controller_type,
-            stopped_message=controller_type.value.default_hint_if_non_existent)
-        backend: backends.CloudVmRayBackend = (
-            backend_utils.get_backend_from_handle(handle))
-
-    requested_components: Set[serve_utils.ServiceComponent] = set()
-    if not targets:
-        # No targets specified -> request all components
-        requested_components = {
-            serve_utils.ServiceComponent.CONTROLLER,
-            serve_utils.ServiceComponent.LOAD_BALANCER,
-            serve_utils.ServiceComponent.REPLICA
-        }
-    else:
-        # Parse provided targets
-        if isinstance(targets, (str, serve_utils.ServiceComponent)):
-            requested_components = {serve_utils.ServiceComponent(targets)}
-        else:  # list
-            requested_components = {
-                serve_utils.ServiceComponent(t) for t in targets
-            }
-
-    normalized_targets: Set[serve_utils.ServiceComponentTarget] = set()
-    if serve_utils.ServiceComponent.CONTROLLER in requested_components:
-        normalized_targets.add(
-            serve_utils.ServiceComponentTarget(
-                serve_utils.ServiceComponent.CONTROLLER))
-    if serve_utils.ServiceComponent.LOAD_BALANCER in requested_components:
-        normalized_targets.add(
-            serve_utils.ServiceComponentTarget(
-                serve_utils.ServiceComponent.LOAD_BALANCER))
-    if serve_utils.ServiceComponent.REPLICA in requested_components:
-        with rich_utils.safe_status(
-                ux_utils.spinner_message('Getting live replica infos...')):
-            replica_targets = _get_all_replica_targets(service_name, backend,
-                                                       handle)
-        if not replica_ids:
-            # Replica target requested but no specific IDs
-            # -> Get all replica logs
-            normalized_targets.update(replica_targets)
-        else:
-            # Replica target requested with specific IDs
-            requested_replica_targets = [
-                serve_utils.ServiceComponentTarget(
-                    serve_utils.ServiceComponent.REPLICA, rid)
-                for rid in replica_ids
-            ]
-            for target in requested_replica_targets:
-                if target not in replica_targets:
-                    logger.warning(f'Replica ID {target.replica_id} not found '
-                                   f'for {service_name}. Skipping...')
-                else:
-                    normalized_targets.add(target)
-
-    def sync_down_logs_by_target(target: serve_utils.ServiceComponentTarget):
-        component = target.component
-        # We need to set one side of the pipe to a logs stream, and the other
-        # side to a file.
-        log_path = str(pathlib.Path(local_dir) / f'{target}.log')
-        stream_logs_code: str
-
-        if component == serve_utils.ServiceComponent.CONTROLLER:
-            stream_logs_code = (
-                serve_utils.ServeCodeGen.stream_serve_process_logs(
-                    service_name,
-                    stream_controller=True,
-                    follow=False,
-                    tail=tail))
-        elif component == serve_utils.ServiceComponent.LOAD_BALANCER:
-            stream_logs_code = (
-                serve_utils.ServeCodeGen.stream_serve_process_logs(
-                    service_name,
-                    stream_controller=False,
-                    follow=False,
-                    tail=tail))
-        elif component == serve_utils.ServiceComponent.REPLICA:
-            replica_id = target.replica_id
-            assert replica_id is not None, service_name
-            stream_logs_code = serve_utils.ServeCodeGen.stream_replica_logs(
-                service_name, replica_id, follow=False, tail=tail)
-        else:
-            assert False, component
-
-        # Refer to the notes in
-        # sky/backends/cloud_vm_ray_backend.py::CloudVmRayBackend::tail_logs.
-        backend.run_on_head(handle,
-                            stream_logs_code,
-                            stream_logs=False,
-                            process_stream=False,
-                            ssh_mode=command_runner.SshMode.INTERACTIVE,
-                            log_path=log_path)
-
-    subprocess_utils.run_in_parallel(sync_down_logs_by_target,
-                                     list(normalized_targets))
-
-    return local_dir
+    return impl.sync_down_logs(service_name,
+                               local_dir=local_dir,
+                               targets=targets,
+                               replica_ids=replica_ids,
+                               tail=tail,
+                               pool=False)