Repository: skypilot-org/skypilot
Pull request #3478 (2024-04-25 03:37:24)

[UX] Fix task schema check for env
<!-- Describe the changes in this PR -->

To reproduce:
```yaml
envs:
  - my_env1: abc
  - my_env2: abc
  - my_env3: abc
```

`sky launch -c test task.yaml`

```
AttributeError: 'list' object has no attribute 'update'
```

With this PR:

```
ValueError: Invalid task YAML: [{'my_env1': 'abc'}, {'my_env2': 'abc'}, {'my_env3': 'abc'}] is not of type 'dict'. Check problematic field(s): $.envs
```

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [ ] Code formatting: `bash format.sh`
- [ ] Any manual or new tests for this PR (please specify below)
- [ ] All smoke tests: `pytest tests/test_smoke.py` 
- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
- [ ] Backward compatibility tests: `bash tests/backward_comaptibility_tests.sh`



Patch:
diff --git a/sky/task.py b/sky/task.py
index 9350e2610..33c74797e 100644
--- a/sky/task.py
+++ b/sky/task.py
@@ -348,6 +348,15 @@ class Task:
         config: Dict[str, Any],
         env_overrides: Optional[List[Tuple[str, str]]] = None,
     ) -> 'Task':
+        # More robust handling for 'envs': explicitly convert keys and values to
+        # str, since users may pass '123' as keys/values which will get parsed
+        # as int causing validate_schema() to fail.
+        envs = config.get('envs')
+        if envs is not None and isinstance(envs, dict):
+            config['envs'] = {str(k): str(v) for k, v in envs.items()}
+
+        common_utils.validate_schema(config, schemas.get_task_schema(),
+                                     'Invalid task YAML: ')
         if env_overrides is not None:
             # We must override env vars before constructing the Task, because
             # the Storage object creation is eager and it (its name/source
@@ -359,16 +368,6 @@ class Task:
             new_envs.update(env_overrides)
             config['envs'] = new_envs
 
-        # More robust handling for 'envs': explicitly convert keys and values to
-        # str, since users may pass '123' as keys/values which will get parsed
-        # as int causing validate_schema() to fail.
-        envs = config.get('envs')
-        if envs is not None and isinstance(envs, dict):
-            config['envs'] = {str(k): str(v) for k, v in envs.items()}
-
-        common_utils.validate_schema(config, schemas.get_task_schema(),
-                                     'Invalid task YAML: ')
-
         # Fill in any Task.envs into file_mounts (src/dst paths, storage
         # name/source).
         if config.get('file_mounts') is not None:
diff --git a/sky/utils/common_utils.py b/sky/utils/common_utils.py
index 14a79a4d7..0cef1a8db 100644
--- a/sky/utils/common_utils.py
+++ b/sky/utils/common_utils.py
@@ -567,20 +567,24 @@ def validate_schema(obj, schema, err_msg_prefix='', skip_none=True):
                            'The `envs` field contains invalid keys:\n' +
                            e.message)
             else:
-                err_msg = err_msg_prefix + 'The following fields are invalid:'
+                err_msg = err_msg_prefix
                 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!r}, did you mean '
+                            err_msg += (f'Instead of {field!r}, did you mean '
                                         f'{most_similar_field[0]!r}?')
                         else:
-                            err_msg += f'\nFound unsupported field {field!r}.'
+                            err_msg += f'Found unsupported field {field!r}.'
         else:
+            message = e.message
+            # Object in jsonschema is represented as dict in Python. Replace
+            # 'object' with 'dict' for better readability.
+            message = message.replace('type \'object\'', 'type \'dict\'')
             # Example e.json_path value: '$.resources'
-            err_msg = (err_msg_prefix + e.message +
+            err_msg = (err_msg_prefix + message +
                        f'. Check problematic field(s): {e.json_path}')
 
     if err_msg:


Test patch:
diff --git a/tests/test_yaml_parser.py b/tests/test_yaml_parser.py
index c02d77b8b..033808492 100644
--- a/tests/test_yaml_parser.py
+++ b/tests/test_yaml_parser.py
@@ -111,3 +111,26 @@ def test_invalid_fields_storage(tmp_path):
     with pytest.raises(AssertionError) as e:
         Task.from_yaml(config_path)
     assert 'Invalid storage args' in e.value.args[0]
+
+
+def test_invalid_envs_key(tmp_path):
+    config_path = _create_config_file(
+        textwrap.dedent(f"""\
+            envs:
+                invalid_env_$_key: value
+            """), tmp_path)
+    with pytest.raises(ValueError) as e:
+        Task.from_yaml(config_path)
+    assert 'does not match any of the regexes:' in e.value.args[0]
+
+
+def test_invalid_envs_type(tmp_path):
+    config_path = _create_config_file(
+        textwrap.dedent(f"""\
+            envs:
+                - env_key1: abc
+                - env_key2: abc
+            """), tmp_path)
+    with pytest.raises(ValueError) as e:
+        Task.from_yaml(config_path)
+    assert 'is not of type \'dict\'' in e.value.args[0]