From eace607d18df91ff2bb35e4786c2edc3230ca788 Mon Sep 17 00:00:00 2001 From: Dominik Safaric Date: Fri, 13 Feb 2026 08:28:29 +0100 Subject: [PATCH 1/3] Further performance improvements --- pyproject.toml | 1 + .../telemetry/core/config.py | 13 +-- .../telemetry/core/events.py | 91 +++++++++++++++---- uv.lock | 11 +++ 4 files changed, 89 insertions(+), 27 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 87306fc..7730faf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,6 +29,7 @@ dependencies = [ "filelock>=3.19.1", "requests>=2.32.5", "ruff>=0.11.6", + "nvidia-ml-py>=13.590.48", ] [project.optional-dependencies] diff --git a/src/tabpfn_common_utils/telemetry/core/config.py b/src/tabpfn_common_utils/telemetry/core/config.py index 4150d26..a680284 100644 --- a/src/tabpfn_common_utils/telemetry/core/config.py +++ b/src/tabpfn_common_utils/telemetry/core/config.py @@ -10,7 +10,6 @@ import logging import os -from datetime import datetime, timezone from typing import Any, Dict import requests @@ -20,20 +19,15 @@ logger = logging.getLogger(__name__) -@ttl_cache(ttl_seconds=60 * 5) +@ttl_cache(ttl_seconds=60 * 60) def download_config() -> Dict[str, Any]: """Download the configuration from server. Returns: Dict[str, Any]: The configuration. """ - # Bust the cache - params = { - "timestamp": datetime.now(tz=timezone.utc).isoformat(), - } - # The default configuration - default = {"enabled": False} + default = {"enabled": True} # This is a public URL anyone can and should read from url = os.environ.get( @@ -41,7 +35,8 @@ def download_config() -> Dict[str, Any]: "https://storage.googleapis.com/prior-labs-tabpfn-public/config/telemetry.json", ) try: - resp = requests.get(url, params=params) + # We use a very short timeout to avoid blocking the main thread + resp = requests.get(url, timeout=0.25) except Exception: logger.debug(f"Failed to download telemetry config: {url}") return default diff --git a/src/tabpfn_common_utils/telemetry/core/events.py b/src/tabpfn_common_utils/telemetry/core/events.py index db0367b..b1489ab 100644 --- a/src/tabpfn_common_utils/telemetry/core/events.py +++ b/src/tabpfn_common_utils/telemetry/core/events.py @@ -1,11 +1,20 @@ +import importlib import os import platform import sys import uuid from dataclasses import dataclass, asdict, field from datetime import datetime, timezone +from importlib.metadata import version, PackageNotFoundError from functools import lru_cache from typing import Any, Dict, Literal, Optional +from pynvml import ( + nvmlInit, + nvmlDeviceGetCount, + nvmlDeviceGetHandleByIndex, + nvmlDeviceGetName, + nvmlShutdown, +) from .runtime import get_execution_context from .state import get_property, set_property @@ -79,7 +88,7 @@ def _get_sklearn_version() -> str: Returns: str: Version string if scikit-learn is installed. """ - return _get_package_version("sklearn") + return _get_package_version("scikit-learn", "sklearn") @lru_cache(maxsize=1) @@ -112,7 +121,6 @@ def _get_autogluon_version() -> str: return _get_package_version("autogluon.core") -@lru_cache(maxsize=None) def _get_gpu_type() -> Optional[str]: """Detect a local GPU using PyTorch (the TabPFN dependency) and return its human-readable name. @@ -120,15 +128,57 @@ def _get_gpu_type() -> Optional[str]: Returns: Optional[str]: Human-readable name of the GPU if available. """ + # First, we try to use the NVML library to get the GPU names + # This is the preferred method as it is faster than using PyTorch try: - import torch # type: ignore[import] - except ImportError: - return None + nvmlInit() + + counts = nvmlDeviceGetCount() + if counts == 0: + return None + + # Retrieve the names of the devices + devices: list[str] = [ + nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(i)) for i in range(counts) + ] + + # Shutdown the NVML library + nvmlShutdown() + + # Because NVML runs very fast, we just return the device name + # without caching it. We return the first device name and assume + # that the VM has the same GPU type for all devices. + return devices[0] + except Exception: + pass + + # Only then, as an alternative, we try to use PyTorch to get the GPU name + # This is the slowest method as it requires importing the PyTorch library + # so we do not prefer this over the previous methods + return _get_torch_gpu_type() + + +@lru_cache(maxsize=1) +def _get_torch_gpu_type() -> Optional[str]: + """Get the type of GPU using PyTorch. + + Returns: + Optional[str]: Type of GPU if available. + """ + # First, we try to load the torch module from the sys.modules + # because we can assume it is already loaded and we avoid + # re-initializing it. + torch = sys.modules.get("torch") + if torch is None: + try: + import torch # type: ignore[import] + except ImportError: + return None try: if torch.cuda.is_available(): - name = torch.cuda.get_device_name(0) - return name or "unknown" + name = torch.cuda.get_device_name(0) or "unknown" + return name except Exception: # noqa: BLE001 pass @@ -137,29 +187,34 @@ def _get_gpu_type() -> Optional[str]: try: if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): # torch doesn't expose an MPS "model string" - return "Apple M-series GPU (MPS)" + name = "Apple M-series GPU (MPS)" + return name except Exception: pass return None -def _get_package_version(package_name: str) -> str: - """Get the version of a package if it's installed. +def _get_package_version(dist_name: str, module_name: Optional[str] = None) -> str: + """Read the package distribution metadata and return the version. Args: - package_name: Name of the package to import (e.g., "torch", "tabpfn"). + dist_name: Name of the package to read the metadata for. + module_name: Name of the module to read the version for, if different + from the distribution name. Returns: - str: Version string if the package is installed, "unknown" otherwise. + str: Version string if the package is installed, None otherwise. """ try: - import importlib - - module = importlib.import_module(package_name) # type: ignore[import] - return getattr(module, "__version__", "unknown") - except ImportError: - return "unknown" + return version(dist_name) + except PackageNotFoundError: + # Fallback to importing the package and getting the version + try: + module = importlib.import_module(module_name or dist_name) + return getattr(module, "__version__", "unknown") + except ImportError: + return "unknown" @lru_cache(maxsize=1) diff --git a/uv.lock b/uv.lock index 7703701..fb97f85 100644 --- a/uv.lock +++ b/uv.lock @@ -726,6 +726,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/68/67/1175790323026d3337cc285cc9c50eca637d70472b5e622529df74bb8f37/numpy-2.2.5-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:d2e3bdadaba0e040d1e7ab39db73e0afe2c74ae277f5614dad53eadbecbbb169", size = 12859001, upload-time = "2025-04-19T22:48:57.665Z" }, ] +[[package]] +name = "nvidia-ml-py" +version = "13.590.48" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/af/a0/f4fc18cf72f06821a9a665085435b901449986855519d5b3843532db35e9/nvidia_ml_py-13.590.48.tar.gz", hash = "sha256:8184d1be52914ac7f0991cd1c0d946c65dc88a840c754cd12c274b77b88760dd", size = 49732, upload-time = "2026-01-22T01:14:56.456Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fd/72/fb2af0d259a651affdce65fd6a495f0e07a685a0136baf585c5065204ee7/nvidia_ml_py-13.590.48-py3-none-any.whl", hash = "sha256:fd43d30ee9cd0b7940f5f9f9220b68d42722975e3992b6c21d14144c48760e43", size = 50680, upload-time = "2026-01-22T01:14:55.281Z" }, +] + [[package]] name = "packaging" version = "25.0" @@ -1222,6 +1231,7 @@ dependencies = [ { name = "filelock" }, { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, { name = "numpy", version = "2.2.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "nvidia-ml-py" }, { name = "pandas" }, { name = "platformdirs" }, { name = "posthog" }, @@ -1261,6 +1271,7 @@ requires-dist = [ { name = "filelock", specifier = ">=3.19.1" }, { name = "hatchling", marker = "extra == 'build'", specifier = ">=1.25" }, { name = "numpy", specifier = ">=1.21.6" }, + { name = "nvidia-ml-py", specifier = ">=13.590.48" }, { name = "pandas", specifier = ">=1.4.0,<3" }, { name = "platformdirs", specifier = ">=4" }, { name = "posthog", specifier = "~=6.7" }, From b7d702a47c48b1ab9fcb5301a1b574d2dc58e396 Mon Sep 17 00:00:00 2001 From: Dominik Safaric Date: Fri, 13 Feb 2026 08:45:46 +0100 Subject: [PATCH 2/3] Ensure proper nvml shutdown --- src/tabpfn_common_utils/telemetry/core/events.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/tabpfn_common_utils/telemetry/core/events.py b/src/tabpfn_common_utils/telemetry/core/events.py index b1489ab..604215e 100644 --- a/src/tabpfn_common_utils/telemetry/core/events.py +++ b/src/tabpfn_common_utils/telemetry/core/events.py @@ -142,15 +142,15 @@ def _get_gpu_type() -> Optional[str]: nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(i)) for i in range(counts) ] - # Shutdown the NVML library - nvmlShutdown() - # Because NVML runs very fast, we just return the device name # without caching it. We return the first device name and assume # that the VM has the same GPU type for all devices. return devices[0] except Exception: pass + finally: + # Shutdown the NVML library + nvmlShutdown() # Only then, as an alternative, we try to use PyTorch to get the GPU name # This is the slowest method as it requires importing the PyTorch library From ce2b18cf0c253838415ab05cd515cf357c159aba Mon Sep 17 00:00:00 2001 From: Dominik Safaric Date: Fri, 13 Feb 2026 08:51:26 +0100 Subject: [PATCH 3/3] nvml fix --- src/tabpfn_common_utils/telemetry/core/events.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/tabpfn_common_utils/telemetry/core/events.py b/src/tabpfn_common_utils/telemetry/core/events.py index 604215e..c01750b 100644 --- a/src/tabpfn_common_utils/telemetry/core/events.py +++ b/src/tabpfn_common_utils/telemetry/core/events.py @@ -130,9 +130,15 @@ def _get_gpu_type() -> Optional[str]: """ # First, we try to use the NVML library to get the GPU names # This is the preferred method as it is faster than using PyTorch + nvml_initialized = False + try: nvmlInit() + # Set a flag to indicate that the NVML library was initialized + # if running on CPU, this part of the code will not be reached + nvml_initialized = True + counts = nvmlDeviceGetCount() if counts == 0: return None @@ -150,7 +156,8 @@ def _get_gpu_type() -> Optional[str]: pass finally: # Shutdown the NVML library - nvmlShutdown() + if nvml_initialized: + nvmlShutdown() # Only then, as an alternative, we try to use PyTorch to get the GPU name # This is the slowest method as it requires importing the PyTorch library