Allow to disable postmortem on fail in binaries

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/356

Attaching PDB on failure is not working when running in distributed environment. This change allows to disable this behavior by passing a command line argument.

Reviewed By: miqueljubert

Differential Revision: D38514736

fbshipit-source-id: 2e0008d6fbc6a4518a605debe67d76f8354364fc

d2go/setup.py

@@ -68,6 +68,11 @@ def basic_argument_parser(
         default=None,
         type=str,
     )
+    parser.add_argument(
+        "--disable-post-mortem",
+        action="store_true",
+        help="whether to NOT connect pdb on failure, which only works locally",
+    )
 
     if distributed:
         parser.add_argument(
@@ -98,6 +103,7 @@ def build_basic_cli_args(
     machine_rank: Optional[Union[int, str]] = None,
     dist_url: Optional[str] = None,
     dist_backend: Optional[str] = None,
+    disable_post_mortem: bool = False,
 ) -> List[str]:
     """
     Returns parameters in the form of CLI arguments for the binary using
@@ -114,6 +120,8 @@ def build_basic_cli_args(
         args += ["--runner", runner_name]
     if save_return_file is not None:
         args += ["--save-return-file", str(save_return_file)]
+    if disable_post_mortem:
+        args += ["--disable-post-mortem"]
     if num_processes is not None:
         args += ["--num-processes", str(num_processes)]
     if num_machines is not None:

tools/benchmark_data.py

@@ -144,8 +144,9 @@ def main(
 
 def run_with_cmdline_args(args):
     cfg, output_dir, runner_name = prepare_for_launch(args)
+    main_func = main if args.disable_post_mortem else post_mortem_if_fail_for_main(main)
     launch(
-        post_mortem_if_fail_for_main(main),
+        main_func,
         num_processes_per_machine=args.num_processes,
         num_machines=args.num_machines,
         machine_rank=args.machine_rank,

tools/evaluator.py

@@ -24,7 +24,6 @@ from d2go.setup import (
     setup_after_launch,
 )
 from d2go.utils.misc import print_metrics_table
-from mobile_cv.common.misc.py import post_mortem_if_fail
 from mobile_cv.predictor.api import create_predictor
 
 logger = logging.getLogger("d2go.tools.caffe2_evaluator")
@@ -61,11 +60,11 @@ def main(
     )
 
 
-@post_mortem_if_fail()
 def run_with_cmdline_args(args):
     cfg, output_dir, runner_name = prepare_for_launch(args)
+    main_func = main if args.disable_post_mortem else post_mortem_if_fail_for_main(main)
     launch(
-        post_mortem_if_fail_for_main(main),
+        main_func,
         args.num_processes,
         num_machines=args.num_machines,
         machine_rank=args.machine_rank,

tools/exporter.py

@@ -16,8 +16,12 @@ import mobile_cv.lut.lib.pt.flops_utils as flops_utils
 from d2go.config import CfgNode, temp_defrost
 from d2go.export.exporter import convert_and_export_predictor
 from d2go.runner import BaseRunner
-from d2go.setup import basic_argument_parser, prepare_for_launch, setup_after_launch
-from mobile_cv.common.misc.py import post_mortem_if_fail
+from d2go.setup import (
+    basic_argument_parser,
+    post_mortem_if_fail_for_main,
+    prepare_for_launch,
+    setup_after_launch,
+)
 
 
 logger = logging.getLogger("d2go.tools.export")
@@ -89,10 +93,10 @@ def main(
     )
 
 
-@post_mortem_if_fail()
 def run_with_cmdline_args(args):
     cfg, output_dir, runner_name = prepare_for_launch(args)
-    return main(
+    main_func = main if args.disable_post_mortem else post_mortem_if_fail_for_main(main)
+    return main_func(
         cfg,
         output_dir,
         runner_name,

tools/train_net.py

@@ -106,8 +106,9 @@ def main(
 def run_with_cmdline_args(args):
     cfg, output_dir, runner_name = prepare_for_launch(args)
 
+    main_func = main if args.disable_post_mortem else post_mortem_if_fail_for_main(main)
     outputs = launch(
-        post_mortem_if_fail_for_main(main),
+        main_func,
         num_processes_per_machine=args.num_processes,
         num_machines=args.num_machines,
         machine_rank=args.machine_rank,