diff --git a/caliban/cli.py b/caliban/cli.py index f90894b..5ea0dac 100644 --- a/caliban/cli.py +++ b/caliban/cli.py @@ -22,7 +22,6 @@ import google.auth._cloud_sdk as csdk from absl.flags import argparse_flags -from blessings import Terminal import caliban.config as conf import caliban.config.experiment as ce @@ -34,10 +33,9 @@ import caliban.platform.gke.util as gke_u import caliban.util as u import caliban.util.argparse as ua +import caliban.util.schema as us from caliban import __version__ -t = Terminal() - def _job_mode(use_gpu: bool, gpu_spec: Optional[ct.GPUSpec], tpu_spec: Optional[ct.TPUSpec]) -> conf.JobMode: @@ -159,7 +157,7 @@ def extra_dirs(parser): "-d", "--dir", action="append", - type=ua.validated_directory, + type=ua.argparse_schema(us.Directory), help="Extra directories to include. List these from large to small " "to take full advantage of Docker's build cache.") @@ -191,7 +189,7 @@ def region_arg(parser): def cloud_key_arg(parser): parser.add_argument("--cloud_key", - type=ua.validated_file, + type=ua.argparse_schema(us.File), help="Path to GCloud service account key. " "(Defaults to $GOOGLE_APPLICATION_CREDENTIALS.)") diff --git a/caliban/config/__init__.py b/caliban/config/__init__.py index ed3a05e..f85340c 100644 --- a/caliban/config/__init__.py +++ b/caliban/config/__init__.py @@ -17,27 +17,26 @@ Utilities for our job runner, for working with configs. """ -from __future__ import absolute_import, division, print_function - -import argparse import os import sys from enum import Enum from typing import Any, Dict, List, Optional -import commentjson -import yaml +import schema as s import caliban.platform.cloud.types as ct +import caliban.util.schema as us class JobMode(str, Enum): + """Represents the two modes that you can use to execute a Caliban job.""" CPU = 'CPU' GPU = 'GPU' + @staticmethod + def parse(label): + return JobMode(label.upper()) -# Special config for Caliban. -CalibanConfig = Dict[str, Any] DRY_RUN_FLAG = "--dry_run" CALIBAN_CONFIG = ".calibanconfig.json" @@ -59,6 +58,37 @@ class JobMode(str, Enum): "type": "ACCELERATOR_TYPE_UNSPECIFIED" } +# Schema for Caliban Config + +AptPackages = s.Or( + [str], { + s.Optional("gpu", default=list): [str], + s.Optional("cpu", default=list): [str] + }, + error=""""apt_packages" entry must be a dictionary or list, not '{}'""") + +CalibanConfig = s.Schema({ + s.Optional("build_time_credentials", default=False): + bool, + s.Optional("default_mode", default=JobMode.CPU): + s.Use(JobMode.parse), + s.Optional("project_id"): + s.And(str, len), + s.Optional("cloud_key"): + s.And(str, len), + s.Optional("base_image"): + str, + s.Optional("apt_packages", default=dict): + AptPackages, + + # Allow extra entries without killing the schema to allow for backwards + # compatibility. + s.Optional(str): + str, +}) + +# Accessors + def gpu(job_mode: JobMode) -> bool: """Returns True if the supplied JobMode is JobMode.GPU, False otherwise. @@ -67,41 +97,6 @@ def gpu(job_mode: JobMode) -> bool: return job_mode == JobMode.GPU -def load_yaml_config(path): - """returns the config parsed based on the info in the flags. - - Grabs the config file, written in yaml, slurps it in. - """ - with open(path) as f: - config = yaml.load(f, Loader=yaml.FullLoader) - - return config - - -def load_config(path, mode='yaml'): - """Load a JSON or YAML config. - - """ - if mode == 'json': - with open(path) as f: - return commentjson.load(f) - - return load_yaml_config(path) - - -def valid_json(path: str) -> Dict[str, Any]: - """Loads JSON if the path points to a valid JSON file; otherwise, throws an - exception that's picked up by argparse. - - """ - try: - return load_config(path, mode='json') - except commentjson.JSONLibraryException: - raise argparse.ArgumentTypeError( - """File '{}' doesn't seem to contain valid JSON. Try again!""".format( - path)) - - def extract_script_args(m: Dict[str, Any]) -> List[str]: """Strip off the "--" argument if it was passed in as a separator.""" script_args = m.get("script_args") @@ -166,28 +161,25 @@ def apt_packages(conf: CalibanConfig, mode: JobMode) -> List[str]: the requests in the config. """ - packages = conf.get("apt_packages") or {} + packages = conf["apt_packages"] if isinstance(packages, dict): k = "gpu" if gpu(mode) else "cpu" - return packages.get(k, []) - - elif isinstance(packages, list): - return packages + return packages[k] - else: - raise argparse.ArgumentTypeError( - """{}'s "apt_packages" entry must be a dictionary or list, not '{}'""". - format(CALIBAN_CONFIG, packages)) + return packages -def caliban_config() -> CalibanConfig: +def caliban_config(conf_path: str = CALIBAN_CONFIG) -> CalibanConfig: """Returns a dict that represents a `.calibanconfig.json` file if present, empty dictionary otherwise. + + If the supplied conf_path is present, but doesn't pass the supplied schema, + errors and kills the program. + """ - if not os.path.isfile(CALIBAN_CONFIG): + if not os.path.isfile(conf_path): return {} - with open(CALIBAN_CONFIG) as f: - conf = commentjson.load(f) - return conf + with us.error_schema(conf_path): + return s.And(us.Json, CalibanConfig).validate(conf_path) diff --git a/caliban/config/experiment.py b/caliban/config/experiment.py index c0fecd5..4305d2d 100644 --- a/caliban/config/experiment.py +++ b/caliban/config/experiment.py @@ -29,6 +29,8 @@ import commentjson import caliban.util as u +import caliban.util.argparse as ua +import caliban.util.schema as us # int, str and bool are allowed in a final experiment; lists are markers for # expansion. @@ -263,11 +265,10 @@ def validate_experiment_config(items: ExpConf) -> ExpConf: def load_experiment_config(s): - if s.lower() == 'stdin': + if isinstance(s, str) and s.lower() == 'stdin': json = commentjson.load(sys.stdin) else: - with open(u.validated_file(s)) as f: - json = commentjson.load(f) + json = ua.argparse_schema(us.Json)(s) return validate_experiment_config(json) diff --git a/caliban/main.py b/caliban/main.py index d4a52f5..3dada18 100644 --- a/caliban/main.py +++ b/caliban/main.py @@ -34,6 +34,7 @@ import caliban.platform.notebook as pn import caliban.platform.run as pr import caliban.platform.shell as ps +import caliban.util.schema as cs ll.getLogger('caliban.main').setLevel(logging.ERROR) t = Terminal() @@ -154,7 +155,8 @@ def run_app(arg_input): def main(): logging.use_python_logging() try: - app.run(run_app, flags_parser=cli.parse_flags) + with cs.fatal_errors(): + app.run(run_app, flags_parser=cli.parse_flags) except KeyboardInterrupt: logging.info('Shutting down.') sys.exit(0) diff --git a/caliban/util/argparse.py b/caliban/util/argparse.py index efd6e34..79f24bc 100644 --- a/caliban/util/argparse.py +++ b/caliban/util/argparse.py @@ -25,6 +25,7 @@ import caliban.util as u import caliban.util.fs as ufs +import schema as s t = Terminal() @@ -39,6 +40,25 @@ def expand_args(items: Dict[str, str]) -> List[str]: return list(it.chain.from_iterable(pairs)) +def argparse_schema(schema): + """Wrapper that performs validation and converts SchemaErrors into + ArgumentTypeErrors for better argument error reporting. + + """ + + def check(x): + try: + return schema.validate(x) + except s.SchemaError as e: + raise argparse.ArgumentTypeError(e.code) from None + + return check + + +# TODO: Now that we use schema, validated_package and parse_kv_pair should be +# converted to schema instances. + + def validated_package(path: str) -> u.Package: """similar to generate_package but runs argparse validation on packages that don't actually exist in the filesystem. @@ -89,26 +109,3 @@ def is_key(k: Optional[str]) -> bool: """ return k is not None and len(k) > 0 and k[0] == "-" - - -def validated_directory(path: str) -> str: - """This validates that the supplied directory exists locally. - - """ - if not os.path.isdir(path): - raise argparse.ArgumentTypeError( - """Directory '{}' doesn't exist in this directory. Check yourself!""". - format(path)) - return path - - -def validated_file(path: str) -> str: - """This validates that the supplied file exists. Tilde expansion is supported. - - """ - expanded = os.path.expanduser(path) - if not os.path.isfile(expanded): - raise argparse.ArgumentTypeError( - """File '{}' isn't a valid file on your system. Try again!""".format( - path)) - return path diff --git a/caliban/util/schema.py b/caliban/util/schema.py new file mode 100644 index 0000000..c15b839 --- /dev/null +++ b/caliban/util/schema.py @@ -0,0 +1,93 @@ +#!/usr/bin/python +# +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Useful shared schemas. +""" +import os +import sys +from contextlib import contextmanager +from typing import Optional + +import commentjson + +import caliban.util as u +import schema as s + + +class FatalSchemaError(Exception): + """Wrapper for an exception that can bubble itself up to the top level of the +program.""" + + def __init__(self, message, context): + self.message = message + self.context = context + super().__init__(self.message) + + +@contextmanager +def error_schema(context: Optional[str] = None): + """Wrap functions that check schemas in this context manager to throw an + appropriate error with a nice message. + + """ + prefix = "" + if context is not None: + prefix = f"\nValidation error while parsing {context}:\n" + + try: + yield + except s.SchemaError as e: + raise FatalSchemaError(e.code, prefix) + + +@contextmanager +def fatal_errors(): + """Context manager meant to wrap an entire program and present schema errors in + an easy-to-read way. + + """ + try: + yield + except FatalSchemaError as e: + u.err(f"{e.context}\n{e.message}\n\n") + sys.exit(1) + except s.SchemaError as e: + u.err(f"\n{e.code}\n\n") + sys.exit(1) + + +def load_json(path): + with open(path) as f: + return commentjson.load(f) + + +# TODO Once a release with this patch happens: +# https://github.com/keleshev/schema/pull/238,, Change `Or` to `Schema`. This +# problem only occurs for callable validators. + +Directory = s.Or( + os.path.isdir, + False, + error="""Directory '{}' doesn't exist in this directory. Check yourself!""") + +File = s.Or(lambda path: os.path.isfile(os.path.expanduser(path)), + False, + error="""File '{}' isn't a valid file on your system. Try again!""") + +Json = s.And( + File, + s.Use(load_json, + error="""File '{}' doesn't seem to contain valid JSON. Try again!""")) diff --git a/setup.py b/setup.py index a72392d..a66b300 100644 --- a/setup.py +++ b/setup.py @@ -49,6 +49,7 @@ def readme(): 'google-cloud-core>=1.0.3', 'google-cloud-container>=0.3.0', 'psycopg2-binary==2.8.5', + 'schema==0.7.2', 'urllib3>=1.25.7', 'yaspin>=0.16.0', # This is not a real dependency of ours, but we need it to override the diff --git a/tests/caliban/config/test_config.py b/tests/caliban/config/test_config.py index 126c21d..3f94229 100644 --- a/tests/caliban/config/test_config.py +++ b/tests/caliban/config/test_config.py @@ -14,14 +14,22 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json import os from argparse import ArgumentTypeError -import caliban.platform.cloud.types as ct import caliban.config as c +import caliban.platform.cloud.types as ct +import caliban.util.schema as us import pytest +def test_gpu(): + assert c.gpu(c.JobMode.GPU) + assert not c.gpu(c.JobMode.CPU) + assert not c.gpu("face") + + def test_extract_region(monkeypatch): if os.environ.get('REGION'): monkeypatch.delenv('REGION') @@ -42,3 +50,33 @@ def test_extract_region(monkeypatch): assert c.extract_region({}) == c.DEFAULT_REGION assert c.extract_region({"region": "us-west1"}) == ct.US.west1 + + +def test_caliban_config(tmpdir): + """Tests validation of the CalibanConfig schema and the method that returns the + parsed config. + + """ + valid = {"apt_packages": {"cpu": ["face"]}, "random": "entry"} + valid_path = tmpdir.join('valid.json') + + with open(valid_path, 'w') as f: + json.dump(valid, f) + + invalid = {"apt_packages": "face"} + invalid_path = tmpdir.join('invalid.json') + + with open(invalid_path, 'w') as f: + json.dump(invalid, f) + + # Failing the schema raises an error. + with pytest.raises(us.FatalSchemaError): + c.caliban_config(invalid_path) + + # paths that don't exist return an empty map: + assert c.caliban_config('random_path') == {} + + # If the config is valid, c.apt_packages can fetch the packages we specified. + config = c.caliban_config(valid_path) + assert c.apt_packages(config, c.JobMode.GPU) == [] + assert c.apt_packages(config, c.JobMode.CPU) == ["face"] diff --git a/tests/caliban/config/test_experiment.py b/tests/caliban/config/test_experiment.py index 5ca1e3d..ad4860c 100644 --- a/tests/caliban/config/test_experiment.py +++ b/tests/caliban/config/test_experiment.py @@ -14,10 +14,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json from argparse import ArgumentTypeError import caliban.config.experiment as c import caliban.util as u +import caliban.util.schema as us import pytest @@ -238,3 +240,27 @@ def test_compound_key_handling(): assert test['after_expansion'] == c.expand_experiment_config(test['input']) assert test['after_dictproduct'] == list( u.dict_product(test['after_tupleization'])) + + +def test_load_experiment_config(tmpdir): + valid = {"key": ['a', 'b'], "random": [True, False]} + valid_path = tmpdir.join('valid.json') + + with open(valid_path, 'w') as f: + json.dump(valid, f) + + invalid_path = tmpdir.join('invalid.json') + + with open(invalid_path, 'w') as f: + f.write("{{{I am not JSON!\n") + + # Failing the schema with invalid json raises an ARGPARSE error, not a schema + # error. We haven't converted experimentconfig to schema yet. + # + # We use schema to validate, but the ua.argparse_schema wrapper converts the + # error internally. + with pytest.raises(ArgumentTypeError): + c.load_experiment_config(invalid_path) + + # A valid config should round trip. + assert c.load_experiment_config(valid_path) == valid diff --git a/tests/caliban/util/test_schema.py b/tests/caliban/util/test_schema.py new file mode 100644 index 0000000..af9bb56 --- /dev/null +++ b/tests/caliban/util/test_schema.py @@ -0,0 +1,41 @@ +#!/usr/bin/python +# +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import tempfile + +import schema as s + +import caliban.util.schema as us +import pytest + + +def test_directory(tmpdir): + # Proper directories pass validation. + assert us.Directory.validate(tmpdir) == tmpdir + + # random dirs that I made up dont! + with pytest.raises(s.SchemaError): + assert us.Directory.validate('random') + + +def test_file(): + with tempfile.NamedTemporaryFile() as tmp: + # Existing files pass validation. + assert us.File.validate(tmp.name) == tmp.name + + # random paths that I made up dont! + with pytest.raises(s.SchemaError): + assert us.File.validate('random') diff --git a/tutorials/basic/.calibanconfig.json b/tutorials/basic/.calibanconfig.json new file mode 100644 index 0000000..e324ca9 --- /dev/null +++ b/tutorials/basic/.calibanconfig.json @@ -0,0 +1 @@ +{"apt_packages": ["cake"]}