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

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

sky/provision/kubernetes/config.py

@@ -3,13 +3,11 @@ import copy
 import logging
 import math
 import os
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 from sky.adaptors import kubernetes
 from sky.provision import common
-from sky.provision.kubernetes import network_utils
 from sky.provision.kubernetes import utils as kubernetes_utils
-from sky.utils import kubernetes_enums
 from sky.utils import yaml_utils
 
 logger = logging.getLogger(__name__)
@@ -28,11 +26,6 @@ def bootstrap_instances(
 
     _configure_services(namespace, context, config.provider_config)
 
-    networking_mode = network_utils.get_networking_mode(
-        config.provider_config.get('networking_mode'), context)
-    if networking_mode == kubernetes_enums.KubernetesNetworkingMode.NODEPORT:
-        config = _configure_ssh_jump(namespace, context, config)
-
     requested_service_account = config.node_config['spec']['serviceAccountName']
     if (requested_service_account ==
             kubernetes_utils.DEFAULT_SERVICE_ACCOUNT_NAME):
@@ -481,41 +474,6 @@ def _configure_autoscaler_cluster_role_binding(
                 f'{created_msg(binding_field, name)}')
 
 
-def _configure_ssh_jump(namespace, context, config: common.ProvisionConfig):
-    """Creates a SSH jump pod to connect to the cluster.
-
-    Also updates config['auth']['ssh_proxy_command'] to use the newly created
-    jump pod.
-    """
-    provider_config = config.provider_config
-    pod_cfg = config.node_config
-
-    ssh_jump_name = pod_cfg['metadata']['labels']['skypilot-ssh-jump']
-    ssh_jump_image = provider_config['ssh_jump_image']
-
-    volumes = pod_cfg['spec']['volumes']
-    # find 'secret-volume' and get the secret name
-    secret_volume = next(filter(lambda x: x['name'] == 'secret-volume',
-                                volumes))
-    ssh_key_secret_name = secret_volume['secret']['secretName']
-
-    # TODO(romilb): We currently split SSH jump pod and svc creation. Service
-    #  is first created in authentication.py::setup_kubernetes_authentication
-    #  and then SSH jump pod creation happens here. This is because we need to
-    #  set the ssh_proxy_command in the ray YAML before we pass it to the
-    #  autoscaler. If in the future if we can write the ssh_proxy_command to the
-    #  cluster yaml through this method, then we should move the service
-    #  creation here.
-
-    # TODO(romilb): We should add a check here to make sure the service is up
-    #  and available before we create the SSH jump pod. If for any reason the
-    #  service is missing, we should raise an error.
-
-    kubernetes_utils.setup_ssh_jump_pod(ssh_jump_name, ssh_jump_image,
-                                        ssh_key_secret_name, namespace, context)
-    return config
-
-
 def _configure_skypilot_system_namespace(
         provider_config: Dict[str, Any]) -> None:
     """Creates the namespace for skypilot-system mounting if it does not exist.
@@ -666,4 +624,9 @@ def _configure_services(namespace: str, context: Optional[str],
 
 
 class KubernetesError(Exception):
-    pass
+
+    def __init__(self,
+                 *args,
+                 insufficent_resources: Optional[List[str]] = None):
+        self.insufficent_resources = insufficent_resources
+        super().__init__(*args)

sky/provision/kubernetes/constants.py

@@ -18,7 +18,6 @@ SKY_K8S_EXEC_AUTH_KUBECONFIG_CACHE = '~/.sky/generated/kubeconfigs'
 
 # Labels for the Pods created by SkyPilot
 TAG_RAY_CLUSTER_NAME = 'ray-cluster-name'
-TAG_SKYPILOT_CLUSTER_NAME = 'skypilot-cluster-name'
 TAG_POD_INITIALIZED = 'skypilot-initialized'
 TAG_SKYPILOT_DEPLOYMENT_NAME = 'skypilot-deployment-name'