# 手动将上面代码 VLOG(10) 改成了 VLOG(7)，因为 GLOG_v=10 PT 会挂，暂不分析
# 当然也可以用 GLOG_vmodule=auto_growth_best_fit_allocator=10 来只收集显存信息，前期只分析显存可以这样，后期因为需要分析周围 log，所以都需要 dump
# 设置 FLAGS_new_executor_sequential_run 确保顺序是可比的
GLOG_v=7 FLAGS_new_executor_sequential_run=true SOT_LOG_LEVEL=0 ENABLE_FALL_BACK=True MIN_GRAPH_SIZE=0 FLAGS_enable_pir_api=1 python tools/train.py -c ppcls/configs/ImageNet/CSWinTransformer/CSWinTransformer_base_384.yaml \
 -o Global.epochs=1 \
 -o Global.save_interval=1 \
 -o Global.eval_interval=1 \
 -o Global.seed=1234 \
 -o DataLoader.Train.dataset.image_root=/workspace/PaddleClas/dataset/ILSVRC2012/ \
 -o DataLoader.Train.dataset.cls_label_path=/workspace/PaddleClas/dataset/ILSVRC2012/train_list.txt \
 -o DataLoader.Train.sampler.batch_size=8 \
 -o DataLoader.Eval.dataset.image_root=/workspace/PaddleClas/dataset/ILSVRC2012/ \
 -o DataLoader.Eval.dataset.cls_label_path=/workspace/PaddleClas/dataset/ILSVRC2012/val_list.txt \
 -o DataLoader.Eval.sampler.batch_size=8 \
 -o DataLoader.Train.loader.num_workers=0 \
 -o DataLoader.Train.sampler.shuffle=False \
 -o Global.output_dir=output/ppcls/configs/ImageNet/CSWinTransformer/CSWinTransformer_base_384 \
 -o Global.to_static=True >! oom-pir-alloc.log 2>&1
 GLOG_v=7 FLAGS_new_executor_sequential_run=true SOT_LOG_LEVEL=0 ENABLE_FALL_BACK=True MIN_GRAPH_SIZE=0 FLAGS_enable_pir_api=0 python tools/train.py -c ppcls/configs/ImageNet/CSWinTransformer/CSWinTransformer_base_384.yaml \
 -o Global.epochs=1 \
 -o Global.save_interval=1 \
 -o Global.eval_interval=1 \
 -o Global.seed=1234 \
 -o DataLoader.Train.dataset.image_root=/workspace/PaddleClas/dataset/ILSVRC2012/ \
 -o DataLoader.Train.dataset.cls_label_path=/workspace/PaddleClas/dataset/ILSVRC2012/train_list.txt \
 -o DataLoader.Train.sampler.batch_size=8 \
 -o DataLoader.Eval.dataset.image_root=/workspace/PaddleClas/dataset/ILSVRC2012/ \
 -o DataLoader.Eval.dataset.cls_label_path=/workspace/PaddleClas/dataset/ILSVRC2012/val_list.txt \
 -o DataLoader.Eval.sampler.batch_size=8 \
 -o DataLoader.Train.loader.num_workers=0 \
 -o DataLoader.Train.sampler.shuffle=False \
 -o Global.output_dir=output/ppcls/configs/ImageNet/CSWinTransformer/CSWinTransformer_base_384 \
 -o Global.to_static=True >! oom-pt-alloc.log 2>&1

from __future__ import annotations
from dataclasses import dataclass
from typing import Callable
import sys
import re
LOG_PATH_TEMPLATE = "oom-{}-alloc.log"
# LOG_PATH = "oom-pt-alloc.log"
# LOG_PATH = "oom-pir-alloc.log"
class MemoryInfoItem:
    def __init__(self, ptr: str, size: int):
        self.ptr = ptr
        self.size = size
def format_bytes(size: int) -> str:
    units = ["B", "KB", "MB", "GB", "TB"]
    unit_idx = 0
    sign = "" if size >= 0 else "-"
    size = abs(size)
    while size >= 1024 and unit_idx < len(units):
        size /= 1024
        unit_idx += 1
    return f"{sign}{size:.2f} {units[unit_idx]}"
class AllocMemoryInfoItem(MemoryInfoItem):
    LOG_MATCH_REGEX = re.compile(r".+Alloc (?P<size>\d+) bytes, ptr = (?P<ptr>0x\w+)")
    def __init__(self, ptr: str, size: int):
        super().__init__(ptr, size)
    def __repr__(self):
        return f"AllocMemoryInfoItem(ptr={self.ptr}, size={self.size})"
class FreeMemoryInfoItem(MemoryInfoItem):
    LOG_MATCH_REGEX = re.compile(r".+Free (?P<size>\d+) bytes, ptr = (?P<ptr>0x\w+)")
    def __init__(self, ptr: str, size: int):
        super().__init__(ptr, size)
    def __repr__(self):
        return f"FreeMemoryInfoItem(ptr={self.ptr}, size={self.size})"
def extract_memory_info(logs: list[str], start_fn: Callable[[str], bool], end_fn:  Callable[[str], bool]) -> list[MemoryInfoItem]:
    memory_info = []
    started = False
    for line in logs:
        if not started and not start_fn(line):
            continue
        started = True
        if end_fn(line):
            break
        if match_obj := AllocMemoryInfoItem.LOG_MATCH_REGEX.match(line):
            memory_info.append(AllocMemoryInfoItem(match_obj.group("ptr"), int(match_obj.group("size"))))
        elif match_obj := FreeMemoryInfoItem.LOG_MATCH_REGEX.match(line):
            memory_info.append(FreeMemoryInfoItem(match_obj.group("ptr"), int(match_obj.group("size"))))
    return memory_info
def read_log(log_path: str) -> list[str]:
    with open(log_path, "r") as f:
        lines = f.readlines()
    return lines
class MemoryAnalyzer:
    def __call__(self, memory_info: list[MemoryInfoItem]):
        ...
class RemainingMemoryAnalyzer(MemoryAnalyzer):
    def __init__(self):
        self.remaining_memory = 0
        self.max_memory = 0
    def __call__(self, memory_info: list[MemoryInfoItem]):
        remaining_memory = 0
        for item in memory_info:
            if isinstance(item, AllocMemoryInfoItem):
                remaining_memory += item.size
            elif isinstance(item, FreeMemoryInfoItem):
                remaining_memory -= item.size
            if remaining_memory > self.max_memory:
                self.max_memory = remaining_memory
        self.remaining_memory = remaining_memory
    
    def summary(self):
        print(f"Remaining memory: {format_bytes(self.remaining_memory)}")
        print(f"Max memory: {format_bytes(self.max_memory)}")
class AllocFreeStatAnalyzer(MemoryAnalyzer):
    def __init__(self):
        self.alloc_size = 0
        self.free_size = 0
        self.alloc_count = 0
        self.free_count = 0
    def __call__(self, memory_info: list[MemoryInfoItem]):
        alloc_size = 0
        free_size = 0
        alloc_count = 0
        free_count = 0
        for item in memory_info:
            if isinstance(item, AllocMemoryInfoItem):
                alloc_size += item.size
                alloc_count += 1
            elif isinstance(item, FreeMemoryInfoItem):
                free_size += item.size
                free_count += 1
        self.alloc_size = alloc_size
        self.free_size = free_size
        self.alloc_count = alloc_count
        self.free_count = free_count
    def summary(self):
        print(f"Allocated memory: {format_bytes(self.alloc_size)}")
        print(f"Allocated count: {self.alloc_count}")
        print(f"Freed memory: {format_bytes(self.free_size)}")
        print(f"Freed count: {self.free_count}")
def analyse_memory_info(memory_info: list[MemoryInfoItem]):
    # remaining_memory
    remaining_memory_analyzer = RemainingMemoryAnalyzer()
    remaining_memory_analyzer(memory_info)
    remaining_memory_analyzer.summary()
    
    # alloc free stat
    alloc_free_stat_analyzer = AllocFreeStatAnalyzer()
    alloc_free_stat_analyzer(memory_info)
    alloc_free_stat_analyzer.summary()
logs = read_log(LOG_PATH_TEMPLATE.format(sys.argv[1]))
memory_info = extract_memory_info(
    logs,
    lambda line: "START SOT_CALL_0" in line,# or "EPOCH START" in line,
    lambda line: "EPOCH END" in line or "Traceback" in line or "START SOT_CALL_429" in line
)
# print(memory_info)
analyse_memory_info(memory_info)