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

sky/clouds/utils/gcp_utils.py

@@ -10,7 +10,7 @@ import dataclasses
 import json
 import time
 import typing
-from typing import List, Optional, Set
+from typing import Any, Dict, List, Optional, Set, Tuple
 
 import cachetools
 
@@ -18,6 +18,7 @@ from sky import sky_logging
 from sky import skypilot_config
 from sky.provision.gcp import constants
 from sky.provision.kubernetes import utils as kubernetes_utils
+from sky.utils import resources_utils
 from sky.utils import subprocess_utils
 
 if typing.TYPE_CHECKING:
@@ -36,9 +37,10 @@ def is_tpu(resources: Optional['resources_lib.Resources']) -> bool:
 def is_tpu_vm(resources: Optional['resources_lib.Resources']) -> bool:
     if not is_tpu(resources):
         return False
-    assert (resources is not None and len(resources.accelerators) == 1)
+    assert (resources is not None and resources.accelerators is not None and
+            len(resources.accelerators) == 1)
     acc, _ = list(resources.accelerators.items())[0]
-    if kubernetes_utils.is_tpu_on_gke(acc):
+    if kubernetes_utils.is_tpu_on_gke(acc, normalize=False):
         return False
     if resources.accelerator_args is None:
         return True
@@ -48,7 +50,7 @@ def is_tpu_vm(resources: Optional['resources_lib.Resources']) -> bool:
 def is_tpu_vm_pod(resources: Optional['resources_lib.Resources']) -> bool:
     if not is_tpu_vm(resources):
         return False
-    assert resources is not None
+    assert resources is not None and resources.accelerators is not None
     acc, _ = list(resources.accelerators.items())[0]
     return not acc.endswith('-8')
 
@@ -136,10 +138,16 @@ def _list_reservations_for_instance_type(
     For example, if we have a specific reservation with n1-highmem-8
     in us-central1-c. `sky launch --gpus V100` will fail.
     """
-    prioritize_reservations = skypilot_config.get_nested(
-        ('gcp', 'prioritize_reservations'), False)
-    specific_reservations = skypilot_config.get_nested(
-        ('gcp', 'specific_reservations'), [])
+    prioritize_reservations = skypilot_config.get_effective_region_config(
+        cloud='gcp',
+        region=None,
+        keys=('prioritize_reservations',),
+        default_value=False)
+    specific_reservations = skypilot_config.get_effective_region_config(
+        cloud='gcp',
+        region=None,
+        keys=('specific_reservations',),
+        default_value=[])
     if not prioritize_reservations and not specific_reservations:
         return []
     logger.debug(f'Querying GCP reservations for instance {instance_type!r}')
@@ -169,14 +177,23 @@ def _list_reservations_for_instance_type(
 
 def get_minimal_compute_permissions() -> List[str]:
     permissions = copy.copy(constants.VM_MINIMAL_PERMISSIONS)
-    if skypilot_config.get_nested(('gcp', 'vpc_name'), None) is None:
+    if skypilot_config.get_effective_region_config(
+            cloud='gcp', region=None, keys=('vpc_name',),
+            default_value=None) is None:
         # If custom VPC is not specified, permissions to modify network are
         # required to ensure SkyPilot to be able to setup the network, and
         # allow opening ports (e.g., via `resources.ports`).
         permissions += constants.FIREWALL_PERMISSIONS
 
-    if (skypilot_config.get_nested(('gcp', 'prioritize_reservations'), False) or
-            skypilot_config.get_nested(('gcp', 'specific_reservations'), [])):
+    if (skypilot_config.get_effective_region_config(
+            cloud='gcp',
+            region=None,
+            keys=('prioritize_reservations',),
+            default_value=False) or skypilot_config.get_effective_region_config(
+                cloud='gcp',
+                region=None,
+                keys=('specific_reservations',),
+                default_value=[])):
         permissions += constants.RESERVATION_PERMISSIONS
 
     permissions += constants.GCP_MINIMAL_PERMISSIONS
@@ -190,3 +207,62 @@ def get_minimal_storage_permissions() -> List[str]:
     permissions += constants.GCP_MINIMAL_PERMISSIONS
 
     return permissions
+
+
+# Get the DWS configuration for the given context in GKE.
+def get_dws_config(
+    context: str,
+    k8s_kueue_local_queue_name: Optional[str],
+    cluster_config_overrides: Optional[Dict[str, Any]] = None,
+) -> Tuple[bool, bool, Optional[int]]:
+    """Get the DWS configuration for the given context.
+
+        Args:
+            context: The context to get the DWS configuration for.
+            k8s_kueue_local_queue_name: The name of the Kueue local queue.
+            cluster_config_overrides: The cluster config overrides.
+
+        Returns:
+            A tuple of (enable_flex_start,
+                        enable_flex_start_queued_provisioning,
+                        max_run_duration_seconds).
+
+        Raises:
+            ValueError: If k8s_kueue_local_queue_name is missing to enable
+                        flex start queued provisioning for the given context.
+        """
+    dws_config = skypilot_config.get_effective_region_config(
+        cloud='kubernetes',
+        region=context,
+        keys=('dws',),
+        default_value={},
+        override_configs=cluster_config_overrides)
+    if not dws_config:
+        return False, False, None
+
+    enabled = dws_config.get('enabled', False)
+    if not enabled:
+        return False, False, None
+
+    enable_flex_start = False
+    enable_flex_start_queued_provisioning = False
+    max_run_duration_seconds = None
+    # If users already use Kueue, use the flex start with queued
+    # provisioning mode.
+    if k8s_kueue_local_queue_name:
+        enable_flex_start_queued_provisioning = True
+    else:
+        enable_flex_start = True
+
+    if not enable_flex_start_queued_provisioning:
+        return (enable_flex_start, enable_flex_start_queued_provisioning,
+                max_run_duration_seconds)
+
+    # Max run duration is only used in the flex start with queued
+    # provisioning mode.
+    max_run_duration = dws_config.get('max_run_duration', None)
+    if max_run_duration:
+        max_run_duration_seconds = resources_utils.parse_time_minutes(
+            max_run_duration) * 60
+    return (enable_flex_start, enable_flex_start_queued_provisioning,
+            max_run_duration_seconds)

sky/clouds/utils/oci_utils.py

@@ -105,22 +105,34 @@ class OCIConfig:
     @classmethod
     def get_compartment(cls, region):
         # Allow task(cluster)-specific compartment/VCN parameters.
-        default_compartment_ocid = skypilot_config.get_nested(
-            ('oci', 'default', 'compartment_ocid'), None)
-        compartment = skypilot_config.get_nested(
-            ('oci', region, 'compartment_ocid'), default_compartment_ocid)
+        default_compartment_ocid = skypilot_config.get_effective_region_config(
+            cloud='oci',
+            region='default',
+            keys=('compartment_ocid',),
+            default_value=None)
+        compartment = skypilot_config.get_effective_region_config(
+            cloud='oci',
+            region=region,
+            keys=('compartment_ocid',),
+            default_value=default_compartment_ocid)
         return compartment
 
     @classmethod
     def get_vcn_ocid(cls, region):
         # Will reuse the regional VCN if specified.
-        vcn = skypilot_config.get_nested(('oci', region, 'vcn_ocid'), None)
+        vcn = skypilot_config.get_effective_region_config(cloud='oci',
+                                                          region=region,
+                                                          keys=('vcn_ocid',),
+                                                          default_value=None)
         return vcn
 
     @classmethod
     def get_vcn_subnet(cls, region):
         # Will reuse the subnet if specified.
-        vcn = skypilot_config.get_nested(('oci', region, 'vcn_subnet'), None)
+        vcn = skypilot_config.get_effective_region_config(cloud='oci',
+                                                          region=region,
+                                                          keys=('vcn_subnet',),
+                                                          default_value=None)
         return vcn
 
     @classmethod
@@ -129,16 +141,22 @@ class OCIConfig:
         # we give a choice to set the default image tag (for gpu instances) in
         # the sky's user-config file (if not specified, use the hardcode one at
         # last)
-        return skypilot_config.get_nested(('oci', 'default', 'image_tag_gpu'),
-                                          'skypilot:gpu-ubuntu-2204')
+        return skypilot_config.get_effective_region_config(
+            cloud='oci',
+            region='default',
+            keys=('image_tag_gpu',),
+            default_value='skypilot:gpu-ubuntu-2204')
 
     @classmethod
     def get_default_image_tag(cls) -> str:
         # Get the default image tag. Instead of hardcoding, we give a choice to
         # set the default image tag in the sky's user-config file. (if not
         # specified, use the hardcode one at last)
-        return skypilot_config.get_nested(
-            ('oci', 'default', 'image_tag_general'), 'skypilot:cpu-ubuntu-2204')
+        return skypilot_config.get_effective_region_config(
+            cloud='oci',
+            region='default',
+            keys=('image_tag_general',),
+            default_value='skypilot:cpu-ubuntu-2204')
 
     @classmethod
     def get_sky_user_config_file(cls) -> str:
@@ -152,16 +170,22 @@ class OCIConfig:
 
     @classmethod
     def get_profile(cls) -> str:
-        return skypilot_config.get_nested(
-            ('oci', 'default', 'oci_config_profile'), 'DEFAULT')
+        return skypilot_config.get_effective_region_config(
+            cloud='oci',
+            region='default',
+            keys=('oci_config_profile',),
+            default_value='DEFAULT')
 
     @classmethod
     def get_default_image_os(cls) -> str:
         # Get the default image OS. Instead of hardcoding, we give a choice to
         # set the default image OS type in the sky's user-config file. (if not
         # specified, use the hardcode one at last)
-        return skypilot_config.get_nested(('oci', 'default', 'image_os_type'),
-                                          'ubuntu')
+        return skypilot_config.get_effective_region_config(
+            cloud='oci',
+            region='default',
+            keys=('image_os_type',),
+            default_value='ubuntu')
 
 
 oci_config = OCIConfig()