Source code for mipcandy.profiler

from dataclasses import dataclass
from inspect import stack
from os import PathLike
from time import time
from typing import Sequence, override

import torch
from psutil import cpu_percent, virtual_memory

from mipcandy.data import dump_allocated_tensors
from mipcandy.types import Device, AmbiguousShape




@dataclass
class ProfilerFrame(object):
    stack: str
    cpu: float
    mem: float
    gpu: list[float] | None = None
    gpu_mem: list[float] | None = None



    @override
    def __str__(self) -> str:
        r = f"[{self.stack}] CPU: {self.cpu:.2f}% @ Memory: {self.mem:.2f}%\n"
        if self.gpu and self.gpu_mem:
            for i, gpu in enumerate(self.gpu):
                r += f"\t\tGPU {i}: {gpu:.2f}% @ Memory: {self.gpu_mem[i]:.2f}%\n"
        return r




    def export(self, duration: float) -> str:
        return f"{duration:.2f}s\t{self}"






class _LineBreak(object):
    def __init__(self, message: str) -> None:
        self.message: str = message



    @override
    def __str__(self) -> str:
        return f"<{self.message}>\n"




    def export(self, duration: float) -> str:
        return f"{duration:.2f}s\t{self}"






class Profiler(object):
    def __init__(self, title: str, save_as: str | PathLike[str], *, gpus: Sequence[Device] = ()) -> None:
        self.title: str = title
        self.save_as: str = save_as
        self.total_mem: float = self.get_total_mem()
        self.has_gpu: bool = len(gpus) > 0
        self._gpus: Sequence[Device] = gpus
        self.total_gpu_mem: list[float] = [self.get_total_gpu_mem(device) for device in gpus]
        with open(save_as, "w") as f:
            f.write(f"# {title}\nTotal memory: {self.total_mem}, Total GPU memory: {self.total_gpu_mem}\n\n")
        self._t0: float = time()
        self._allocated_tensors: tuple[float, list[tuple[
            float, AmbiguousShape, torch.dtype, torch.device, bool, str]]] = (0, [])



    @staticmethod
    def get_cpu_usage() -> float:
        return cpu_percent()




    def get_mem_usage(self) -> float:
        return 100 * virtual_memory().used / self.total_mem




    @staticmethod
    def get_total_mem() -> float:
        return virtual_memory().total




    @staticmethod
    def get_gpu_usage(device: Device) -> float:
        return torch.cuda.utilization(device)




    def get_gpu_mem_usage(self, device: Device) -> float:
        return 100 * torch.cuda.device_memory_used(device) / self.total_gpu_mem[self._gpus.index(device)]




    @staticmethod
    def get_total_gpu_mem(device: Device) -> float:
        return torch.cuda.get_device_properties(device).total_memory




    def _save(self, obj: ProfilerFrame | _LineBreak | str) -> None:
        with open(self.save_as, "a") as f:
            t = time()
            f.write(obj + "\n" if isinstance(obj, str) else f"{obj.export(t - self._t0)}\n")
            self._t0 = t




    def record_allocated_tensors(self, *, limit: int = 10) -> str:
        allocated_tensors = dump_allocated_tensors()
        counted_tensors = []
        added_tensors = []
        removed_tensors = []
        tensors, prev_tensors = allocated_tensors[1], self._allocated_tensors[1]
        for tensor in tensors + prev_tensors:
            if tensor in counted_tensors:
                continue
            sz, shape, dtype, device, requires_grad, grad_fn = tensor
            t = f"{sz:8.1f} MB | {shape} | {dtype} | {device} | grad={requires_grad} | {grad_fn}"
            if tensor in prev_tensors:
                if tensor in tensors:
                    num_diff = tensors.count(tensor) - prev_tensors.count(tensor)
                    if num_diff > 0:
                        added_tensors.append(f"{num_diff} x {t}")
                    if num_diff < 0:
                        removed_tensors.append(f"{num_diff} x {t}")
                else:
                    removed_tensors.append(f"{-prev_tensors.count(tensor)} x {t}")
            else:
                added_tensors.append(f"{tensors.count(tensor)} x {t}")
            counted_tensors.append(tensor)
        if len(added_tensors) > limit:
            added_tensors = added_tensors[:limit]
        if len(removed_tensors) > limit:
            removed_tensors = removed_tensors[:limit]
        r = (f"Total size diff: {allocated_tensors[0] - self._allocated_tensors[0]} MB\n"
             f"Added tensors:\n{"\t\n".join(added_tensors)}\nRemoved tensors:\n{"\t\n".join(removed_tensors)}\n")
        self._save(r)
        self._allocated_tensors = allocated_tensors
        return r




    def record(self, *, stack_trace_offset: int = 1) -> ProfilerFrame:
        frame = ProfilerFrame(" -> ".join([f"{f.function}:{f.lineno}" for f in reversed(stack()[stack_trace_offset:])]),
                              self.get_cpu_usage(), self.get_mem_usage())
        if self.has_gpu:
            frame.gpu = [torch.cuda.utilization(device) for device in self._gpus]
            frame.gpu_mem = [self.get_gpu_mem_usage(device) for device in self._gpus]
        self._save(frame)
        return frame




    def line_break(self, message: str) -> _LineBreak:
        r = _LineBreak(message)
        self._save(r)
        return r