/ludwig/benchmarking/profiler.py

import contextlib
import glob
import logging
import os
import shutil
import threading
import time
from queue import Empty as EmptyQueueException
from queue import Queue
from subprocess import PIPE, Popen
from typing import Any
from xml.etree.ElementTree import fromstring

import psutil
import torch
from cpuinfo import get_cpu_info
from gpustat.core import GPUStatCollection

from ludwig.benchmarking.profiler_dataclasses import profiler_dataclass_to_flat_dict, TorchProfilerMetrics
from ludwig.benchmarking.reporting import get_metrics_from_system_usage_profiler, get_metrics_from_torch_profiler
from ludwig.constants import LUDWIG_TAG
from ludwig.globals import LUDWIG_VERSION
from ludwig.utils.data_utils import save_json

STOP_MESSAGE = "stop"
logger = logging.getLogger()


def get_gpu_info():
    """Gathers general hardware information about an nvidia GPU.

    This function was copied from `experiment_impact_tracker` to get around a Pandas 2.0 breaking change impacting the
    package. https://github.com/Breakend/experiment-impact-
    tracker/blob/master/experiment_impact_tracker/gpu/nvidia.py#L48-L73
    """
    p = Popen(["nvidia-smi", "-q", "-x"], stdout=PIPE)
    outs, errors = p.communicate()
    xml = fromstring(outs)
    data = []
    driver_version = xml.findall("driver_version")[0].text
    cuda_version = xml.findall("cuda_version")[0].text

    for _gpu_id, gpu in enumerate(xml.getiterator("gpu")):
        gpu_data = {}
        name = list(gpu.getiterator("product_name"))[0].text
        memory_usage = gpu.findall("fb_memory_usage")[0]
        total_memory = memory_usage.findall("total")[0].text

        gpu_data["name"] = name
        gpu_data["total_memory"] = total_memory
        gpu_data["driver_version"] = driver_version
        gpu_data["cuda_version"] = cuda_version
        data.append(gpu_data)
    return data


def monitor(queue: Queue, info: dict[str, Any], logging_interval: int, cuda_is_available: bool) -> None:
    """Monitors hardware resource use.

    Collects system specific metrics (CPU/CUDA, CPU/CUDA memory) at a `logging_interval` interval and pushes
    results back to the parent process.

    Args:
        queue: queue from which we can push and retrieve messages sent to the function targeted by the thread.
        info: dictionary containing system resource usage information about the running process.
        logging_interval: time interval at which we will poll the system for usage metrics.
        cuda_is_available: stores torch.cuda.is_available().
    """
    info["global_cpu_memory_available"] = [psutil.virtual_memory().available]
    info["global_cpu_utilization"] = [psutil.cpu_percent()]
    # get the pid of the parent process.
    tracked_process = psutil.Process(os.getpid())

    # will return a meaningless 0 value on the first call because `interval` arg is set to None.
    tracked_process.cpu_percent(interval=logging_interval)
    with tracked_process.oneshot():
        info["cpu_utilization"] = [tracked_process.cpu_percent() / info["num_cpu"]]
        info["cpu_memory_usage"] = [tracked_process.memory_full_info().uss]
        try:
            info["num_accessible_cpus"] = len(tracked_process.cpu_affinity())
        except (AttributeError, NotImplementedError):
            pass

    while True:
        try:
            message = queue.get(block=False)
            if isinstance(message, str):
                if message == STOP_MESSAGE:
                    # synchronize CUDA to get accurate timing for jobs running on GPU.
                    if cuda_is_available:
                        torch.cuda.synchronize()
                    queue.put(info)
                    return
            else:
                queue.put(message)
        except EmptyQueueException:
            pass
        if cuda_is_available:
            gpu_infos = GPUStatCollection.new_query()
            for i, gpu_info in enumerate(gpu_infos):
                gpu_key = f"cuda_{i}"
                info[f"{gpu_key}_memory_used"].append(gpu_info.memory_used)
        with tracked_process.oneshot():
            info["cpu_utilization"].append(tracked_process.cpu_percent() / info["num_cpu"])
            info["cpu_memory_usage"].append(tracked_process.memory_full_info().uss)
        info["global_cpu_memory_available"].append(psutil.virtual_memory().available)
        info["global_cpu_utilization"].append(psutil.cpu_percent())
        time.sleep(logging_interval)


class LudwigProfiler(contextlib.ContextDecorator):
    """Track system resource (hardware and software) usage.

    Warning: If `use_torch_profiler=True` while profiling on CUDA, it's not possible to benchmark DataLoaders
    with `num_workers > 0` due to CUDA multiprocessing limitations. See warning under `profile` class
    definition: https://github.com/pytorch/pytorch/blob/master/torch/autograd/profiler.py

    Attributes:
        tag: a string tag describing the code block/function that we're tracking.
            (e.g trainer.train, preprocessing, etc.)
        output_dir: path where metrics are saved.
        logging_interval: time interval in seconds at which system is polled for resource usage.
    """

    def __init__(self, tag: str, use_torch_profiler: bool, output_dir: str, logging_interval: float = 0.1) -> None:
        self.tag = tag
        self._tag = LUDWIG_TAG + self.tag
        self.use_torch_profiler = use_torch_profiler
        self.output_dir = output_dir
        self.logging_interval = logging_interval
        self.cuda_is_available = torch.cuda.is_available()
        self.launched = False
        if self.use_torch_profiler:
            self.profiler_activities = [torch.profiler.ProfilerActivity.CPU]
            if self.cuda_is_available:
                self.profiler_activities.append(torch.profiler.ProfilerActivity.CUDA)
        os.makedirs(os.path.join(self.output_dir), exist_ok=True)

    def _init_tracker_info(self):
        """Initialize new self.info, self.torch_profiler, and self.torch_record_function instances.

        Important to call this in __enter__ if the user decides not to create a new class instance and therefore
        __init__ wouldn't be called.
        """
        self.info = {"code_block_tag": self.tag}
        if self.use_torch_profiler:
            self.torch_profiler = torch.profiler.profile(activities=self.profiler_activities, profile_memory=True)
            self.torch_record_function = torch.profiler.record_function(self._tag)

    def _populate_static_information(self) -> None:
        """Populate the report with static software and hardware information."""
        self.info["ludwig_version"] = LUDWIG_VERSION
        self.info["start_disk_usage"] = shutil.disk_usage(os.path.expanduser("~")).used

        # CPU information
        cpu_info = get_cpu_info()
        self.info["cpu_architecture"] = cpu_info["arch"]
        self.info["num_cpu"] = psutil.cpu_count()
        self.info["cpu_name"] = cpu_info.get("brand_raw", "unknown")
        self.info["total_cpu_memory_size"] = psutil.virtual_memory().total

        # GPU information
        if self.cuda_is_available:
            gpu_infos = get_gpu_info()
            gpu_usage = GPUStatCollection.new_query()
            for i, gpu_info in enumerate(gpu_infos):
                gpu_key = f"cuda_{i}"
                self.info[f"{gpu_key}_memory_used"] = [gpu_usage[i].memory_used]
                self.info[f"{gpu_key}_name"] = gpu_info["name"]
                self.info[f"{gpu_key}_total_memory"] = gpu_info["total_memory"]
                self.info[f"{gpu_key}_driver_version"] = gpu_info["driver_version"]
                self.info[f"{gpu_key}_cuda_version"] = gpu_info["cuda_version"]

        # recording in microseconds to be in line with torch profiler time recording.
        self.info["start_time"] = time.perf_counter_ns() / 1000

    def __enter__(self):
        """Populate static information and monitors resource usage."""
        if self.launched:
            raise RuntimeError("LudwigProfiler already launched. You can't use the same instance.")

        self._init_tracker_info()
        self._populate_static_information()

        if self.use_torch_profiler:
            # contextlib.ExitStack gracefully handles situations where __enter__ or __exit__ calls throw exceptions.
            with contextlib.ExitStack() as ctx_exit_stack:
                try:
                    # Launch torch.profiler to track PyTorch operators.
                    ctx_exit_stack.enter_context(self.torch_profiler)
                except RuntimeError:
                    # PyTorch profiler is already enabled on this thread.
                    # Using the running PyTorch profiler to track events.
                    self.torch_profiler = None

                ctx_exit_stack.enter_context(self.torch_record_function)
                self._ctx_exit_stack = ctx_exit_stack.pop_all()
        try:
            # Starting thread to monitor system resource usage.
            self.queue = Queue()
            self.t = threading.Thread(
                target=monitor,
                args=(
                    self.queue,
                    self.info,
                    self.logging_interval,
                    self.cuda_is_available,
                ),
            )
            self.t.start()
            self.launched = True
        except Exception:
            self.launched = False
            logger.exception("Encountered exception when launching tracker thread.")

        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        """Stop profiling, postprocess and export resource usage metrics."""
        try:
            self.queue.put(STOP_MESSAGE)
            self.t.join()
            result = self.queue.get()
            # If monitor thread crashed, result may be a string instead of dict
            if isinstance(result, dict):
                self.info = result
            # recording in microseconds to be in line with torch profiler time recording.
            self.info["end_time"] = time.perf_counter_ns() / 1000
            self.info["end_disk_usage"] = shutil.disk_usage(os.path.expanduser("~")).used
            self.launched = False
        except Exception:
            logger.exception("Encountered exception when joining tracker thread.")
        finally:
            if self.use_torch_profiler:
                self._ctx_exit_stack.close()
                self._export_torch_metrics()
            self._export_system_usage_metrics()

    def _export_system_usage_metrics(self):
        """Export system resource usage metrics (no torch operators)."""
        system_usage_metrics = get_metrics_from_system_usage_profiler(self.info)
        output_subdir = os.path.join(self.output_dir, "system_resource_usage", system_usage_metrics.code_block_tag)
        os.makedirs(output_subdir, exist_ok=True)
        num_prev_runs = len(glob.glob(os.path.join(output_subdir, "run_*.json")))
        file_name = os.path.join(output_subdir, f"run_{num_prev_runs}.json")
        save_json(file_name, profiler_dataclass_to_flat_dict(system_usage_metrics))

    def _reformat_torch_usage_metrics_tags(
        self, torch_usage_metrics: dict[str, Any]
    ) -> dict[str, list[TorchProfilerMetrics]]:
        reformatted_dict = {}
        for key, value in torch_usage_metrics.items():
            assert key.startswith(LUDWIG_TAG)
            reformatted_key = key[len(LUDWIG_TAG) :]
            reformatted_dict[reformatted_key] = value
        return reformatted_dict

    def _export_torch_metrics(self):
        """Export resource usage metrics of torch operators."""
        if self.torch_profiler:
            torch_usage_metrics = get_metrics_from_torch_profiler(self.torch_profiler)
            torch_usage_metrics = self._reformat_torch_usage_metrics_tags(torch_usage_metrics)
            for tag, runs in torch_usage_metrics.items():
                temp_dir = os.path.join(self.output_dir, "torch_ops_resource_usage", tag)
                os.makedirs(temp_dir, exist_ok=True)
                for run in runs:
                    num_prev_runs = len(glob.glob(os.path.join(temp_dir, "run_*.json")))
                    save_json(os.path.join(temp_dir, f"run_{num_prev_runs}.json"), profiler_dataclass_to_flat_dict(run))