Repository: Lightning-AI/pytorch-lightning
File: src/lightning/pytorch/callbacks/device_stats_monitor.py (Python)

Code before:
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Device Stats Monitor
====================

Monitors and logs device stats during training.

"""

from typing import Any, Optional

from typing_extensions import override

import lightning.pytorch as pl
from lightning.pytorch.accelerators.cpu import _PSUTIL_AVAILABLE
from lightning.pytorch.callbacks.callback import Callback
from lightning.pytorch.utilities.exceptions import MisconfigurationException
from lightning.pytorch.utilities.types import STEP_OUTPUT


class DeviceStatsMonitor(Callback):
    r"""Automatically monitors and logs device stats during training, validation and testing stage.
    ``DeviceStatsMonitor`` is a special callback as it requires a ``logger`` to passed as argument to the ``Trainer``.

    Args:
        cpu_stats: if ``None``, it will log CPU stats only if the accelerator is CPU.
            If ``True``, it will log CPU stats regardless of the accelerator.
            If ``False``, it will not log CPU stats regardless of the accelerator.

    Raises:
        MisconfigurationException:
            If ``Trainer`` has no logger.
        ModuleNotFoundError:
            If ``psutil`` is not installed and CPU stats are monitored.

    Logged Metrics:
        Device statistics are logged with keys prefixed as
        ``DeviceStatsMonitor.{hook_name}/{base_metric_name}`` (e.g.,
        ``DeviceStatsMonitor.on_train_batch_start/cpu_percent``). The source of these
        metrics depends on the active :class:`~lightning.pytorch.accelerators.accelerator.Accelerator`
        and the ``cpu_stats`` flag.

        CPU (via ``psutil``): Logs ``cpu_percent``, ``cpu_vm_percent``, ``cpu_swap_percent``.
        All are percentages (%).
        CUDA GPU (via :func:`torch.cuda.memory_stats`): Logs detailed memory statistics from
        PyTorch's allocator (e.g., ``allocated_bytes.all.current``, ``num_ooms``; all in Bytes).
        GPU compute utilization is not logged by default.
        Other Accelerators (e.g., TPU, MPS): Logs device-specific stats.
        - TPU example: ``avg. free memory (MB)``.
        - MPS example: ``mps.current_allocated_bytes``.

Review comment (suggestion):
`Raises:` or `Args:` are Sphinx-specific keywords compared to `Logged Metrics:`, so pls let's move it just to the top of this docstring

Code after:
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Device Stats Monitor
====================

Monitors and logs device stats during training.

"""

from typing import Any, Optional

from typing_extensions import override

import lightning.pytorch as pl
from lightning.pytorch.accelerators.cpu import _PSUTIL_AVAILABLE
from lightning.pytorch.callbacks.callback import Callback
from lightning.pytorch.utilities.exceptions import MisconfigurationException
from lightning.pytorch.utilities.types import STEP_OUTPUT


class DeviceStatsMonitor(Callback):
    r"""Automatically monitors and logs device stats during training, validation and testing stage.
    ``DeviceStatsMonitor`` is a special callback as it requires a ``logger`` to passed as argument to the ``Trainer``.

    **Logged Metrics**

    Logs device statistics with keys prefixed as ``DeviceStatsMonitor.{hook_name}/{base_metric_name}``.
    The actual metrics depend on the active accelerator and the ``cpu_stats`` flag. Below are an overview of the
    possible available metrics and their meaning.

    - CPU (via ``psutil``)

        - ``cpu_percent`` — System-wide CPU utilization (%)
        - ``cpu_vm_percent`` — System-wide virtual memory (RAM) utilization (%)
        - ``cpu_swap_percent`` — System-wide swap memory utilization (%)

    - CUDA GPU (via ``torch.cuda.memory_stats``)

        Logs memory statistics from PyTorch caching allocator (all in bytes).
        GPU compute utilization is not logged by default.

        - General Memory Usage:

            - ``allocated_bytes.all.current`` — Current allocated GPU memory
            - ``allocated_bytes.all.peak`` — Peak allocated GPU memory
            - ``reserved_bytes.all.current`` — Current reserved GPU memory (allocated + cached)
            - ``reserved_bytes.all.peak`` — Peak reserved GPU memory
            - ``active_bytes.all.current`` — Current GPU memory in active use
            - ``active_bytes.all.peak`` — Peak GPU memory in active use
            - ``inactive_split_bytes.all.current`` — Memory in inactive, splittable blocks