skypilot-nightly 1.0.0.dev20250710__py3-none-any.whl → 1.0.0.dev20250712__py3-none-any.whl

sky/global_user_state.py

@@ -220,17 +220,16 @@ def _glob_to_similar(glob_pattern):
     return like_pattern
 
 
-def create_table():
+def create_table(engine: sqlalchemy.engine.Engine):
     # Enable WAL mode to avoid locking issues.
     # See: issue #1441 and PR #1509
     # https://github.com/microsoft/WSL/issues/2395
     # TODO(romilb): We do not enable WAL for WSL because of known issue in WSL.
     #  This may cause the database locked problem from WSL issue #1441.
-    if (_SQLALCHEMY_ENGINE.dialect.name
-            == db_utils.SQLAlchemyDialect.SQLITE.value and
+    if (engine.dialect.name == db_utils.SQLAlchemyDialect.SQLITE.value and
             not common_utils.is_wsl()):
         try:
-            with orm.Session(_SQLALCHEMY_ENGINE) as session:
+            with orm.Session(engine) as session:
                 session.execute(sqlalchemy.text('PRAGMA journal_mode=WAL'))
                 session.commit()
         except sqlalchemy_exc.OperationalError as e:
@@ -240,12 +239,12 @@ def create_table():
             # is not critical and is likely to be enabled by other processes.
 
     # Create tables if they don't exist
-    db_utils.add_tables_to_db_sqlalchemy(Base.metadata, _SQLALCHEMY_ENGINE)
+    db_utils.add_tables_to_db_sqlalchemy(Base.metadata, engine)
 
     # For backward compatibility.
     # TODO(zhwu): Remove this function after all users have migrated to
     # the latest version of SkyPilot.
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(engine) as session:
         # Add autostop column to clusters table
         db_utils.add_column_to_table_sqlalchemy(session,
                                                 'clusters',
@@ -391,15 +390,15 @@ def initialize_and_get_db() -> sqlalchemy.engine.Engine:
                 conn_string = skypilot_config.get_nested(('db',), None)
             if conn_string:
                 logger.debug(f'using db URI from {conn_string}')
-                _SQLALCHEMY_ENGINE = sqlalchemy.create_engine(
-                    conn_string, poolclass=sqlalchemy.NullPool)
+                engine = sqlalchemy.create_engine(conn_string,
+                                                  poolclass=sqlalchemy.NullPool)
             else:
                 db_path = os.path.expanduser('~/.sky/state.db')
                 pathlib.Path(db_path).parents[0].mkdir(parents=True,
                                                        exist_ok=True)
-                _SQLALCHEMY_ENGINE = sqlalchemy.create_engine('sqlite:///' +
-                                                              db_path)
-            create_table()
+                engine = sqlalchemy.create_engine('sqlite:///' + db_path)
+            create_table(engine)
+            _SQLALCHEMY_ENGINE = engine
     return _SQLALCHEMY_ENGINE
 
 

sky/jobs/state.py

@@ -112,17 +112,16 @@ ha_recovery_script_table = sqlalchemy.Table(
 )
 
 
-def create_table():
+def create_table(engine: sqlalchemy.engine.Engine):
     # Enable WAL mode to avoid locking issues.
     # See: issue #3863, #1441 and PR #1509
     # https://github.com/microsoft/WSL/issues/2395
     # TODO(romilb): We do not enable WAL for WSL because of known issue in WSL.
     #  This may cause the database locked problem from WSL issue #1441.
-    if (_SQLALCHEMY_ENGINE.dialect.name
-            == db_utils.SQLAlchemyDialect.SQLITE.value and
+    if (engine.dialect.name == db_utils.SQLAlchemyDialect.SQLITE.value and
             not common_utils.is_wsl()):
         try:
-            with orm.Session(_SQLALCHEMY_ENGINE) as session:
+            with orm.Session(engine) as session:
                 session.execute(sqlalchemy.text('PRAGMA journal_mode=WAL'))
                 session.commit()
         except sqlalchemy_exc.OperationalError as e:
@@ -132,10 +131,10 @@ def create_table():
             # is not critical and is likely to be enabled by other processes.
 
     # Create tables if they don't exist
-    db_utils.add_tables_to_db_sqlalchemy(Base.metadata, _SQLALCHEMY_ENGINE)
+    db_utils.add_tables_to_db_sqlalchemy(Base.metadata, engine)
 
     # Backward compatibility: add columns that not exist in older databases
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(engine) as session:
         db_utils.add_column_to_table_sqlalchemy(session, 'spot',
                                                 'failure_reason',
                                                 sqlalchemy.Text())
@@ -228,15 +227,15 @@ def initialize_and_get_db() -> sqlalchemy.engine.Engine:
                 conn_string = skypilot_config.get_nested(('db',), None)
             if conn_string:
                 logger.debug(f'using db URI from {conn_string}')
-                _SQLALCHEMY_ENGINE = sqlalchemy.create_engine(
-                    conn_string, poolclass=sqlalchemy.NullPool)
+                engine = sqlalchemy.create_engine(conn_string,
+                                                  poolclass=sqlalchemy.NullPool)
             else:
                 db_path = os.path.expanduser('~/.sky/spot_jobs.db')
                 pathlib.Path(db_path).parents[0].mkdir(parents=True,
                                                        exist_ok=True)
-                _SQLALCHEMY_ENGINE = sqlalchemy.create_engine('sqlite:///' +
-                                                              db_path)
-            create_table()
+                engine = sqlalchemy.create_engine('sqlite:///' + db_path)
+            create_table(engine)
+            _SQLALCHEMY_ENGINE = engine
     return _SQLALCHEMY_ENGINE
 
 

sky/jobs/utils.py

@@ -30,6 +30,7 @@ from sky.backends import backend_utils
 from sky.jobs import constants as managed_job_constants
 from sky.jobs import scheduler
 from sky.jobs import state as managed_job_state
+from sky.server import common as server_common
 from sky.skylet import constants
 from sky.skylet import job_lib
 from sky.skylet import log_lib
@@ -38,6 +39,7 @@ from sky.utils import annotations
 from sky.utils import command_runner
 from sky.utils import common_utils
 from sky.utils import controller_utils
+from sky.utils import env_options
 from sky.utils import infra_utils
 from sky.utils import log_utils
 from sky.utils import message_utils
@@ -128,9 +130,15 @@ def terminate_cluster(cluster_name: str, max_retry: int = 6) -> None:
             time.sleep(backoff.current_backoff())
 
 
-def _check_consolidation_mode_consistency(
+def _validate_consolidation_mode_config(
         current_is_consolidation_mode: bool) -> None:
-    """Check the consistency of the consolidation mode."""
+    """Validate the consolidation mode config."""
+    if (current_is_consolidation_mode and
+            not env_options.Options.IS_DEVELOPER.get() and
+            server_common.is_api_server_local()):
+        with ux_utils.print_exception_no_traceback():
+            raise exceptions.NotSupportedError(
+                'Consolidation mode is not supported when running locally.')
     # Check whether the consolidation mode config is changed.
     if current_is_consolidation_mode:
         controller_cn = (
@@ -176,7 +184,7 @@ def _check_consolidation_mode_consistency(
 def is_consolidation_mode() -> bool:
     consolidation_mode = skypilot_config.get_nested(
         ('jobs', 'controller', 'consolidation_mode'), default_value=False)
-    _check_consolidation_mode_consistency(consolidation_mode)
+    _validate_consolidation_mode_config(consolidation_mode)
     return consolidation_mode
 
 

sky/optimizer.py

@@ -997,23 +997,29 @@ class Optimizer:
     @staticmethod
     def _print_candidates(node_to_candidate_map: _TaskToPerCloudCandidates):
         for node, candidate_set in node_to_candidate_map.items():
-            if node.best_resources:
-                accelerator = node.best_resources.accelerators
-            else:
-                accelerator = list(node.resources)[0].accelerators
+            best_resources = node.best_resources
+            if best_resources is None:
+                best_resources = list(node.resources)[0]
             is_multi_instances = False
-            if accelerator:
-                acc_name, acc_count = list(accelerator.items())[0]
+            if best_resources.accelerators:
+                acc_name, acc_count = list(
+                    best_resources.accelerators.items())[0]
                 for cloud, candidate_list in candidate_set.items():
-                    if len(candidate_list) > 1:
+                    # Filter only the candidates matching the best
+                    # resources chosen by the optimizer.
+                    best_resources_candidates = [
+                        res for res in candidate_list if
+                        res.get_accelerators_str() == f'{acc_name}:{acc_count}'
+                    ]
+                    if len(best_resources_candidates) > 1:
                         is_multi_instances = True
-                        instance_list = [
+                        instance_list = set([
                             res.instance_type
-                            for res in candidate_list
+                            for res in best_resources_candidates
                             if res.instance_type is not None
-                        ]
+                        ])
                         candidate_str = resources_utils.format_resource(
-                            candidate_list[0], simplify=True)
+                            best_resources, simplify=True)
 
                         logger.info(
                             f'{colorama.Style.DIM}🔍 Multiple {cloud} instances '
@@ -1327,8 +1333,7 @@ def _fill_in_launchable_resources(
     launchable: Dict[resources_lib.Resources, List[resources_lib.Resources]] = (
         collections.defaultdict(list))
     all_fuzzy_candidates = set()
-    cloud_candidates: _PerCloudCandidates = collections.defaultdict(
-        List[resources_lib.Resources])
+    cloud_candidates: _PerCloudCandidates = collections.defaultdict(list)
     resource_hints: Dict[resources_lib.Resources,
                          List[str]] = collections.defaultdict(list)
     if blocked_resources is None:
@@ -1365,7 +1370,10 @@ def _fill_in_launchable_resources(
                 launchable[resources].extend(
                     resources_utils.make_launchables_for_valid_region_zones(
                         cheapest))
-                cloud_candidates[cloud] = feasible_resources.resources_list
+                # Each cloud can occur multiple times in feasible_list,
+                # for different region/zone.
+                cloud_candidates[cloud].extend(
+                    feasible_resources.resources_list)
             else:
                 all_fuzzy_candidates.update(
                     feasible_resources.fuzzy_candidate_list)

sky/provision/kubernetes/utils.py

@@ -1,5 +1,6 @@
 """Kubernetes utilities for SkyPilot."""
 import dataclasses
+import enum
 import functools
 import hashlib
 import json
@@ -57,6 +58,69 @@ HIGH_AVAILABILITY_DEPLOYMENT_VOLUME_MOUNT_NAME = 'sky-data'
 # and store all data that needs to be persisted in future.
 HIGH_AVAILABILITY_DEPLOYMENT_VOLUME_MOUNT_PATH = '/home/sky'
 
+
+class KubernetesHighPerformanceNetworkType(enum.Enum):
+    """Enum for different Kubernetes cluster types with high performance
+    network configurations.
+
+    This enum defines cluster types that support optimized networking for
+    distributed ML workloads:
+    - GCP_TCPX: GKE clusters with GPUDirect-TCPX support
+      (A3 High instances: a3-highgpu-8g)
+    - GCP_TCPXO: GKE clusters with GPUDirect-TCPXO support
+      (A3 Mega instances: a3-megagpu-8g)
+    - GCP_GPUDIRECT_RDMA: GKE clusters with GPUDirect-RDMA support
+      (A4/A3 Ultra instances)
+    - NEBIUS: Nebius clusters with InfiniBand support for high-throughput,
+      low-latency networking
+    - NONE: Standard clusters without specialized networking optimizations
+
+    The network configurations align with corresponding VM-based
+    implementations:
+    - GCP settings match
+      sky.provision.gcp.constants.GPU_DIRECT_TCPX_SPECIFIC_OPTIONS
+    - Nebius settings match the InfiniBand configuration used in Nebius VMs
+    """
+
+    GCP_TCPX = 'gcp_tcpx'
+    GCP_TCPXO = 'gcp_tcpxo'
+    GCP_GPUDIRECT_RDMA = 'gcp_gpudirect_rdma'
+    NEBIUS = 'nebius'
+    NONE = 'none'
+
+    def get_network_env_vars(self) -> Dict[str, str]:
+        """Get network environment variables for this cluster type."""
+        if self == KubernetesHighPerformanceNetworkType.NEBIUS:
+            # Nebius cluster with InfiniBand - use InfiniBand optimizations
+            return {
+                'NCCL_IB_HCA': 'mlx5',
+                'UCX_NET_DEVICES': ('mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1,'
+                                    'mlx5_4:1,mlx5_5:1,mlx5_6:1,mlx5_7:1')
+            }
+        else:
+            # GCP clusters and generic clusters - environment variables are
+            # handled directly in the template
+            return {}
+
+    def supports_high_performance_networking(self) -> bool:
+        """Check if this cluster type supports high performance networking."""
+        return self is not KubernetesHighPerformanceNetworkType.NONE
+
+    def supports_gpu_direct(self) -> bool:
+        """Check if this cluster type supports GPUDirect networking."""
+        return self in (KubernetesHighPerformanceNetworkType.GCP_TCPX,
+                        KubernetesHighPerformanceNetworkType.GCP_TCPXO,
+                        KubernetesHighPerformanceNetworkType.GCP_GPUDIRECT_RDMA)
+
+    def requires_ipc_lock_capability(self) -> bool:
+        """Check if this cluster type requires IPC_LOCK capability."""
+        return self.supports_high_performance_networking()
+
+    def requires_tcpxo_daemon(self) -> bool:
+        """Check if this cluster type requires TCPXO daemon."""
+        return self == KubernetesHighPerformanceNetworkType.GCP_TCPXO
+
+
 # TODO(romilb): Move constants to constants.py
 DEFAULT_NAMESPACE = 'default'
 
@@ -758,6 +822,74 @@ class GKEAutoscaler(Autoscaler):
                     return True
         return False
 
+    @classmethod
+    @annotations.lru_cache(scope='request', maxsize=10)
+    def get_available_machine_types(cls, context: str) -> List[str]:
+        """Returns the list of machine types that are available in the cluster.
+        """
+        # Assume context naming convention of
+        # gke_PROJECT-ID_LOCATION_CLUSTER-NAME
+        valid, project_id, location, cluster_name = cls._validate_context_name(
+            context)
+        if not valid:
+            # Context name is not in the format of
+            # gke_PROJECT-ID_LOCATION_CLUSTER-NAME.
+            # Cannot determine if the context can autoscale.
+            # Return empty list.
+            logger.debug(f'Context {context} is not in the format of '
+                         f'gke_PROJECT-ID_LOCATION_CLUSTER-NAME. '
+                         'Returning empty machine type list.')
+            return []
+        try:
+            logger.debug(
+                f'Attempting to get information about cluster {cluster_name}')
+            container_service = gcp.build('container',
+                                          'v1',
+                                          credentials=None,
+                                          cache_discovery=False)
+            cluster = container_service.projects().locations().clusters().get(
+                name=f'projects/{project_id}'
+                f'/locations/{location}'
+                f'/clusters/{cluster_name}').execute()
+        except ImportError:
+            # If the gcp module is not installed, return empty list.
+            # Remind the user once per day to install the gcp module for better
+            # pod scheduling with GKE autoscaler.
+            if time.time() - cls._pip_install_gcp_hint_last_sent > 60 * 60 * 24:
+                logger.info(
+                    'Could not fetch autoscaler information from GKE. '
+                    'Run pip install "skypilot[gcp]" for more intelligent pod '
+                    'scheduling with GKE autoscaler.')
+                cls._pip_install_gcp_hint_last_sent = time.time()
+            return []
+        except gcp.http_error_exception() as e:
+            # Cluster information is not available.
+            # Return empty list.
+            logger.debug(f'{e.message}', exc_info=True)
+            return []
+
+        machine_types = []
+        # Get the list of machine types that are available in the cluster.
+        node_pools = cluster.get('nodePools', [])
+        for node_pool in node_pools:
+            name = node_pool.get('name', '')
+            logger.debug(f'Checking if node pool {name} '
+                         'has autoscaling enabled.')
+            autoscaling_enabled = (node_pool.get('autoscaling',
+                                                 {}).get('enabled', False))
+            if autoscaling_enabled:
+                logger.debug(f'Node pool {name} has autoscaling enabled.')
+                try:
+                    machine_type = node_pool.get('config',
+                                                 {}).get('machineType', '')
+                    if machine_type:
+                        machine_types.append(machine_type)
+                except KeyError:
+                    logger.debug(f'Encountered KeyError while checking machine '
+                                 f'type of node pool {name}.')
+                    continue
+        return machine_types
+
     @classmethod
     def _validate_context_name(cls, context: str) -> Tuple[bool, str, str, str]:
         """Validates the context name is in the format of

sky/setup_files/dependencies.py

@@ -129,6 +129,7 @@ extras_require: Dict[str, List[str]] = {
         'azure-mgmt-compute>=33.0.0',
         'azure-storage-blob>=12.23.1',
         'msgraph-sdk',
+        'msrestazure',
     ] + local_ray,
     # We need google-api-python-client>=2.69.0 to enable 'discardLocalSsd'
     # parameter for stopping instances. Reference:

sky/skypilot_config.py

@@ -52,6 +52,7 @@ import contextlib
 import copy
 import json
 import os
+import pathlib
 import tempfile
 import threading
 import typing
@@ -848,7 +849,9 @@ def update_api_server_config_no_lock(config: config_utils.Config) -> None:
 
     global_config_path = _resolve_server_config_path()
     if global_config_path is None:
-        global_config_path = get_user_config_path()
+        # Fallback to ~/.sky/config.yaml, and make sure it exists.
+        global_config_path = os.path.expanduser(get_user_config_path())
+        pathlib.Path(global_config_path).touch(exist_ok=True)
 
     db_updated = False
     if os.environ.get(constants.ENV_VAR_IS_SKYPILOT_SERVER) is not None: