skypilot-nightly 1.0.0.dev20250910__py3-none-any.whl → 1.0.0.dev20250912__py3-none-any.whl

sky/server/requests/executor.py

@@ -31,6 +31,7 @@ import time
 import typing
 from typing import Any, Callable, Generator, List, Optional, TextIO, Tuple
 
+import psutil
 import setproctitle
 
 from sky import exceptions
@@ -130,8 +131,9 @@ queue_backend = server_config.QueueBackend.MULTIPROCESSING
 def executor_initializer(proc_group: str):
     setproctitle.setproctitle(f'SkyPilot:executor:{proc_group}:'
                               f'{multiprocessing.current_process().pid}')
+    # Executor never stops, unless the whole process is killed.
     threading.Thread(target=metrics_lib.process_monitor,
-                     args=(f'worker:{proc_group}',),
+                     args=(f'worker:{proc_group}', threading.Event()),
                      daemon=True).start()
 
 
@@ -373,11 +375,13 @@ def _request_execution_wrapper(request_id: str,
     4. Handle the SIGTERM signal to abort the request gracefully.
     5. Maintain the lifecycle of the temp dir used by the request.
     """
+    pid = multiprocessing.current_process().pid
+    proc = psutil.Process(pid)
+    rss_begin = proc.memory_info().rss
     db_utils.set_max_connections(num_db_connections_per_worker)
     # Handle the SIGTERM signal to abort the request processing gracefully.
     signal.signal(signal.SIGTERM, _sigterm_handler)
 
-    pid = multiprocessing.current_process().pid
     logger.info(f'Running request {request_id} with pid {pid}')
     with api_requests.update_request(request_id) as request_task:
         assert request_task is not None, request_id
@@ -443,8 +447,41 @@ def _request_execution_wrapper(request_id: str,
             _restore_output(original_stdout, original_stderr)
             logger.info(f'Request {request_id} finished')
         finally:
-            with metrics_lib.time_it(name='release_memory', group='internal'):
-                common_utils.release_memory()
+            try:
+                # Capture the peak RSS before GC.
+                peak_rss = max(proc.memory_info().rss,
+                               metrics_lib.peak_rss_bytes)
+                with metrics_lib.time_it(name='release_memory',
+                                         group='internal'):
+                    common_utils.release_memory()
+                _record_memory_metrics(request_name, proc, rss_begin, peak_rss)
+            except Exception as e:  # pylint: disable=broad-except
+                logger.error(f'Failed to record memory metrics: '
+                             f'{common_utils.format_exception(e)}')
+
+
+_first_request = True
+
+
+def _record_memory_metrics(request_name: str, proc: psutil.Process,
+                           rss_begin: int, peak_rss: int) -> None:
+    """Record the memory metrics for a request."""
+    # Do not record full memory delta for the first request as it
+    # will loads the sky core modules and make the memory usage
+    # estimation inaccurate.
+    global _first_request
+    if _first_request:
+        _first_request = False
+        return
+    rss_end = proc.memory_info().rss
+
+    # Answer "how much RSS this request contributed?"
+    metrics_lib.SKY_APISERVER_REQUEST_RSS_INCR_BYTES.labels(
+        name=request_name).observe(max(rss_end - rss_begin, 0))
+    # Estimate the memory usage by the request by capturing the
+    # peak memory delta during the request execution.
+    metrics_lib.SKY_APISERVER_REQUEST_MEMORY_USAGE_BYTES.labels(
+        name=request_name).observe(max(peak_rss - rss_begin, 0))
 
 
 async def execute_request_coroutine(request: api_requests.Request):

sky/server/server.py

@@ -1214,6 +1214,7 @@ async def logs(
         request_body=cluster_job_body,
         func=core.tail_logs,
         schedule_type=requests_lib.ScheduleType.SHORT,
+        request_cluster_name=cluster_job_body.cluster_name,
     )
     task = asyncio.create_task(executor.execute_request_coroutine(request_task))
 

sky/server/uvicorn.py

@@ -213,11 +213,17 @@ class Server(uvicorn.Server):
             # Same as set PYTHONASYNCIODEBUG=1, but with custom threshold.
             event_loop.set_debug(True)
             event_loop.slow_callback_duration = lag_threshold
-        threading.Thread(target=metrics_lib.process_monitor,
-                         args=('server',),
-                         daemon=True).start()
-        with self.capture_signals():
-            asyncio.run(self.serve(*args, **kwargs))
+        stop_monitor = threading.Event()
+        monitor = threading.Thread(target=metrics_lib.process_monitor,
+                                   args=('server', stop_monitor),
+                                   daemon=True)
+        monitor.start()
+        try:
+            with self.capture_signals():
+                asyncio.run(self.serve(*args, **kwargs))
+        finally:
+            stop_monitor.set()
+            monitor.join()
 
 
 def run(config: uvicorn.Config, max_db_connections: Optional[int] = None):

sky/skylet/constants.py

@@ -62,11 +62,14 @@ SKY_UV_INSTALL_CMD = (f'{SKY_UV_CMD} -V >/dev/null 2>&1 || '
                       'curl -LsSf https://astral.sh/uv/install.sh '
                       f'| UV_INSTALL_DIR={SKY_UV_INSTALL_DIR} sh')
 SKY_UV_PIP_CMD: str = (f'VIRTUAL_ENV={SKY_REMOTE_PYTHON_ENV} {SKY_UV_CMD} pip')
-# Deleting the SKY_REMOTE_PYTHON_ENV_NAME from the PATH to deactivate the
-# environment. `deactivate` command does not work when conda is used.
+SKY_UV_RUN_CMD: str = (f'VIRTUAL_ENV={SKY_REMOTE_PYTHON_ENV} {SKY_UV_CMD} run')
+# Deleting the SKY_REMOTE_PYTHON_ENV_NAME from the PATH and unsetting relevant
+# VIRTUAL_ENV envvars to deactivate the environment. `deactivate` command does
+# not work when conda is used.
 DEACTIVATE_SKY_REMOTE_PYTHON_ENV = (
     'export PATH='
-    f'$(echo $PATH | sed "s|$(echo ~)/{SKY_REMOTE_PYTHON_ENV_NAME}/bin:||")')
+    f'$(echo $PATH | sed "s|$(echo ~)/{SKY_REMOTE_PYTHON_ENV_NAME}/bin:||") && '
+    'unset VIRTUAL_ENV && unset VIRTUAL_ENV_PROMPT')
 
 # Prefix for SkyPilot environment variables
 SKYPILOT_ENV_VAR_PREFIX = 'SKYPILOT_'
@@ -98,7 +101,7 @@ SKYLET_VERSION = '18'
 SKYLET_LIB_VERSION = 4
 SKYLET_VERSION_FILE = '~/.sky/skylet_version'
 SKYLET_GRPC_PORT = 46590
-SKYLET_GRPC_TIMEOUT_SECONDS = 5
+SKYLET_GRPC_TIMEOUT_SECONDS = 10
 
 # Docker default options
 DEFAULT_DOCKER_CONTAINER_NAME = 'sky_container'
@@ -229,7 +232,7 @@ RAY_INSTALLATION_COMMANDS = (
     'export PATH=$PATH:$HOME/.local/bin; '
     # Writes ray path to file if it does not exist or the file is empty.
     f'[ -s {SKY_RAY_PATH_FILE} ] || '
-    f'{{ {ACTIVATE_SKY_REMOTE_PYTHON_ENV} && '
+    f'{{ {SKY_UV_RUN_CMD} '
     f'which ray > {SKY_RAY_PATH_FILE} || exit 1; }}; ')
 
 SKYPILOT_WHEEL_INSTALLATION_COMMANDS = (
@@ -374,7 +377,6 @@ OVERRIDEABLE_CONFIG_KEYS_IN_TASK: List[Tuple[str, ...]] = [
     ('ssh', 'pod_config'),
     ('kubernetes', 'custom_metadata'),
     ('kubernetes', 'pod_config'),
-    ('kubernetes', 'context_configs'),
     ('kubernetes', 'provision_timeout'),
     ('kubernetes', 'dws'),
     ('kubernetes', 'kueue'),
@@ -449,7 +451,7 @@ SKYPILOT_DEFAULT_WORKSPACE = 'default'
 # BEGIN constants used for service catalog.
 HOSTED_CATALOG_DIR_URL = 'https://raw.githubusercontent.com/skypilot-org/skypilot-catalog/master/catalogs'  # pylint: disable=line-too-long
 HOSTED_CATALOG_DIR_URL_S3_MIRROR = 'https://skypilot-catalog.s3.us-east-1.amazonaws.com/catalogs'  # pylint: disable=line-too-long
-CATALOG_SCHEMA_VERSION = 'v7'
+CATALOG_SCHEMA_VERSION = 'v8'
 CATALOG_DIR = '~/.sky/catalogs'
 ALL_CLOUDS = ('aws', 'azure', 'gcp', 'ibm', 'lambda', 'scp', 'oci',
               'kubernetes', 'runpod', 'vast', 'vsphere', 'cudo', 'fluidstack',
@@ -510,3 +512,6 @@ SKY_LOCKS_DIR = os.path.expanduser('~/.sky/locks')
 
 ENV_VAR_LOOP_LAG_THRESHOLD_MS = (SKYPILOT_ENV_VAR_PREFIX +
                                  'DEBUG_LOOP_LAG_THRESHOLD_MS')
+
+ARM64_ARCH = 'arm64'
+X86_64_ARCH = 'x86_64'

sky/skylet/log_lib.py

@@ -354,6 +354,17 @@ def run_bash_command_with_log(bash_command: str,
                             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):
+    return_code = run_bash_command_with_log(bash_command, log_path, env_vars,
+                                            stream_logs, with_ray)
+    return {'return_code': return_code, 'pid': os.getpid()}
+
+
 def _follow_job_logs(file,
                      job_id: int,
                      start_streaming: bool,

sky/skylet/log_lib.pyi

@@ -129,6 +129,15 @@ def run_bash_command_with_log(bash_command: str,
     ...
 
 
+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 = ...):
+    ...
+
+
 def tail_logs(job_id: int,
               log_dir: Optional[str],
               managed_job_id: Optional[int] = ...,

sky/task.py

@@ -20,6 +20,7 @@ from sky.provision import docker_utils
 from sky.serve import service_spec
 from sky.skylet import constants
 from sky.utils import common_utils
+from sky.utils import git
 from sky.utils import registry
 from sky.utils import schemas
 from sky.utils import ux_utils
@@ -1596,6 +1597,67 @@ class Task:
                 d[k] = v
         return d
 
+    def update_workdir(self, workdir: Optional[str], git_url: Optional[str],
+                       git_ref: Optional[str]) -> 'Task':
+        """Updates the task workdir.
+
+        Args:
+            workdir: The workdir to update.
+            git_url: The git url to update.
+            git_ref: The git ref to update.
+        """
+        if self.workdir is None or isinstance(self.workdir, str):
+            if workdir is not None:
+                self.workdir = workdir
+                return self
+            if git_url is not None:
+                self.workdir = {}
+                self.workdir['url'] = git_url
+                if git_ref is not None:
+                    self.workdir['ref'] = git_ref
+                return self
+            return self
+        if git_url is not None:
+            self.workdir['url'] = git_url
+        if git_ref is not None:
+            self.workdir['ref'] = git_ref
+        return self
+
+    def update_envs_and_secrets_from_workdir(self) -> 'Task':
+        """Updates the task envs and secrets from the workdir."""
+        if self.workdir is None:
+            return self
+        if not isinstance(self.workdir, dict):
+            return self
+        url = self.workdir['url']
+        ref = self.workdir.get('ref', '')
+        token = os.environ.get(git.GIT_TOKEN_ENV_VAR)
+        ssh_key_path = os.environ.get(git.GIT_SSH_KEY_PATH_ENV_VAR)
+        try:
+            git_repo = git.GitRepo(url, ref, token, ssh_key_path)
+            clone_info = git_repo.get_repo_clone_info()
+            if clone_info is None:
+                return self
+            self.envs[git.GIT_URL_ENV_VAR] = clone_info.url
+            if ref:
+                ref_type = git_repo.get_ref_type()
+                if ref_type == git.GitRefType.COMMIT:
+                    self.envs[git.GIT_COMMIT_HASH_ENV_VAR] = ref
+                elif ref_type == git.GitRefType.BRANCH:
+                    self.envs[git.GIT_BRANCH_ENV_VAR] = ref
+                elif ref_type == git.GitRefType.TAG:
+                    self.envs[git.GIT_TAG_ENV_VAR] = ref
+            if clone_info.token is None and clone_info.ssh_key is None:
+                return self
+            if clone_info.token is not None:
+                self.secrets[git.GIT_TOKEN_ENV_VAR] = clone_info.token
+            if clone_info.ssh_key is not None:
+                self.secrets[git.GIT_SSH_KEY_ENV_VAR] = clone_info.ssh_key
+        except exceptions.GitError as e:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(f'{str(e)}') from None
+        return self
+
     def to_yaml_config(self,
                        use_user_specified_yaml: bool = False) -> Dict[str, Any]:
         """Returns a yaml-style dict representation of the task.

sky/templates/kubernetes-ray.yml.j2

@@ -654,8 +654,125 @@ available_node_types:
                 # after v0.11.0 release.
                 touch /tmp/apt_ssh_setup_started
 
-                DEBIAN_FRONTEND=noninteractive $(prefix_cmd) apt-get update > /tmp/apt-update.log 2>&1 || \
-                echo "Warning: apt-get update failed. Continuing anyway..." >> /tmp/apt-update.log
+                # Helper: run apt-get update with retries
+                apt_update_with_retries() {
+                  # do not fail the whole shell; we handle return codes
+                  set +e
+                  local log=/tmp/apt-update.log
+                  local tries=3
+                  local delay=1
+                  local i
+                  for i in $(seq 1 $tries); do
+                    DEBIAN_FRONTEND=noninteractive $(prefix_cmd) apt-get update >> "$log" 2>&1 && { set -e; return 0; }
+                    echo "apt-get update attempt $i/$tries failed; retrying in ${delay}s" >> "$log"
+                    sleep $delay
+                    delay=$((delay * 2))
+                  done
+                  set -e
+                  return 1
+                }
+                apt_install_with_retries() {
+                  local packages="$@"
+                  [ -z "$packages" ] && return 0
+                  set +e
+                  local log=/tmp/apt-update.log
+                  local tries=3
+                  local delay=1
+                  local i
+                  for i in $(seq 1 $tries); do
+                    DEBIAN_FRONTEND=noninteractive $(prefix_cmd) apt-get install -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" $packages && { set -e; return 0; }
+                    echo "apt-get install failed for: $packages (attempt $i/$tries). Running -f install and retrying..." >> "$log"
+                    DEBIAN_FRONTEND=noninteractive $(prefix_cmd) apt-get -f install -y >> "$log" 2>&1 || true
+                    DEBIAN_FRONTEND=noninteractive $(prefix_cmd) apt-get clean >> "$log" 2>&1 || true
+                    sleep $delay
+                    delay=$((delay * 2))
+                  done
+                  set -e
+                  return 1
+                }
+                apt_update_install_with_retries() {
+                  apt_update_with_retries
+                  apt_install_with_retries "$@"
+                }
+                backup_dir=/etc/apt/sources.list.backup_skypilot
+                backup_source() {
+                  $(prefix_cmd) mkdir -p "$backup_dir"
+                  if [ -f /etc/apt/sources.list ] && [ ! -f "$backup_dir/sources.list" ]; then
+                    $(prefix_cmd) cp -a /etc/apt/sources.list "$backup_dir/sources.list" || true
+                  fi
+                }
+                restore_source() {
+                  if [ -f "$backup_dir/sources.list" ]; then
+                    $(prefix_cmd) cp -a "$backup_dir/sources.list" /etc/apt/sources.list || true
+                  fi
+                }
+                update_apt_sources() {
+                  local host=$1
+                  local apt_file=$2
+                  $(prefix_cmd) sed -i -E "s|https?://[a-zA-Z0-9.-]+\.ubuntu\.com/ubuntu|http://$host/ubuntu|g" $apt_file
+                }
+                # Helper: install packages across mirrors with retries
+                apt_install_with_mirrors() {
+                  local required=$1; shift
+                  local packages="$@"
+                  [ -z "$packages" ] && return 0
+                  set +e
+                  # Install packages with default sources first
+                  local log=/tmp/apt-update.log
+                  echo "$(date +%Y-%m-%d\ %H:%M:%S) Installing packages: $packages" >> "$log"
+                  restore_source
+                  apt_update_install_with_retries $packages >> "$log" 2>&1 && { set -e; return 0; }
+                  echo "Install failed with default sources: $packages" >> "$log"
+                  # Detect distro (ubuntu/debian)
+                  local APT_OS="unknown"
+                  if [ -f /etc/os-release ]; then
+                    . /etc/os-release
+                    case "$ID" in
+                      debian) APT_OS="debian" ;;
+                      ubuntu) APT_OS="ubuntu" ;;
+                      *)
+                        if [ -n "$ID_LIKE" ]; then
+                          case " $ID $ID_LIKE " in
+                            *ubuntu*) APT_OS="ubuntu" ;;
+                            *debian*) APT_OS="debian" ;;
+                          esac
+                        fi
+                        ;;
+                    esac
+                  fi
+                  # Build mirror candidates
+                  # deb.debian.org is a CDN endpoint, if one backend goes down,
+                  # the CDN automatically fails over to another mirror,
+                  # so we only retry for ubuntu here.
+                  if [ "$APT_OS" = "ubuntu" ]; then
+                    # Backup current sources once
+                    backup_source
+                    # Selected from https://launchpad.net/ubuntu/+archivemirrors
+                    # and results from apt-select
+                    local MIRROR_CANDIDATES="mirrors.wikimedia.org mirror.umd.edu"
+                    for host in $MIRROR_CANDIDATES; do
+                      echo "Trying APT mirror ($APT_OS): $host" >> "$log"
+                      if [ -f /etc/apt/sources.list ]; then
+                        update_apt_sources $host /etc/apt/sources.list
+                      else
+                        echo "Error: /etc/apt/sources.list not found" >> "$log"
+                        break
+                      fi
+                      apt_update_install_with_retries $packages >> "$log" 2>&1 && { set -e; return 0; }
+                      echo "Install failed with mirror ($APT_OS): $host" >> "$log"
+                      # Restore to default sources
+                      restore_source
+                    done
+                  fi
+                  set -e
+                  if [ "$required" = "1" ]; then
+                    echo "Error: required package install failed across all mirrors: $packages" >> "$log"
+                    return 1
+                  else
+                    echo "Optional package install failed across all mirrors: $packages; skipping." >> "$log"
+                    return 0
+                  fi
+                }
                 # Install both fuse2 and fuse3 for compatibility for all possible fuse adapters in advance,
                 # so that both fusemount and fusermount3 can be masked before enabling SSH access.
                 PACKAGES="rsync curl wget netcat gcc patch pciutils fuse fuse3 openssh-server";
@@ -682,7 +799,7 @@ available_node_types:
                 done;
                 if [ ! -z "$INSTALL_FIRST" ]; then
                   echo "Installing core packages: $INSTALL_FIRST";
-                  DEBIAN_FRONTEND=noninteractive $(prefix_cmd) apt-get install -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" $INSTALL_FIRST;
+                  apt_install_with_mirrors 1 $INSTALL_FIRST || { echo "Error: core package installation failed." >> /tmp/apt-update.log; exit 1; }
                 fi;
                 # SSH and other packages are not necessary, so we disable set -e
                 set +e

sky/utils/accelerator_registry.py

@@ -107,10 +107,12 @@ def canonicalize_accelerator_name(accelerator: str,
     if not names and cloud_str in ['Kubernetes', None]:
         with rich_utils.safe_status(
                 ux_utils.spinner_message('Listing accelerators on Kubernetes')):
+            # Only search for Kubernetes to reduce the lookup cost.
+            # For other clouds, the catalog has been searched in previous steps.
             searched = catalog.list_accelerators(
                 name_filter=accelerator,
                 case_sensitive=False,
-                clouds=cloud_str,
+                clouds='Kubernetes',
             )
         names = list(searched.keys())
         if accelerator in names:

sky/utils/command_runner.py

@@ -469,15 +469,19 @@ class CommandRunner:
         """Close the cached connection to the remote machine."""
         pass
 
-    def port_forward_command(self,
-                             port_forward: List[Tuple[int, int]],
-                             connect_timeout: int = 1) -> List[str]:
+    def port_forward_command(
+            self,
+            port_forward: List[Tuple[int, int]],
+            connect_timeout: int = 1,
+            ssh_mode: SshMode = SshMode.INTERACTIVE) -> List[str]:
         """Command for forwarding ports from localhost to the remote machine.
 
         Args:
             port_forward: A list of ports to forward from the localhost to the
                 remote host.
             connect_timeout: The timeout for the connection.
+            ssh_mode: The mode to use for ssh.
+                See SSHMode for more details.
         """
         raise NotImplementedError
 
@@ -592,6 +596,7 @@ class SSHCommandRunner(CommandRunner):
         ssh_proxy_command: Optional[str] = None,
         docker_user: Optional[str] = None,
         disable_control_master: Optional[bool] = False,
+        port_forward_execute_remote_command: Optional[bool] = False,
     ):
         """Initialize SSHCommandRunner.
 
@@ -618,6 +623,10 @@ class SSHCommandRunner(CommandRunner):
             disable_control_master: bool; specifies either or not the ssh
                 command will utilize ControlMaster. We currently disable
                 it for k8s instance.
+            port_forward_execute_remote_command: bool; specifies whether to
+                add -N to the port forwarding command. This is useful if you
+                want to run a command on the remote machine to make sure the
+                SSH tunnel is established.
         """
         super().__init__(node)
         ip, port = node
@@ -646,22 +655,28 @@ class SSHCommandRunner(CommandRunner):
             self.ssh_user = ssh_user
             self.port = port
             self._docker_ssh_proxy_command = None
+        self.port_forward_execute_remote_command = (
+            port_forward_execute_remote_command)
 
-    def port_forward_command(self,
-                             port_forward: List[Tuple[int, int]],
-                             connect_timeout: int = 1) -> List[str]:
+    def port_forward_command(
+            self,
+            port_forward: List[Tuple[int, int]],
+            connect_timeout: int = 1,
+            ssh_mode: SshMode = SshMode.INTERACTIVE) -> List[str]:
         """Command for forwarding ports from localhost to the remote machine.
 
         Args:
             port_forward: A list of ports to forward from the local port to the
                 remote port.
             connect_timeout: The timeout for the ssh connection.
+            ssh_mode: The mode to use for ssh.
+                See SSHMode for more details.
 
         Returns:
             The command for forwarding ports from localhost to the remote
             machine.
         """
-        return self.ssh_base_command(ssh_mode=SshMode.INTERACTIVE,
+        return self.ssh_base_command(ssh_mode=ssh_mode,
                                      port_forward=port_forward,
                                      connect_timeout=connect_timeout)
 
@@ -680,7 +695,11 @@ class SSHCommandRunner(CommandRunner):
             for local, remote in port_forward:
                 logger.debug(
                     f'Forwarding local port {local} to remote port {remote}.')
-                ssh += ['-NL', f'{local}:localhost:{remote}']
+                if self.port_forward_execute_remote_command:
+                    ssh += ['-L']
+                else:
+                    ssh += ['-NL']
+                ssh += [f'{local}:localhost:{remote}']
         if self._docker_ssh_proxy_command is not None:
             docker_ssh_proxy_command = self._docker_ssh_proxy_command(ssh)
         else:
@@ -894,9 +913,11 @@ class KubernetesCommandRunner(CommandRunner):
         else:
             return f'pod/{self.pod_name}'
 
-    def port_forward_command(self,
-                             port_forward: List[Tuple[int, int]],
-                             connect_timeout: int = 1) -> List[str]:
+    def port_forward_command(
+            self,
+            port_forward: List[Tuple[int, int]],
+            connect_timeout: int = 1,
+            ssh_mode: SshMode = SshMode.INTERACTIVE) -> List[str]:
         """Command for forwarding ports from localhost to the remote machine.
 
         Args:
@@ -904,7 +925,10 @@ class KubernetesCommandRunner(CommandRunner):
                 remote port. Currently, only one port is supported, i.e. the
                 list should have only one element.
             connect_timeout: The timeout for the ssh connection.
+            ssh_mode: The mode to use for ssh.
+                See SSHMode for more details.
         """
+        del ssh_mode  # unused
         assert port_forward and len(port_forward) == 1, (
             'Only one port is supported for Kubernetes port-forward.')
         kubectl_args = [

sky/utils/command_runner.pyi

@@ -106,6 +106,13 @@ class CommandRunner:
               max_retry: int = ...) -> None:
         ...
 
+    def port_forward_command(
+            self,
+            port_forward: List[Tuple[int, int]],
+            connect_timeout: int = 1,
+            ssh_mode: SshMode = SshMode.INTERACTIVE) -> List[str]:
+        ...
+
     @classmethod
     def make_runner_list(cls: typing.Type[CommandRunner],
                          node_list: Iterable[Tuple[Any, ...]],
@@ -127,6 +134,7 @@ class SSHCommandRunner(CommandRunner):
     ssh_control_name: Optional[str]
     docker_user: str
     disable_control_master: Optional[bool]
+    port_forward_execute_remote_command: Optional[bool]
 
     def __init__(
         self,
@@ -200,6 +208,13 @@ class SSHCommandRunner(CommandRunner):
               max_retry: int = ...) -> None:
         ...
 
+    def port_forward_command(
+            self,
+            port_forward: List[Tuple[int, int]],
+            connect_timeout: int = 1,
+            ssh_mode: SshMode = SshMode.INTERACTIVE) -> List[str]:
+        ...
+
 
 class KubernetesCommandRunner(CommandRunner):
 
@@ -272,6 +287,13 @@ class KubernetesCommandRunner(CommandRunner):
               max_retry: int = ...) -> None:
         ...
 
+    def port_forward_command(
+            self,
+            port_forward: List[Tuple[int, int]],
+            connect_timeout: int = 1,
+            ssh_mode: SshMode = SshMode.INTERACTIVE) -> List[str]:
+        ...
+
 
 class LocalProcessCommandRunner(CommandRunner):
 

sky/utils/context_utils.py

@@ -10,6 +10,8 @@ import sys
 import typing
 from typing import Any, Callable, IO, Optional, Tuple, TypeVar
 
+from typing_extensions import ParamSpec
+
 from sky import sky_logging
 from sky.utils import context
 from sky.utils import subprocess_utils
@@ -173,9 +175,14 @@ def cancellation_guard(func: F) -> F:
     return typing.cast(F, wrapper)
 
 
+P = ParamSpec('P')
+T = TypeVar('T')
+
+
 # TODO(aylei): replace this with asyncio.to_thread once we drop support for
 # python 3.8
-def to_thread(func, /, *args, **kwargs):
+def to_thread(func: Callable[P, T], /, *args: P.args,
+              **kwargs: P.kwargs) -> 'asyncio.Future[T]':
     """Asynchronously run function *func* in a separate thread.
 
     This is same as asyncio.to_thread added in python 3.9
@@ -183,5 +190,11 @@ def to_thread(func, /, *args, **kwargs):
     loop = asyncio.get_running_loop()
     # This is critical to pass the current coroutine context to the new thread
     pyctx = contextvars.copy_context()
-    func_call = functools.partial(pyctx.run, func, *args, **kwargs)
+    func_call: Callable[..., T] = functools.partial(
+        # partial deletes arguments type and thus can't figure out the return
+        # type of pyctx.run
+        pyctx.run,  # type: ignore
+        func,
+        *args,
+        **kwargs)
     return loop.run_in_executor(None, func_call)

sky/utils/db/migration_utils.py

@@ -17,7 +17,7 @@ logger = sky_logging.init_logger(__name__)
 DB_INIT_LOCK_TIMEOUT_SECONDS = 10
 
 GLOBAL_USER_STATE_DB_NAME = 'state_db'
-GLOBAL_USER_STATE_VERSION = '007'
+GLOBAL_USER_STATE_VERSION = '008'
 GLOBAL_USER_STATE_LOCK_PATH = '~/.sky/locks/.state_db.lock'
 
 SPOT_JOBS_DB_NAME = 'spot_jobs_db'