From 24dfd0b13f64c965f7588732a65f0ad5ad9713ae Mon Sep 17 00:00:00 2001 From: Francesc Alted Date: Thu, 12 Dec 2024 11:27:51 +0100 Subject: [PATCH] Reuse the previous get_cpu_info function --- src/blosc2/core.py | 25 ++++++++++++++++++++----- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/src/blosc2/core.py b/src/blosc2/core.py index 3dcd6f6e..1bc9db7c 100644 --- a/src/blosc2/core.py +++ b/src/blosc2/core.py @@ -23,6 +23,7 @@ from functools import lru_cache from typing import TYPE_CHECKING +import cpuinfo import numpy as np import blosc2 @@ -1151,6 +1152,24 @@ def linux_cache_size(cache_level: int, default_size: int) -> int: return cache_size +def _get_cpu_info(): + cpu_info = cpuinfo.get_cpu_info() + # cpuinfo does not correctly retrieve the cache sizes for Apple Silicon, so do it manually + if platform.system() == "Darwin": + cpu_info["l1_data_cache_size"] = apple_silicon_cache_size(1) + cpu_info["l2_cache_size"] = apple_silicon_cache_size(2) + cpu_info["l3_cache_size"] = apple_silicon_cache_size(3) + # cpuinfo does not correctly retrieve the cache sizes for all CPUs on Linux, so ask the kernel + if platform.system() == "Linux": + l1_data_cache_size = cpu_info.get("l1_data_cache_size", 32 * 1024) + cpu_info["l1_data_cache_size"] = linux_cache_size(1, l1_data_cache_size) + l2_cache_size = cpu_info.get("l2_cache_size", 256 * 1024) + cpu_info["l2_cache_size"] = linux_cache_size(2, l2_cache_size) + l3_cache_size = cpu_info.get("l3_cache_size", 1024 * 1024) + cpu_info["l3_cache_size"] = linux_cache_size(3, l3_cache_size) + return cpu_info + + def write_cached_cpu_info(cpu_info_dict: dict[str, any]) -> None: with open(pathlib.Path.home() / ".blosc2-cpuinfo.json", "w") as f: json.dump(cpu_info_dict, f, indent=4) @@ -1170,11 +1189,7 @@ def get_cpu_info() -> dict: if cached_info: return cached_info - try: - import cpuinfo - except ImportError: - return {} - cpu_info_dict = cpuinfo.get_cpu_info() + cpu_info_dict = _get_cpu_info() with contextlib.suppress(OSError): # In case cpu info cannot be stored, will need to be recomputed in the next process write_cached_cpu_info(cpu_info_dict)