Skip to content

Commit

Permalink
wrap cli exec to gracefully catch keyboard exit signal (#603)
Browse files Browse the repository at this point in the history 
Signed-off-by: Jeffrey Martin <[email protected]>
  • Loading branch information
@jmartin-tech
jmartin-tech authored Apr 16, 2024
1 parent 4127ae5 commit 88d048f
Showing 1 changed file with 156 additions and 154 deletions.
310 changes: 156 additions & 154 deletions garak/cli.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -328,168 +328,170 @@ def main(arguments=[]) -> None:

import garak.evaluators

if not args.version and not args.report:
command.start_run()

# do a special thing for CLIprobe options, generator options
if "probe_options" in args or "probe_option_file" in args:
if "probe_options" in args:
try:
probe_cli_config = json.loads(args.probe_options)
except json.JSONDecodeError as e:
logging.warning("Failed to parse JSON probe_options: %s", e.args[0])

elif "probe_option_file" in args:
with open(args.probe_option_file, encoding="utf-8") as f:
probe_options_json = f.read().strip()
try:
probe_cli_config = json.loads(probe_options_json)
except json.decoder.JSONDecodeError as e:
logging.warning("Failed to parse JSON probe_options: %s", {e.args[0]})
raise e
try:
if not args.version and not args.report:
command.start_run()

# do a special thing for CLIprobe options, generator options
if "probe_options" in args or "probe_option_file" in args:
if "probe_options" in args:
try:
probe_cli_config = json.loads(args.probe_options)
except json.JSONDecodeError as e:
logging.warning("Failed to parse JSON probe_options: %s", e.args[0])

elif "probe_option_file" in args:
with open(args.probe_option_file, encoding="utf-8") as f:
probe_options_json = f.read().strip()
try:
probe_cli_config = json.loads(probe_options_json)
except json.decoder.JSONDecodeError as e:
logging.warning("Failed to parse JSON probe_options: %s", {e.args[0]})
raise e

_config.plugins.probes = _config._combine_into(
probe_cli_config, _config.plugins.probes
)

_config.plugins.probes = _config._combine_into(
probe_cli_config, _config.plugins.probes
)
if "generator_options" in args or "generator_option_file" in args:
if "generator_options" in args:
try:
generator_cli_config = json.loads(args.generator_options)
except json.JSONDecodeError as e:
logging.warning("Failed to parse JSON generator_options: %s", e.args[0])

elif "generator_option_file" in args:
with open(args.generator_option_file, encoding="utf-8") as f:
generator_options_json = f.read().strip()
try:
generator_cli_config = json.loads(generator_options_json)
except json.decoder.JSONDecodeError as e:
logging.warning(
"Failed to parse JSON generator_options: %s", {e.args[0]}
)
raise e

_config.plugins.generators = _config._combine_into(
generator_cli_config, _config.plugins.generators
)

if "generator_options" in args or "generator_option_file" in args:
if "generator_options" in args:
try:
generator_cli_config = json.loads(args.generator_options)
except json.JSONDecodeError as e:
logging.warning("Failed to parse JSON generator_options: %s", e.args[0])
# process commands
if args.interactive:
from garak.interactive import interactive_mode

elif "generator_option_file" in args:
with open(args.generator_option_file, encoding="utf-8") as f:
generator_options_json = f.read().strip()
try:
generator_cli_config = json.loads(generator_options_json)
except json.decoder.JSONDecodeError as e:
logging.warning(
"Failed to parse JSON generator_options: %s", {e.args[0]}
)
raise e

_config.plugins.generators = _config._combine_into(
generator_cli_config, _config.plugins.generators
)

# process commands
if args.interactive:
from garak.interactive import interactive_mode

try:
interactive_mode()
except Exception as e:
logging.error(e)
print(e)
sys.exit(1)

if args.version:
pass

elif args.plugin_info:
command.plugin_info(args.plugin_info)

elif args.list_probes:
command.print_probes()

elif args.list_detectors:
command.print_detectors()

elif args.list_buffs:
command.print_buffs()

elif args.list_generators:
command.print_generators()

elif args.list_config:
print("cli args:\n ", args)
command.list_config()

elif args.report:
from garak.report import Report

report_location = args.report
print(f"📜 Converting garak reports {report_location}")
report = Report(args.report).load().get_evaluations()
report.export()
print(f"📜 AVID reports generated at {report.write_location}")

# model is specified, we're doing something
elif _config.plugins.model_type:
if (
_config.plugins.model_type
in ("openai", "replicate", "ggml", "huggingface", "litellm")
and not _config.plugins.model_name
):
message = f"⚠️ Model type '{_config.plugins.model_type}' also needs a model name\n You can set one with e.g. --model_name \"billwurtz/gpt-1.0\""
logging.error(message)
raise ValueError(message)
print(f"📜 reporting to {_config.transient.report_filename}")

generator_module_name = _config.plugins.model_type.split(".")[0]
generator_mod = importlib.import_module(
"garak.generators." + generator_module_name
)
if "." not in _config.plugins.model_type:
if generator_mod.default_class:
generator_class_name = generator_mod.default_class
interactive_mode()
except Exception as e:
logging.error(e)
print(e)
sys.exit(1)

if args.version:
pass

elif args.plugin_info:
command.plugin_info(args.plugin_info)

elif args.list_probes:
command.print_probes()

elif args.list_detectors:
command.print_detectors()

elif args.list_buffs:
command.print_buffs()

elif args.list_generators:
command.print_generators()

elif args.list_config:
print("cli args:\n ", args)
command.list_config()

elif args.report:
from garak.report import Report

report_location = args.report
print(f"📜 Converting garak reports {report_location}")
report = Report(args.report).load().get_evaluations()
report.export()
print(f"📜 AVID reports generated at {report.write_location}")

# model is specified, we're doing something
elif _config.plugins.model_type:
if (
_config.plugins.model_type
in ("openai", "replicate", "ggml", "huggingface", "litellm")
and not _config.plugins.model_name
):
message = f"⚠️ Model type '{_config.plugins.model_type}' also needs a model name\n You can set one with e.g. --model_name \"billwurtz/gpt-1.0\""
logging.error(message)
raise ValueError(message)
print(f"📜 reporting to {_config.transient.report_filename}")

generator_module_name = _config.plugins.model_type.split(".")[0]
generator_mod = importlib.import_module(
"garak.generators." + generator_module_name
)
if "." not in _config.plugins.model_type:
if generator_mod.default_class:
generator_class_name = generator_mod.default_class
else:
raise ValueError(
"module {generator_module_name} has no default class; pass module.ClassName to --model_type"
)
else:
raise ValueError(
"module {generator_module_name} has no default class; pass module.ClassName to --model_type"
)
else:
generator_class_name = _config.plugins.model_type.split(".")[1]

# if 'model_name' not in args:
# generator = getattr(generator_mod, generator_class_name)()
# else:
generator = getattr(generator_mod, generator_class_name)(
_config.plugins.model_name
)
if hasattr(_config.run, "generations") and _config.run.generations is not None:
generator.generations = _config.run.generations
if hasattr(_config.run, "seed") and _config.run.seed is not None:
generator.seed = _config.run.seed

if "generate_autodan" in args and args.generate_autodan:
from garak.resources.autodan import autodan_generate
generator_class_name = _config.plugins.model_type.split(".")[1]

try:
prompt = _config.probe_options["prompt"]
target = _config.probe_options["target"]
except Exception as e:
print(
"AutoDAN generation requires --probe_options with a .json containing a `prompt` and `target` "
"string"
)
autodan_generate(generator=generator, prompt=prompt, target=target)
# if 'model_name' not in args:
# generator = getattr(generator_mod, generator_class_name)()
# else:
generator = getattr(generator_mod, generator_class_name)(
_config.plugins.model_name
)
if hasattr(_config.run, "generations") and _config.run.generations is not None:
generator.generations = _config.run.generations
if hasattr(_config.run, "seed") and _config.run.seed is not None:
generator.seed = _config.run.seed

if "generate_autodan" in args and args.generate_autodan:
from garak.resources.autodan import autodan_generate

try:
prompt = _config.probe_options["prompt"]
target = _config.probe_options["target"]
except Exception as e:
print(
"AutoDAN generation requires --probe_options with a .json containing a `prompt` and `target` "
"string"
)
autodan_generate(generator=generator, prompt=prompt, target=target)

probe_names = _config.parse_plugin_spec(
_config.plugins.probe_spec, "probes", _config.run.probe_tags
)
detector_names = _config.parse_plugin_spec(
_config.plugins.detector_spec, "detectors"
)
buff_names = _config.parse_plugin_spec(_config.plugins.buff_spec, "buffs")

probe_names = _config.parse_plugin_spec(
_config.plugins.probe_spec, "probes", _config.run.probe_tags
)
detector_names = _config.parse_plugin_spec(
_config.plugins.detector_spec, "detectors"
)
buff_names = _config.parse_plugin_spec(_config.plugins.buff_spec, "buffs")
evaluator = garak.evaluators.ThresholdEvaluator(_config.run.eval_threshold)

evaluator = garak.evaluators.ThresholdEvaluator(_config.run.eval_threshold)
if detector_names == []:
command.probewise_run(generator, probe_names, evaluator, buff_names)

if detector_names == []:
command.probewise_run(generator, probe_names, evaluator, buff_names)
else:
command.pxd_run(
generator, probe_names, detector_names, evaluator, buff_names
)

command.end_run()
else:
command.pxd_run(
generator, probe_names, detector_names, evaluator, buff_names
)

command.end_run()

else:
print("nothing to do 🤷 try --help")
if _config.plugins.model_name and not _config.plugins.model_type:
print(
"💡 try setting --model_type (--model_name is currently set but not --model_type)"
)
logging.info("nothing to do 🤷")
print("nothing to do 🤷 try --help")
if _config.plugins.model_name and not _config.plugins.model_type:
print(
"💡 try setting --model_type (--model_name is currently set but not --model_type)"
)
logging.info("nothing to do 🤷")
except KeyboardInterrupt:
print("User cancel received, terminating all runs")

0 comments on commit 88d048f

Please sign in to comment.