skypilot-nightly 1.0.0.dev20251009__py3-none-any.whl → 1.0.0.dev20251107__py3-none-any.whl

sky/backends/backend_utils.py

@@ -48,6 +48,7 @@ from sky.server.requests import requests as requests_lib
 from sky.skylet import autostop_lib
 from sky.skylet import constants
 from sky.usage import usage_lib
+from sky.utils import auth_utils
 from sky.utils import cluster_utils
 from sky.utils import command_runner
 from sky.utils import common
@@ -755,7 +756,7 @@ def write_cluster_config(
             assert k not in credentials, f'{k} already in credentials'
             credentials[k] = v
 
-    private_key_path, _ = auth.get_or_generate_keys()
+    private_key_path, _ = auth_utils.get_or_generate_keys()
     auth_config = {'ssh_private_key': private_key_path}
     region_name = resources_vars.get('region')
 
@@ -1124,6 +1125,8 @@ def _add_auth_to_cluster_config(cloud: clouds.Cloud, tmp_yaml_path: str):
         config = auth.setup_fluidstack_authentication(config)
     elif isinstance(cloud, clouds.Hyperbolic):
         config = auth.setup_hyperbolic_authentication(config)
+    elif isinstance(cloud, clouds.Shadeform):
+        config = auth.setup_shadeform_authentication(config)
     elif isinstance(cloud, clouds.PrimeIntellect):
         config = auth.setup_primeintellect_authentication(config)
     elif isinstance(cloud, clouds.Seeweb):
@@ -1855,6 +1858,13 @@ def check_owner_identity(cluster_name: str) -> None:
                                                      summary_response=True)
     if record is None:
         return
+    _check_owner_identity_with_record(cluster_name, record)
+
+
+def _check_owner_identity_with_record(cluster_name: str,
+                                      record: Dict[str, Any]) -> None:
+    if env_options.Options.SKIP_CLOUD_IDENTITY_CHECK.get():
+        return
     handle = record['handle']
     if not isinstance(handle, backends.CloudVmRayResourceHandle):
         return
@@ -1941,7 +1951,8 @@ def tag_filter_for_cluster(cluster_name: str) -> Dict[str, str]:
 
 @context_utils.cancellation_guard
 def _query_cluster_status_via_cloud_api(
-    handle: 'cloud_vm_ray_backend.CloudVmRayResourceHandle'
+    handle: 'cloud_vm_ray_backend.CloudVmRayResourceHandle',
+    retry_if_missing: bool,
 ) -> List[Tuple[status_lib.ClusterStatus, Optional[str]]]:
     """Returns the status of the cluster as a list of tuples corresponding
     to the node status and an optional reason string for said status.
@@ -1968,8 +1979,11 @@ def _query_cluster_status_via_cloud_api(
         cloud_name = repr(handle.launched_resources.cloud)
         try:
             node_status_dict = provision_lib.query_instances(
-                cloud_name, cluster_name, cluster_name_on_cloud,
-                provider_config)
+                cloud_name,
+                cluster_name,
+                cluster_name_on_cloud,
+                provider_config,
+                retry_if_missing=retry_if_missing)
             logger.debug(f'Querying {cloud_name} cluster '
                          f'{cluster_name_in_hint} '
                          f'status:\n{pprint.pformat(node_status_dict)}')
@@ -2149,6 +2163,8 @@ def check_can_clone_disk_and_override_task(
 
 def _update_cluster_status(
         cluster_name: str,
+        record: Dict[str, Any],
+        retry_if_missing: bool,
         include_user_info: bool = True,
         summary_response: bool = False) -> Optional[Dict[str, Any]]:
     """Update the cluster status.
@@ -2177,12 +2193,6 @@ def _update_cluster_status(
           fetched from the cloud provider or there are leaked nodes causing
           the node number larger than expected.
     """
-    record = global_user_state.get_cluster_from_name(
-        cluster_name,
-        include_user_info=include_user_info,
-        summary_response=summary_response)
-    if record is None:
-        return None
     handle = record['handle']
     if handle.cluster_yaml is None:
         # Remove cluster from db since this cluster does not have a config file
@@ -2201,7 +2211,8 @@ def _update_cluster_status(
         return record
     cluster_name = handle.cluster_name
 
-    node_statuses = _query_cluster_status_via_cloud_api(handle)
+    node_statuses = _query_cluster_status_via_cloud_api(
+        handle, retry_if_missing=retry_if_missing)
 
     all_nodes_up = (all(status[0] == status_lib.ClusterStatus.UP
                         for status in node_statuses) and
@@ -2376,7 +2387,8 @@ def _update_cluster_status(
             # and check again. This is a best-effort leak prevention check.
             # See https://github.com/skypilot-org/skypilot/issues/4431.
             time.sleep(_LAUNCH_DOUBLE_CHECK_DELAY)
-            node_statuses = _query_cluster_status_via_cloud_api(handle)
+            node_statuses = _query_cluster_status_via_cloud_api(
+                handle, retry_if_missing=False)
             # Note: even if all the node_statuses are UP now, we will still
             # consider this cluster abnormal, and its status will be INIT.
 
@@ -2620,7 +2632,8 @@ def refresh_cluster_record(
         cluster_lock_already_held: bool = False,
         cluster_status_lock_timeout: int = CLUSTER_STATUS_LOCK_TIMEOUT_SECONDS,
         include_user_info: bool = True,
-        summary_response: bool = False) -> Optional[Dict[str, Any]]:
+        summary_response: bool = False,
+        retry_if_missing: bool = True) -> Optional[Dict[str, Any]]:
     """Refresh the cluster, and return the possibly updated record.
 
     The function will update the cached cluster status in the global state. For
@@ -2649,6 +2662,8 @@ def refresh_cluster_record(
           value is <0, do not timeout (wait for the lock indefinitely). By
           default, this is set to CLUSTER_STATUS_LOCK_TIMEOUT_SECONDS. Warning:
           if correctness is required, you must set this to -1.
+        retry_if_missing: Whether to retry the call to the cloud api if the
+          cluster is not found when querying the live status on the cloud.
 
     Returns:
         If the cluster is terminated or does not exist, return None.
@@ -2675,10 +2690,9 @@ def refresh_cluster_record(
     # using the correct cloud credentials.
     workspace = record.get('workspace', constants.SKYPILOT_DEFAULT_WORKSPACE)
     with skypilot_config.local_active_workspace_ctx(workspace):
-        check_owner_identity(cluster_name)
-
-        if not isinstance(record['handle'], backends.CloudVmRayResourceHandle):
-            return record
+        # check_owner_identity returns if the record handle is
+        # not a CloudVmRayResourceHandle
+        _check_owner_identity_with_record(cluster_name, record)
 
         # The loop logic allows us to notice if the status was updated in the
         # global_user_state by another process and stop trying to get the lock.
@@ -2695,7 +2709,9 @@ def refresh_cluster_record(
                 return record
 
             if cluster_lock_already_held:
-                return _update_cluster_status(cluster_name, include_user_info,
+                return _update_cluster_status(cluster_name, record,
+                                              retry_if_missing,
+                                              include_user_info,
                                               summary_response)
 
             # Try to acquire the lock so we can fetch the status.
@@ -2711,7 +2727,8 @@ def refresh_cluster_record(
                             record, force_refresh_statuses):
                         return record
                     # Update and return the cluster status.
-                    return _update_cluster_status(cluster_name,
+                    return _update_cluster_status(cluster_name, record,
+                                                  retry_if_missing,
                                                   include_user_info,
                                                   summary_response)
 
@@ -2749,7 +2766,8 @@ def refresh_cluster_status_handle(
     *,
     force_refresh_statuses: Optional[Set[status_lib.ClusterStatus]] = None,
     cluster_lock_already_held: bool = False,
-    cluster_status_lock_timeout: int = CLUSTER_STATUS_LOCK_TIMEOUT_SECONDS
+    cluster_status_lock_timeout: int = CLUSTER_STATUS_LOCK_TIMEOUT_SECONDS,
+    retry_if_missing: bool = True,
 ) -> Tuple[Optional[status_lib.ClusterStatus],
            Optional[backends.ResourceHandle]]:
     """Refresh the cluster, and return the possibly updated status and handle.
@@ -2764,7 +2782,8 @@ def refresh_cluster_status_handle(
         cluster_lock_already_held=cluster_lock_already_held,
         cluster_status_lock_timeout=cluster_status_lock_timeout,
         include_user_info=False,
-        summary_response=True)
+        summary_response=True,
+        retry_if_missing=retry_if_missing)
     if record is None:
         return None, None
     return record['status'], record['handle']
@@ -3115,25 +3134,23 @@ def refresh_cluster_records() -> None:
     exclude_managed_clusters = True
     if env_options.Options.SHOW_DEBUG_INFO.get():
         exclude_managed_clusters = False
-    cluster_names = global_user_state.get_cluster_names(
-        exclude_managed_clusters=exclude_managed_clusters,)
+    cluster_names = set(
+        global_user_state.get_cluster_names(
+            exclude_managed_clusters=exclude_managed_clusters,))
 
     # TODO(syang): we should try not to leak
     # request info in backend_utils.py.
     # Refactor this to use some other info to
     # determine if a launch is in progress.
-    request = requests_lib.get_request_tasks(
-        req_filter=requests_lib.RequestTaskFilter(
-            status=[requests_lib.RequestStatus.RUNNING],
-            cluster_names=cluster_names,
-            include_request_names=['sky.launch']))
     cluster_names_with_launch_request = {
-        request.cluster_name for request in request
+        request.cluster_name for request in requests_lib.get_request_tasks(
+            req_filter=requests_lib.RequestTaskFilter(
+                status=[requests_lib.RequestStatus.RUNNING],
+                include_request_names=['sky.launch'],
+                fields=['cluster_name']))
     }
-    cluster_names_without_launch_request = [
-        cluster_name for cluster_name in cluster_names
-        if cluster_name not in cluster_names_with_launch_request
-    ]
+    cluster_names_without_launch_request = (cluster_names -
+                                            cluster_names_with_launch_request)
 
     def _refresh_cluster_record(cluster_name):
         return _refresh_cluster(cluster_name,
@@ -3142,7 +3159,7 @@ def refresh_cluster_records() -> None:
                                 include_user_info=False,
                                 summary_response=True)
 
-    if len(cluster_names) > 0:
+    if len(cluster_names_without_launch_request) > 0:
         # Do not refresh the clusters that have an active launch request.
         subprocess_utils.run_in_parallel(_refresh_cluster_record,
                                          cluster_names_without_launch_request)
@@ -3154,6 +3171,7 @@ def get_clusters(
     all_users: bool = True,
     include_credentials: bool = False,
     summary_response: bool = False,
+    include_handle: bool = True,
     # Internal only:
     # pylint: disable=invalid-name
     _include_is_managed: bool = False,
@@ -3237,12 +3255,11 @@ def get_clusters(
         """Add resource str to record"""
         for record in _get_records_with_handle(records):
             handle = record['handle']
-            record[
-                'resources_str'] = resources_utils.get_readable_resources_repr(
-                    handle, simplify=True)
-            record[
-                'resources_str_full'] = resources_utils.get_readable_resources_repr(
-                    handle, simplify=False)
+            resource_str_simple, resource_str_full = (
+                resources_utils.get_readable_resources_repr(
+                    handle, simplified_only=False))
+            record['resources_str'] = resource_str_simple
+            record['resources_str_full'] = resource_str_full
             if not summary_response:
                 record['cluster_name_on_cloud'] = handle.cluster_name_on_cloud
 
@@ -3268,9 +3285,17 @@ def get_clusters(
                 expanded_private_key_path = os.path.expanduser(
                     ssh_private_key_path)
                 if not os.path.exists(expanded_private_key_path):
-                    auth.create_ssh_key_files_from_db(ssh_private_key_path)
+                    success = auth_utils.create_ssh_key_files_from_db(
+                        ssh_private_key_path)
+                    if not success:
+                        # If the ssh key files are not found, we do not
+                        # update the record with credentials.
+                        logger.debug(
+                            f'SSH keys not found for cluster {record["name"]} '
+                            f'at key path {ssh_private_key_path}')
+                        continue
             else:
-                private_key_path, _ = auth.get_or_generate_keys()
+                private_key_path, _ = auth_utils.get_or_generate_keys()
                 expanded_private_key_path = os.path.expanduser(private_key_path)
             if expanded_private_key_path in cached_private_keys:
                 credential['ssh_private_key_content'] = cached_private_keys[
@@ -3302,6 +3327,8 @@ def get_clusters(
             record['accelerators'] = (
                 f'{handle.launched_resources.accelerators}'
                 if handle.launched_resources.accelerators else None)
+            if not include_handle:
+                record.pop('handle', None)
 
     # Add handle info to the records
     _update_records_with_handle_info(records)
@@ -3330,7 +3357,10 @@ def get_clusters(
                                   force_refresh_statuses=force_refresh_statuses,
                                   include_user_info=True,
                                   summary_response=summary_response)
-        if 'error' not in record:
+        # record may be None if the cluster is deleted during refresh,
+        # e.g. all the Pods of a cluster on Kubernetes have been
+        # deleted before refresh.
+        if record is not None and 'error' not in record:
             _update_records_with_handle_info([record])
             if include_credentials:
                 _update_records_with_credentials([record])
@@ -3342,45 +3372,56 @@ def get_clusters(
     # request info in backend_utils.py.
     # Refactor this to use some other info to
     # determine if a launch is in progress.
-    request = requests_lib.get_request_tasks(
-        req_filter=requests_lib.RequestTaskFilter(
-            status=[requests_lib.RequestStatus.RUNNING],
-            cluster_names=cluster_names,
-            include_request_names=['sky.launch']))
     cluster_names_with_launch_request = {
-        request.cluster_name for request in request
+        request.cluster_name for request in requests_lib.get_request_tasks(
+            req_filter=requests_lib.RequestTaskFilter(
+                status=[requests_lib.RequestStatus.RUNNING],
+                include_request_names=['sky.launch'],
+                cluster_names=cluster_names,
+                fields=['cluster_name']))
     }
+    # Preserve the index of the cluster name as it appears on "records"
     cluster_names_without_launch_request = [
-        cluster_name for cluster_name in cluster_names
+        (i, cluster_name)
+        for i, cluster_name in enumerate(cluster_names)
         if cluster_name not in cluster_names_with_launch_request
     ]
     # for clusters that have an active launch request, we do not refresh the status
-    updated_records = [
-        record for record in records
-        if record['name'] in cluster_names_with_launch_request
-    ]
+    updated_records = []
     if len(cluster_names_without_launch_request) > 0:
         with progress:
             updated_records = subprocess_utils.run_in_parallel(
-                _refresh_cluster_record, cluster_names_without_launch_request)
-
+                _refresh_cluster_record, [
+                    cluster_name
+                    for _, cluster_name in cluster_names_without_launch_request
+                ])
+    # Preserve the index of the cluster name as it appears on "records"
+    # before filtering for clusters being launched.
+    updated_records_dict: Dict[int, Optional[Dict[str, Any]]] = {
+        cluster_names_without_launch_request[i][0]: updated_records[i]
+        for i in range(len(cluster_names_without_launch_request))
+    }
     # Show information for removed clusters.
     kept_records = []
     autodown_clusters, remaining_clusters, failed_clusters = [], [], []
     for i, record in enumerate(records):
-        if updated_records[i] is None:
+        if i not in updated_records_dict:
+            # record was not refreshed, keep the original record
+            kept_records.append(record)
+            continue
+        updated_record = updated_records_dict[i]
+        if updated_record is None:
             if record['to_down']:
-                autodown_clusters.append(cluster_names[i])
+                autodown_clusters.append(record['name'])
             else:
-                remaining_clusters.append(cluster_names[i])
-        elif updated_records[i]['status'] == 'UNKNOWN':
-            failed_clusters.append(
-                (cluster_names[i], updated_records[i]['error']))
+                remaining_clusters.append(record['name'])
+        elif updated_record['status'] == 'UNKNOWN':
+            failed_clusters.append((record['name'], updated_record['error']))
             # Keep the original record if the status is unknown,
             # so that the user can still see the cluster.
             kept_records.append(record)
         else:
-            kept_records.append(updated_records[i])
+            kept_records.append(updated_record)
 
     if autodown_clusters:
         plural = 's' if len(autodown_clusters) > 1 else ''