skypilot-nightly 1.0.0.dev20250905__py3-none-any.whl → 1.0.0.dev20251210__py3-none-any.whl

sky/skylet/log_lib.py

@@ -8,11 +8,13 @@ import functools
 import io
 import multiprocessing.pool
 import os
+import queue as queue_lib
 import shlex
 import subprocess
 import sys
 import tempfile
 import textwrap
+import threading
 import time
 from typing import (Deque, Dict, Iterable, Iterator, List, Optional, TextIO,
                     Tuple, Union)
@@ -39,6 +41,11 @@ logger = sky_logging.init_logger(__name__)
 
 LOG_FILE_START_STREAMING_AT = 'Waiting for task resources on '
 
+# 16-64KiB seems to be the sweet spot:
+# https://github.com/grpc/grpc.github.io/issues/371
+# TODO(kevin): Benchmark this ourselves and verify.
+DEFAULT_LOG_CHUNK_SIZE = 16 * 1024  # 16KiB
+
 
 class _ProcessingArgs:
     """Arguments for processing logs."""
@@ -165,7 +172,7 @@ def run_with_log(
     streaming_prefix: Optional[str] = None,
     log_cmd: bool = False,
     **kwargs,
-) -> Union[int, Tuple[int, str, str]]:
+) -> Union[int, Tuple[int, str, str], Tuple[int, int]]:
     """Runs a command and logs its output to a file.
 
     Args:
@@ -176,6 +183,8 @@ def run_with_log(
         process_stream: Whether to post-process the stdout/stderr of the
             command, such as replacing or skipping lines on the fly. If
             enabled, lines are printed only when '\r' or '\n' is found.
+        streaming_prefix: Optional prefix for each log line. Can contain {pid}
+            placeholder which will be replaced with the subprocess PID.
 
     Returns the returncode or returncode, stdout and stderr of the command.
       Note that the stdout and stderr is already decoded.
@@ -213,7 +222,21 @@ def run_with_log(
                           stdin=stdin,
                           **kwargs) as proc:
         try:
-            subprocess_utils.kill_process_daemon(proc.pid)
+            if ctx is not None:
+                # When runs in coroutine, use kill_pg if available to avoid
+                # the overhead of refreshing the process tree in the daemon.
+                subprocess_utils.kill_process_daemon(proc.pid, use_kill_pg=True)
+            else:
+                # For backward compatibility, do not specify use_kill_pg by
+                # default.
+                subprocess_utils.kill_process_daemon(proc.pid)
+
+            # Format streaming_prefix with subprocess PID if it contains {pid}
+            formatted_streaming_prefix = streaming_prefix
+            if streaming_prefix and '{pid}' in streaming_prefix:
+                formatted_streaming_prefix = streaming_prefix.format(
+                    pid=proc.pid)
+
             stdout = ''
             stderr = ''
             stdout_stream_handler = None
@@ -242,7 +265,7 @@ def run_with_log(
                     line_processor=line_processor,
                     # Replace CRLF when the output is logged to driver by ray.
                     replace_crlf=with_ray,
-                    streaming_prefix=streaming_prefix,
+                    streaming_prefix=formatted_streaming_prefix,
                 )
                 stdout_stream_handler = functools.partial(
                     _handle_io_stream,
@@ -264,7 +287,6 @@ def run_with_log(
                 stdout, stderr = context_utils.pipe_and_wait_process(
                     ctx,
                     proc,
-                    cancel_callback=subprocess_utils.kill_children_processes,
                     stdout_stream_handler=stdout_stream_handler,
                     stderr_stream_handler=stderr_stream_handler)
             elif process_stream:
@@ -336,7 +358,8 @@ def run_bash_command_with_log(bash_command: str,
                               log_path: str,
                               env_vars: Optional[Dict[str, str]] = None,
                               stream_logs: bool = False,
-                              with_ray: bool = False):
+                              with_ray: bool = False,
+                              streaming_prefix: Optional[str] = None):
     with tempfile.NamedTemporaryFile('w', prefix='sky_app_',
                                      delete=False) as fp:
         bash_command = make_task_bash_script(bash_command, env_vars=env_vars)
@@ -351,9 +374,26 @@ def run_bash_command_with_log(bash_command: str,
                             log_path,
                             stream_logs=stream_logs,
                             with_ray=with_ray,
+                            streaming_prefix=streaming_prefix,
                             shell=True)
 
 
+def run_bash_command_with_log_and_return_pid(
+        bash_command: str,
+        log_path: str,
+        env_vars: Optional[Dict[str, str]] = None,
+        stream_logs: bool = False,
+        with_ray: bool = False,
+        streaming_prefix: Optional[str] = None):
+    return_code = run_bash_command_with_log(bash_command,
+                                            log_path,
+                                            env_vars,
+                                            stream_logs,
+                                            with_ray,
+                                            streaming_prefix=streaming_prefix)
+    return {'return_code': return_code, 'pid': os.getpid()}
+
+
 def _follow_job_logs(file,
                      job_id: int,
                      start_streaming: bool,
@@ -395,9 +435,9 @@ def _follow_job_logs(file,
                     wait_last_logs = False
                     continue
                 status_str = status.value if status is not None else 'None'
-                print(ux_utils.finishing_message(
-                    f'Job finished (status: {status_str}).'),
-                      flush=True)
+                finish = ux_utils.finishing_message(
+                    f'Job finished (status: {status_str}).')
+                yield finish + '\n'
                 return
 
             time.sleep(SKY_LOG_TAILING_GAP_SECONDS)
@@ -552,3 +592,207 @@ def tail_logs(job_id: Optional[int],
         except FileNotFoundError:
             print(f'{colorama.Fore.RED}ERROR: Logs for job {job_id} (status:'
                   f' {status.value}) does not exist.{colorama.Style.RESET_ALL}')
+
+
+def tail_logs_iter(job_id: Optional[int],
+                   log_dir: Optional[str],
+                   managed_job_id: Optional[int] = None,
+                   follow: bool = True,
+                   tail: int = 0) -> Iterator[str]:
+    """Tail the logs of a job. This is mostly the same as tail_logs, but
+    returns an iterator instead of printing to stdout/stderr."""
+    if job_id is None:
+        # This only happens when job_lib.get_latest_job_id() returns None,
+        # which means no job has been submitted to this cluster. See
+        # sky.skylet.job_lib.JobLibCodeGen.tail_logs for more details.
+        logger.info('Skip streaming logs as no job has been submitted.')
+        return
+    job_str = f'job {job_id}'
+    if managed_job_id is not None:
+        job_str = f'managed job {managed_job_id}'
+    if log_dir is None:
+        msg = f'{job_str.capitalize()} not found (see `sky queue`).'
+        yield msg + '\n'
+        return
+    logger.debug(f'Tailing logs for job, real job_id {job_id}, managed_job_id '
+                 f'{managed_job_id}.')
+    log_path = os.path.join(log_dir, 'run.log')
+    log_path = os.path.expanduser(log_path)
+
+    status = job_lib.update_job_status([job_id], silent=True)[0]
+
+    # Wait for the log to be written. This is needed due to the `ray submit`
+    # will take some time to start the job and write the log.
+    retry_cnt = 0
+    while status is not None and not status.is_terminal():
+        retry_cnt += 1
+        if os.path.exists(log_path) and status != job_lib.JobStatus.INIT:
+            break
+        if retry_cnt >= SKY_LOG_WAITING_MAX_RETRY:
+            err = (f'{colorama.Fore.RED}ERROR: Logs for '
+                   f'{job_str} (status: {status.value}) does not exist '
+                   f'after retrying {retry_cnt} times.'
+                   f'{colorama.Style.RESET_ALL}')
+            yield err + '\n'
+            return
+        waiting = (f'INFO: Waiting {SKY_LOG_WAITING_GAP_SECONDS}s for the logs '
+                   'to be written...')
+        yield waiting + '\n'
+        time.sleep(SKY_LOG_WAITING_GAP_SECONDS)
+        status = job_lib.update_job_status([job_id], silent=True)[0]
+
+    start_stream_at = LOG_FILE_START_STREAMING_AT
+    # Explicitly declare the type to avoid mypy warning.
+    lines: Iterable[str] = []
+    if follow and status in [
+            job_lib.JobStatus.SETTING_UP,
+            job_lib.JobStatus.PENDING,
+            job_lib.JobStatus.RUNNING,
+    ]:
+        # Not using `ray job logs` because it will put progress bar in
+        # multiple lines.
+        with open(log_path, 'r', newline='', encoding='utf-8') as log_file:
+            # Using `_follow` instead of `tail -f` to streaming the whole
+            # log and creating a new process for tail.
+            start_streaming = False
+            if tail > 0:
+                head_lines_of_log_file = _peek_head_lines(log_file)
+                lines = collections.deque(log_file, maxlen=tail)
+                start_streaming = _should_stream_the_whole_tail_lines(
+                    head_lines_of_log_file, lines, start_stream_at)
+                for line in lines:
+                    if start_stream_at in line:
+                        start_streaming = True
+                    if start_streaming:
+                        yield line
+            # Now, the cursor is at the end of the last lines
+            # if tail > 0
+            for line in _follow_job_logs(log_file,
+                                         job_id=job_id,
+                                         start_streaming=start_streaming,
+                                         start_streaming_at=start_stream_at):
+                yield line
+    else:
+        try:
+            start_streaming = False
+            with open(log_path, 'r', encoding='utf-8') as log_file:
+                if tail > 0:
+                    # If tail > 0, we need to read the last n lines.
+                    # We use double ended queue to rotate the last n lines.
+                    head_lines_of_log_file = _peek_head_lines(log_file)
+                    lines = collections.deque(log_file, maxlen=tail)
+                    start_streaming = _should_stream_the_whole_tail_lines(
+                        head_lines_of_log_file, lines, start_stream_at)
+                else:
+                    lines = log_file
+                for line in lines:
+                    if start_stream_at in line:
+                        start_streaming = True
+                    if start_streaming:
+                        yield line
+                status_str = status.value if status is not None else 'None'
+                # Only show "Job finished" for actually terminal states
+                if status is not None and status.is_terminal():
+                    finish = ux_utils.finishing_message(
+                        f'Job finished (status: {status_str}).')
+                    yield finish + '\n'
+            return
+        except FileNotFoundError:
+            err = (
+                f'{colorama.Fore.RED}ERROR: Logs for job {job_id} (status:'
+                f' {status.value}) does not exist.{colorama.Style.RESET_ALL}')
+            yield err + '\n'
+
+
+class LogBuffer:
+    """In-memory buffer for chunking log lines for streaming."""
+
+    def __init__(self, max_chars: int = DEFAULT_LOG_CHUNK_SIZE):
+        """Initialize the log buffer.
+
+        Args:
+            max_chars: Maximum buffer size (in characters, not bytes) before
+                       flushing. The actual amount of bytes (UTF-8 encoding)
+                       could be more than this, depending on the characters,
+                       i.e. ASCII characters take 1 byte, while others
+                       may take 2-4 bytes. But this is fine as our default
+                       chunk size is well below the default value of
+                       grpc.max_receive_message_length which is 4MB.
+        """
+        self.max_chars = max_chars
+        self._buffer = io.StringIO()
+
+    def _should_flush(self) -> bool:
+        return self._buffer.tell() >= self.max_chars
+
+    def flush(self) -> str:
+        """Get the current buffered content and clear the buffer.
+
+        Returns:
+            The buffered log lines as a single string
+        """
+        if not self._buffer.tell():
+            return ''
+        chunk = self._buffer.getvalue()
+        self._buffer.truncate(0)
+        self._buffer.seek(0)
+        return chunk
+
+    def write(self, line: str) -> bool:
+        """Add a line to the buffer.
+
+        Args:
+            line: The log line to add
+
+        Returns:
+            True if buffer should be flushed after adding the line
+        """
+        self._buffer.write(line)
+        return self._should_flush()
+
+    def close(self):
+        self._buffer.close()
+
+
+def buffered_iter_with_timeout(buffer: LogBuffer, iterable: Iterable[str],
+                               timeout: float) -> Iterable[str]:
+    """Iterates over an iterable, writing each item to a buffer,
+        and flushing the buffer when it is full or no item is
+        yielded within the timeout duration."""
+    # TODO(kevin): Simplify this using asyncio.timeout, once we move
+    # the skylet event loop and gRPC server to asyncio.
+    # https://docs.python.org/3/library/asyncio-task.html#timeouts
+
+    queue: queue_lib.Queue = queue_lib.Queue()
+    sentinel = object()
+
+    def producer():
+        try:
+            for item in iterable:
+                queue.put(item)
+        finally:
+            queue.put(sentinel)
+
+    thread = threading.Thread(target=producer, daemon=True)
+    thread.start()
+
+    while True:
+        try:
+            item = queue.get(timeout=timeout)
+        except queue_lib.Empty:
+            out = buffer.flush()
+            if out:
+                yield out
+            continue
+
+        if item is sentinel:
+            thread.join()
+            out = buffer.flush()
+            if out:
+                yield out
+            return
+
+        if buffer.write(item):
+            out = buffer.flush()
+            if out:
+                yield out

sky/skylet/log_lib.pyi

@@ -4,7 +4,7 @@ overloaded type hints for run_with_log(), as we need to determine
 the return type based on the value of require_outputs.
 """
 import typing
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, Iterable, Iterator, List, Optional, Tuple, Union
 
 from typing_extensions import Literal
 
@@ -42,7 +42,7 @@ class _ProcessingArgs:
         ...
 
 
-def _get_context() -> Optional[context.Context]:
+def _get_context() -> Optional[context.SkyPilotContext]:
     ...
 
 
@@ -68,7 +68,7 @@ def run_with_log(cmd: Union[List[str], str],
                  process_stream: bool = ...,
                  line_processor: Optional[log_utils.LineProcessor] = ...,
                  streaming_prefix: Optional[str] = ...,
-                 ray_job_id: Optional[str] = ...,
+                 log_cmd: bool = ...,
                  **kwargs) -> int:
     ...
 
@@ -87,7 +87,7 @@ def run_with_log(cmd: Union[List[str], str],
                  process_stream: bool = ...,
                  line_processor: Optional[log_utils.LineProcessor] = ...,
                  streaming_prefix: Optional[str] = ...,
-                 ray_job_id: Optional[str] = ...,
+                 log_cmd: bool = ...,
                  **kwargs) -> Tuple[int, str, str]:
     ...
 
@@ -106,8 +106,8 @@ def run_with_log(cmd: Union[List[str], str],
                  process_stream: bool = ...,
                  line_processor: Optional[log_utils.LineProcessor] = ...,
                  streaming_prefix: Optional[str] = ...,
-                 ray_job_id: Optional[str] = ...,
-                 **kwargs) -> Union[int, Tuple[int, str, str]]:
+                 log_cmd: bool = ...,
+                 **kwargs) -> Tuple[int, int]:
     ...
 
 
@@ -125,7 +125,18 @@ def run_bash_command_with_log(bash_command: str,
                               log_path: str,
                               env_vars: Optional[Dict[str, str]] = ...,
                               stream_logs: bool = ...,
-                              with_ray: bool = ...):
+                              with_ray: bool = ...,
+                              streaming_prefix: Optional[str] = ...) -> int:
+    ...
+
+
+def run_bash_command_with_log_and_return_pid(
+        bash_command: str,
+        log_path: str,
+        env_vars: Optional[Dict[str, str]] = ...,
+        stream_logs: bool = ...,
+        with_ray: bool = ...,
+        streaming_prefix: Optional[str] = ...) -> Dict[str, Union[int, str]]:
     ...
 
 
@@ -134,3 +145,32 @@ def tail_logs(job_id: int,
               managed_job_id: Optional[int] = ...,
               follow: bool = ...) -> None:
     ...
+
+
+def tail_logs_iter(job_id: Optional[int],
+                   log_dir: Optional[str],
+                   managed_job_id: Optional[int] = ...,
+                   follow: bool = ...,
+                   tail: int = ...) -> Iterator[str]:
+    ...
+
+
+class LogBuffer:
+    max_chars: int
+
+    def __init__(self, max_chars: int = ...):
+        ...
+
+    def flush(self) -> str:
+        ...
+
+    def write(self, line: str) -> bool:
+        ...
+
+    def close(self):
+        ...
+
+
+def buffered_iter_with_timeout(buffer: LogBuffer, iterable: Iterable[str],
+                               timeout: float) -> Iterable[str]:
+    ...

sky/skylet/providers/ibm/node_provider.py

@@ -24,7 +24,7 @@ import socket
 import threading
 import time
 from pathlib import Path
-from pprint import pprint
+from pprint import pformat, pprint
 from typing import Any, Dict, List, Optional
 from uuid import uuid4
 
@@ -67,13 +67,13 @@ def log_in_out(func):
         logger.debug(
             f"\n\nEnter {name} from {inspect.stack()[0][3]} "
             f"{inspect.stack()[1][3]} {inspect.stack()[2][3]} with args: "
-            f"entered with args:\n{pprint(args)} and kwargs {pprint(kwargs)}"
+            f"entered with args:\n{pformat(args)} and kwargs {pformat(kwargs)}"
         )
         try:
             result = func(*args, **kwargs)
             logger.debug(
                 f"Leave {name} from {inspect.stack()[1][3]} with result "
-                f"Func Result:{pprint(result)}\n\n"
+                f"Func Result:{pformat(result)}\n\n"
             )
         except Exception:
             cli_logger.error(f"Error in {name}")
@@ -445,7 +445,7 @@ class IBMVPCNodeProvider(NodeProvider):
         """returns the worker's node private ip address"""
         node = self._get_cached_node(node_id)
 
-        # if a bug ocurred, or node data was fetched before primary_ip
+        # if a bug occurred, or node data was fetched before primary_ip
         # was assigned, refetch node data from cloud.
         try:
             primary_ip = node["network_interfaces"][0].get("primary_ip")["address"]
@@ -502,8 +502,12 @@ class IBMVPCNodeProvider(NodeProvider):
 
         logger.info(f"Creating new VM instance {name}")
 
-        security_group_identity_model = {"id": self.vpc_tags["security_group_id"]}
-        subnet_identity_model = {"id": self.vpc_tags["subnet_id"]}
+        if self.vpc_tags is None:
+            raise ValueError("vpc_tags must be initialized before creating instances")
+        vpc_tags = self.vpc_tags  # Help mypy with type narrowing
+
+        security_group_identity_model = {"id": vpc_tags["security_group_id"]}
+        subnet_identity_model = {"id": vpc_tags["subnet_id"]}
         primary_network_interface = {
             "name": "eth0",
             "subnet": subnet_identity_model,
@@ -536,7 +540,7 @@ class IBMVPCNodeProvider(NodeProvider):
         instance_prototype["keys"] = [key_identity_model]
         instance_prototype["profile"] = {"name": profile_name}
         instance_prototype["resource_group"] = {"id": self.resource_group_id}
-        instance_prototype["vpc"] = {"id": self.vpc_tags["vpc_id"]}
+        instance_prototype["vpc"] = {"id": vpc_tags["vpc_id"]}
         instance_prototype["image"] = {"id": base_config["image_id"]}
 
         instance_prototype["zone"] = {"name": self.zone}
@@ -584,7 +588,7 @@ class IBMVPCNodeProvider(NodeProvider):
         floating_ip_name = f"{RAY_RECYCLABLE}-{uuid4().hex[:4]}"
         # create a new floating ip
         logger.debug(f"Creating floating IP {floating_ip_name}")
-        floating_ip_prototype = {}
+        floating_ip_prototype: Dict[str, Any] = {}
         floating_ip_prototype["name"] = floating_ip_name
         floating_ip_prototype["zone"] = {"name": self.zone}
         floating_ip_prototype["resource_group"] = {"id": self.resource_group_id}

sky/skylet/providers/ibm/vpc_provider.py

@@ -10,6 +10,7 @@ import textwrap
 import time
 import uuid
 from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Dict
 
 import requests
 
@@ -173,7 +174,7 @@ class IBMVPCProvider:
                 "a subnet"
             )
 
-        subnet_prototype = {}
+        subnet_prototype: Dict[str, Any] = {}
         subnet_prototype["zone"] = {"name": zone_name}
         subnet_prototype["ip_version"] = "ipv4"
         subnet_prototype["name"] = subnet_name
@@ -186,7 +187,7 @@ class IBMVPCProvider:
 
     def create_public_gateway(self, vpc_id, zone_name, subnet_data):
 
-        gateway_prototype = {}
+        gateway_prototype: Dict[str, Any] = {}
         gateway_prototype["vpc"] = {"id": vpc_id}
         gateway_prototype["zone"] = {"name": zone_name}
         gateway_prototype["name"] = f"{subnet_data['name']}-gw"
@@ -345,7 +346,7 @@ class IBMVPCProvider:
                     return True
                 tries -= 1
                 time.sleep(sleep_interval)
-            logger.error("Failed to delete instance within the alloted time\n")
+            logger.error("Failed to delete instance within the allotted time\n")
             return False
 
         for subnet_id in self.get_vpc_subnets(vpc_data, region, field="id"):
@@ -522,7 +523,7 @@ class ClusterCleaner:
             if e.code == 404:
                 print(("VPC doesn't exist."))
                 return None
-            else: raise 
+            else: raise
 
     def delete_subnets(vpc_data):
         def _poll_subnet_exists(subnet_id):
@@ -560,12 +561,12 @@ class ClusterCleaner:
                     deleting_resource = False
                 except ibm_cloud_sdk_core.ApiException as e:
                     if e.code == 404:
-                        print("gateway doesn't exist.") 
+                        print("gateway doesn't exist.")
                         deleting_resource = False
                     if e.code == 409:
                         print("gateway still in use.")
-                        # will retry until cloud functions timeout. 
-                        time.sleep(5) 
+                        # will retry until cloud functions timeout.
+                        time.sleep(5)
 
     def delete_vms(vpc_id):
         def _poll_vpc_contains_vms(vpc_id):
@@ -586,7 +587,7 @@ class ClusterCleaner:
             )
 
         def _del_instance(vm_data):
-            # first delete ips created by node_provider 
+            # first delete ips created by node_provider
             nic_id = vm_data["network_interfaces"][0]["id"]
             res = ibm_vpc_client.list_instance_network_interface_floating_ips(
                 vm_data["id"], nic_id
@@ -598,7 +599,7 @@ class ClusterCleaner:
                     ibm_vpc_client.delete_floating_ip(ip["id"])
             print(f"Deleting VM: {vm_data['id']}")
             ibm_vpc_client.delete_instance(id=vm_data["id"])
-            
+
         res = ibm_vpc_client.list_instances(vpc_id=vpc_id).get_result()
         num_instances = res["total_count"]
 
@@ -619,12 +620,12 @@ class ClusterCleaner:
                 deleting_resource = False
             except ibm_cloud_sdk_core.ApiException as e:
                 if e.code == 404:
-                    print("VPC doesn't exist.") 
+                    print("VPC doesn't exist.")
                     deleting_resource = False
                 if e.code == 409:
                     print("VPC still in use.")
-                    # will retry until cloud functions timeout. 
-                    time.sleep(5) 
+                    # will retry until cloud functions timeout.
+                    time.sleep(5)
 
     def delete_vpc(vpc_id):
         vpc_data = get_vpc_data(vpc_id)

sky/skylet/runtime_utils.py

@@ -0,0 +1,21 @@
+"""Runtime utilities for SkyPilot."""
+import os
+
+from sky.skylet import constants
+
+
+def get_runtime_dir_path(path_suffix: str = '') -> str:
+    """Get an expanded path within the SkyPilot runtime directory.
+
+    Args:
+        path_suffix: Path suffix to join with the runtime dir
+        (e.g., '.sky/jobs.db').
+
+    Returns:
+        The full expanded path.
+    """
+    runtime_dir = os.path.expanduser(
+        os.environ.get(constants.SKY_RUNTIME_DIR_ENV_VAR_KEY, '~'))
+    if path_suffix:
+        return os.path.join(runtime_dir, path_suffix)
+    return runtime_dir