()
| 10 | |
| 11 | |
| 12 | def main(): |
| 13 | parser = argparse.ArgumentParser( |
| 14 | prog="megengine.tools.profiler", description="Profiling megengine program" |
| 15 | ) |
| 16 | parser.add_argument( |
| 17 | "-m", "--module", action="store_true", help="whether launch program as module" |
| 18 | ) |
| 19 | parser.add_argument("-o", "--output", type=str, help="output file location") |
| 20 | parser.add_argument( |
| 21 | "-f", |
| 22 | "--format", |
| 23 | action="append", |
| 24 | type=str, |
| 25 | help="output file format", |
| 26 | choices=Profiler.valid_formats, |
| 27 | ) |
| 28 | parser.add_argument( |
| 29 | "--merge_trace_events", action="store_true", |
| 30 | ) |
| 31 | parser.add_argument( |
| 32 | "--clean", action="store_true", |
| 33 | ) |
| 34 | for opt in Profiler.valid_options: |
| 35 | parser.add_argument("--" + opt, type=int, default=None) |
| 36 | args, extras = parser.parse_known_args(sys.argv[1:]) |
| 37 | prof_args = {} |
| 38 | for opt in Profiler.valid_options: |
| 39 | optval = getattr(args, opt, None) |
| 40 | if optval is not None: |
| 41 | prof_args[opt] = optval |
| 42 | |
| 43 | if args.output is not None: |
| 44 | prof_args["path"] = args.output |
| 45 | |
| 46 | if args.format: |
| 47 | prof_args["formats"] = args.format |
| 48 | |
| 49 | if len(extras) == 0: |
| 50 | if not args.merge_trace_events: |
| 51 | parser.print_usage() |
| 52 | exit(1) |
| 53 | else: |
| 54 | filename = extras[0] |
| 55 | if not args.module: |
| 56 | if not os.path.exists(filename): |
| 57 | get_logger().fatal("cannot find file {}".format(filename)) |
| 58 | exit(1) |
| 59 | filename = os.path.realpath(filename) |
| 60 | # Replace profiler's dir with script's dir in front of module search path. |
| 61 | sys.path[0] = os.path.dirname(filename) |
| 62 | |
| 63 | sys.argv[:] = [filename, *extras[1:]] |
| 64 | |
| 65 | profiler = Profiler(**prof_args) |
| 66 | |
| 67 | if args.clean: |
| 68 | for file in os.listdir(profiler.directory): |
| 69 | os.remove(os.path.join(profiler.directory, file)) |
no test coverage detected