chore: add a dedicated exception for cli errors

harness/determined/cli/__init__.py

@@ -5,7 +5,6 @@
     setup_session,
     require_feature_flag,
     login_sdk_client,
-    report_cli_error,
     print_warnings,
 )
 from determined.cli import (

harness/determined/cli/_util.py

@@ -1,7 +1,7 @@
 import argparse
 import functools
 import sys
-from typing import Any, Callable, Dict, List, Sequence, Union
+from typing import Any, Callable, Dict, List, Sequence
 
 import termcolor
 
@@ -115,11 +115,6 @@ def wrapper(args: argparse.Namespace) -> None:
     return decorator
 
 
-def report_cli_error(msg: Union[str, Exception]) -> None:
-    print(termcolor.colored(f"Error: {msg}", "red"), file=sys.stderr)
-    sys.exit(1)
-
-
 def print_warnings(warnings: Sequence[bindings.v1LaunchWarning]) -> None:
     for warning in warnings:
         print(termcolor.colored(api.WARNING_MESSAGE_MAP[warning], "yellow"), file=sys.stderr)

harness/determined/cli/agent.py

@@ -9,6 +9,7 @@
 from determined import cli
 from determined.cli import render
 from determined.cli import task as cli_task
+from determined.cli.errors import CliError
 from determined.common import api
 from determined.common.api import authentication, bindings
 from determined.common.check import check_false
@@ -147,8 +148,7 @@ def patch(args: argparse.Namespace) -> None:
         check_false(args.all and args.agent_id)
 
         if not (args.all or args.agent_id):
-            print("Error: must specify exactly one of `--all` or agent_id", file=sys.stderr)
-            sys.exit(1)
+            raise CliError("Must specify exactly on of --all or --agent-id")
 
         if args.agent_id:
             agent_ids = [args.agent_id]

harness/determined/cli/cli.py

@@ -52,7 +52,7 @@
 )
 from determined.errors import EnterpriseOnlyError
 
-from .errors import FeatureFlagDisabled
+from .errors import CliError, FeatureFlagDisabled
 
 
 @authentication.required
@@ -192,13 +192,13 @@ def main(
 
         parsed_args = parser.parse_args(args)
 
-        def die(message: str, always_print_traceback: bool = False) -> None:
+        def die(message: str, always_print_traceback: bool = False, exit_code: int = 1) -> None:
             if always_print_traceback or debug_mode():
                 import traceback
 
                 traceback.print_exc(file=sys.stderr)
 
-            parser.exit(1, colored(message + "\n", "red"))
+            parser.exit(exit_code, colored(message + "\n", "red"))
 
         v = vars(parsed_args)
         if not v.get("func"):
@@ -280,6 +280,10 @@ def die(message: str, always_print_traceback: bool = False) -> None:
             die(f"Determined Enterprise Edition is required for this functionality: {e}")
         except FeatureFlagDisabled as e:
             die(f"Master does not support this operation: {e}")
+        except CliError as e:
+            if e.e_stack:
+                print(colored(f"Error: {e}", "yellow"), file=sys.stderr)
+            die(f"{e.name}: {e.message}", exit_code=e.exit_code)
         except Exception:
             die("Failed to {}".format(parsed_args.func.__name__), always_print_traceback=True)
     except KeyboardInterrupt:

harness/determined/cli/command.py

@@ -10,6 +10,7 @@
 
 from determined import cli
 from determined.cli import render
+from determined.cli.errors import CliArgError
 from determined.common import api, context, util, yaml
 from determined.common.api import authentication
 
@@ -199,7 +200,7 @@ def list_tasks(args: Namespace) -> None:
             cli.setup_session(args), args.workspace_name
         )
         if workspace is None:
-            return cli.report_cli_error(f'Workspace "{args.workspace_name}" not found.')
+            raise CliArgError(f'Workspace "{args.workspace_name}" not found.')
 
         params["workspaceId"] = workspace.id
 

harness/determined/cli/errors.py

@@ -1,7 +1,43 @@
+from typing import Any, Optional
+
+
 class FeatureFlagDisabled(Exception):
     """
     Exception indicating that there is a currently disabled feature flag
     that is required to use a feature
     """
 
     pass
+
+
+class CliError(Exception):
+    """
+    Base class for all CLI errors.
+    """
+
+    name: str
+
+    def __init__(
+        self, message: str, e_stack: Optional[Exception] = None, exit_code: int = 1
+    ) -> None:
+        """
+        Args:
+        - e_stack: The exception that caused this error.
+        - exit_code: The exit code to use when exiting the CLI.
+        """
+        super().__init__(message)
+        self.name = "Error"
+        self.exit_code = exit_code
+        self.e_stack = e_stack
+        self.message = message
+
+
+class CliArgError(CliError):
+    """
+    Exception indicating that there was a problem with the arguments
+    passed to the CLI.
+    """
+
+    def __init__(self, message: str, *args: Any, **kwargs: Any) -> None:
+        super().__init__(message, *args, **kwargs)
+        self.name = "Input Error"

harness/determined/cli/experiment.py

@@ -19,6 +19,7 @@
 from determined import cli
 from determined.cli import checkpoint, render
 from determined.cli.command import CONFIG_DESC, parse_config_overrides
+from determined.cli.errors import CliArgError, CliError
 from determined.common import api, context, set_logger, util, yaml
 from determined.common.api import authentication, bindings, logs
 from determined.common.declarative_argparse import Arg, Cmd, Group
@@ -65,36 +66,32 @@ def read_git_metadata(model_def_path: pathlib.Path) -> Tuple[str, str, str, str]
     try:
         from git import Repo
     except ImportError as e:  # pragma: no cover
-        print(f"Error: Please verify that git is installed correctly: {e}")
-        sys.exit(1)
+        raise CliError(f"Error: Please verify that git is installed correctly: {e}")
 
     if model_def_path.is_dir():
         repo_path = model_def_path.resolve()
     else:
         repo_path = model_def_path.parent.resolve()
 
     if not repo_path.joinpath(".git").is_dir():
-        print(
+        raise CliError(
             f"Error: No git directory found at {repo_path}. Please "
             "initialize a git repository or refrain from "
             "using the --git feature."
         )
-        sys.exit(1)
 
     try:
         repo = Repo(str(repo_path))
     except Exception as e:
-        print(f"Failed to initialize git repository at {repo_path}: {e}")
-        sys.exit(1)
+        raise CliError(f"Failed to initialize git repository at {repo_path}: {e}")
 
     if repo.is_dirty():
-        print(
+        raise CliError(
             "Git working directory is dirty. Please commit the "
             "following changes before creating an experiment "
             "with the --git feature:\n"
+            f"\n{repo.git.status()}"
         )
-        print(repo.git.status())
-        sys.exit(1)
 
     commit = repo.commit()
     commit_hash = commit.hexsha
@@ -113,8 +110,7 @@ def read_git_metadata(model_def_path: pathlib.Path) -> Tuple[str, str, str, str]
         remote_url = repo.git.config(f"remote.{remote_name}.url", get=True)
         print(f"Using remote URL '{remote_url}' from upstream branch '{upstream_branch}'")
     except Exception as e:
-        print("Failed to find the upstream branch: ", e)
-        sys.exit(1)
+        raise CliError("Failed to find the upstream branch: ", e)
 
     return (remote_url, commit_hash, committer, commit_date)
 
@@ -125,8 +121,7 @@ def _parse_config_text_or_exit(
     experiment_config = util.safe_load_yaml_with_exceptions(config_text)
 
     if not experiment_config or not isinstance(experiment_config, dict):
-        print(f"Error: invalid experiment config file {path}", path)
-        sys.exit(1)
+        raise CliArgError(f"Error: invalid experiment config file {path}")
 
     parse_config_overrides(experiment_config, config_overrides)
 

harness/determined/cli/tensorboard.py

@@ -1,4 +1,3 @@
-import sys
 from argparse import ONE_OR_MORE, FileType, Namespace
 from functools import partial
 from pathlib import Path
@@ -8,6 +7,7 @@
 
 from determined import cli
 from determined.cli import command, task
+from determined.cli.errors import CliArgError
 from determined.common import api, context
 from determined.common.api import authentication, bindings, request
 from determined.common.check import check_eq
@@ -17,8 +17,7 @@
 @authentication.required
 def start_tensorboard(args: Namespace) -> None:
     if not (args.trial_ids or args.experiment_ids):
-        print("Either experiment_ids or trial_ids must be specified.")
-        sys.exit(1)
+        raise CliArgError("Either experiment_ids or trial_ids must be specified.")
 
     config = command.parse_config(args.config_file, None, args.config, [])
 

harness/determined/cli/workspace.py

@@ -4,6 +4,7 @@
 from typing import Any, Dict, List, Optional, Sequence
 
 from determined import cli
+from determined.cli.errors import CliArgError
 from determined.cli.user import AGENT_USER_GROUP_ARGS
 from determined.common import api, util
 from determined.common.api import authentication, bindings, errors
@@ -32,11 +33,9 @@ def get_workspace_id_from_args(args: Namespace) -> Optional[int]:
             cli.setup_session(args), args.workspace_name
         )
         if workspace is None:
-            cli.report_cli_error(f'Workspace "{args.workspace_name}" not found')
-            return None
+            raise CliArgError(f'Workspace "{args.workspace_name}" not found')
         if workspace.archived:
-            cli.report_cli_error(f'Workspace "{args.workspace_name}" is archived')
-            return None
+            raise CliArgError(f'Workspace "{args.workspace_name}" is archived')
         workspace_id = workspace.id
     return workspace_id
 

harness/determined/deploy/aws/cli.py

@@ -2,7 +2,6 @@
 import base64
 import json
 import re
-import sys
 from pathlib import Path
 from typing import Callable, Dict, Tuple, Type
 
@@ -12,6 +11,7 @@
 
 from determined.common.declarative_argparse import Arg, ArgGroup, BoolOptArg, Cmd
 from determined.deploy.errors import MasterTimeoutExpired
+from harness.determined.cli.errors import CliError
 
 from . import aws, constants
 from .deployment_types import base, govcloud, secure, simple, vpc
@@ -55,11 +55,10 @@ def error_no_credentials() -> None:
         colored("Unable to locate AWS credentials.", "red"),
         "Did you run %s?" % colored("aws configure", "yellow"),
     )
-    print(
+    raise CliError(
         "See the AWS Documentation for information on how to use AWS credentials:",
         "https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-files.html",
     )
-    sys.exit(1)
 
 
 def get_deployment_class(deployment_type: str) -> Type[base.DeterminedDeployment]:
@@ -84,18 +83,18 @@ def deploy_aws(command: str, args: argparse.Namespace) -> None:
             f"`det deploy` is only supported in {constants.misc.SUPPORTED_REGIONS} - "
             f"tried to deploy to {boto3_session.region_name}"
         )
-        print("use the --region argument to deploy to a supported region")
-        sys.exit(1)
+        raise CliError("use the --region argument to deploy to a supported region")
 
     if command == "list":
         try:
             output = aws.list_stacks(boto3_session)
         except NoCredentialsError:
             error_no_credentials()
         except Exception as e:
-            print(e)
-            print("Listing stacks failed. Check the AWS CloudFormation Console for details.")
-            sys.exit(1)
+            raise CliError(
+                "Listing stacks failed. Check the AWS CloudFormation Console for details.",
+                e_stack=e,
+            )
         for item in output:
             print(item["StackName"])
         return
@@ -109,8 +108,7 @@ def deploy_aws(command: str, args: argparse.Namespace) -> None:
     #     sys.exit(1)
 
     if not re.match(constants.misc.CLOUDFORMATION_REGEX, args.cluster_id):
-        print("Deployment Failed - cluster-id much match ^[a-zA-Z][-a-zA-Z0-9]*$")
-        sys.exit(1)
+        raise CliError("Deployment Failed - cluster-id much match ^[a-zA-Z][-a-zA-Z0-9]*$")
 
     if command == "down":
         if not args.yes:
@@ -128,18 +126,18 @@ def deploy_aws(command: str, args: argparse.Namespace) -> None:
         except NoCredentialsError:
             error_no_credentials()
         except Exception as e:
-            print(e)
-            print("Stack Deletion Failed. Check the AWS CloudFormation Console for details.")
-            sys.exit(1)
+            raise CliError(
+                "Stack Deletion Failed. Check the AWS CloudFormation Console for details.",
+                e_stack=e,
+            )
 
         print("Delete Successful")
         return
 
     if (args.cpu_env_image and not args.gpu_env_image) or (
         args.gpu_env_image and not args.cpu_env_image
     ):
-        print("If a CPU or GPU environment image is specified, both should be.")
-        sys.exit(1)
+        raise CliError("If a CPU or GPU environment image is specified, both should be.")
 
     if args.deployment_type != constants.deployment_types.SIMPLE:
         if args.agent_subnet_id is not None:
@@ -247,13 +245,9 @@ def deploy_aws(command: str, args: argparse.Namespace) -> None:
     except NoCredentialsError:
         error_no_credentials()
     except Exception as e:
-        print(e)
-        print(
-            colored(
-                "Stack Deployment Failed. Check the AWS CloudFormation Console for details.", "red"
-            )
+        raise CliError(
+            "Stack Deployment Failed. Check the AWS CloudFormation Console for details.", e_stack=e
         )
-        sys.exit(1)
 
     if not args.no_wait_for_master:
         try:
@@ -265,8 +259,9 @@ def deploy_aws(command: str, args: argparse.Namespace) -> None:
                     "red",
                 )
             )
-            print("For details, SSH to master instance and check /var/log/cloud-init-output.log.")
-            sys.exit(1)
+            raise CliError(
+                "For details, SSH to master instance and check /var/log/cloud-init-output.log."
+            )
 
     print("Determined Deployment Successful")