wandb 0.20.1__py3-none-win32.whl → 0.20.2rc20250616__py3-none-win32.whl

wandb/sdk/internal/system/assets/gpu.py

@@ -1,416 +0,0 @@
-import logging
-import threading
-from collections import deque
-from typing import TYPE_CHECKING, List
-
-try:
-    import psutil
-except ImportError:
-    psutil = None
-
-from wandb.vendor.pynvml import pynvml
-
-from .aggregators import aggregate_mean
-from .asset_registry import asset_registry
-from .interfaces import Interface, Metric, MetricsMonitor
-
-if TYPE_CHECKING:
-    from typing import Deque
-
-    from wandb.sdk.internal.settings_static import SettingsStatic
-
-    GPUHandle = object
-
-
-logger = logging.getLogger(__name__)
-
-
-def gpu_in_use_by_this_process(gpu_handle: "GPUHandle", pid: int) -> bool:
-    if psutil is None:
-        return False
-
-    try:
-        base_process = psutil.Process(pid=pid)
-    except psutil.NoSuchProcess:
-        # do not report any gpu metrics if the base process can't be found
-        return False
-
-    our_processes = base_process.children(recursive=True)
-    our_processes.append(base_process)
-
-    our_pids = {process.pid for process in our_processes}
-
-    compute_pids = {
-        process.pid
-        for process in pynvml.nvmlDeviceGetComputeRunningProcesses(gpu_handle)  # type: ignore
-    }
-    graphics_pids = {
-        process.pid
-        for process in pynvml.nvmlDeviceGetGraphicsRunningProcesses(gpu_handle)  # type: ignore
-    }
-
-    pids_using_device = compute_pids | graphics_pids
-
-    return len(pids_using_device & our_pids) > 0
-
-
-class GPUMemoryUtilization:
-    """GPU memory utilization in percent for each GPU."""
-
-    # name = "memory_utilization"
-    name = "gpu.{}.memory"
-    # samples: Deque[Tuple[datetime.datetime, float]]
-    samples: "Deque[List[float]]"
-
-    def __init__(self, pid: int) -> None:
-        self.pid = pid
-        self.samples = deque([])
-
-    def sample(self) -> None:
-        memory_utilization_rate = []
-        device_count = pynvml.nvmlDeviceGetCount()  # type: ignore
-        for i in range(device_count):
-            handle = pynvml.nvmlDeviceGetHandleByIndex(i)  # type: ignore
-            memory_utilization_rate.append(
-                pynvml.nvmlDeviceGetUtilizationRates(handle).memory  # type: ignore
-            )
-        self.samples.append(memory_utilization_rate)
-
-    def clear(self) -> None:
-        self.samples.clear()
-
-    def aggregate(self) -> dict:
-        if not self.samples:
-            return {}
-        stats = {}
-        device_count = pynvml.nvmlDeviceGetCount()  # type: ignore
-        for i in range(device_count):
-            samples = [sample[i] for sample in self.samples]
-            aggregate = aggregate_mean(samples)
-            stats[self.name.format(i)] = aggregate
-
-            handle = pynvml.nvmlDeviceGetHandleByIndex(i)  # type: ignore
-            if gpu_in_use_by_this_process(handle, self.pid):
-                stats[self.name.format(f"process.{i}")] = aggregate
-
-        return stats
-
-
-class GPUMemoryAllocated:
-    """GPU memory allocated in percent for each GPU."""
-
-    # name = "memory_allocated"
-    name = "gpu.{}.memoryAllocated"
-    # samples: Deque[Tuple[datetime.datetime, float]]
-    samples: "Deque[List[float]]"
-
-    def __init__(self, pid: int) -> None:
-        self.pid = pid
-        self.samples = deque([])
-
-    def sample(self) -> None:
-        memory_allocated = []
-        device_count = pynvml.nvmlDeviceGetCount()  # type: ignore
-        for i in range(device_count):
-            handle = pynvml.nvmlDeviceGetHandleByIndex(i)  # type: ignore
-            memory_info = pynvml.nvmlDeviceGetMemoryInfo(handle)  # type: ignore
-            memory_allocated.append(memory_info.used / memory_info.total * 100)
-        self.samples.append(memory_allocated)
-
-    def clear(self) -> None:
-        self.samples.clear()
-
-    def aggregate(self) -> dict:
-        if not self.samples:
-            return {}
-        stats = {}
-        device_count = pynvml.nvmlDeviceGetCount()  # type: ignore
-        for i in range(device_count):
-            samples = [sample[i] for sample in self.samples]
-            aggregate = aggregate_mean(samples)
-            stats[self.name.format(i)] = aggregate
-
-            handle = pynvml.nvmlDeviceGetHandleByIndex(i)  # type: ignore
-            if gpu_in_use_by_this_process(handle, self.pid):
-                stats[self.name.format(f"process.{i}")] = aggregate
-
-        return stats
-
-
-class GPUMemoryAllocatedBytes:
-    """GPU memory allocated in bytes for each GPU."""
-
-    # name = "memory_allocated"
-    name = "gpu.{}.memoryAllocatedBytes"
-    # samples: Deque[Tuple[datetime.datetime, float]]
-    samples: "Deque[List[float]]"
-
-    def __init__(self, pid: int) -> None:
-        self.pid = pid
-        self.samples = deque([])
-
-    def sample(self) -> None:
-        memory_allocated = []
-        device_count = pynvml.nvmlDeviceGetCount()  # type: ignore
-        for i in range(device_count):
-            handle = pynvml.nvmlDeviceGetHandleByIndex(i)  # type: ignore
-            memory_info = pynvml.nvmlDeviceGetMemoryInfo(handle)  # type: ignore
-            memory_allocated.append(memory_info.used)
-        self.samples.append(memory_allocated)
-
-    def clear(self) -> None:
-        self.samples.clear()
-
-    def aggregate(self) -> dict:
-        if not self.samples:
-            return {}
-        stats = {}
-        device_count = pynvml.nvmlDeviceGetCount()  # type: ignore
-        for i in range(device_count):
-            samples = [sample[i] for sample in self.samples]
-            aggregate = aggregate_mean(samples)
-            stats[self.name.format(i)] = aggregate
-
-            handle = pynvml.nvmlDeviceGetHandleByIndex(i)  # type: ignore
-            if gpu_in_use_by_this_process(handle, self.pid):
-                stats[self.name.format(f"process.{i}")] = aggregate
-
-        return stats
-
-
-class GPUUtilization:
-    """GPU utilization in percent for each GPU."""
-
-    # name = "gpu_utilization"
-    name = "gpu.{}.gpu"
-    # samples: Deque[Tuple[datetime.datetime, float]]
-    samples: "Deque[List[float]]"
-
-    def __init__(self, pid: int) -> None:
-        self.pid = pid
-        self.samples = deque([])
-
-    def sample(self) -> None:
-        gpu_utilization_rate = []
-        device_count = pynvml.nvmlDeviceGetCount()  # type: ignore
-        for i in range(device_count):
-            handle = pynvml.nvmlDeviceGetHandleByIndex(i)  # type: ignore
-            gpu_utilization_rate.append(
-                pynvml.nvmlDeviceGetUtilizationRates(handle).gpu  # type: ignore
-            )
-        self.samples.append(gpu_utilization_rate)
-
-    def clear(self) -> None:
-        self.samples.clear()
-
-    def aggregate(self) -> dict:
-        if not self.samples:
-            return {}
-        stats = {}
-        device_count = pynvml.nvmlDeviceGetCount()  # type: ignore
-        for i in range(device_count):
-            samples = [sample[i] for sample in self.samples]
-            aggregate = aggregate_mean(samples)
-            stats[self.name.format(i)] = aggregate
-
-            handle = pynvml.nvmlDeviceGetHandleByIndex(i)  # type: ignore
-            if gpu_in_use_by_this_process(handle, self.pid):
-                stats[self.name.format(f"process.{i}")] = aggregate
-
-        return stats
-
-
-class GPUTemperature:
-    """GPU temperature in Celsius for each GPU."""
-
-    # name = "gpu_temperature"
-    name = "gpu.{}.temp"
-    # samples: Deque[Tuple[datetime.datetime, float]]
-    samples: "Deque[List[float]]"
-
-    def __init__(self, pid: int) -> None:
-        self.pid = pid
-        self.samples = deque([])
-
-    def sample(self) -> None:
-        temperature = []
-        device_count = pynvml.nvmlDeviceGetCount()  # type: ignore
-        for i in range(device_count):
-            handle = pynvml.nvmlDeviceGetHandleByIndex(i)  # type: ignore
-            temperature.append(
-                pynvml.nvmlDeviceGetTemperature(  # type: ignore
-                    handle,
-                    pynvml.NVML_TEMPERATURE_GPU,
-                )
-            )
-        self.samples.append(temperature)
-
-    def clear(self) -> None:
-        self.samples.clear()
-
-    def aggregate(self) -> dict:
-        if not self.samples:
-            return {}
-        stats = {}
-        device_count = pynvml.nvmlDeviceGetCount()  # type: ignore
-        for i in range(device_count):
-            samples = [sample[i] for sample in self.samples]
-            aggregate = aggregate_mean(samples)
-            stats[self.name.format(i)] = aggregate
-
-            handle = pynvml.nvmlDeviceGetHandleByIndex(i)  # type: ignore
-            if gpu_in_use_by_this_process(handle, self.pid):
-                stats[self.name.format(f"process.{i}")] = aggregate
-
-        return stats
-
-
-class GPUPowerUsageWatts:
-    """GPU power usage in Watts for each GPU."""
-
-    name = "gpu.{}.powerWatts"
-    # samples: Deque[Tuple[datetime.datetime, float]]
-    samples: "Deque[List[float]]"
-
-    def __init__(self, pid: int) -> None:
-        self.pid = pid
-        self.samples = deque([])
-
-    def sample(self) -> None:
-        power_usage = []
-        device_count = pynvml.nvmlDeviceGetCount()  # type: ignore
-        for i in range(device_count):
-            handle = pynvml.nvmlDeviceGetHandleByIndex(i)  # type: ignore
-            power_watts = pynvml.nvmlDeviceGetPowerUsage(handle) / 1000  # type: ignore
-            power_usage.append(power_watts)
-        self.samples.append(power_usage)
-
-    def clear(self) -> None:
-        self.samples.clear()
-
-    def aggregate(self) -> dict:
-        stats = {}
-        device_count = pynvml.nvmlDeviceGetCount()  # type: ignore
-        for i in range(device_count):
-            samples = [sample[i] for sample in self.samples]
-            aggregate = aggregate_mean(samples)
-            stats[self.name.format(i)] = aggregate
-
-            handle = pynvml.nvmlDeviceGetHandleByIndex(i)  # type: ignore
-            if gpu_in_use_by_this_process(handle, self.pid):
-                stats[self.name.format(f"process.{i}")] = aggregate
-
-        return stats
-
-
-class GPUPowerUsagePercent:
-    """GPU power usage in percent for each GPU."""
-
-    name = "gpu.{}.powerPercent"
-    # samples: Deque[Tuple[datetime.datetime, float]]
-    samples: "Deque[List[float]]"
-
-    def __init__(self, pid: int) -> None:
-        self.pid = pid
-        self.samples = deque([])
-
-    def sample(self) -> None:
-        power_usage = []
-        device_count = pynvml.nvmlDeviceGetCount()  # type: ignore
-        for i in range(device_count):
-            handle = pynvml.nvmlDeviceGetHandleByIndex(i)  # type: ignore
-            power_watts = pynvml.nvmlDeviceGetPowerUsage(handle)  # type: ignore
-            power_capacity_watts = pynvml.nvmlDeviceGetEnforcedPowerLimit(handle)  # type: ignore
-            power_usage.append((power_watts / power_capacity_watts) * 100)
-        self.samples.append(power_usage)
-
-    def clear(self) -> None:
-        self.samples.clear()
-
-    def aggregate(self) -> dict:
-        if not self.samples:
-            return {}
-        stats = {}
-        device_count = pynvml.nvmlDeviceGetCount()  # type: ignore
-        for i in range(device_count):
-            samples = [sample[i] for sample in self.samples]
-            aggregate = aggregate_mean(samples)
-            stats[self.name.format(i)] = aggregate
-
-            handle = pynvml.nvmlDeviceGetHandleByIndex(i)  # type: ignore
-            if gpu_in_use_by_this_process(handle, self.pid):
-                stats[self.name.format(f"process.{i}")] = aggregate
-
-        return stats
-
-
-@asset_registry.register
-class GPU:
-    def __init__(
-        self,
-        interface: "Interface",
-        settings: "SettingsStatic",
-        shutdown_event: threading.Event,
-    ) -> None:
-        self.name = self.__class__.__name__.lower()
-        self.metrics: List[Metric] = [
-            GPUMemoryAllocated(settings.x_stats_pid),
-            GPUMemoryAllocatedBytes(settings.x_stats_pid),
-            GPUMemoryUtilization(settings.x_stats_pid),
-            GPUUtilization(settings.x_stats_pid),
-            GPUTemperature(settings.x_stats_pid),
-            GPUPowerUsageWatts(settings.x_stats_pid),
-            GPUPowerUsagePercent(settings.x_stats_pid),
-        ]
-        self.metrics_monitor = MetricsMonitor(
-            self.name,
-            self.metrics,
-            interface,
-            settings,
-            shutdown_event,
-        )
-
-    @classmethod
-    def is_available(cls) -> bool:
-        try:
-            pynvml.nvmlInit()  # type: ignore
-            return True
-        except pynvml.NVMLError_LibraryNotFound:  # type: ignore
-            return False
-        except Exception:
-            logger.exception("Error initializing NVML.")
-            return False
-
-    def start(self) -> None:
-        self.metrics_monitor.start()
-
-    def finish(self) -> None:
-        self.metrics_monitor.finish()
-
-    def probe(self) -> dict:
-        info = {}
-        try:
-            pynvml.nvmlInit()  # type: ignore
-            # todo: this is an adapter for the legacy stats system:
-            info["gpu"] = pynvml.nvmlDeviceGetName(pynvml.nvmlDeviceGetHandleByIndex(0))  # type: ignore
-            info["gpu_count"] = pynvml.nvmlDeviceGetCount()  # type: ignore
-
-            device_count = pynvml.nvmlDeviceGetCount()  # type: ignore
-            devices = []
-            for i in range(device_count):
-                handle = pynvml.nvmlDeviceGetHandleByIndex(i)  # type: ignore
-                gpu_info = pynvml.nvmlDeviceGetMemoryInfo(handle)  # type: ignore
-                devices.append(
-                    {
-                        "name": pynvml.nvmlDeviceGetName(handle),
-                        "memory_total": gpu_info.total,
-                    }
-                )
-            info["gpu_devices"] = devices
-
-        except pynvml.NVMLError:
-            pass
-        except Exception:
-            logger.exception("Error Probing GPU.")
-
-        return info

wandb/sdk/internal/system/assets/gpu_amd.py

@@ -1,233 +0,0 @@
-import json
-import logging
-import shutil
-import subprocess
-import threading
-from collections import deque
-from typing import TYPE_CHECKING, Any, Dict, Final, List, Literal, Union
-
-from wandb.sdk.lib import telemetry
-
-from .aggregators import aggregate_mean
-from .asset_registry import asset_registry
-from .interfaces import Interface, Metric, MetricsMonitor
-
-if TYPE_CHECKING:
-    from typing import Deque
-
-    from wandb.sdk.internal.settings_static import SettingsStatic
-
-
-logger = logging.getLogger(__name__)
-ROCM_SMI_CMD: Final[str] = shutil.which("rocm-smi") or "/usr/bin/rocm-smi"
-
-
-_StatsKeys = Literal[
-    "gpu",
-    "memoryAllocated",
-    "temp",
-    "powerWatts",
-    "powerPercent",
-]
-_Stats = Dict[_StatsKeys, float]
-
-
-_InfoDict = Dict[str, Union[int, List[Dict[str, Any]]]]
-
-
-def get_rocm_smi_stats() -> Dict[str, Any]:
-    command = [str(ROCM_SMI_CMD), "-a", "--json"]
-    output = subprocess.check_output(command, universal_newlines=True).strip()
-    if "No AMD GPUs specified" in output:
-        return {}
-    return json.loads(output.split("\n")[0])  # type: ignore
-
-
-def parse_stats(stats: Dict[str, str]) -> _Stats:
-    """Parse stats from rocm-smi output."""
-    parsed_stats: _Stats = {}
-
-    try:
-        parsed_stats["gpu"] = float(stats.get("GPU use (%)"))  # type: ignore
-    except (TypeError, ValueError):
-        logger.warning("Could not parse GPU usage as float")
-    try:
-        parsed_stats["memoryAllocated"] = float(stats.get("GPU memory use (%)"))  # type: ignore
-    except (TypeError, ValueError):
-        logger.warning("Could not parse GPU memory allocation as float")
-    try:
-        parsed_stats["temp"] = float(stats.get("Temperature (Sensor memory) (C)"))  # type: ignore
-    except (TypeError, ValueError):
-        logger.warning("Could not parse GPU temperature as float")
-    try:
-        parsed_stats["powerWatts"] = float(
-            stats.get("Average Graphics Package Power (W)")  # type: ignore
-        )
-    except (TypeError, ValueError):
-        logger.warning("Could not parse GPU power as float")
-    try:
-        parsed_stats["powerPercent"] = (
-            float(stats.get("Average Graphics Package Power (W)"))  # type: ignore
-            / float(stats.get("Max Graphics Package Power (W)"))  # type: ignore
-            * 100
-        )
-    except (TypeError, ValueError):
-        logger.warning("Could not parse GPU average/max power as float")
-
-    return parsed_stats
-
-
-class GPUAMDStats:
-    """Stats for AMD GPU devices."""
-
-    name = "gpu.{gpu_id}.{key}"
-    samples: "Deque[List[_Stats]]"
-
-    def __init__(self) -> None:
-        self.samples = deque()
-
-    def sample(self) -> None:
-        try:
-            raw_stats = get_rocm_smi_stats()
-            cards = []
-
-            card_keys = [
-                key for key in sorted(raw_stats.keys()) if key.startswith("card")
-            ]
-
-            for card_key in card_keys:
-                card_stats = raw_stats[card_key]
-                stats = parse_stats(card_stats)
-                if stats:
-                    cards.append(stats)
-
-            if cards:
-                self.samples.append(cards)
-
-        except (OSError, ValueError, TypeError, subprocess.CalledProcessError):
-            logger.exception("GPU stats error")
-
-    def clear(self) -> None:
-        self.samples.clear()
-
-    def aggregate(self) -> dict:
-        if not self.samples:
-            return {}
-        stats = {}
-        device_count = len(self.samples[0])
-
-        for i in range(device_count):
-            samples = [sample[i] for sample in self.samples]
-
-            for key in samples[0].keys():
-                samples_key = [s[key] for s in samples]
-                aggregate = aggregate_mean(samples_key)
-                stats[self.name.format(gpu_id=i, key=key)] = aggregate
-
-        return stats
-
-
-@asset_registry.register
-class GPUAMD:
-    """GPUAMD is a class for monitoring AMD GPU devices.
-
-    Uses AMD's rocm_smi tool to get GPU stats.
-    For the list of supported environments and devices, see
-    https://github.com/RadeonOpenCompute/ROCm/blob/develop/docs/deploy/
-    """
-
-    def __init__(
-        self,
-        interface: "Interface",
-        settings: "SettingsStatic",
-        shutdown_event: threading.Event,
-    ) -> None:
-        self.name = self.__class__.__name__.lower()
-        self.metrics: List[Metric] = [
-            GPUAMDStats(),
-        ]
-        self.metrics_monitor = MetricsMonitor(
-            self.name,
-            self.metrics,
-            interface,
-            settings,
-            shutdown_event,
-        )
-        telemetry_record = telemetry.TelemetryRecord()
-        telemetry_record.env.amd_gpu = True
-        interface._publish_telemetry(telemetry_record)
-
-    @classmethod
-    def is_available(cls) -> bool:
-        rocm_smi_available = shutil.which(ROCM_SMI_CMD) is not None
-        if not rocm_smi_available:
-            # If rocm-smi is not available, we can't monitor AMD GPUs
-            return False
-
-        is_driver_initialized = False
-
-        try:
-            # inspired by https://github.com/ROCm/rocm_smi_lib/blob/5d2cd0c2715ae45b8f9cfe1e777c6c2cd52fb601/python_smi_tools/rocm_smi.py#L71C1-L81C17
-            with open("/sys/module/amdgpu/initstate") as file:
-                file_content = file.read()
-                if "live" in file_content:
-                    is_driver_initialized = True
-        except FileNotFoundError:
-            pass
-
-        can_read_rocm_smi = False
-        try:
-            # try to read stats from rocm-smi and parse them
-            raw_stats = get_rocm_smi_stats()
-            card_keys = [
-                key for key in sorted(raw_stats.keys()) if key.startswith("card")
-            ]
-
-            for card_key in card_keys:
-                card_stats = raw_stats[card_key]
-                parse_stats(card_stats)
-
-            can_read_rocm_smi = True
-        except Exception:
-            pass
-
-        return is_driver_initialized and can_read_rocm_smi
-
-    def start(self) -> None:
-        self.metrics_monitor.start()
-
-    def finish(self) -> None:
-        self.metrics_monitor.finish()
-
-    def probe(self) -> dict:
-        info: _InfoDict = {}
-        try:
-            stats = get_rocm_smi_stats()
-
-            info["gpu_count"] = len(
-                [key for key in stats.keys() if key.startswith("card")]
-            )
-            key_mapping = {
-                "id": "GPU ID",
-                "unique_id": "Unique ID",
-                "vbios_version": "VBIOS version",
-                "performance_level": "Performance Level",
-                "gpu_overdrive": "GPU OverDrive value (%)",
-                "gpu_memory_overdrive": "GPU Memory OverDrive value (%)",
-                "max_power": "Max Graphics Package Power (W)",
-                "series": "Card series",
-                "model": "Card model",
-                "vendor": "Card vendor",
-                "sku": "Card SKU",
-                "sclk_range": "Valid sclk range",
-                "mclk_range": "Valid mclk range",
-            }
-
-            info["gpu_devices"] = [
-                {k: stats[key][v] for k, v in key_mapping.items() if stats[key].get(v)}
-                for key in stats.keys()
-                if key.startswith("card")
-            ]
-        except Exception:
-            logger.exception("GPUAMD probe error")
-        return info