memory ¶

def add_memory_limit_args(
    parser: argparse.ArgumentParser,
    *,
    mlx_memory_limit_gib: float | None = None,
    mlx_cache_limit_gib: float | None = None,
    mlx_disable_cache: bool = False,
    mlx_wired_limit_gib: float | None = None,
    torch_mps_high_watermark_ratio: float | None = None,
    torch_mps_low_watermark_ratio: float | None = None,
) -> None:
    """
    Add configurable Apple Silicon memory-limit options to an argument parser.

    Parameters:
        parser (argparse.ArgumentParser): Parser to which the options are added.
        mlx_memory_limit_gib (float | None): Default MLX memory limit in GiB.
        mlx_cache_limit_gib (float | None): Default MLX cache limit in GiB.
        mlx_disable_cache (bool): Whether the cache limit defaults to zero.
        mlx_wired_limit_gib (float | None): Default MLX wired-memory limit in GiB.
        torch_mps_high_watermark_ratio (float | None): Default PyTorch MPS high-watermark ratio.
        torch_mps_low_watermark_ratio (float | None): Default PyTorch MPS low-watermark ratio.
    """
    parser.add_argument("--mlx-memory-limit-gib",
                        type=float,
                        default=mlx_memory_limit_gib,
                        help="Set MLX memory limit in GiB for memory-tier testing (DiT path).")
    parser.add_argument("--mlx-cache-limit-gib",
                        type=float,
                        default=mlx_cache_limit_gib,
                        help="Set MLX cache limit in GiB. Use --mlx-disable-cache to force 0.")
    parser.add_argument("--mlx-disable-cache",
                        action="store_true",
                        default=mlx_disable_cache,
                        help="Set MLX cache limit to 0 for stricter memory-tier tests.")
    parser.add_argument("--mlx-wired-limit-gib",
                        type=float,
                        default=mlx_wired_limit_gib,
                        help="Set MLX wired-memory limit in GiB where supported by macOS/MLX.")
    parser.add_argument("--torch-mps-high-watermark-ratio",
                        type=float,
                        default=torch_mps_high_watermark_ratio,
                        help="Set PYTORCH_MPS_HIGH_WATERMARK_RATIO before importing torch.")
    parser.add_argument("--torch-mps-low-watermark-ratio",
                        type=float,
                        default=torch_mps_low_watermark_ratio,
                        help="Set PYTORCH_MPS_LOW_WATERMARK_RATIO before importing torch.")

def apply_memory_limits(
    *,
    mlx_memory_limit_gib: float | None = None,
    mlx_cache_limit_gib: float | None = None,
    mlx_disable_cache: bool = False,
    mlx_wired_limit_gib: float | None = None,
    torch_mps_high_watermark_ratio: float | None = None,
    torch_mps_low_watermark_ratio: float | None = None,
    mx_module: Any | None = None,
) -> AppliedMemoryLimits:
    """Apply optional MLX allocator limits and PyTorch MPS watermarks.

    PyTorch reads MPS watermark variables when the MPS backend initializes, so
    call this before importing PyTorch. Specifying only a high watermark sets the
    low watermark to ``0.0``. MLX limit-setting failures are recorded in the
    result and do not prevent other limits from being applied.

    Parameters:
        mlx_memory_limit_gib (float | None): Maximum MLX memory in GiB.
        mlx_cache_limit_gib (float | None): Maximum MLX cache size in GiB.
        mlx_disable_cache (bool): Whether to disable the MLX cache.
        mlx_wired_limit_gib (float | None): Maximum MLX wired memory in GiB.
        torch_mps_high_watermark_ratio (float | None): PyTorch MPS high watermark
            ratio.
        torch_mps_low_watermark_ratio (float | None): PyTorch MPS low watermark
            ratio.

    Returns:
        AppliedMemoryLimits: Configured values, applied and previous MLX byte
            limits, MPS watermark values, and per-limit errors.
    """
    if torch_mps_high_watermark_ratio is not None and torch_mps_low_watermark_ratio is None:
        torch_mps_low_watermark_ratio = 0.0

    high = _set_mps_env("PYTORCH_MPS_HIGH_WATERMARK_RATIO", torch_mps_high_watermark_ratio)
    low = _set_mps_env("PYTORCH_MPS_LOW_WATERMARK_RATIO", torch_mps_low_watermark_ratio)

    memory_bytes = gib_to_bytes(mlx_memory_limit_gib)
    cache_bytes = 0 if mlx_disable_cache else gib_to_bytes(mlx_cache_limit_gib)
    wired_bytes = gib_to_bytes(mlx_wired_limit_gib)

    applied: dict[str, int] = {}
    previous: dict[str, int] = {}
    errors: dict[str, str] = {}
    if memory_bytes is not None or cache_bytes is not None or wired_bytes is not None:
        if mx_module is None:
            import mlx.core as mx

            mx_module = mx

        # Apply each limit independently; record failures without stopping.
        limits = [
            ("mlx_memory_limit", memory_bytes, mx_module.set_memory_limit),
            ("mlx_cache_limit", cache_bytes, mx_module.set_cache_limit),
            ("mlx_wired_limit", wired_bytes, mx_module.set_wired_limit),
        ]
        for name, value, setter in limits:
            if value is not None:
                try:
                    previous[name] = int(setter(value))
                    applied[name] = value
                except Exception as exc:  # noqa: BLE001 - macOS/system-limit dependent.
                    errors[name] = f"{type(exc).__name__}: {exc}"

    return AppliedMemoryLimits(
        mlx_memory_limit_gib=mlx_memory_limit_gib,
        mlx_cache_limit_gib=mlx_cache_limit_gib,
        mlx_disable_cache=mlx_disable_cache,
        mlx_wired_limit_gib=mlx_wired_limit_gib,
        torch_mps_high_watermark_ratio=high,
        torch_mps_low_watermark_ratio=low,
        applied_bytes=applied,
        previous_bytes=previous,
        errors=errors,
    )

def cleanup_mlx(mx_module: Any | None = None) -> None:
    """Collect unreachable MLX objects, then release their allocator cache."""
    if mx_module is None:
        import mlx.core as mx

        mx_module = mx
    gc.collect()
    mx_module.clear_cache()

def cleanup_torch_mps(torch_module: Any | None = None) -> None:
    """Collect unreachable Torch objects, then release the MPS allocator cache."""
    if torch_module is None:
        import torch

        torch_module = torch
    gc.collect()
    if torch_module.backends.mps.is_available():
        torch_module.mps.empty_cache()

def gib_to_bytes(value: float | None) -> int | None:
    """
    Convert a positive memory limit from GiB to bytes.

    Parameters:
        value (float | None): Memory limit in GiB, or `None` when unset.

    Returns:
        int | None: The memory limit in bytes, or `None` when no limit is provided.

    Raises:
        ValueError: If `value` is zero or negative.
    """
    if value is None:
        return None
    if value <= 0:
        raise ValueError(f"Memory limit must be positive GiB, got {value}")
    return int(value * GIB)