1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
|
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
import argparse
import logging
from typing import Optional, Sequence
from tools.flight_recorder.components.fr_logger import FlightRecorderLogger
logger: FlightRecorderLogger = FlightRecorderLogger()
class JobConfig:
"""
A helper class to manage the script configuration.
"""
def __init__(self: "JobConfig"):
self.parser = argparse.ArgumentParser(
description="PyTorch Flight recorder analyzing script."
)
self.parser.add_argument(
"trace_dir",
nargs="?",
help="Directory containing one trace file per rank, named with <prefix>_<rank>.",
)
self.parser.add_argument(
"--selected-ranks",
default=None,
nargs="+",
type=int,
help="List of ranks we want to show traces for.",
)
self.parser.add_argument(
"--allow-incomplete-ranks",
action="store_true",
help=(
"FR trace require all ranks to have dumps for analysis. "
"This flag allows best-effort partial analysis of results "
"and printing of collected data."
),
)
self.parser.add_argument(
"--pg-filters",
default=None,
nargs="+",
type=str,
help=(
"List of filter strings, it could be pg name or pg desc. "
"If specified, only show traces for the given pg."
),
)
self.parser.add_argument("-o", "--output", default=None)
self.parser.add_argument(
"-p",
"--prefix",
help=(
"Common filename prefix to strip such that rank can be extracted. "
"If not specified, will attempt to infer a common prefix."
),
default=None,
)
self.parser.add_argument("-j", "--just_print_entries", action="store_true")
self.parser.add_argument("-v", "--verbose", action="store_true")
def parse_args(
self: "JobConfig", args: Optional[Sequence[str]]
) -> argparse.Namespace:
args = self.parser.parse_args(args)
if args.selected_ranks is not None:
assert (
args.just_print_entries
), "Not support selecting ranks without printing entries"
if args.pg_filters is not None:
assert (
args.just_print_entries
), "Not support selecting pg filters without printing entries"
if args.verbose:
logger.set_log_level(logging.DEBUG)
return args
|
