Repository: skypilot-org/skypilot
Pull request: Enhance YAML schema validation
Fixes #674.

Currently, the YAML schema validation only throws errors for unknown fields. It does not have support for required / optional fields, enum checks and type checks. 

I propose using the [JSON Schema](https://json-schema.org/) format to define and validate schemas for the YAML configuration file.  Although there are a number of mature Python libraries for data validation ([schema](https://github.com/keleshev/schema), [pydantic](https://pydantic-docs.helpmanual.io/), [ceberus](https://docs.python-cerberus.org/en/stable/) etc.), the main benefit of using JSON Schema is that it is an open specification and so we are not bound to any one implementation.

This PR replaces the current YAML validation with a validation based on JSON Schema that allows for handling unknown fields, enum checks, type checks and potentially a lot more complex validation as required. Similar to the current implementation, it also suggests the most similar field when an unknown field is encountered.

Tested:
- [Full configuration](https://sky-proj-sky.readthedocs-hosted.com/en/latest/reference/yaml-spec.html)
- Failing enum checks (e.g. `cloud: gcs` instead of `cloud: gcp` )
- Failing type checks (e.g. `num_nodes: one` instead of `num_nodes: 1`)
- Unknown fields (e.g. `num_workers` instead of `num_nodes`)

From what I see from the docs, there are no required fields so the schema does not specify any.

Problem statement:
Enforce strict YAML schema compliance
Sometimes I accidentally type `num_workers: 2` instead of `num_nodes: 2` in a sky YAML. `num_workers: 2` silently succeeds, and the user is left wondering why only 1 node was spun up.

We should enforce strict checking against the YAML schema, and raise errors if any extra keys are supplied in the sky YAML.

Hints from issue discussion:
This should be partially fixed by #759, as we will error out for any unknown field.
Reopening this because in order to achieve full schema compliance, we also need to check values of YAML fields. For example, I accidentally typed `cloud: gcp` instead of `gcs`, and it silently succeeds and the optimizer picked AWS for me and ran the job. Instead, it should have raised an error saying something like `Invalid cloud "gcs", please pick from [azure, aws, gcp]`. 

In other words, we should be checking not only for unrecognized fields, but also for invalid values in those fields.

Commit message: Use older JSON schema

The version of jsonschema that supports Python 3.6 does not contain a validator for the 202012 schema, so we have to use an older schema to support Python 3.6

Patch to review:
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index c010c3bd4ed..c1cc5cb8fbb 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -24,6 +24,7 @@
 import colorama
 import filelock
 import jinja2
+import jsonschema
 import psutil
 import requests
 from requests import adapters
@@ -43,6 +44,7 @@
 from sky import spot as spot_lib
 from sky.skylet import log_lib
 from sky.utils import timeline
+from sky.utils import validator
 
 if typing.TYPE_CHECKING:
     from sky import resources
@@ -1755,26 +1757,24 @@ def current_backoff(self) -> float:
         return self._backoff
 
 
-def check_fields(provided_fields, known_fields):
-    known_fields = set(known_fields)
-    unknown_fields = []
-    for field in provided_fields:
-        if field not in known_fields:
-            unknown_fields.append(field)
-
-    if len(unknown_fields) > 0:
-        invalid_keys = 'The following fields are invalid:\n'
-        for unknown_key in unknown_fields:
-            similar_keys = difflib.get_close_matches(unknown_key, known_fields)
-            key_invalid = f'    Unknown field \'{unknown_key}\'.'
-            if len(similar_keys) == 1:
-                key_invalid += f' Did you mean \'{similar_keys[0]}\'?'
-            if len(similar_keys) > 1:
-                key_invalid += f' Did you mean one of {similar_keys}?'
-            key_invalid += '\n'
-            invalid_keys += key_invalid
+def validate_schema(obj, schema, err_msg_prefix=''):
+    err_msg = None
+    try:
+        validator.SchemaValidator(schema).validate(obj)
+    except jsonschema.ValidationError as e:
+        err_msg = err_msg_prefix + e.message
+        if e.validator == 'additionalProperties':
+            known_fields = set(e.schema.get('properties', {}).keys())
+            for field in e.instance:
+                if field not in known_fields:
+                    most_similar_field = difflib.get_close_matches(
+                        field, known_fields, 1)
+                    if most_similar_field:
+                        err_msg += f'\nInstead of {field}, did you mean {most_similar_field[0]}?'
+
+    if err_msg:
         with print_exception_no_traceback():
-            raise ValueError(invalid_keys)
+            raise ValueError(err_msg)
 
 
 @contextlib.contextmanager
diff --git a/sky/clouds/cloud.py b/sky/clouds/cloud.py
index f7a21459030..cbf534acd60 100644
--- a/sky/clouds/cloud.py
+++ b/sky/clouds/cloud.py
@@ -27,6 +27,9 @@ class _CloudRegistry(dict):
     def from_str(self, name: Optional[str]) -> Optional['Cloud']:
         if name is None:
             return None
+        if name.lower() not in self:
+            raise ValueError(f'Cloud {name} is not a valid cloud among '
+                             '{list(self.keys())}')
         return self.get(name.lower())
 
     def register(self, cloud_cls: 'Cloud') -> None:
diff --git a/sky/data/storage.py b/sky/data/storage.py
index 99f7a30e59f..9fac9c8ebc4 100644
--- a/sky/data/storage.py
+++ b/sky/data/storage.py
@@ -11,6 +11,7 @@
 from sky.adaptors import aws
 from sky.adaptors import gcp
 from sky.backends import backend_utils
+from sky.utils import schemas
 from sky.data import data_transfer
 from sky.data import data_utils
 from sky import exceptions
@@ -240,9 +241,6 @@ class Storage(object):
         storage.delete()
     """
 
-    # Update the key list when a new field is added.
-    _YAML_KEYS = ['name', 'source', 'store', 'mode', 'persistent']
-
     class StorageMetadata(object):
         """A pickle-able tuple of:
 
@@ -622,7 +620,8 @@ def _sync_store(self, store: AbstractStore):
 
     @classmethod
     def from_yaml_config(cls, config: Dict[str, str]) -> 'Storage':
-        backend_utils.check_fields(config.keys(), cls._YAML_KEYS)
+        backend_utils.validate_schema(config, schemas.get_storage_schema(),
+                                      'Invalid storage YAML: ')
 
         name = config.pop('name', None)
         source = config.pop('source', None)
diff --git a/sky/resources.py b/sky/resources.py
index 6e7a628f4df..ace7d0cc193 100644
--- a/sky/resources.py
+++ b/sky/resources.py
@@ -6,6 +6,7 @@
 from sky import sky_logging
 from sky import spot
 from sky.backends import backend_utils
+from sky.utils import schemas
 
 logger = sky_logging.init_logger(__name__)
 
@@ -43,12 +44,6 @@ class Resources:
     # 3. Modify the to_config method to handle the new fields.
     _VERSION = 3
 
-    # Update the key list when a new field is added.
-    _YAML_KEYS = [
-        'cloud', 'instance_type', 'accelerators', 'accelerator_args',
-        'use_spot', 'spot_recovery', 'disk_size', 'region'
-    ]
-
     def __init__(
         self,
         cloud: Optional[clouds.Cloud] = None,
@@ -447,7 +442,8 @@ def from_yaml_config(cls, config: Optional[Dict[str, str]]) -> 'Resources':
         if config is None:
             return Resources()
 
-        backend_utils.check_fields(config.keys(), cls._YAML_KEYS)
+        backend_utils.validate_schema(config, schemas.get_resources_schema(),
+                                      'Invalid resources YAML: ')
 
         resources_fields = dict()
         if config.get('cloud') is not None:
diff --git a/sky/setup_files/setup.py b/sky/setup_files/setup.py
index 89de71ded8d..89c28d45c39 100644
--- a/sky/setup_files/setup.py
+++ b/sky/setup_files/setup.py
@@ -28,6 +28,7 @@
     'colorama',
     'cryptography',
     'jinja2',
+    'jsonschema',
     'networkx',
     'oauth2client',
     'pandas',
diff --git a/sky/task.py b/sky/task.py
index 3ba80d3452f..347d4f74835 100644
--- a/sky/task.py
+++ b/sky/task.py
@@ -12,6 +12,7 @@
 from sky.data import storage as storage_lib
 from sky.data import data_transfer as data_transfer_lib
 from sky.data import data_utils
+from sky.utils import schemas
 
 if typing.TYPE_CHECKING:
     from sky import resources as resources_lib
@@ -66,12 +67,6 @@ def _is_valid_env_var(name: str) -> bool:
 class Task:
     """Task: a coarse-grained stage in an application."""
 
-    # Update the key list when a new field is added.
-    _YAML_KEYS = [
-        'name', 'run', 'workdir', 'setup', 'num_nodes', 'envs', 'file_mounts',
-        'inputs', 'outputs', 'resources'
-    ]
-
     def __init__(
         self,
         name: Optional[str] = None,
@@ -214,7 +209,8 @@ def from_yaml(yaml_path):
         if config is None:
             config = {}
 
-        backend_utils.check_fields(config.keys(), Task._YAML_KEYS)
+        backend_utils.validate_schema(config, schemas.get_task_schema(),
+                                      'Invalid task YAML: ')
 
         task = Task(
             config.pop('name', None),
diff --git a/sky/utils/schemas.py b/sky/utils/schemas.py
new file mode 100644
index 00000000000..eaf81ad3d2a
--- /dev/null
+++ b/sky/utils/schemas.py
@@ -0,0 +1,120 @@
+"""This module contains schemas used to validate objects.
+
+Schemas conform to the JSON Schema specification as defined at
+https://json-schema.org/
+"""
+
+from sky.clouds import cloud
+from sky.data import storage
+
+
+def get_resources_schema():
+    return {
+        '$schema': 'http://json-schema.org/draft-07/schema#',
+        'type': 'object',
+        'required': [],
+        'additionalProperties': False,
+        'properties': {
+            'cloud': {
+                'type': 'string',
+                'case_insensitive_enum': list(cloud.CLOUD_REGISTRY.keys())
+            },
+            'region': {
+                'type': 'string',
+            },
+            'accelerators': {
+                'type': 'string',
+            },
+            'instance_type': {
+                'type': 'string',
+            },
+            'use_spot': {
+                'type': 'boolean',
+            },
+            'spot_recovery': {
+                'type': 'string',
+            },
+            'disk_size': {
+                'type': 'integer',
+            },
+            'accelerator_args': {
+                'type': 'object',
+                'required': [],
+                'additionalProperties': False,
+                'properties': {
+                    'tf_version': {
+                        'type': 'string',
+                    },
+                    'tpu_name': {
+                        'type': 'string',
+                    }
+                }
+            }
+        }
+    }
+
+
+def get_storage_schema():
+    return {
+        '$schema': 'http://json-schema.org/draft-07/schema#',
+        'type': 'object',
+        'required': [],
+        'additionalProperties': False,
+        'properties': {
+            'name': {
+                'type': 'string',
+            },
+            'source': {
+                'type': 'string',
+            },
+            'store': {
+                'type': 'string',
+                'case_insensitive_enum': [
+                    type.value for type in storage.StoreType
+                ]
+            },
+            'persistent': {
+                'type': 'boolean',
+            },
+            'mode': {
+                'type': 'string',
+                'case_insensitive_enum': [
+                    mode.value for mode in storage.StorageMode
+                ]
+            }
+        }
+    }
+
+
+def get_task_schema():
+    return {
+        '$schema': 'https://json-schema.org/draft/2020-12/schema',
+        'type': 'object',
+        'required': [],
+        'additionalProperties': False,
+        'properties': {
+            'name': {
+                'type': 'string',
+            },
+            'workdir': {
+                'type': 'string',
+            },
+            'num_nodes': {
+                'type': 'integer',
+            },
+            # Resources config is validated separately using RESOURCES_SCHEMA
+            'resources': {
+                'type': 'object',
+            },
+            # Storage config is validated separately using STORAGE_SCHEMA
+            'file_mounts': {
+                'type': 'object',
+            },
+            'setup': {
+                'type': 'string',
+            },
+            'run': {
+                'type': 'string',
+            }
+        }
+    }
diff --git a/sky/utils/validator.py b/sky/utils/validator.py
new file mode 100644
index 00000000000..c1e9cbe5d19
--- /dev/null
+++ b/sky/utils/validator.py
@@ -0,0 +1,18 @@
+"""This module contains a custom validator for the JSON Schema specification.
+
+The main motivation behind extending the existing JSON Schema validator is to
+allow for case-insensitive enum matching since this is currently not supported
+by the latest JSON Schema specification.
+"""
+import jsonschema
+
+
+def case_insensitive_enum(validator, enums, instance, schema):  # pylint: disable=unused-argument
+    if instance.lower() not in [enum.lower() for enum in enums]:
+        yield jsonschema.ValidationError(
+            f'{instance!r} is not one of {enums!r}')
+
+
+SchemaValidator = jsonschema.validators.extend(
+    jsonschema.Draft7Validator,
+    validators={'case_insensitive_enum': case_insensitive_enum})


Code review comments:
diff_hunk: @@ -0,0 +1,18 @@
+"""This module contains a custom validator for the JSON Schema specification.
+
+The main motivation behind extending the existing JSON Schema validator is to
+allow for case-insensitive enum matching since this is currently not supported
+by the latest JSON Schema specification.
+"""
+import jsonschema
+
+
+def case_insensitive_enum(validator, enums, instance, schema):  # pylint: disable=unused-argument
path: sky/utils/validator.py
text: @author:
As explained at https://python-jsonschema.readthedocs.io/en/stable/creating/, validation functions should take in 4 arguments. Even though there are 2 unused arguments here, I've decided to leave them in for clarity.

@user1:
For unused argument, please follow the style guide [here](https://google.github.io/styleguide/pyguide.html#2-python-language-rules), using `del instance, schema  # unused` at the beginning of this function.

@author:
Done, thanks for pointing this out