adds strict_db flag to run/cloud/cluster commands

caliban/cli.py

@@ -316,6 +316,7 @@ def local_run_parser(base):
   image_id_arg(parser)
   docker_run_arg(parser)
   xgroup_submit_arg(parser)
+  strict_db_arg(parser)
 
 
 def gpu_spec_arg(parser, validate_count: bool = False):
@@ -386,6 +387,7 @@ def container_parser(parser):
   job_name_arg(parser)
   label_arg(parser)
   xgroup_submit_arg(parser)
+  strict_db_arg(parser)
 
 
 def cloud_parser(base):
@@ -687,6 +689,7 @@ def cluster_job_submit_cmd(base):
   dry_run_arg(parser)
   job_export_arg(parser)
   xgroup_submit_arg(parser)
+  strict_db_arg(parser)
 
   require_module(parser)
   add_script_args(parser)
@@ -947,3 +950,12 @@ def max_jobs_arg(parser):
             f'then this specifies the total number of jobs to return, ordered '
             f'by creation date, or all jobs if max_jobs==0.'),
   )
+
+
+# ----------------------------------------------------------------------------
+def strict_db_arg(parser):
+  parser.add_argument(
+      '--strict_db',
+      action='store_true',
+      help=(f'if the caliban database specified by the CALIBAN_DB_URL '
+            f'environment variable cannot be reached, then fail and exit'))

caliban/cloud/core.py

@@ -559,7 +559,8 @@ def submit_ml_job(job_mode: conf.JobMode,
                   experiment_config: Optional[conf.ExpConf] = None,
                   script_args: Optional[List[str]] = None,
                   request_retries: Optional[int] = None,
-                  xgroup: Optional[str] = None) -> None:
+                  xgroup: Optional[str] = None,
+                  strict_db: bool = False) -> None:
   """Top level function in the module. This function:
 
   - builds an image using the supplied docker_args, in either CPU or GPU mode
@@ -606,6 +607,7 @@ def submit_ml_job(job_mode: conf.JobMode,
     a timeout or a rate limiting request.
   - xgroup: experiment group for this submission, if None a new group will
     be created
+  - strict_db: if database specified by the CALIBAN_DB_URL is not found, exit
   """
   if script_args is None:
     script_args = []
@@ -628,7 +630,7 @@ def submit_ml_job(job_mode: conf.JobMode,
   if request_retries is None:
     request_retries = 10
 
-  engine = get_mem_engine() if dry_run else get_sql_engine()
+  engine = get_mem_engine() if dry_run else get_sql_engine(strict=strict_db)
 
   with session_scope(engine) as session:
     container_spec = generate_container_spec(session, docker_args, image_tag)

caliban/docker.py

@@ -810,6 +810,7 @@ def run_experiments(job_mode: c.JobMode,
                     dry_run: bool = False,
                     experiment_config: Optional[c.ExpConf] = None,
                     xgroup: Optional[str] = None,
+                    strict_db: bool = False,
                     **build_image_kwargs) -> None:
   """Builds an image using the supplied **build_image_kwargs and calls `docker
   run` on the resulting image using sensible defaults.
@@ -830,6 +831,8 @@ def run_experiments(job_mode: c.JobMode,
   - dry_run: if True, no actual jobs will be executed and docker won't
     actually build; logging side effects will show the user what will happen
     without dry_run=True.
+  - xgroup: experiment group for this command
+  - strict_db: if database specified by the CALIBAN_DB_URL is not found, exit
 
   any extra kwargs supplied are passed through to build_image.
   """
@@ -845,7 +848,7 @@ def run_experiments(job_mode: c.JobMode,
   docker_args = {k: v for k, v in build_image_kwargs.items()}
   docker_args['job_mode'] = job_mode
 
-  engine = get_mem_engine() if dry_run else get_sql_engine()
+  engine = get_mem_engine() if dry_run else get_sql_engine(strict=strict_db)
 
   with session_scope(engine) as session:
     container_spec = generate_container_spec(session, docker_args, image_id)

caliban/gke/cli.py

@@ -72,7 +72,13 @@ def wrapper(args: dict,
                           zone=zone,
                           creds=creds)
 
-    return fn(args, cluster=cluster) if cluster else None
+    if cluster is None:
+      logging.error(f'unable to resolve cluster')
+      logging.error(f'you can see your available clusters using the command:')
+      logging.error(f'caliban cluster ls')
+      return
+
+    return fn(args, cluster=cluster)
 
   return wrapper
 
@@ -359,6 +365,7 @@ def _job_submit(args: dict, cluster: Cluster) -> None:
   xgroup = args.get('xgroup')
   image_tag = args.get('image_tag')
   export = args.get('export', None)
+  strict_db = args.get('strict_db', False)
 
   labels = args.get('label')
   if labels is not None:
@@ -415,7 +422,7 @@ def _job_submit(args: dict, cluster: Cluster) -> None:
   accel, accel_count = accel_spec
 
   # --------------------------------------------------------------------------
-  engine = get_mem_engine() if dry_run else get_sql_engine()
+  engine = get_mem_engine() if dry_run else get_sql_engine(strict=strict_db)
 
   with session_scope(engine) as session:
     container_spec = generate_container_spec(session, docker_m, image_tag)

caliban/history/utils.py

@@ -25,7 +25,7 @@
 
 from sqlalchemy import create_engine
 from sqlalchemy.engine.base import Engine
-from sqlalchemy.exc import OperationalError
+from sqlalchemy.exc import OperationalError, ArgumentError
 from sqlalchemy.orm import Session, sessionmaker
 
 import caliban.config as conf
@@ -108,7 +108,7 @@ def get_sql_engine(
   try:
     return _create_sqa_engine(url=url, echo=echo)
 
-  except (OperationalError, OSError) as e:
+  except (OperationalError, OSError, ArgumentError) as e:
     logging.error("")
     logging.error(
         t.red(

caliban/main.py

@@ -97,6 +97,7 @@ def run_app(arg_input):
     image_id = args.get("image_id")
     exp_config = args.get("experiment_config")
     xgroup = args.get('xgroup')
+    strict_db = args.get('strict_db')
 
     docker.run_experiments(job_mode,
                            run_args=docker_run_args,
@@ -106,6 +107,7 @@ def run_app(arg_input):
                            dry_run=dry_run,
                            package=package,
                            xgroup=xgroup,
+                           strict_db=strict_db,
                            **docker_args)
 
   elif command == "cloud":
@@ -123,6 +125,7 @@ def run_app(arg_input):
     exp_config = args.get("experiment_config")
     labels = u.sanitize_labels(args.get("label") or [])
     xgroup = args.get('xgroup')
+    strict_db = args.get('strict_db')
 
     # Arguments to internally build the image required to submit to Cloud.
     docker_m = {"job_mode": job_mode, "package": package, **docker_args}
@@ -143,6 +146,7 @@ def run_app(arg_input):
         script_args=script_args,
         experiment_config=exp_config,
         xgroup=xgroup,
+        strict_db=strict_db,
     )
   else:
     logging.info("Unknown command: {}".format(command))