Add Schema for Calibanconfig

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.)")
 

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)

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)
 

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)

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

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!"""))

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