metrics.py

from dataclasses import dataclass
import time
import tracemalloc


@dataclass
class Metrics:
    elapsed_time: float
    cpu_time: float
    peak_memory_kb: int
    nb_visited: int | None = None

    def __str__(self) -> str:
        line = f"Elapsed time: {self.elapsed_time:.6f} s" + " | " + f"CPU time: {self.cpu_time:.6f} s" + " | " + f"Peak memory: {self.peak_memory_kb} KB"

        if self.nb_visited is not None:
            line += " | " + f"Number of visits: {self.nb_visited}"

        return line + "\n"


def get_metrics(solver, parameters):
    """
    solver: function(parameters) -> (solution, stats)
    stats: optional dict containing algorithm counters
    """

    tracemalloc.start()

    start_time = time.perf_counter()
    start_cpu = time.process_time()

    solution, capacity_left, stats = solver(parameters)

    end_time = time.perf_counter()
    end_cpu = time.process_time()

    _, peak = tracemalloc.get_traced_memory()
    tracemalloc.stop()

    metrics = Metrics(
        elapsed_time=end_time - start_time,
        cpu_time=end_cpu - start_cpu,
        peak_memory_kb=peak // 1024,
        nb_visited=stats.get("nb_visits") if stats is not None else None
    )

    return solution, capacity_left, metrics