wandb

1
#!/usr/bin/env python
2
"""Generate dbt files for telemetry.
3

4
Data directory for telemetry records:
5
    https://github.com/wandb/analytics/tree/master/dbt/data
6

7
Usage:
8
    ./wandb/tools/telemetry-tool.py --output-dir analytics/dbt/seeds/
9
"""
10

11
import argparse
12
import csv
13
import os
14
from typing import Any, List
15

16
from wandb.proto import wandb_telemetry_pb2 as tpb
17

18
DEFAULT_DIR: str = "analytics/dbt/seeds"
19
parser = argparse.ArgumentParser()
20
parser.add_argument(
21
    "--output-dir", default=DEFAULT_DIR if os.path.exists(DEFAULT_DIR) else ""
22
)
23
parser.add_argument(
24
    "--output-telemetry-record-types", default="map_run_cli_telemetry_record_types.csv"
25
)
26
parser.add_argument("--output-imports", default="map_run_cli_imports.csv")
27
parser.add_argument("--output-features", default="map_run_cli_features.csv")
28
parser.add_argument("--output-environments", default="map_run_cli_environments.csv")
29
parser.add_argument("--output-labels", default="map_run_cli_labels.csv")
30
parser.add_argument(
31
    "--output-deprecated-features", default="map_run_cli_deprecated.csv"
32
)
33
parser.add_argument("--output-issues", default="map_run_cli_issues.csv")
34
args = parser.parse_args()
35

36

37
def write_csv(record: str, fields: List[Any]):
38
    record_arg = f"output_{record}s"
39
    fname = os.path.join(args.output_dir, getattr(args, record_arg))
40
    print("Writing:", fname)
41
    with open(fname, "w") as fp:
42
        writer = csv.DictWriter(fp, fieldnames=[record, "key"], lineterminator="\n")
43
        writer.writeheader()
44
        for f in fields:
45
            # let's skip private fields
46
            if f.name.startswith("_"):
47
                continue
48
            writer.writerow({record: f.name, "key": f.number})
49

50

51
def main():
52
    telemetry_records = list(tpb.TelemetryRecord.DESCRIPTOR.fields)
53
    write_csv(record="telemetry_record_type", fields=telemetry_records)
54

55
    import_records = list(tpb.Imports.DESCRIPTOR.fields)
56
    write_csv(record="import", fields=import_records)
57

58
    feature_records = list(tpb.Feature.DESCRIPTOR.fields)
59
    write_csv(record="feature", fields=feature_records)
60

61
    env_records = list(tpb.Env.DESCRIPTOR.fields)
62
    write_csv(record="environment", fields=env_records)
63

64
    label_records = list(tpb.Labels.DESCRIPTOR.fields)
65
    write_csv(record="label", fields=label_records)
66

67
    deprecated_records = list(tpb.Deprecated.DESCRIPTOR.fields)
68
    write_csv(record="deprecated_feature", fields=deprecated_records)
69

70
    issue_records = list(tpb.Issues.DESCRIPTOR.fields)
71
    write_csv(record="issue", fields=issue_records)
72

73

74
if __name__ == "__main__":
75
    main()
76