Repository: skypilot-org/skypilot
Pull request: Add YAML schema check
Fixes #674.

### Tests
- [x] Full YAML from documentation
```
num_node: 1
# returns error
========== 
resource:
   cloud: aws
   disk_space: 256
# returns error
```

Also tested `file_mounts`:
```
# returns useful error with suggestion:
file_mounts:
  ~/resnet_ddp:
        name: tmp-resnet-test
        sources: ./examples/spot/resnet_ddp
        persistent: false
        mode: COPY
```

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.

Commit message: remove bad yaml

Patch to review:
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index c45ca820b07..33ed35f85d3 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -1,6 +1,7 @@
 """Util constants/functions for the backends."""
 import colorama
 import datetime
+import difflib
 import enum
 import hashlib
 import http.client as httplib
@@ -1476,3 +1477,24 @@ def current_backoff(self) -> float:
         self._backoff += random.uniform(-self.JITTER * self._backoff,
                                         self.JITTER * self._backoff)
         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
+        raise ValueError(invalid_keys)
diff --git a/sky/data/storage.py b/sky/data/storage.py
index d663acbb775..37fb6a69294 100644
--- a/sky/data/storage.py
+++ b/sky/data/storage.py
@@ -236,6 +236,9 @@ 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:
 
@@ -607,10 +610,13 @@ 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)
+
         name = config.pop('name', None)
         source = config.pop('source', None)
         store = config.pop('store', None)
         mode_str = config.pop('mode', None)
+
         if isinstance(mode_str, str):
             # Make mode case insensitive, if specified
             mode = StorageMode(mode_str.upper())
@@ -618,6 +624,9 @@ def from_yaml_config(cls, config: Dict[str, str]) -> 'Storage':
             # Make sure this keeps the same as the default mode in __init__
             mode = StorageMode.MOUNT
         persistent = config.pop('persistent', True)
+
+        assert not config, f'Invalid storage args: {config.keys()}'
+
         # Validation of the config object happens on instantiation.
         storage_obj = cls(name=name,
                           source=source,
@@ -625,9 +634,6 @@ def from_yaml_config(cls, config: Dict[str, str]) -> 'Storage':
                           mode=mode)
         if store is not None:
             storage_obj.add_store(StoreType(store.upper()))
-        if config:
-            raise exceptions.StorageSpecError(
-                f'Invalid storage spec. Unknown fields: {config.keys()}')
         return storage_obj
 
     def to_yaml_config(self) -> Dict[str, str]:
diff --git a/sky/resources.py b/sky/resources.py
index cbb3f61b86d..b73b4b14691 100644
--- a/sky/resources.py
+++ b/sky/resources.py
@@ -5,6 +5,7 @@
 from sky import global_user_state
 from sky import sky_logging
 from sky import spot
+from sky.backends import backend_utils
 
 logger = sky_logging.init_logger(__name__)
 
@@ -42,6 +43,12 @@ 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,
@@ -438,6 +445,9 @@ def copy(self, **override) -> 'Resources':
     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)
+
         resources_fields = dict()
         if config.get('cloud') is not None:
             resources_fields['cloud'] = clouds.CLOUD_REGISTRY.from_str(
@@ -458,8 +468,7 @@ def from_yaml_config(cls, config: Optional[Dict[str, str]]) -> 'Resources':
         if config.get('region') is not None:
             resources_fields['region'] = config.pop('region')
 
-        if len(config) > 0:
-            raise ValueError(f'Unknown fields in resources config: {config}')
+        assert not config, f'Invalid resource args: {config.keys()}'
         return Resources(**resources_fields)
 
     def to_yaml_config(self) -> Dict[str, Union[str, int]]:
diff --git a/sky/task.py b/sky/task.py
index ef9ca3e22a7..085041ec3ab 100644
--- a/sky/task.py
+++ b/sky/task.py
@@ -8,6 +8,7 @@
 
 import sky
 from sky import clouds
+from sky.backends import backend_utils
 from sky.data import storage as storage_lib
 from sky.data import data_transfer as data_transfer_lib
 from sky.data import data_utils
@@ -65,6 +66,12 @@ 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,
@@ -202,6 +209,8 @@ def from_yaml(yaml_path):
         if config is None:
             config = {}
 
+        backend_utils.check_fields(config.keys(), Task.yaml_keys)
+
         task = Task(
             config.pop('name', None),
             run=config.pop('run', None),
@@ -263,9 +272,8 @@ def from_yaml(yaml_path):
             if acc.startswith('tpu-') and task.num_nodes > 1:
                 raise ValueError('Multi-node TPU cluster not supported. '
                                  f'Got num_nodes={task.num_nodes}')
-        if len(config) > 0:
-            raise ValueError(f'Unknown fields in in YAML: {config.keys()}')
         task.set_resources({resources})
+        assert not config, f'Invalid task args: {config.keys()}'
         return task
 
     def to_yaml_config(self) -> Dict[str, Any]:


Code review comments:
diff_hunk: @@ -236,6 +236,9 @@ class Storage(object):
         storage.delete()
     """
 
+    # Update the key list when a new field is added.
+    yaml_keys = ['name', 'source', 'store', 'mode', 'persistent']
path: sky/data/storage.py
text: @user1:
use `_yaml_keys` as it's internal to the module?

@user2:
nit: Use upper case for constant.

@user1:
Oh right we should use upper case here.