skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/utils/schemas.py

@@ -4,6 +4,9 @@ Schemas conform to the JSON Schema specification as defined at
 https://json-schema.org/
 """
 import enum
+from typing import Any, Dict, List, Tuple
+
+from sky.skylet import constants
 
 
 def _check_not_both_fields_present(field1: str, field2: str):
@@ -83,13 +86,28 @@ def _get_single_resources_schema():
             'use_spot': {
                 'type': 'boolean',
             },
-            # Deprecated: use 'job_recovery' instead. This is for backward
-            # compatibility, and can be removed in 0.8.0.
-            'spot_recovery': {
-                'type': 'string',
-            },
             'job_recovery': {
-                'type': 'string',
+                # Either a string or a dict.
+                'anyOf': [{
+                    'type': 'string',
+                }, {
+                    'type': 'object',
+                    'required': [],
+                    'additionalProperties': False,
+                    'properties': {
+                        'strategy': {
+                            'anyOf': [{
+                                'type': 'string',
+                            }, {
+                                'type': 'null',
+                            }],
+                        },
+                        'max_restarts_on_errors': {
+                            'type': 'integer',
+                            'minimum': 0,
+                        },
+                    }
+                }],
             },
             'disk_size': {
                 'type': 'integer',
@@ -111,6 +129,8 @@ def _get_single_resources_schema():
                             'type': 'integer',
                         }]
                     }
+                }, {
+                    'type': 'null',
                 }],
             },
             'labels': {
@@ -145,7 +165,8 @@ def _get_single_resources_schema():
                     'type': 'null',
                 }]
             },
-            # The following fields are for internal use only.
+            # The following fields are for internal use only. Should not be
+            # specified in the task config.
             '_docker_login_config': {
                 'type': 'object',
                 'required': ['username', 'password', 'server'],
@@ -168,6 +189,9 @@ def _get_single_resources_schema():
             '_requires_fuse': {
                 'type': 'boolean',
             },
+            '_cluster_config_overrides': {
+                'type': 'object',
+            },
         }
     }
 
@@ -227,8 +251,6 @@ def get_resources_schema():
                 'items': multi_resources_schema,
             }
         },
-        # Avoid job_recovery and spot_recovery being present at the same time.
-        **_check_not_both_fields_present('job_recovery', 'spot_recovery')
     }
 
 
@@ -270,6 +292,12 @@ def get_storage_schema():
                     mode.value for mode in storage.StorageMode
                 ]
             },
+            '_is_sky_managed': {
+                'type': 'boolean',
+            },
+            '_bucket_sub_path': {
+                'type': 'string',
+            },
             '_force_delete': {
                 'type': 'boolean',
             }
@@ -279,6 +307,9 @@ def get_storage_schema():
 
 def get_service_schema():
     """Schema for top-level `service:` field (for SkyServe)."""
+    # To avoid circular imports, only import when needed.
+    # pylint: disable=import-outside-toplevel
+    from sky.serve import load_balancing_policies
     return {
         '$schema': 'https://json-schema.org/draft/2020-12/schema',
         'type': 'object',
@@ -299,6 +330,9 @@ def get_service_schema():
                         'initial_delay_seconds': {
                             'type': 'number',
                         },
+                        'timeout_seconds': {
+                            'type': 'number',
+                        },
                         'post_data': {
                             'anyOf': [{
                                 'type': 'string',
@@ -345,24 +379,100 @@ def get_service_schema():
                     'downscale_delay_seconds': {
                         'type': 'number',
                     },
-                    # TODO(MaoZiming): Fields `qps_upper_threshold`,
-                    # `qps_lower_threshold` and `auto_restart` are deprecated.
-                    # Temporarily keep these fields for backward compatibility.
-                    # Remove after 2 minor release, i.e., 0.6.0.
-                    'auto_restart': {
-                        'type': 'boolean',
-                    },
-                    'qps_upper_threshold': {
-                        'type': 'number',
-                    },
-                    'qps_lower_threshold': {
-                        'type': 'number',
-                    },
                 }
             },
+            'ports': {
+                'type': 'integer',
+            },
             'replicas': {
                 'type': 'integer',
             },
+            'load_balancing_policy': {
+                'type': 'string',
+                'case_insensitive_enum': list(
+                    load_balancing_policies.LB_POLICIES.keys())
+            },
+            'tls': {
+                'type': 'object',
+                'required': ['keyfile', 'certfile'],
+                'additionalProperties': False,
+                'properties': {
+                    'keyfile': {
+                        'type': 'string',
+                    },
+                    'certfile': {
+                        'type': 'string',
+                    },
+                },
+            },
+        }
+    }
+
+
+def _filter_schema(schema: dict, keys_to_keep: List[Tuple[str, ...]]) -> dict:
+    """Recursively filter a schema to include only certain keys.
+
+    Args:
+        schema: The original schema dictionary.
+        keys_to_keep: List of tuples with the path of keys to retain.
+
+    Returns:
+        The filtered schema.
+    """
+    # Convert list of tuples to a dictionary for easier access
+    paths_dict: Dict[str, Any] = {}
+    for path in keys_to_keep:
+        current = paths_dict
+        for step in path:
+            if step not in current:
+                current[step] = {}
+            current = current[step]
+
+    def keep_keys(current_schema: dict, current_path_dict: dict,
+                  new_schema: dict) -> dict:
+        # Base case: if we reach a leaf in the path_dict, we stop.
+        if (not current_path_dict or not isinstance(current_schema, dict) or
+                not current_schema.get('properties')):
+            return current_schema
+
+        if 'properties' not in new_schema:
+            new_schema = {
+                key: current_schema[key]
+                for key in current_schema
+                # We do not support the handling of `oneOf`, `anyOf`, `allOf`,
+                # `required` for now.
+                if key not in
+                {'properties', 'oneOf', 'anyOf', 'allOf', 'required'}
+            }
+            new_schema['properties'] = {}
+        for key, sub_schema in current_schema['properties'].items():
+            if key in current_path_dict:
+                # Recursively keep keys if further path dict exists
+                new_schema['properties'][key] = {}
+                current_path_value = current_path_dict.pop(key)
+                new_schema['properties'][key] = keep_keys(
+                    sub_schema, current_path_value,
+                    new_schema['properties'][key])
+
+        return new_schema
+
+    # Start the recursive filtering
+    new_schema = keep_keys(schema, paths_dict, {})
+    assert not paths_dict, f'Unprocessed keys: {paths_dict}'
+    return new_schema
+
+
+def _experimental_task_schema() -> dict:
+    config_override_schema = _filter_schema(
+        get_config_schema(), constants.OVERRIDEABLE_CONFIG_KEYS_IN_TASK)
+    return {
+        'experimental': {
+            'type': 'object',
+            'required': [],
+            'additionalProperties': False,
+            'properties': {
+                'config_overrides': config_override_schema,
+            }
         }
     }
 
@@ -432,6 +542,10 @@ def get_task_schema():
                     'type': 'number'
                 }
             },
+            'file_mounts_mapping': {
+                'type': 'object',
+            },
+            **_experimental_task_schema(),
         }
     }
 
@@ -513,15 +627,6 @@ _NETWORK_CONFIG_SCHEMA = {
 }
 
 _LABELS_SCHEMA = {
-    # Deprecated: 'instance_tags' is replaced by 'labels'. Keeping for backward
-    # compatibility. Will be removed after 0.7.0.
-    'instance_tags': {
-        'type': 'object',
-        'required': [],
-        'additionalProperties': {
-            'type': 'string',
-        },
-    },
     'labels': {
         'type': 'object',
         'required': [],
@@ -531,6 +636,32 @@ _LABELS_SCHEMA = {
     }
 }
 
+_PRORPERTY_NAME_OR_CLUSTER_NAME_TO_PROPERTY = {
+    'oneOf': [
+        {
+            'type': 'string'
+        },
+        {
+            # A list of single-element dict to pretain the
+            # order.
+            # Example:
+            #  property_name:
+            #    - my-cluster1-*: my-property-1
+            #    - my-cluster2-*: my-property-2
+            #    - "*"": my-property-3
+            'type': 'array',
+            'items': {
+                'type': 'object',
+                'additionalProperties': {
+                    'type': 'string'
+                },
+                'maxProperties': 1,
+                'minProperties': 1,
+            },
+        }
+    ]
+}
+
 
 class RemoteIdentityOptions(enum.Enum):
     """Enum for remote identity types.
@@ -541,6 +672,7 @@ class RemoteIdentityOptions(enum.Enum):
     """
     LOCAL_CREDENTIALS = 'LOCAL_CREDENTIALS'
     SERVICE_ACCOUNT = 'SERVICE_ACCOUNT'
+    NO_UPLOAD = 'NO_UPLOAD'
 
 
 def get_default_remote_identity(cloud: str) -> str:
@@ -559,36 +691,16 @@ _REMOTE_IDENTITY_SCHEMA = {
     }
 }
 
-_REMOTE_IDENTITY_SCHEMA_AWS = {
+_REMOTE_IDENTITY_SCHEMA_KUBERNETES = {
     'remote_identity': {
-        'oneOf': [
-            {
+        'anyOf': [{
+            'type': 'string'
+        }, {
+            'type': 'object',
+            'additionalProperties': {
                 'type': 'string'
-            },
-            {
-                # A list of single-element dict to pretain the order.
-                # Example:
-                #  remote_identity:
-                #    - my-cluster1-*: my-iam-role-1
-                #    - my-cluster2-*: my-iam-role-2
-                #    - "*"": my-iam-role-3
-                'type': 'array',
-                'items': {
-                    'type': 'object',
-                    'additionalProperties': {
-                        'type': 'string'
-                    },
-                    'maxProperties': 1,
-                    'minProperties': 1,
-                },
             }
-        ]
-    }
-}
-
-_REMOTE_IDENTITY_SCHEMA_KUBERNETES = {
-    'remote_identity': {
-        'type': 'string'
+        }]
     },
 }
 
@@ -618,6 +730,11 @@ def get_config_schema():
                     'resources': resources_schema,
                 }
             },
+            'bucket': {
+                'type': 'string',
+                'pattern': '^(https|s3|gs|r2|cos)://.+',
+                'required': [],
+            }
         }
     }
     cloud_configs = {
@@ -626,9 +743,20 @@ def get_config_schema():
             'required': [],
             'additionalProperties': False,
             'properties': {
-                'security_group_name': {
-                    'type': 'string'
+                'prioritize_reservations': {
+                    'type': 'boolean',
                 },
+                'specific_reservations': {
+                    'type': 'array',
+                    'items': {
+                        'type': 'string',
+                    },
+                },
+                'disk_encrypted': {
+                    'type': 'boolean',
+                },
+                'security_group_name':
+                    (_PRORPERTY_NAME_OR_CLUSTER_NAME_TO_PROPERTY),
                 **_LABELS_SCHEMA,
                 **_NETWORK_CONFIG_SCHEMA,
             },
@@ -648,16 +776,54 @@ def get_config_schema():
                         'type': 'string',
                     },
                 },
+                'managed_instance_group': {
+                    'type': 'object',
+                    'required': ['run_duration'],
+                    'additionalProperties': False,
+                    'properties': {
+                        'run_duration': {
+                            'type': 'integer',
+                        },
+                        'provision_timeout': {
+                            'type': 'integer',
+                        }
+                    }
+                },
+                'force_enable_external_ips': {
+                    'type': 'boolean'
+                },
+                'enable_gvnic': {
+                    'type': 'boolean'
+                },
                 **_LABELS_SCHEMA,
                 **_NETWORK_CONFIG_SCHEMA,
             },
             **_check_not_both_fields_present('instance_tags', 'labels')
         },
+        'azure': {
+            'type': 'object',
+            'required': [],
+            'additionalProperties': False,
+            'properties': {
+                'storage_account': {
+                    'type': 'string',
+                },
+                'resource_group_vm': {
+                    'type': 'string',
+                },
+            }
+        },
         'kubernetes': {
             'type': 'object',
             'required': [],
             'additionalProperties': False,
             'properties': {
+                'allowed_contexts': {
+                    'type': 'array',
+                    'items': {
+                        'type': 'string',
+                    },
+                },
                 'networking': {
                     'type': 'string',
                     'case_insensitive_enum': [
@@ -723,6 +889,9 @@ def get_config_schema():
                     'image_tag_gpu': {
                         'type': 'string',
                     },
+                    'vcn_ocid': {
+                        'type': 'string',
+                    },
                     'vcn_subnet': {
                         'type': 'string',
                     },
@@ -731,6 +900,13 @@ def get_config_schema():
         },
     }
 
+    admin_policy_schema = {
+        'type': 'string',
+        # Check regex to be a valid python module path
+        'pattern': (r'^[a-zA-Z_][a-zA-Z0-9_]*'
+                    r'(\.[a-zA-Z_][a-zA-Z0-9_]*)+$'),
+    }
+
     allowed_clouds = {
         # A list of cloud names that are allowed to be used
         'type': 'array',
@@ -741,9 +917,52 @@ def get_config_schema():
         }
     }
 
+    docker_configs = {
+        'type': 'object',
+        'required': [],
+        'additionalProperties': False,
+        'properties': {
+            'run_options': {
+                'anyOf': [{
+                    'type': 'string',
+                }, {
+                    'type': 'array',
+                    'items': {
+                        'type': 'string',
+                    }
+                }]
+            }
+        }
+    }
+    gpu_configs = {
+        'type': 'object',
+        'required': [],
+        'additionalProperties': False,
+        'properties': {
+            'disable_ecc': {
+                'type': 'boolean',
+            },
+        }
+    }
+
+    api_server = {
+        'type': 'object',
+        'required': [],
+        'additionalProperties': False,
+        'properties': {
+            'endpoint': {
+                'type': 'string',
+                # Apply validation for URL
+                'pattern': r'^https?://.*$',
+            },
+        }
+    }
+
     for cloud, config in cloud_configs.items():
         if cloud == 'aws':
-            config['properties'].update(_REMOTE_IDENTITY_SCHEMA_AWS)
+            config['properties'].update({
+                'remote_identity': _PRORPERTY_NAME_OR_CLUSTER_NAME_TO_PROPERTY
+            })
         elif cloud == 'kubernetes':
             config['properties'].update(_REMOTE_IDENTITY_SCHEMA_KUBERNETES)
         else:
@@ -755,11 +974,12 @@ def get_config_schema():
         'additionalProperties': False,
         'properties': {
             'jobs': controller_resources_schema,
-            'spot': controller_resources_schema,
             'serve': controller_resources_schema,
             'allowed_clouds': allowed_clouds,
+            'admin_policy': admin_policy_schema,
+            'docker': docker_configs,
+            'nvidia_gpus': gpu_configs,
+            'api_server': api_server,
             **cloud_configs,
         },
-        # Avoid spot and jobs being present at the same time.
-        **_check_not_both_fields_present('spot', 'jobs')
     }

sky/{status_lib.py → utils/status_lib.py}

@@ -6,22 +6,27 @@ import colorama
 
 
 class ClusterStatus(enum.Enum):
-    """Cluster status as recorded in table 'clusters'."""
+    """Cluster status as recorded in local cache.
+
+    This can be different from the actual cluster status, and can be refreshed
+    by running ``sky status --refresh``.
+    """
     # NOTE: these statuses are as recorded in our local cache, the table
     # 'clusters'.  The actual cluster state may be different (e.g., an UP
     # cluster getting killed manually by the user or the cloud provider).
 
-    # Initializing.  This means a backend.provision() call has started but has
-    # not successfully finished. The cluster may be undergoing setup, may have
-    # failed setup, may be live or down.
     INIT = 'INIT'
+    """Initializing.
+
+    This means a provisioning has started but has not successfully finished. The
+    cluster may be undergoing setup, may have failed setup, may be live or down.
+    """
 
-    # The cluster is recorded as up.  This means a backend.provision() has
-    # previously succeeded.
     UP = 'UP'
+    """The cluster is up. This means a provisioning has previously succeeded."""
 
-    # Stopped.  This means a `sky stop` call has previously succeeded.
     STOPPED = 'STOPPED'
+    """The cluster is stopped."""
 
     def colored_str(self):
         color = _STATUS_TO_COLOR[self]