skypilot-nightly 1.0.0.dev20250427__py3-none-any.whl → 1.0.0.dev20250429__py3-none-any.whl

sky/__init__.py

@@ -5,7 +5,7 @@ from typing import Optional
 import urllib.request
 
 # Replaced with the current commit when building the wheels.
-_SKYPILOT_COMMIT_SHA = '966b4b5f431d8912d7fecd151b5e069bca0b9f13'
+_SKYPILOT_COMMIT_SHA = '7b804dafe2f6b775f8a357ac6e147b83e792af93'
 
 
 def _get_git_commit():
@@ -35,7 +35,7 @@ def _get_git_commit():
 
 
 __commit__ = _get_git_commit()
-__version__ = '1.0.0.dev20250427'
+__version__ = '1.0.0.dev20250429'
 __root_dir__ = os.path.dirname(os.path.abspath(__file__))
 
 

sky/adaptors/nebius.py

@@ -29,11 +29,6 @@ MAX_RETRIES_TO_INSTANCE_WAIT = 120  # Maximum number of retries
 
 POLL_INTERVAL = 5
 
-_iam_token = None
-_sdk = None
-_tenant_id = None
-_project_id = None
-
 _IMPORT_ERROR_MESSAGE = ('Failed to import dependencies for Nebius AI Cloud.'
                          'Try pip install "skypilot[nebius]"')
 
@@ -81,56 +76,49 @@ def vpc():
     return vpc_v1
 
 
+@annotations.lru_cache(scope='request')
 def get_iam_token():
-    global _iam_token
-    if _iam_token is None:
-        try:
-            with open(os.path.expanduser(NEBIUS_IAM_TOKEN_PATH),
-                      encoding='utf-8') as file:
-                _iam_token = file.read().strip()
-        except FileNotFoundError:
-            return None
-    return _iam_token
+    try:
+        with open(os.path.expanduser(NEBIUS_IAM_TOKEN_PATH),
+                  encoding='utf-8') as file:
+            return file.read().strip()
+    except FileNotFoundError:
+        return None
 
 
+@annotations.lru_cache(scope='request')
 def is_token_or_cred_file_exist():
     return (os.path.exists(os.path.expanduser(NEBIUS_IAM_TOKEN_PATH)) or
             os.path.exists(os.path.expanduser(NEBIUS_CREDENTIALS_PATH)))
 
 
+@annotations.lru_cache(scope='request')
 def get_project_id():
-    global _project_id
-    if _project_id is None:
-        try:
-            with open(os.path.expanduser(NEBIUS_PROJECT_ID_PATH),
-                      encoding='utf-8') as file:
-                _project_id = file.read().strip()
-        except FileNotFoundError:
-            return None
-    return _project_id
+    try:
+        with open(os.path.expanduser(NEBIUS_PROJECT_ID_PATH),
+                  encoding='utf-8') as file:
+            return file.read().strip()
+    except FileNotFoundError:
+        return None
 
 
+@annotations.lru_cache(scope='request')
 def get_tenant_id():
-    global _tenant_id
-    if _tenant_id is None:
-        try:
-            with open(os.path.expanduser(NEBIUS_TENANT_ID_PATH),
-                      encoding='utf-8') as file:
-                _tenant_id = file.read().strip()
-        except FileNotFoundError:
-            return None
-    return _tenant_id
+    try:
+        with open(os.path.expanduser(NEBIUS_TENANT_ID_PATH),
+                  encoding='utf-8') as file:
+            return file.read().strip()
+    except FileNotFoundError:
+        return None
 
 
+@annotations.lru_cache(scope='request')
 def sdk():
-    global _sdk
-    if _sdk is None:
-        if get_iam_token() is not None:
-            _sdk = nebius.sdk.SDK(credentials=get_iam_token())
-            return _sdk
-        _sdk = nebius.sdk.SDK(
-            credentials_file_name=os.path.expanduser(NEBIUS_CREDENTIALS_PATH))
-    return _sdk
+    token = get_iam_token()
+    if token is not None:
+        return nebius.sdk.SDK(credentials=token)
+    return nebius.sdk.SDK(
+        credentials_file_name=os.path.expanduser(NEBIUS_CREDENTIALS_PATH))
 
 
 def get_nebius_credentials(boto3_session):

sky/backends/backend_utils.py

@@ -179,6 +179,9 @@ _RAY_YAML_KEYS_TO_RESTORE_EXCEPTIONS = [
     ('available_node_types', 'ray.head.default', 'node_config', 'UserData'),
     ('available_node_types', 'ray.head.default', 'node_config',
      'azure_arm_parameters', 'cloudInitSetupCommands'),
+    ('available_node_types', 'ray_head_default', 'node_config', 'pvc_spec'),
+    ('available_node_types', 'ray_head_default', 'node_config',
+     'deployment_spec'),
 ]
 # These keys are expected to change when provisioning on an existing cluster,
 # but they don't actually represent a change that requires re-provisioning the
@@ -705,6 +708,13 @@ def write_cluster_config(
     is_custom_docker = ('true' if to_provision.extract_docker_image()
                         is not None else 'false')
 
+    # Here, if users specify the controller to be high availability, we will
+    # provision a high availability controller. Whether the cloud supports
+    # this feature has been checked by
+    # CloudImplementationFeatures.HIGH_AVAILABILITY_CONTROLLERS
+    high_availability_specified = controller_utils.high_availability_specified(
+        cluster_name_on_cloud)
+
     # Use a tmp file path to avoid incomplete YAML file being re-used in the
     # future.
     tmp_yaml_path = yaml_path + '.tmp'
@@ -790,6 +800,9 @@ def write_cluster_config(
                 'sky_wheel_hash': wheel_hash,
                 # Authentication (optional).
                 **auth_config,
+
+                # High availability
+                'high_availability': high_availability_specified,
             }),
         output_path=tmp_yaml_path)
     config_dict['cluster_name'] = cluster_name
@@ -802,8 +815,12 @@ def write_cluster_config(
             cluster_config_overrides=to_provision.cluster_config_overrides)
         kubernetes_utils.combine_metadata_fields(tmp_yaml_path)
         yaml_obj = common_utils.read_yaml(tmp_yaml_path)
-        pod_config = yaml_obj['available_node_types']['ray_head_default'][
-            'node_config']
+        pod_config: Dict[str, Any] = yaml_obj['available_node_types'][
+            'ray_head_default']['node_config']
+
+        # Check pod spec only. For high availability controllers, we deploy pvc & deployment for the controller. Read kubernetes-ray.yml.j2 for more details.
+        pod_config.pop('deployment_spec', None)
+        pod_config.pop('pvc_spec', None)
         valid, message = kubernetes_utils.check_pod_config(pod_config)
         if not valid:
             raise exceptions.InvalidCloudConfigs(

sky/backends/cloud_vm_ray_backend.py

@@ -1583,6 +1583,10 @@ class RetryingVmProvisioner(object):
                 except provision_common.StopFailoverError:
                     with ux_utils.print_exception_no_traceback():
                         raise
+                except exceptions.InconsistentHighAvailabilityError:
+                    # No teardown happens for this error.
+                    with ux_utils.print_exception_no_traceback():
+                        raise
                 except Exception as e:  # pylint: disable=broad-except
                     # NOTE: We try to cleanup the cluster even if the previous
                     # cluster does not exist. Also we are fast at
@@ -2032,6 +2036,7 @@ class RetryingVmProvisioner(object):
                 # Recheck cluster name as the 'except:' block below may
                 # change the cloud assignment.
                 common_utils.check_cluster_name_is_valid(cluster_name)
+
                 if dryrun:
                     cloud_user = None
                 else:
@@ -2459,6 +2464,7 @@ class CloudVmRayResourceHandle(backends.backend.ResourceHandle):
                     'Tried to use cached cluster info, but it\'s missing for '
                     f'cluster "{self.cluster_name}"')
             self._update_cluster_info()
+
         assert self.cached_cluster_info is not None, self
         runners = provision_lib.get_command_runners(
             self.cached_cluster_info.provider_name, self.cached_cluster_info,
@@ -2689,6 +2695,7 @@ class CloudVmRayBackend(backends.Backend['CloudVmRayResourceHandle']):
             self._optimize_target) or common.OptimizeTarget.COST
         self._requested_features = kwargs.pop('requested_features',
                                               self._requested_features)
+        self._dump_final_script = kwargs.pop('dump_final_script', False)
         assert not kwargs, f'Unexpected kwargs: {kwargs}'
 
     def check_resources_fit_cluster(
@@ -3272,18 +3279,30 @@ class CloudVmRayBackend(backends.Backend['CloudVmRayResourceHandle']):
                                                          env_vars=setup_envs)
             encoded_script = shlex.quote(setup_script)
 
-            def _dump_setup_script(setup_script: str) -> None:
+            def _dump_final_script(
+                    setup_script: str,
+                    target_dir: str = remote_setup_file_name) -> None:
                 with tempfile.NamedTemporaryFile('w', prefix='sky_setup_') as f:
                     f.write(setup_script)
                     f.flush()
                     setup_sh_path = f.name
                     runner.rsync(source=setup_sh_path,
-                                 target=remote_setup_file_name,
+                                 target=target_dir,
                                  up=True,
                                  stream_logs=False)
 
+            # Always dump the full setup script to the persistent path first
+            # In high availability mode, we need to dump the full setup script
+            # to a persistent path BEFORE any other operations. This ensures
+            # that if the pod restarts, it can find and execute the complete
+            # setup script, rather than a reference to a temporary file that
+            # would no longer exist after restart.
+            if self._dump_final_script:
+                _dump_final_script(setup_script,
+                                   constants.PERSISTENT_SETUP_SCRIPT_PATH)
+
             if detach_setup or _is_command_length_over_limit(encoded_script):
-                _dump_setup_script(setup_script)
+                _dump_final_script(setup_script)
                 create_script_code = 'true'
             else:
                 create_script_code = (f'{{ echo {encoded_script} > '
@@ -3335,7 +3354,7 @@ class CloudVmRayBackend(backends.Backend['CloudVmRayResourceHandle']):
                         'Failed to run setup command inline due to '
                         'command length limit. Dumping setup script to '
                         'file and running it with SSH.')
-                    _dump_setup_script(setup_script)
+                    _dump_final_script(setup_script)
                     returncode = _run_setup(setup_cmd)
 
             def error_message() -> str:
@@ -3426,14 +3445,14 @@ class CloudVmRayBackend(backends.Backend['CloudVmRayResourceHandle']):
         code = job_lib.JobLibCodeGen.queue_job(job_id, job_submit_cmd)
         job_submit_cmd = ' && '.join([mkdir_code, create_script_code, code])
 
-        def _dump_code_to_file(codegen: str) -> None:
+        def _dump_code_to_file(codegen: str,
+                               target_dir: str = SKY_REMOTE_APP_DIR) -> None:
             runners = handle.get_command_runners()
             head_runner = runners[0]
             with tempfile.NamedTemporaryFile('w', prefix='sky_app_') as fp:
                 fp.write(codegen)
                 fp.flush()
-                script_path = os.path.join(SKY_REMOTE_APP_DIR,
-                                           f'sky_job_{job_id}')
+                script_path = os.path.join(target_dir, f'sky_job_{job_id}')
                 # We choose to sync code + exec, because the alternative of 'ray
                 # submit' may not work as it may use system python (python2) to
                 # execute the script. Happens for AWS.
@@ -3442,6 +3461,12 @@ class CloudVmRayBackend(backends.Backend['CloudVmRayResourceHandle']):
                                   up=True,
                                   stream_logs=False)
 
+        # Should also be ealier than _is_command_length_over_limit
+        # Same reason as in _setup
+        if self._dump_final_script:
+            _dump_code_to_file(job_submit_cmd,
+                               constants.PERSISTENT_RUN_SCRIPT_DIR)
+
         if _is_command_length_over_limit(job_submit_cmd):
             _dump_code_to_file(codegen)
             job_submit_cmd = f'{mkdir_code} && {code}'
@@ -3457,7 +3482,7 @@ class CloudVmRayBackend(backends.Backend['CloudVmRayResourceHandle']):
             # We cannot set the managed job to PENDING state in the job template
             # (jobs-controller.yaml.j2), as it may need to wait for the run
             # commands to be scheduled on the job controller in high-load cases.
-            job_submit_cmd = job_submit_cmd + ' && ' + managed_job_code
+            job_submit_cmd += ' && ' + managed_job_code
 
         returncode, stdout, stderr = self.run_on_head(handle,
                                                       job_submit_cmd,

sky/backends/local_docker_backend.py

@@ -276,7 +276,6 @@ class LocalDockerBackend(backends.Backend['LocalDockerResourceHandle']):
                  detach_run: bool,
                  dryrun: bool = False) -> None:
         """ Launches the container."""
-
         if detach_run:
             raise NotImplementedError('detach_run=True is not supported in '
                                       'LocalDockerBackend.')
@@ -364,7 +363,7 @@ class LocalDockerBackend(backends.Backend['LocalDockerResourceHandle']):
                 if k.startswith(_DOCKER_LABEL_PREFIX):
                     # Remove 'skymeta_' from key
                     metadata[k[len(_DOCKER_LABEL_PREFIX):]] = v
-            self.images[c.name] = [c.image, metadata]
+            self.images[c.name] = (c.image, metadata)
             self.containers[c.name] = c
 
     def _execute_task_one_node(self, handle: LocalDockerResourceHandle,

sky/cli.py

@@ -23,6 +23,7 @@ NOTE: the order of command definitions in this file corresponds to how they are
 listed in "sky --help".  Take care to put logically connected commands close to
 each other.
 """
+import collections
 import copy
 import datetime
 import functools
@@ -162,7 +163,7 @@ def _get_cluster_records_and_set_ssh_config(
                 '-o StrictHostKeyChecking=no '
                 '-o UserKnownHostsFile=/dev/null '
                 '-o IdentitiesOnly=yes '
-                '-W %h:%p '
+                '-W \'[%h]:%p\' '
                 f'{handle.ssh_user}@127.0.0.1 '
                 '-o ProxyCommand='
                 # TODO(zhwu): write the template to a temp file, don't use
@@ -3413,7 +3414,7 @@ def show_gpus(
 
     # TODO(zhwu,romilb): We should move most of these kubernetes related
     # queries into the backend, especially behind the server.
-    def _get_kubernetes_realtime_gpu_table(
+    def _get_kubernetes_realtime_gpu_tables(
             context: Optional[str] = None,
             name_filter: Optional[str] = None,
             quantity_filter: Optional[int] = None):
@@ -3423,15 +3424,14 @@ def show_gpus(
         else:
             qty_header = 'REQUESTABLE_QTY_PER_NODE'
             free_header = 'TOTAL_FREE_GPUS'
-        realtime_gpu_table = log_utils.create_table(
-            ['GPU', qty_header, 'TOTAL_GPUS', free_header])
-        realtime_gpu_availability_list = sdk.stream_and_get(
+
+        realtime_gpu_availability_lists = sdk.stream_and_get(
             sdk.realtime_kubernetes_gpu_availability(
                 context=context,
                 name_filter=name_filter,
                 quantity_filter=quantity_filter))
-        if not realtime_gpu_availability_list:
-            err_msg = 'No GPUs found in Kubernetes cluster. '
+        if not realtime_gpu_availability_lists:
+            err_msg = 'No GPUs found in any allowed Kubernetes cluster. '
             debug_msg = 'To further debug, run: sky check '
             if name_filter is not None:
                 gpu_info_msg = f' {name_filter!r}'
@@ -3439,26 +3439,52 @@ def show_gpus(
                     gpu_info_msg += (' with requested quantity'
                                      f' {quantity_filter}')
                 err_msg = (f'Resources{gpu_info_msg} not found '
-                           'in Kubernetes cluster. ')
+                           'in any allowed Kubernetes cluster. ')
                 debug_msg = ('To show available accelerators on kubernetes,'
                              ' run: sky show-gpus --cloud kubernetes ')
             full_err_msg = (err_msg + kubernetes_constants.NO_GPU_HELP_MESSAGE +
                             debug_msg)
             raise ValueError(full_err_msg)
         no_permissions_str = '<no permissions>'
-        for realtime_gpu_availability in sorted(realtime_gpu_availability_list):
-            gpu_availability = models.RealtimeGpuAvailability(
-                *realtime_gpu_availability)
-            available_qty = (gpu_availability.available
-                             if gpu_availability.available != -1 else
-                             no_permissions_str)
-            realtime_gpu_table.add_row([
-                gpu_availability.gpu,
-                _list_to_str(gpu_availability.counts),
-                gpu_availability.capacity,
-                available_qty,
-            ])
-        return realtime_gpu_table
+        realtime_gpu_infos = []
+        total_gpu_info: Dict[str, List[int]] = collections.defaultdict(
+            lambda: [0, 0])
+
+        for (ctx, availability_list) in realtime_gpu_availability_lists:
+            realtime_gpu_table = log_utils.create_table(
+                ['GPU', qty_header, 'TOTAL_GPUS', free_header])
+            for realtime_gpu_availability in sorted(availability_list):
+                gpu_availability = models.RealtimeGpuAvailability(
+                    *realtime_gpu_availability)
+                available_qty = (gpu_availability.available
+                                 if gpu_availability.available != -1 else
+                                 no_permissions_str)
+                realtime_gpu_table.add_row([
+                    gpu_availability.gpu,
+                    _list_to_str(gpu_availability.counts),
+                    gpu_availability.capacity,
+                    available_qty,
+                ])
+                gpu = gpu_availability.gpu
+                capacity = gpu_availability.capacity
+                # we want total, so skip permission denied.
+                available = max(gpu_availability.available, 0)
+                if capacity > 0:
+                    total_gpu_info[gpu][0] += capacity
+                    total_gpu_info[gpu][1] += available
+            realtime_gpu_infos.append((ctx, realtime_gpu_table))
+
+        # display an aggregated table for all contexts
+        # if there are more than one contexts with GPUs
+        if len(realtime_gpu_infos) > 1:
+            total_realtime_gpu_table = log_utils.create_table(
+                ['GPU', 'TOTAL_GPUS', free_header])
+            for gpu, stats in total_gpu_info.items():
+                total_realtime_gpu_table.add_row([gpu, stats[0], stats[1]])
+        else:
+            total_realtime_gpu_table = None
+
+        return realtime_gpu_infos, total_realtime_gpu_table
 
     def _format_kubernetes_node_info(context: Optional[str]):
         node_table = log_utils.create_table(
@@ -3479,7 +3505,7 @@ def show_gpus(
             'Kubernetes per node accelerator availability ')
         if nodes_info.hint:
             k8s_per_node_acc_message += nodes_info.hint
-        return (f'{colorama.Fore.CYAN}{colorama.Style.BRIGHT}'
+        return (f'{colorama.Fore.LIGHTMAGENTA_EX}{colorama.Style.NORMAL}'
                 f'{k8s_per_node_acc_message}'
                 f'{colorama.Style.RESET_ALL}\n'
                 f'{node_table.get_string()}')
@@ -3516,8 +3542,7 @@ def show_gpus(
                     # If --cloud kubernetes is not specified, we want to catch
                     # the case where no GPUs are available on the cluster and
                     # print the warning at the end.
-                    k8s_realtime_table = _get_kubernetes_realtime_gpu_table(
-                        context)
+                    k8s_realtime_infos, total_table = _get_kubernetes_realtime_gpu_tables(context)  # pylint: disable=line-too-long
                 except ValueError as e:
                     if not cloud_is_kubernetes:
                         # Make it a note if cloud is not kubernetes
@@ -3525,13 +3550,24 @@ def show_gpus(
                     k8s_messages += str(e)
                 else:
                     print_section_titles = True
-                    context_str = f'(Context: {context})' if context else ''
-                    yield (f'{colorama.Fore.CYAN}{colorama.Style.BRIGHT}'
-                           f'Kubernetes GPUs {context_str}'
-                           f'{colorama.Style.RESET_ALL}\n')
-                    yield from k8s_realtime_table.get_string()
-                    yield '\n\n'
-                    yield _format_kubernetes_node_info(context)
+
+                    # print total table
+                    if total_table is not None:
+                        yield (f'{colorama.Fore.GREEN}{colorama.Style.BRIGHT}'
+                               'Total Kubernetes GPUs'
+                               f'{colorama.Style.RESET_ALL}\n')
+                        yield from total_table.get_string()
+                        yield '\n-----\n\n'
+
+                    # print individual infos.
+                    for (ctx, k8s_realtime_table) in k8s_realtime_infos:
+                        context_str = f'(Context: {ctx})' if ctx else ''
+                        yield (f'{colorama.Fore.CYAN}{colorama.Style.BRIGHT}'
+                               f'Kubernetes GPUs {context_str}'
+                               f'{colorama.Style.RESET_ALL}\n')
+                        yield from k8s_realtime_table.get_string()
+                        yield '\n\n'
+                        yield _format_kubernetes_node_info(ctx) + '\n-----\n\n'
                 if kubernetes_autoscaling:
                     k8s_messages += (
                         '\n' + kubernetes_utils.KUBERNETES_AUTOSCALER_NOTE)
@@ -3620,13 +3656,29 @@ def show_gpus(
             # Print section title if not showing all and instead a specific
             # accelerator is requested
             print_section_titles = True
-            yield (f'{colorama.Fore.CYAN}{colorama.Style.BRIGHT}'
-                   f'Kubernetes GPUs{colorama.Style.RESET_ALL}\n')
             # TODO(romilb): Show filtered per node GPU availability here as well
             try:
-                k8s_realtime_table = _get_kubernetes_realtime_gpu_table(
-                    name_filter=name, quantity_filter=quantity)
-                yield from k8s_realtime_table.get_string()
+                k8s_realtime_infos, total_table = _get_kubernetes_realtime_gpu_tables(  # pylint: disable=line-too-long
+                    context=region,
+                    name_filter=name,
+                    quantity_filter=quantity)
+
+                # print total table
+                if total_table is not None:
+                    yield (f'{colorama.Fore.GREEN}{colorama.Style.BRIGHT}'
+                           'Total Kubernetes GPUs'
+                           f'{colorama.Style.RESET_ALL}\n')
+                    yield from total_table.get_string()
+                    yield '\n-----\n\n'
+
+                # print individual tables
+                for (ctx, k8s_realtime_table) in k8s_realtime_infos:
+                    context_str = f'(Context: {ctx})' if ctx else ''
+                    yield (f'{colorama.Fore.CYAN}{colorama.Style.BRIGHT}'
+                           f'Kubernetes GPUs {context_str}'
+                           f'{colorama.Style.RESET_ALL}\n')
+                    yield from k8s_realtime_table.get_string()
+                    yield '\n\n'
             except ValueError as e:
                 # In the case of a specific accelerator, show the error message
                 # immediately (e.g., "Resources H100 not found ...")
@@ -5911,11 +5963,12 @@ def api_info():
     user_name = os.getenv(constants.USER_ENV_VAR, getpass.getuser())
     user_hash = common_utils.get_user_hash()
     dashboard_url = server_common.get_dashboard_url(url)
-    click.echo(f'Using SkyPilot API server: {url} Dashboard: {dashboard_url}\n'
+    click.echo(f'Using SkyPilot API server: {url}\n'
                f'{ux_utils.INDENT_SYMBOL}Status: {api_server_info["status"]}, '
                f'commit: {api_server_info["commit"]}, '
                f'version: {api_server_info["version"]}\n'
-               f'{ux_utils.INDENT_LAST_SYMBOL}User: {user_name} ({user_hash})')
+               f'{ux_utils.INDENT_SYMBOL}User: {user_name} ({user_hash})\n'
+               f'{ux_utils.INDENT_LAST_SYMBOL}Dashboard: {dashboard_url}')
 
 
 def main():