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

sky/provision/aws/instance.py

@@ -311,9 +311,10 @@ def _get_head_instance_id(instances: List) -> Optional[str]:
     return head_instance_id
 
 
-def run_instances(region: str, cluster_name_on_cloud: str,
+def run_instances(region: str, cluster_name: str, cluster_name_on_cloud: str,
                   config: common.ProvisionConfig) -> common.ProvisionRecord:
     """See sky/provision/__init__.py"""
+    del cluster_name  # unused
     ec2 = _default_ec2_resource(region)
     # NOTE: We set max_attempts=0 for fast failing when the resource is not
     # available (although the doc says it will only retry for network
@@ -629,9 +630,10 @@ def query_instances(
     cluster_name_on_cloud: str,
     provider_config: Optional[Dict[str, Any]] = None,
     non_terminated_only: bool = True,
+    retry_if_missing: bool = False,
 ) -> Dict[str, Tuple[Optional['status_lib.ClusterStatus'], Optional[str]]]:
     """See sky/provision/__init__.py"""
-    del cluster_name  # unused
+    del cluster_name, retry_if_missing  # unused
     assert provider_config is not None, (cluster_name_on_cloud, provider_config)
     region = provider_config['region']
     ec2 = _default_ec2_resource(region)
@@ -743,6 +745,7 @@ def terminate_instances(
 
         # Make this multithreaded: modify all instances' SGs in parallel.
         def modify_instance_sg(instance):
+            assert default_sg is not None  # Type narrowing for mypy
             instance.modify_attribute(Groups=[default_sg.id])
             logger.debug(f'Instance {instance.id} modified to use default SG:'
                          f'{default_sg.id} for quick deletion.')

sky/provision/azure/instance.py

@@ -214,7 +214,7 @@ def _create_network_interface(
             location=provider_config['location'],
             public_ip_allocation_method='Static',
             public_ip_address_version='IPv4',
-            sku=network.PublicIPAddressSku(name='Basic', tier='Regional'))
+            sku=network.PublicIPAddressSku(name='Standard', tier='Regional'))
         ip_poller = network_client.public_ip_addresses.begin_create_or_update(
             resource_group_name=provider_config['resource_group'],
             public_ip_address_name=f'{vm_name}-ip',
@@ -362,9 +362,10 @@ def _create_instances(compute_client: 'azure_compute.ComputeManagementClient',
     return instances
 
 
-def run_instances(region: str, cluster_name_on_cloud: str,
+def run_instances(region: str, cluster_name: str, cluster_name_on_cloud: str,
                   config: common.ProvisionConfig) -> common.ProvisionRecord:
     """See sky/provision/__init__.py"""
+    del cluster_name  # unused
     # TODO(zhwu): This function is too long. We should refactor it.
     provider_config = config.provider_config
     resource_group = provider_config['resource_group']
@@ -956,9 +957,10 @@ def query_instances(
     cluster_name_on_cloud: str,
     provider_config: Optional[Dict[str, Any]] = None,
     non_terminated_only: bool = True,
+    retry_if_missing: bool = False,
 ) -> Dict[str, Tuple[Optional['status_lib.ClusterStatus'], Optional[str]]]:
     """See sky/provision/__init__.py"""
-    del cluster_name  # unused
+    del cluster_name, retry_if_missing  # unused
     assert provider_config is not None, cluster_name_on_cloud
 
     subscription_id = provider_config['subscription_id']

sky/provision/common.py

@@ -97,6 +97,8 @@ class InstanceInfo:
     external_ip: Optional[str]
     tags: Dict[str, str]
     ssh_port: int = 22
+    # The internal service address of the instance on Kubernetes.
+    internal_svc: Optional[str] = None
 
     def get_feasible_ip(self) -> str:
         """Get the most feasible IPs of the instance. This function returns

sky/provision/cudo/instance.py

@@ -40,10 +40,10 @@ def _get_head_instance_id(instances: Dict[str, Any]) -> Optional[str]:
     return head_instance_id
 
 
-def run_instances(region: str, cluster_name_on_cloud: str,
+def run_instances(region: str, cluster_name: str, cluster_name_on_cloud: str,
                   config: common.ProvisionConfig) -> common.ProvisionRecord:
     """Runs instances for the given cluster."""
-
+    del cluster_name  # unused
     pending_status = ['pend', 'init', 'prol', 'boot']
 
     while True:
@@ -195,9 +195,10 @@ def query_instances(
     cluster_name_on_cloud: str,
     provider_config: Optional[Dict[str, Any]] = None,
     non_terminated_only: bool = True,
+    retry_if_missing: bool = False,
 ) -> Dict[str, Tuple[Optional['status_lib.ClusterStatus'], Optional[str]]]:
     """See sky/provision/__init__.py"""
-    del cluster_name  # unused
+    del cluster_name, retry_if_missing  # unused
     assert provider_config is not None, (cluster_name_on_cloud, provider_config)
     instances = _filter_instances(cluster_name_on_cloud, None)
 

sky/provision/do/instance.py

@@ -26,10 +26,10 @@ def _get_head_instance(
     return None
 
 
-def run_instances(region: str, cluster_name_on_cloud: str,
+def run_instances(region: str, cluster_name: str, cluster_name_on_cloud: str,
                   config: common.ProvisionConfig) -> common.ProvisionRecord:
     """Runs instances for the given cluster."""
-
+    del cluster_name  # unused
     pending_status = ['new']
     newly_started_instances = utils.filter_instances(cluster_name_on_cloud,
                                                      pending_status + ['off'])
@@ -246,9 +246,10 @@ def query_instances(
     cluster_name_on_cloud: str,
     provider_config: Optional[Dict[str, Any]] = None,
     non_terminated_only: bool = True,
+    retry_if_missing: bool = False,
 ) -> Dict[str, Tuple[Optional['status_lib.ClusterStatus'], Optional[str]]]:
     """See sky/provision/__init__.py"""
-    del cluster_name  # unused
+    del cluster_name, retry_if_missing  # unused
     # terminated instances are not retrieved by the
     # API making `non_terminated_only` argument moot.
     del non_terminated_only

sky/provision/docker_utils.py

@@ -3,7 +3,7 @@
 import dataclasses
 import shlex
 import time
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from sky import sky_logging
 from sky.skylet import constants
@@ -15,23 +15,52 @@ logger = sky_logging.init_logger(__name__)
 # Configure environment variables. A docker image can have environment variables
 # set in the Dockerfile with `ENV``. We need to export these variables to the
 # shell environment, so that our ssh session can access them.
+# Filter out RAY_RUNTIME_ENV_HOOK to prevent Ray version conflicts.
+# Docker images with Ray 2.48.0+ set this for UV package manager support,
+# but it causes FAILED_DRIVER errors with SkyPilot's Ray 2.9.3.
+# See: https://github.com/skypilot-org/skypilot/pull/7181
 SETUP_ENV_VARS_CMD = (
     'prefix_cmd() '
     '{ if [ $(id -u) -ne 0 ]; then echo "sudo"; else echo ""; fi; } && '
-    'export -p > ~/container_env_var.sh && '
+    'export -p | grep -v RAY_RUNTIME_ENV_HOOK > ~/container_env_var.sh && '
     '$(prefix_cmd) '
     'mv ~/container_env_var.sh /etc/profile.d/container_env_var.sh;')
 
 # Docker daemon may not be ready when the machine is firstly started. The error
 # message starts with the following string. We should wait for a while and retry
 # the command.
-DOCKER_PERMISSION_DENIED_STR = ('permission denied while trying to connect to '
-                                'the Docker daemon socket')
+DOCKER_PERMISSION_DENIED_STR = ('permission denied while trying to connect to ')
 
 DOCKER_SOCKET_NOT_READY_STR = ('Is the docker daemon running?')
+DOCKER_SOCKET_NOT_READY_STR_2 = (
+    'check if the path is correct and if the daemon is running')
 
 _DOCKER_SOCKET_WAIT_TIMEOUT_SECONDS = 30
 
+# Install AWS CLI v2 (not v1 from pip) as it's required for ECR authentication
+# AWS CLI v2 is installed as a standalone binary, not a Python package. See:
+# https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html
+INSTALL_AWS_CLI_CMD = (
+    'which aws || ((command -v unzip >/dev/null 2>&1 || '
+    '(sudo apt-get update && sudo apt-get install -y unzip)) && '
+    'curl -fsSL "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" '
+    '-o "/tmp/awscliv2.zip" && '
+    'unzip -q /tmp/awscliv2.zip -d /tmp && sudo /tmp/aws/install '
+    '&& rm -rf /tmp/awscliv2.zip /tmp/aws)')
+
+
+def _extract_region_from_ecr_server(server: str) -> str:
+    """Extract AWS region from ECR server URL.
+
+    ECR server format: <account-id>.dkr.ecr.<region>.amazonaws.com
+    Returns the region part from the URL.
+    """
+    # Split: ['<account-id>', 'dkr', 'ecr', '<region>', 'amazonaws', 'com']
+    parts = server.split('.')
+    if len(parts) >= 6 and parts[1] == 'dkr' and parts[2] == 'ecr':
+        return parts[3]
+    raise ValueError(f'Invalid ECR server format: {server}')
+
 
 @dataclasses.dataclass
 class DockerLoginConfig:
@@ -157,19 +186,23 @@ class DockerInitializer:
         self.docker_config = docker_config
         self.container_name = docker_config['container_name']
         self.runner = runner
-        self.home_dir = None
+        self.home_dir: Optional[str] = None
         self.initialized = False
         # podman is not fully tested yet.
         use_podman = docker_config.get('use_podman', False)
         self.docker_cmd = 'podman' if use_podman else 'docker'
         self.log_path = log_path
 
-    def _run(self,
-             cmd,
-             run_env='host',
-             wait_for_docker_daemon: bool = False,
-             separate_stderr: bool = False,
-             log_err_when_fail: bool = True) -> str:
+    def _run(
+        self,
+        cmd,
+        run_env='host',
+        wait_for_docker_daemon: bool = False,
+        separate_stderr: bool = False,
+        log_err_when_fail: bool = True,
+        flock_name: Optional[str] = None,
+        flock_args: Optional[str] = None,
+    ) -> str:
 
         if run_env == 'docker':
             cmd = self._docker_expand_user(cmd, any_char=True)
@@ -178,8 +211,13 @@ class DockerInitializer:
             # an error: `the input device is not a TTY`, and it works without
             # `-it` flag.
             # TODO(zhwu): ray use the `-it` flag, we need to check why.
-            cmd = (f'{self.docker_cmd} exec {self.container_name} /bin/bash -c'
-                   f' {shlex.quote(cmd)} ')
+            cmd = (f'{self.docker_cmd} exec -u 0 {self.container_name}'
+                   f' /bin/bash -c {shlex.quote(cmd)} ')
+
+        if flock_name is not None:
+            flock_args = flock_args or ''
+            cmd = (f'flock {flock_args} /tmp/{flock_name} '
+                   f'-c {shlex.quote(cmd)}')
 
         logger.debug(f'+ {cmd}')
         start = time.time()
@@ -191,7 +229,8 @@ class DockerInitializer:
                 separate_stderr=separate_stderr,
                 log_path=self.log_path)
             if (DOCKER_PERMISSION_DENIED_STR in stdout + stderr or
-                    DOCKER_SOCKET_NOT_READY_STR in stdout + stderr):
+                    DOCKER_SOCKET_NOT_READY_STR in stdout + stderr or
+                    DOCKER_SOCKET_NOT_READY_STR_2 in stdout + stderr):
                 if wait_for_docker_daemon:
                     if time.time(
                     ) - start > _DOCKER_SOCKET_WAIT_TIMEOUT_SECONDS:
@@ -231,14 +270,17 @@ class DockerInitializer:
         if self._check_container_exited():
             self.initialized = True
             self._run(f'{self.docker_cmd} start {self.container_name}')
-            self._run('sudo service ssh start', run_env='docker')
+            self._run('sudo service ssh start',
+                      run_env='docker',
+                      flock_name=f'{self.container_name}.sky.lifecycle.lock',
+                      flock_args='-s -w 1')
             return self._run('whoami', run_env='docker')
 
         # SkyPilot: Docker login if user specified a private docker registry.
         if 'docker_login_config' in self.docker_config:
-            # TODO(tian): Maybe support a command to get the login password?
             docker_login_config = DockerLoginConfig(
                 **self.docker_config['docker_login_config'])
+
             if docker_login_config.password:
                 # Password is allowed to be empty, in that case, we will not run
                 # the login command, and assume that the image pulling is
@@ -249,6 +291,25 @@ class DockerInitializer:
                     f'--password {shlex.quote(docker_login_config.password)} '
                     f'{shlex.quote(docker_login_config.server)}',
                     wait_for_docker_daemon=True)
+            elif (docker_login_config.server.endswith('.amazonaws.com') and
+                  '.dkr.ecr.' in docker_login_config.server):
+                # AWS ECR: Use aws ecr get-login-password for authentication
+                # ECR format: <account-id>.dkr.ecr.<region>.amazonaws.com
+                # This command uses the IAM credentials from the EC2 instance
+                # Ref: https://docs.aws.amazon.com/AmazonECR/latest/userguide/registry_auth.html # pylint: disable=line-too-long
+                region = _extract_region_from_ecr_server(
+                    docker_login_config.server)
+
+                # AWS CLI is not pre-installed on AWS instances, unlike gcloud
+                # on GCP instances, so we need to install it first
+                self._run(INSTALL_AWS_CLI_CMD, wait_for_docker_daemon=False)
+
+                self._run(
+                    f'aws ecr get-login-password --region {region} | '
+                    f'{self.docker_cmd} login --username AWS '
+                    f'--password-stdin '
+                    f'{shlex.quote(docker_login_config.server)}',
+                    wait_for_docker_daemon=True)
             elif docker_login_config.server.endswith('-docker.pkg.dev'):
                 # Docker image server is on GCR, we need to do additional setup
                 # to pull the image.
@@ -311,7 +372,9 @@ class DockerInitializer:
                     self._auto_configure_shm(user_docker_run_options)),
                 self.docker_cmd,
             )
-            self._run(f'{remove_container_cmd}; {start_command}')
+            self._run(f'{remove_container_cmd} && {start_command}',
+                      flock_name=f'{self.container_name}.sky.lifecycle.lock',
+                      flock_args='-x -w 10')
 
         # SkyPilot: Setup Commands.
         # TODO(zhwu): the following setups should be aligned with the kubernetes
@@ -329,14 +392,18 @@ class DockerInitializer:
             'echo "export DEBIAN_FRONTEND=noninteractive" >> ~/.bashrc;',
             run_env='docker')
         # Install dependencies.
-        self._run(
-            'sudo apt-get update; '
+        cmd = (
+            'bash -lc \''
+            'exec 200>/var/tmp/sky_apt.lock; '
+            'flock -x -w 120 200 || exit 1; '
+            'export DEBIAN_FRONTEND=noninteractive; '
+            'apt-get -yq update && '
             # Our mount script will install gcsfuse without fuse package.
             # We need to install fuse package first to enable storage mount.
             # The dpkg option is to suppress the prompt for fuse installation.
-            'sudo apt-get -o DPkg::Options::="--force-confnew" install -y '
-            'rsync curl wget patch openssh-server python3-pip fuse;',
-            run_env='docker')
+            'apt-get -o DPkg::Options::=--force-confnew install -y '
+            'rsync curl wget patch openssh-server python3-pip fuse\'')
+        self._run(cmd, run_env='docker')
 
         # Copy local authorized_keys to docker container.
         # Stop and disable jupyter service. This is to avoid port conflict on
@@ -367,7 +434,7 @@ class DockerInitializer:
         # pylint: disable=anomalous-backslash-in-string
         self._run(
             'sudo sed -i "/^Port .*/d" /etc/ssh/sshd_config;'
-            f'sudo echo "Port {port}" >> /etc/ssh/sshd_config;'
+            f'echo "Port {port}" | sudo tee -a /etc/ssh/sshd_config > /dev/null;'
             'mkdir -p ~/.ssh;'
             'cat /tmp/host_ssh_authorized_keys >> ~/.ssh/authorized_keys;'
             'sudo service ssh start;'
@@ -412,9 +479,13 @@ class DockerInitializer:
         user_pos = string.find('~')
         if user_pos > -1:
             if self.home_dir is None:
-                cmd = (f'{self.docker_cmd} exec {self.container_name} '
-                       'printenv HOME')
-                self.home_dir = self._run(cmd, separate_stderr=True)
+                cmd = (f'{self.docker_cmd} exec {self.container_name}'
+                       ' printenv HOME')
+                self.home_dir = self._run(
+                    cmd,
+                    separate_stderr=True,
+                    flock_name=f'{self.container_name}.sky.lifecycle.lock',
+                    flock_args='-s -w 1')
                 # Check for unexpected newline in home directory, which can be
                 # a common issue when the output is mixed with stderr.
                 assert '\n' not in self.home_dir, (

sky/provision/fluidstack/instance.py

@@ -3,11 +3,11 @@ import os
 import time
 from typing import Any, Dict, List, Optional, Tuple
 
-from sky import authentication as auth
 from sky import exceptions
 from sky import sky_logging
 from sky.provision import common
 from sky.provision.fluidstack import fluidstack_utils as utils
+from sky.utils import auth_utils
 from sky.utils import command_runner
 from sky.utils import common_utils
 from sky.utils import status_lib
@@ -27,7 +27,7 @@ logger = sky_logging.init_logger(__name__)
 def get_internal_ip(node_info: Dict[str, Any]) -> None:
     node_info['internal_ip'] = node_info['ip_address']
 
-    private_key_path, _ = auth.get_or_generate_keys()
+    private_key_path, _ = auth_utils.get_or_generate_keys()
     runner = command_runner.SSHCommandRunner(
         (node_info['ip_address'], 22),
         ssh_user='ubuntu',
@@ -78,10 +78,10 @@ def _get_head_instance_id(instances: Dict[str, Any]) -> Optional[str]:
     return head_instance_id
 
 
-def run_instances(region: str, cluster_name_on_cloud: str,
+def run_instances(region: str, cluster_name: str, cluster_name_on_cloud: str,
                   config: common.ProvisionConfig) -> common.ProvisionRecord:
     """Runs instances for the given cluster."""
-
+    del cluster_name  # unused
     pending_status = ['pending', 'provisioning']
     while True:
         instances = _filter_instances(cluster_name_on_cloud, pending_status)
@@ -291,9 +291,10 @@ def query_instances(
     cluster_name_on_cloud: str,
     provider_config: Optional[Dict[str, Any]] = None,
     non_terminated_only: bool = True,
+    retry_if_missing: bool = False,
 ) -> Dict[str, Tuple[Optional['status_lib.ClusterStatus'], Optional[str]]]:
     """See sky/provision/__init__.py"""
-    del cluster_name  # unused
+    del cluster_name, retry_if_missing  # unused
     assert provider_config is not None, (cluster_name_on_cloud, provider_config)
     instances = _filter_instances(cluster_name_on_cloud, None)
     instances = _filter_instances(cluster_name_on_cloud, None)

sky/provision/gcp/config.py

@@ -5,6 +5,8 @@ import time
 import typing
 from typing import Any, Dict, List, Set, Tuple
 
+from typing_extensions import TypedDict
+
 from sky.adaptors import gcp
 from sky.clouds.utils import gcp_utils
 from sky.provision import common
@@ -415,6 +417,9 @@ def _configure_iam_role(config: common.ProvisionConfig, crm, iam) -> dict:
     return iam_role
 
 
+AllowedList = TypedDict('AllowedList', {'IPProtocol': str, 'ports': List[str]})
+
+
 def _check_firewall_rules(cluster_name: str, vpc_name: str, project_id: str,
                           compute):
     """Check if the firewall rules in the VPC are sufficient."""
@@ -466,7 +471,7 @@ def _check_firewall_rules(cluster_name: str, vpc_name: str, project_id: str,
             }
         """
         source2rules: Dict[Tuple[str, str], Dict[str, Set[int]]] = {}
-        source2allowed_list: Dict[Tuple[str, str], List[Dict[str, str]]] = {}
+        source2allowed_list: Dict[Tuple[str, str], List[AllowedList]] = {}
         for rule in rules:
             # Rules applied to specific VM (targetTags) may not work for the
             # current VM, so should be skipped.

sky/provision/gcp/instance.py

@@ -62,9 +62,10 @@ def query_instances(
     cluster_name_on_cloud: str,
     provider_config: Optional[Dict[str, Any]] = None,
     non_terminated_only: bool = True,
+    retry_if_missing: bool = False,
 ) -> Dict[str, Tuple[Optional['status_lib.ClusterStatus'], Optional[str]]]:
     """See sky/provision/__init__.py"""
-    del cluster_name  # unused
+    del cluster_name, retry_if_missing  # unused
     assert provider_config is not None, (cluster_name_on_cloud, provider_config)
     zone = provider_config['availability_zone']
     project_id = provider_config['project_id']
@@ -360,9 +361,10 @@ def _run_instances(region: str, cluster_name_on_cloud: str,
                                   created_instance_ids=created_instance_ids)
 
 
-def run_instances(region: str, cluster_name_on_cloud: str,
+def run_instances(region: str, cluster_name: str, cluster_name_on_cloud: str,
                   config: common.ProvisionConfig) -> common.ProvisionRecord:
     """See sky/provision/__init__.py"""
+    del cluster_name  # unused
     try:
         return _run_instances(region, cluster_name_on_cloud, config)
     except gcp.http_error_exception() as e:

sky/provision/hyperbolic/instance.py

@@ -64,8 +64,9 @@ def _get_head_instance_id(instances: Dict[str, Any]) -> Optional[str]:
     return next(iter(instances.keys()))
 
 
-def run_instances(region: str, cluster_name_on_cloud: str,
+def run_instances(region: str, cluster_name: str, cluster_name_on_cloud: str,
                   config: common.ProvisionConfig) -> common.ProvisionRecord:
+    del cluster_name  # unused
     logger.info(f'Starting run_instances with region={region}, '
                 f'cluster={cluster_name_on_cloud}')
     logger.debug(f'Config: {config}')
@@ -308,9 +309,10 @@ def query_instances(
     cluster_name_on_cloud: str,
     provider_config: Optional[dict] = None,
     non_terminated_only: bool = True,
+    retry_if_missing: bool = False,
 ) -> Dict[str, Tuple[Optional['status_lib.ClusterStatus'], Optional[str]]]:
     """Returns the status of the specified instances for Hyperbolic."""
-    del cluster_name, provider_config  # unused
+    del cluster_name, provider_config, retry_if_missing  # unused
     # Fetch all instances for this cluster
     instances = utils.list_instances(
         metadata={'skypilot': {

sky/provision/instance_setup.py

@@ -10,6 +10,7 @@ from typing import Any, Callable, Dict, List, Optional, Tuple
 from sky import exceptions
 from sky import logs
 from sky import provision
+from sky import resources as resources_lib
 from sky import sky_logging
 from sky.provision import common
 from sky.provision import docker_utils
@@ -38,11 +39,13 @@ _RAY_PRLIMIT = (
     'which prlimit && for id in $(pgrep -f raylet/raylet); '
     'do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;')
 
-_DUMP_RAY_PORTS = (
-    f'{constants.SKY_PYTHON_CMD} -c \'import json, os; '
-    f'json.dump({constants.SKY_REMOTE_RAY_PORT_DICT_STR}, '
-    f'open(os.path.expanduser("{constants.SKY_REMOTE_RAY_PORT_FILE}"), "w", '
-    'encoding="utf-8"))\';')
+DUMP_RAY_PORTS = (f'{constants.SKY_PYTHON_CMD} -c \'import json, os; '
+                  f'runtime_dir = os.path.expanduser(os.environ.get('
+                  f'"{constants.SKY_RUNTIME_DIR_ENV_VAR_KEY}", "~")); '
+                  f'json.dump({constants.SKY_REMOTE_RAY_PORT_DICT_STR}, '
+                  f'open(os.path.join(runtime_dir, '
+                  f'"{constants.SKY_REMOTE_RAY_PORT_FILE}"), "w", '
+                  'encoding="utf-8"))\';')
 
 _RAY_PORT_COMMAND = (
     f'RAY_PORT=$({constants.SKY_PYTHON_CMD} -c '
@@ -84,7 +87,7 @@ def _set_usage_run_id_cmd() -> str:
     latest one when the function is called.
     """
     return (
-        f'cat {usage_constants.USAGE_RUN_ID_FILE} || '
+        f'cat {usage_constants.USAGE_RUN_ID_FILE} 2> /dev/null || '
         # The run id is retrieved locally for the current run, so that the
         # remote cluster will be set with the same run id as the initial
         # launch operation.
@@ -92,12 +95,6 @@ def _set_usage_run_id_cmd() -> str:
         f'{usage_constants.USAGE_RUN_ID_FILE}')
 
 
-def _set_skypilot_env_var_cmd() -> str:
-    """Sets the skypilot environment variables on the remote machine."""
-    env_vars = env_options.Options.all_options()
-    return '; '.join([f'export {k}={v}' for k, v in env_vars.items()])
-
-
 def _auto_retry(should_retry: Callable[[Exception], bool] = lambda _: True):
     """Decorator that retries the function if it fails.
 
@@ -136,6 +133,20 @@ def _hint_worker_log_path(cluster_name: str, cluster_info: common.ClusterInfo,
         logger.info(f'Logs of worker nodes can be found at: {worker_log_path}')
 
 
+class SSHThreadPoolExecutor(futures.ThreadPoolExecutor):
+    """ThreadPoolExecutor that kills children processes on exit."""
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # ssh command runner eventually calls
+        # log_lib.run_with_log, which will spawn
+        # subprocesses. If we are exiting the context
+        # we need to kill the children processes
+        # to avoid leakage.
+        subprocess_utils.kill_children_processes()
+        self.shutdown()
+        return False
+
+
 def _parallel_ssh_with_cache(func,
                              cluster_name: str,
                              stage_name: str,
@@ -148,7 +159,7 @@ def _parallel_ssh_with_cache(func,
         # as 32 is too large for some machines.
         max_workers = subprocess_utils.get_parallel_threads(
             cluster_info.provider_name)
-    with futures.ThreadPoolExecutor(max_workers=max_workers) as pool:
+    with SSHThreadPoolExecutor(max_workers=max_workers) as pool:
         results = []
         runners = provision.get_command_runners(cluster_info.provider_name,
                                                 cluster_info, **ssh_credentials)
@@ -317,7 +328,7 @@ def ray_head_start_command(custom_resource: Optional[str],
         # the warning when the worker count is >12x CPUs.
         'RAY_worker_maximum_startup_concurrency=$(( 3 * $(nproc --all) )) '
         f'{constants.SKY_RAY_CMD} start --head {ray_options} || exit 1;' +
-        _RAY_PRLIMIT + _DUMP_RAY_PORTS + RAY_HEAD_WAIT_INITIALIZED_COMMAND)
+        _RAY_PRLIMIT + DUMP_RAY_PORTS + RAY_HEAD_WAIT_INITIALIZED_COMMAND)
     return cmd
 
 
@@ -425,8 +436,16 @@ def start_ray_on_worker_nodes(cluster_name: str, no_restart: bool,
         # use the external IP of the head node.
         use_external_ip = cluster_info.custom_ray_options.pop(
             'use_external_ip', False)
-    head_ip = (head_instance.internal_ip
-               if not use_external_ip else head_instance.external_ip)
+
+    if use_external_ip:
+        head_ip = head_instance.external_ip
+    else:
+        # For Kubernetes, use the internal service address of the head node.
+        # Keep this consistent with the logic in kubernetes-ray.yml.j2
+        if head_instance.internal_svc:
+            head_ip = head_instance.internal_svc
+        else:
+            head_ip = head_instance.internal_ip
 
     ray_cmd = ray_worker_start_command(custom_resource,
                                        cluster_info.custom_ray_options,
@@ -468,11 +487,38 @@ def start_ray_on_worker_nodes(cluster_name: str, no_restart: bool,
 @common.log_function_start_end
 @_auto_retry()
 @timeline.event
-def start_skylet_on_head_node(cluster_name: str,
-                              cluster_info: common.ClusterInfo,
-                              ssh_credentials: Dict[str, Any]) -> None:
+def start_skylet_on_head_node(
+        cluster_name: resources_utils.ClusterName,
+        cluster_info: common.ClusterInfo, ssh_credentials: Dict[str, Any],
+        launched_resources: resources_lib.Resources) -> None:
     """Start skylet on the head node."""
-    del cluster_name
+    # Avoid circular import.
+    # pylint: disable=import-outside-toplevel
+    from sky.utils import controller_utils
+
+    def _set_skypilot_env_var_cmd() -> str:
+        """Sets the skypilot environment variables on the remote machine."""
+        env_vars = {
+            k: str(v) for (k, v) in env_options.Options.all_options().items()
+        }
+        is_controller = controller_utils.Controllers.from_name(
+            cluster_name.display_name) is not None
+        is_kubernetes = cluster_info.provider_name == 'kubernetes'
+        if is_controller and is_kubernetes:
+            # For jobs/serve controller, we pass in the CPU and memory limits
+            # when starting the skylet to handle cases where these env vars
+            # are not set on the cluster's pod spec. The skylet will read
+            # these env vars when starting (ManagedJobEvent.start()) and write
+            # it to disk.
+            resources = launched_resources.assert_launchable()
+            vcpus, mem = resources.cloud.get_vcpus_mem_from_instance_type(
+                resources.instance_type)
+            if vcpus is not None:
+                env_vars['SKYPILOT_POD_CPU_CORE_LIMIT'] = str(vcpus)
+            if mem is not None:
+                env_vars['SKYPILOT_POD_MEMORY_GB_LIMIT'] = str(mem)
+        return '; '.join([f'export {k}={v}' for k, v in env_vars.items()])
+
     runners = provision.get_command_runners(cluster_info.provider_name,
                                             cluster_info, **ssh_credentials)
     head_runner = runners[0]

sky/provision/kubernetes/__init__.py

@@ -13,4 +13,6 @@ from sky.provision.kubernetes.network import open_ports
 from sky.provision.kubernetes.network import query_ports
 from sky.provision.kubernetes.volume import apply_volume
 from sky.provision.kubernetes.volume import delete_volume
+from sky.provision.kubernetes.volume import get_all_volumes_usedby
 from sky.provision.kubernetes.volume import get_volume_usedby
+from sky.provision.kubernetes.volume import map_all_volumes_usedby