wandb

1
import fire
2
import wandb
3
from wandb.proto import wandb_internal_pb2
4
from wandb.sdk.internal import datastore
5

6

7
def _robust_scan(ds):
8
    """Attempt to scan data, handling incomplete files."""
9
    try:
10
        return ds.scan_data()
11
    except AssertionError as e:
12
        if ds.in_last_block():
13
            wandb.termwarn(
14
                f".wandb file is incomplete ({e}), be sure to sync this run again once it's finished"
15
            )
16
            return None
17
        else:
18
            raise e
19

20

21
def run(wandb_file: str) -> None:
22
    ds = datastore.DataStore()
23
    try:
24
        ds.open_for_scan(wandb_file)
25
    except AssertionError as e:
26
        print(f".wandb file is empty ({e}), skipping: {wandb_file}")
27
        return
28

29
    # save exit for final send
30
    while True:
31
        data = _robust_scan(ds)
32
        if data is None:
33
            break
34
        pb = wandb_internal_pb2.Record()
35
        pb.ParseFromString(data)
36
        record_type = pb.WhichOneof("record_type")
37
        print(f"RECORD TYPE: {record_type}")
38
        print(pb)
39
        print()
40
        input()
41

42

43
if __name__ == "__main__":
44
    fire.Fire(run)
45