skypilot-nightly 1.0.0.dev20250413__py3-none-any.whl → 1.0.0.dev20250421__py3-none-any.whl

sky/utils/controller_utils.py

@@ -6,7 +6,7 @@ import getpass
 import os
 import tempfile
 import typing
-from typing import Any, Dict, Iterable, List, Optional, Set
+from typing import Any, Dict, Iterable, List, Optional, Set, Tuple
 import uuid
 
 import colorama
@@ -46,7 +46,7 @@ logger = sky_logging.init_logger(__name__)
 # controller resources spec.
 CONTROLLER_RESOURCES_NOT_VALID_MESSAGE = (
     '{controller_type} controller resources is not valid, please check '
-    '~/.sky/skyconfig.yaml file and make sure '
+    '~/.sky/config.yaml file and make sure '
     '{controller_type}.controller.resources is a valid resources spec. '
     'Details:\n  {err}')
 
@@ -72,6 +72,7 @@ class _ControllerSpec:
     default_hint_if_non_existent: str
     connection_error_hint: str
     default_resources_config: Dict[str, Any]
+    default_autostop_config: Dict[str, Any]
 
     @property
     def decline_down_when_failed_to_fetch_status_hint(self) -> str:
@@ -118,7 +119,8 @@ class Controllers(enum.Enum):
         default_hint_if_non_existent='No in-progress managed jobs.',
         connection_error_hint=(
             'Failed to connect to jobs controller, please try again later.'),
-        default_resources_config=managed_job_constants.CONTROLLER_RESOURCES)
+        default_resources_config=managed_job_constants.CONTROLLER_RESOURCES,
+        default_autostop_config=managed_job_constants.CONTROLLER_AUTOSTOP)
     SKY_SERVE_CONTROLLER = _ControllerSpec(
         controller_type='serve',
         name='serve controller',
@@ -148,7 +150,8 @@ class Controllers(enum.Enum):
         default_hint_if_non_existent='No live services.',
         connection_error_hint=(
             'Failed to connect to serve controller, please try again later.'),
-        default_resources_config=serve_constants.CONTROLLER_RESOURCES)
+        default_resources_config=serve_constants.CONTROLLER_RESOURCES,
+        default_autostop_config=serve_constants.CONTROLLER_AUTOSTOP)
 
     @classmethod
     def from_name(cls, name: Optional[str]) -> Optional['Controllers']:
@@ -262,8 +265,9 @@ def _get_cloud_dependencies_installation_commands(
                 '  ARCH="amd64"; '
                 'fi && '
                 '(command -v kubectl &>/dev/null || '
-                '("https://dl.k8s.io/release/v1.31.6/bin/linux/$ARCH/kubectl" '
-                '&& sudo install -o root -g root -m 0755 '
+                '(curl -s -LO "https://dl.k8s.io/release/v1.31.6'
+                '/bin/linux/$ARCH/kubectl" && '
+                'sudo install -o root -g root -m 0755 '
                 'kubectl /usr/local/bin/kubectl))')
         elif isinstance(cloud, clouds.Cudo):
             step_prefix = prefix_str.replace('<step>', str(len(commands) + 1))
@@ -389,7 +393,6 @@ def download_and_stream_latest_job_log(
             f'Failed to stream the logs for the user program at '
             f'{log_file}: {common_utils.format_exception(e)}',
             exc_info=True)
-        # Return the log_file anyway.
 
     return log_file
 
@@ -601,6 +604,40 @@ def get_controller_resources(
     return result
 
 
+def get_controller_autostop_config(
+        controller: Controllers) -> Tuple[Optional[int], bool]:
+    """Get the autostop config for the controller.
+
+    Returns:
+      A tuple of (idle_minutes_to_autostop, down), which correspond to the
+      values passed to execution.launch().
+    """
+    controller_autostop_config_copied: Dict[str, Any] = copy.copy(
+        controller.value.default_autostop_config)
+    if skypilot_config.loaded():
+        custom_controller_autostop_config = skypilot_config.get_nested(
+            (controller.value.controller_type, 'controller', 'autostop'), None)
+        if custom_controller_autostop_config is False:
+            # Disabled with `autostop: false` in config.
+            # To indicate autostop is disabled, we return None for
+            # idle_minutes_to_autostop.
+            return None, False
+        elif custom_controller_autostop_config is True:
+            # Enabled with default values. There is no change in behavior, but
+            # this is included by for completeness, since `False` is valid.
+            pass
+        elif custom_controller_autostop_config is not None:
+            # We have specific config values.
+            # Override the controller autostop config with the ones specified in
+            # the config.
+            assert isinstance(custom_controller_autostop_config, dict)
+            controller_autostop_config_copied.update(
+                custom_controller_autostop_config)
+
+    return (controller_autostop_config_copied['idle_minutes'],
+            controller_autostop_config_copied['down'])
+
+
 def _setup_proxy_command_on_controller(
         controller_launched_cloud: 'clouds.Cloud',
         user_config: Dict[str, Any]) -> config_utils.Config:

sky/utils/kubernetes/generate_kubeconfig.sh

@@ -328,9 +328,9 @@ cp kubeconfig ~/.kube/config
 # Verify that you can access the cluster
 kubectl get pods
 
-Also add this to your ~/.sky/skyconfig.yaml to use the new service account:
+Also add this to your ~/.sky/config.yaml to use the new service account:
 
-# ~/.sky/skyconfig.yaml
+# ~/.sky/config.yaml
 kubernetes:
   remote_identity: ${SKYPILOT_SA}
 "

sky/utils/kubernetes/gpu_labeler.py

@@ -3,8 +3,9 @@ import argparse
 import hashlib
 import os
 import subprocess
-from typing import Optional, Tuple
+from typing import Dict, Optional, Tuple
 
+import colorama
 import yaml
 
 import sky
@@ -13,6 +14,10 @@ from sky.provision.kubernetes import utils as kubernetes_utils
 from sky.utils import rich_utils
 
 
+def _format_string(str_to_format: str, colorama_format: str) -> str:
+    return f'{colorama_format}{str_to_format}{colorama.Style.RESET_ALL}'
+
+
 def cleanup(context: Optional[str] = None) -> Tuple[bool, str]:
     """Deletes all Kubernetes resources created by this script
 
@@ -45,7 +50,7 @@ def get_node_hash(node_name: str):
     return md5_hash[:32]
 
 
-def label(context: Optional[str] = None):
+def label(context: Optional[str] = None, wait_for_completion: bool = True):
     deletion_success, reason = cleanup(context=context)
     if not deletion_success:
         print(reason)
@@ -60,8 +65,10 @@ def label(context: Optional[str] = None):
               'in their capacity.')
         return
 
-    print(f'Found {len(unlabeled_gpu_nodes)} '
-          'unlabeled GPU nodes in the cluster')
+    print(
+        _format_string(
+            f'Found {len(unlabeled_gpu_nodes)} '
+            'unlabeled GPU nodes in the cluster', colorama.Fore.YELLOW))
 
     sky_dir = os.path.dirname(sky.__file__)
     manifest_dir = os.path.join(sky_dir, 'utils/kubernetes')
@@ -80,6 +87,7 @@ def label(context: Optional[str] = None):
             print('Error setting up GPU labeling: ' + output)
             return
 
+    jobs_to_node_names: Dict[str, str] = {}
     with rich_utils.client_status('Creating GPU labeler jobs'):
         batch_v1 = kubernetes.batch_api(context=context)
         # Load the job manifest
@@ -113,8 +121,11 @@ def label(context: Optional[str] = None):
             node_name = node.metadata.name
 
             # Modify the job manifest for the current node
-            job_manifest['metadata']['name'] = ('sky-gpu-labeler-'
-                                                f'{get_node_hash(node_name)}')
+            job_name = ('sky-gpu-labeler-'
+                        f'{get_node_hash(node_name)}')
+            jobs_to_node_names[job_name] = node_name
+            job_manifest['metadata']['name'] = job_name
+
             job_manifest['spec']['template']['spec']['nodeSelector'] = {
                 'kubernetes.io/hostname': node_name
             }
@@ -122,17 +133,85 @@ def label(context: Optional[str] = None):
 
             # Create the job for this node`
             batch_v1.create_namespaced_job(namespace, job_manifest)
-            print(f'Created GPU labeler job for node {node_name}')
+            print(
+                _format_string(f'Created GPU labeler job for node {node_name}',
+                               colorama.Style.DIM))
 
     context_str = f' --context {context}' if context else ''
-    print(f'GPU labeling started - this may take 10 min or more to complete.'
-          '\nTo check the status of GPU labeling jobs, run '
-          f'`kubectl get jobs -n kube-system '
-          f'-l job=sky-gpu-labeler{context_str}`'
-          '\nYou can check if nodes have been labeled by running '
-          f'`kubectl describe nodes{context_str}` '
-          'and looking for labels of the format '
-          '`skypilot.co/accelerator: <gpu_name>`. ')
+
+    if wait_for_completion:
+        # Wait for the job to complete
+        with rich_utils.client_status(
+                'Waiting for GPU labeler jobs to complete'):
+            success = wait_for_jobs_completion(jobs_to_node_names,
+                                               'kube-system',
+                                               context=context)
+        if success:
+            print(
+                _format_string('✅ GPU labeling completed successfully',
+                               colorama.Fore.GREEN))
+        else:
+            print(_format_string('❌ GPU labeling failed', colorama.Fore.RED))
+        cleanup(context=context)
+    else:
+        print(
+            f'GPU labeling started - this may take 10 min or more to complete.'
+            '\nTo check the status of GPU labeling jobs, run '
+            f'`kubectl get jobs -n kube-system '
+            f'-l job=sky-gpu-labeler{context_str}`'
+            '\nYou can check if nodes have been labeled by running '
+            f'`kubectl describe nodes{context_str}` '
+            'and looking for labels of the format '
+            '`skypilot.co/accelerator: <gpu_name>`. ')
+
+
+def wait_for_jobs_completion(jobs_to_node_names: Dict[str, str],
+                             namespace: str,
+                             context: Optional[str] = None,
+                             timeout: int = 60 * 20):
+    """Waits for a Kubernetes Job to complete or fail.
+
+    Args:
+        jobs_to_node_names: A dictionary mapping job names to node names.
+        namespace: The namespace the Job is in (default: "default").
+        timeout: Timeout in seconds (default: 1200 seconds = 20 minutes).
+
+    Returns:
+        True if the Job completed successfully, False if it failed or timed out.
+    """
+    batch_v1 = kubernetes.batch_api(context=context)
+    w = kubernetes.watch()
+    completed_jobs = []
+    for event in w.stream(func=batch_v1.list_namespaced_job,
+                          namespace=namespace,
+                          timeout_seconds=timeout):
+        job = event['object']
+        job_name = job.metadata.name
+        if job_name in jobs_to_node_names:
+            node_name = jobs_to_node_names[job_name]
+            if job.status and job.status.completion_time:
+                print(
+                    _format_string(
+                        f'GPU labeler job for node {node_name} '
+                        'completed successfully', colorama.Style.DIM))
+                completed_jobs.append(job_name)
+                num_remaining_jobs = len(jobs_to_node_names) - len(
+                    completed_jobs)
+                if num_remaining_jobs == 0:
+                    w.stop()
+                    return True
+            elif job.status and job.status.failed:
+                print(
+                    _format_string(
+                        f'GPU labeler job for node {node_name} failed',
+                        colorama.Style.DIM))
+                w.stop()
+                return False
+    print(
+        _format_string(
+            f'Timed out after waiting {timeout} seconds '
+            'for job to complete', colorama.Style.DIM))
+    return False  #Timed out
 
 
 def main():
@@ -151,6 +230,10 @@ def main():
     parser.add_argument('--context',
                         type=str,
                         help='the context to use for the Kubernetes cluster.')
+    parser.add_argument('--async',
+                        dest='async_completion',
+                        action='store_true',
+                        help='do not wait for the GPU labeling to complete.')
     args = parser.parse_args()
     context = None
     if args.context:
@@ -165,7 +248,7 @@ def main():
     if args.cleanup:
         cleanup(context=context)
     else:
-        label(context=context)
+        label(context=context, wait_for_completion=not args.async_completion)
 
 
 if __name__ == '__main__':

sky/utils/schemas.py

@@ -725,6 +725,29 @@ def get_config_schema():
         if k != '$schema'
     }
     resources_schema['properties'].pop('ports')
+    autostop_schema = {
+        'anyOf': [
+            {
+                # Use boolean to disable autostop completely, e.g.
+                #   autostop: false
+                'type': 'boolean',
+            },
+            {
+                'type': 'object',
+                'required': [],
+                'additionalProperties': False,
+                'properties': {
+                    'idle_minutes': {
+                        'type': 'integer',
+                        'minimum': 0,
+                    },
+                    'down': {
+                        'type': 'boolean',
+                    },
+                },
+            },
+        ],
+    }
     controller_resources_schema = {
         'type': 'object',
         'required': [],
@@ -736,6 +759,7 @@ def get_config_schema():
                 'additionalProperties': False,
                 'properties': {
                     'resources': resources_schema,
+                    'autostop': autostop_schema,
                 }
             },
             'bucket': {

{skypilot_nightly-1.0.0.dev20250413.dist-info → skypilot_nightly-1.0.0.dev20250421.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: skypilot-nightly
-Version: 1.0.0.dev20250413
+Version: 1.0.0.dev20250421
 Summary: SkyPilot: An intercloud broker for the clouds
 Author: SkyPilot Team
 License: Apache 2.0
@@ -47,6 +47,7 @@ Requires-Dist: python-multipart
 Requires-Dist: aiofiles
 Requires-Dist: httpx
 Requires-Dist: setproctitle
+Requires-Dist: omegaconf<2.5,>=2.4.0dev3
 Provides-Extra: aws
 Requires-Dist: urllib3<2; extra == "aws"
 Requires-Dist: awscli>=1.27.10; extra == "aws"