skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/provision/kubernetes/manifests/smarter-device-manager-daemonset.yaml

@@ -26,6 +26,9 @@ spec:
       hostname: smarter-device-management
       hostNetwork: true
       dnsPolicy: ClusterFirstWithHostNet
+      tolerations:
+      - effect: NoSchedule
+        operator: Exists
       containers:
       - name: smarter-device-manager
         image: us-central1-docker.pkg.dev/skypilot-375900/skypilotk8s/smarter-device-manager:v1.1.2

sky/provision/kubernetes/network.py

@@ -1,6 +1,7 @@
 """Kubernetes network provisioning."""
 from typing import Any, Dict, List, Optional
 
+from sky import sky_logging
 from sky.adaptors import kubernetes
 from sky.provision import common
 from sky.provision.kubernetes import network_utils
@@ -8,6 +9,8 @@ from sky.provision.kubernetes import utils as kubernetes_utils
 from sky.utils import kubernetes_enums
 from sky.utils.resources_utils import port_ranges_to_set
 
+logger = sky_logging.init_logger(__name__)
+
 _PATH_PREFIX = '/skypilot/{namespace}/{cluster_name_on_cloud}/{port}'
 _LOADBALANCER_SERVICE_NAME = '{cluster_name_on_cloud}--skypilot-lb'
 
@@ -55,7 +58,8 @@ def _open_ports_using_loadbalancer(
     kubernetes_utils.merge_custom_metadata(content['service_spec']['metadata'])
 
     network_utils.create_or_replace_namespaced_service(
-        namespace=provider_config.get('namespace', 'default'),
+        namespace=kubernetes_utils.get_namespace_from_config(provider_config),
+        context=kubernetes_utils.get_context_from_config(provider_config),
         service_name=service_name,
         service_spec=content['service_spec'])
 
@@ -65,8 +69,9 @@ def _open_ports_using_ingress(
     ports: List[int],
     provider_config: Dict[str, Any],
 ) -> None:
+    context = kubernetes_utils.get_context_from_config(provider_config)
     # Check if an ingress controller exists
-    if not network_utils.ingress_controller_exists():
+    if not network_utils.ingress_controller_exists(context):
         raise Exception(
             'Ingress controller not found. '
             'Install Nginx ingress controller first: '
@@ -74,13 +79,14 @@ def _open_ports_using_ingress(
         )
 
     # Prepare service names, ports,  for template rendering
-    service_details = [(f'{cluster_name_on_cloud}--skypilot-svc--{port}', port,
-                        _PATH_PREFIX.format(
-                            cluster_name_on_cloud=cluster_name_on_cloud,
-                            port=port,
-                            namespace=kubernetes_utils.
-                            get_current_kube_config_context_namespace()).rstrip(
-                                '/').lstrip('/')) for port in ports]
+    service_details = [
+        (f'{cluster_name_on_cloud}--skypilot-svc--{port}', port,
+         _PATH_PREFIX.format(
+             cluster_name_on_cloud=cluster_name_on_cloud,
+             port=port,
+             namespace=kubernetes_utils.get_kube_config_context_namespace(
+                 context)).rstrip('/').lstrip('/')) for port in ports
+    ]
 
     # Generate ingress and services specs
     # We batch ingress rule creation because each rule triggers a hot reload of
@@ -105,7 +111,9 @@ def _open_ports_using_ingress(
         # Update metadata from config
         kubernetes_utils.merge_custom_metadata(service_spec['metadata'])
         network_utils.create_or_replace_namespaced_service(
-            namespace=provider_config.get('namespace', 'default'),
+            namespace=kubernetes_utils.get_namespace_from_config(
+                provider_config),
+            context=kubernetes_utils.get_context_from_config(provider_config),
             service_name=service_name,
             service_spec=service_spec,
         )
@@ -113,7 +121,8 @@ def _open_ports_using_ingress(
     kubernetes_utils.merge_custom_metadata(content['ingress_spec']['metadata'])
     # Create or update the single ingress for all services
     network_utils.create_or_replace_namespaced_ingress(
-        namespace=provider_config.get('namespace', 'default'),
+        namespace=kubernetes_utils.get_namespace_from_config(provider_config),
+        context=kubernetes_utils.get_context_from_config(provider_config),
         ingress_name=f'{cluster_name_on_cloud}-skypilot-ingress',
         ingress_spec=content['ingress_spec'],
     )
@@ -163,14 +172,16 @@ def _cleanup_ports_for_ingress(
     for port in ports:
         service_name = f'{cluster_name_on_cloud}--skypilot-svc--{port}'
         network_utils.delete_namespaced_service(
-            namespace=provider_config.get('namespace', 'default'),
+            namespace=provider_config.get('namespace',
+                                          kubernetes_utils.DEFAULT_NAMESPACE),
             service_name=service_name,
         )
 
     # Delete the single ingress used for all ports
     ingress_name = f'{cluster_name_on_cloud}-skypilot-ingress'
     network_utils.delete_namespaced_ingress(
-        namespace=provider_config.get('namespace', 'default'),
+        namespace=kubernetes_utils.get_namespace_from_config(provider_config),
+        context=kubernetes_utils.get_context_from_config(provider_config),
         ingress_name=ingress_name,
     )
 
@@ -199,11 +210,13 @@ def query_ports(
             return _query_ports_for_ingress(
                 cluster_name_on_cloud=cluster_name_on_cloud,
                 ports=ports,
+                provider_config=provider_config,
             )
         elif port_mode == kubernetes_enums.KubernetesPortMode.PODIP:
             return _query_ports_for_podip(
                 cluster_name_on_cloud=cluster_name_on_cloud,
                 ports=ports,
+                provider_config=provider_config,
             )
         else:
             return {}
@@ -218,12 +231,23 @@ def _query_ports_for_loadbalancer(
     ports: List[int],
     provider_config: Dict[str, Any],
 ) -> Dict[int, List[common.Endpoint]]:
+    logger.debug(f'Getting loadbalancer IP for cluster {cluster_name_on_cloud}')
     result: Dict[int, List[common.Endpoint]] = {}
     service_name = _LOADBALANCER_SERVICE_NAME.format(
         cluster_name_on_cloud=cluster_name_on_cloud)
+    context = provider_config.get(
+        'context', kubernetes_utils.get_current_kube_config_context_name())
+    namespace = provider_config.get(
+        'namespace',
+        kubernetes_utils.get_kube_config_context_namespace(context))
     external_ip = network_utils.get_loadbalancer_ip(
-        namespace=provider_config.get('namespace', 'default'),
-        service_name=service_name)
+        context=context,
+        namespace=namespace,
+        service_name=service_name,
+        # Timeout is set so that we can retry the query when the
+        # cluster is firstly created and the load balancer is not ready yet.
+        timeout=60,
+    )
 
     if external_ip is None:
         return {}
@@ -237,19 +261,24 @@ def _query_ports_for_loadbalancer(
 def _query_ports_for_ingress(
     cluster_name_on_cloud: str,
     ports: List[int],
+    provider_config: Dict[str, Any],
 ) -> Dict[int, List[common.Endpoint]]:
-    ingress_details = network_utils.get_ingress_external_ip_and_ports()
+    context = provider_config.get(
+        'context', kubernetes_utils.get_current_kube_config_context_name())
+    ingress_details = network_utils.get_ingress_external_ip_and_ports(context)
     external_ip, external_ports = ingress_details
     if external_ip is None:
         return {}
 
+    namespace = provider_config.get(
+        'namespace',
+        kubernetes_utils.get_kube_config_context_namespace(context))
     result: Dict[int, List[common.Endpoint]] = {}
     for port in ports:
         path_prefix = _PATH_PREFIX.format(
             cluster_name_on_cloud=cluster_name_on_cloud,
             port=port,
-            namespace=kubernetes_utils.
-            get_current_kube_config_context_namespace())
+            namespace=namespace)
 
         http_port, https_port = external_ports \
             if external_ports is not None else (None, None)
@@ -268,10 +297,15 @@ def _query_ports_for_ingress(
 def _query_ports_for_podip(
     cluster_name_on_cloud: str,
     ports: List[int],
+    provider_config: Dict[str, Any],
 ) -> Dict[int, List[common.Endpoint]]:
-    namespace = kubernetes_utils.get_current_kube_config_context_namespace()
+    context = provider_config.get(
+        'context', kubernetes_utils.get_current_kube_config_context_name())
+    namespace = provider_config.get(
+        'namespace',
+        kubernetes_utils.get_kube_config_context_namespace(context))
     pod_name = kubernetes_utils.get_head_pod_name(cluster_name_on_cloud)
-    pod_ip = network_utils.get_pod_ip(namespace, pod_name)
+    pod_ip = network_utils.get_pod_ip(context, namespace, pod_name)
 
     result: Dict[int, List[common.Endpoint]] = {}
     if pod_ip is None:

sky/provision/kubernetes/network_utils.py

@@ -1,5 +1,6 @@
 """Kubernetes network provisioning utils."""
 import os
+import time
 from typing import Dict, List, Optional, Tuple, Union
 
 import jinja2
@@ -7,12 +8,15 @@ import yaml
 
 import sky
 from sky import exceptions
+from sky import sky_logging
 from sky import skypilot_config
 from sky.adaptors import kubernetes
 from sky.provision.kubernetes import utils as kubernetes_utils
 from sky.utils import kubernetes_enums
 from sky.utils import ux_utils
 
+logger = sky_logging.init_logger(__name__)
+
 _INGRESS_TEMPLATE_NAME = 'kubernetes-ingress.yml.j2'
 _LOADBALANCER_TEMPLATE_NAME = 'kubernetes-loadbalancer.yml.j2'
 
@@ -43,6 +47,23 @@ def get_port_mode(
     return port_mode
 
 
+def get_networking_mode(
+    mode_str: Optional[str] = None
+) -> kubernetes_enums.KubernetesNetworkingMode:
+    """Get the networking mode from the provider config."""
+    mode_str = mode_str or skypilot_config.get_nested(
+        ('kubernetes', 'networking_mode'),
+        kubernetes_enums.KubernetesNetworkingMode.PORTFORWARD.value)
+    try:
+        networking_mode = kubernetes_enums.KubernetesNetworkingMode.from_str(
+            mode_str)
+    except ValueError as e:
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(str(e) +
+                             ' Please check: ~/.sky/config.yaml.') from None
+    return networking_mode
+
+
 def fill_loadbalancer_template(namespace: str, service_name: str,
                                ports: List[int], selector_key: str,
                                selector_value: str) -> Dict:
@@ -54,6 +75,10 @@ def fill_loadbalancer_template(namespace: str, service_name: str,
 
     with open(template_path, 'r', encoding='utf-8') as fin:
         template = fin.read()
+    annotations = skypilot_config.get_nested(
+        ('kubernetes', 'custom_metadata', 'annotations'), {})
+    labels = skypilot_config.get_nested(
+        ('kubernetes', 'custom_metadata', 'labels'), {})
     j2_template = jinja2.Template(template)
     cont = j2_template.render(
         namespace=namespace,
@@ -61,6 +86,8 @@ def fill_loadbalancer_template(namespace: str, service_name: str,
         ports=ports,
         selector_key=selector_key,
         selector_value=selector_value,
+        annotations=annotations,
+        labels=labels,
     )
     content = yaml.safe_load(cont)
     return content
@@ -77,6 +104,10 @@ def fill_ingress_template(namespace: str, service_details: List[Tuple[str, int,
             f'Template "{_INGRESS_TEMPLATE_NAME}" does not exist.')
     with open(template_path, 'r', encoding='utf-8') as fin:
         template = fin.read()
+    annotations = skypilot_config.get_nested(
+        ('kubernetes', 'custom_metadata', 'annotations'), {})
+    labels = skypilot_config.get_nested(
+        ('kubernetes', 'custom_metadata', 'labels'), {})
     j2_template = jinja2.Template(template)
     cont = j2_template.render(
         namespace=namespace,
@@ -88,6 +119,8 @@ def fill_ingress_template(namespace: str, service_details: List[Tuple[str, int,
         ingress_name=ingress_name,
         selector_key=selector_key,
         selector_value=selector_value,
+        annotations=annotations,
+        labels=labels,
     )
     content = yaml.safe_load(cont)
 
@@ -99,10 +132,10 @@ def fill_ingress_template(namespace: str, service_details: List[Tuple[str, int,
 
 
 def create_or_replace_namespaced_ingress(
-        namespace: str, ingress_name: str,
+        namespace: str, context: Optional[str], ingress_name: str,
         ingress_spec: Dict[str, Union[str, int]]) -> None:
     """Creates an ingress resource for the specified service."""
-    networking_api = kubernetes.networking_api()
+    networking_api = kubernetes.networking_api(context)
 
     try:
         networking_api.read_namespaced_ingress(
@@ -123,9 +156,10 @@ def create_or_replace_namespaced_ingress(
         _request_timeout=kubernetes.API_TIMEOUT)
 
 
-def delete_namespaced_ingress(namespace: str, ingress_name: str) -> None:
+def delete_namespaced_ingress(namespace: str, context: Optional[str],
+                              ingress_name: str) -> None:
     """Deletes an ingress resource."""
-    networking_api = kubernetes.networking_api()
+    networking_api = kubernetes.networking_api(context)
     try:
         networking_api.delete_namespaced_ingress(
             ingress_name, namespace, _request_timeout=kubernetes.API_TIMEOUT)
@@ -137,10 +171,10 @@ def delete_namespaced_ingress(namespace: str, ingress_name: str) -> None:
 
 
 def create_or_replace_namespaced_service(
-        namespace: str, service_name: str,
+        namespace: str, context: Optional[str], service_name: str,
         service_spec: Dict[str, Union[str, int]]) -> None:
     """Creates a service resource for the specified service."""
-    core_api = kubernetes.core_api()
+    core_api = kubernetes.core_api(context)
 
     try:
         core_api.read_namespaced_service(
@@ -174,9 +208,10 @@ def delete_namespaced_service(namespace: str, service_name: str) -> None:
         raise e
 
 
-def ingress_controller_exists(ingress_class_name: str = 'nginx') -> bool:
+def ingress_controller_exists(context: Optional[str],
+                              ingress_class_name: str = 'nginx') -> bool:
     """Checks if an ingress controller exists in the cluster."""
-    networking_api = kubernetes.networking_api()
+    networking_api = kubernetes.networking_api(context)
     ingress_classes = networking_api.list_ingress_class(
         _request_timeout=kubernetes.API_TIMEOUT).items
     return any(
@@ -185,16 +220,17 @@ def ingress_controller_exists(ingress_class_name: str = 'nginx') -> bool:
 
 
 def get_ingress_external_ip_and_ports(
+    context: Optional[str],
     namespace: str = 'ingress-nginx'
 ) -> Tuple[Optional[str], Optional[Tuple[int, int]]]:
     """Returns external ip and ports for the ingress controller."""
-    core_api = kubernetes.core_api()
+    core_api = kubernetes.core_api(context)
     ingress_services = [
         item for item in core_api.list_namespaced_service(
             namespace, _request_timeout=kubernetes.API_TIMEOUT).items
         if item.metadata.name == 'ingress-nginx-controller'
     ]
-    if len(ingress_services) == 0:
+    if not ingress_services:
         return (None, None)
 
     ingress_service = ingress_services[0]
@@ -222,23 +258,36 @@ def get_ingress_external_ip_and_ports(
     return external_ip, None
 
 
-def get_loadbalancer_ip(namespace: str, service_name: str) -> Optional[str]:
+def get_loadbalancer_ip(context: Optional[str],
+                        namespace: str,
+                        service_name: str,
+                        timeout: int = 0) -> Optional[str]:
     """Returns the IP address of the load balancer."""
-    core_api = kubernetes.core_api()
-    service = core_api.read_namespaced_service(
-        service_name, namespace, _request_timeout=kubernetes.API_TIMEOUT)
+    core_api = kubernetes.core_api(context)
 
-    if service.status.load_balancer.ingress is None:
-        return None
+    ip = None
 
-    ip = service.status.load_balancer.ingress[
-        0].ip or service.status.load_balancer.ingress[0].hostname
-    return ip if ip is not None else None
+    start_time = time.time()
+    retry_cnt = 0
+    while ip is None and (retry_cnt == 0 or time.time() - start_time < timeout):
+        service = core_api.read_namespaced_service(
+            service_name, namespace, _request_timeout=kubernetes.API_TIMEOUT)
+        if service.status.load_balancer.ingress is not None:
+            ip = (service.status.load_balancer.ingress[0].ip or
+                  service.status.load_balancer.ingress[0].hostname)
+        if ip is None:
+            retry_cnt += 1
+            if retry_cnt % 5 == 0:
+                logger.debug('Waiting for load balancer IP to be assigned'
+                             '...')
+            time.sleep(1)
+    return ip
 
 
-def get_pod_ip(namespace: str, pod_name: str) -> Optional[str]:
+def get_pod_ip(context: Optional[str], namespace: str,
+               pod_name: str) -> Optional[str]:
     """Returns the IP address of the pod."""
-    core_api = kubernetes.core_api()
+    core_api = kubernetes.core_api(context)
     pod = core_api.read_namespaced_pod(pod_name,
                                        namespace,
                                        _request_timeout=kubernetes.API_TIMEOUT)