be0fd8b822
Address 2 high, 3 medium, and 1 low severity issues identified during an audit of the ROCm support implementation. Also adds 24 new test cases (21 Python + 3 shell) to cover edge cases and prevent regressions on CUDA, CPU, macOS, and Windows pathways. Bug fixes: - install.sh: add validation guard for _rocm_tag to reject malformed output (e.g. "rocm." from garbled amd-smi) that would produce bogus PyTorch index URLs - install.sh: strip Debian epoch prefix (e.g. "2:6.2.0") from dpkg-query output before parsing ROCm version - install_python_stack.py: add timeout=30 to the subprocess probe that checks whether torch is already GPU-enabled, preventing hangs - install_llama_prebuilt.py: log a warning that the ROCm prebuilt is compiled for ROCm 7.2 and may fall back to source build on other ROCm versions - install_llama_prebuilt.py: log a message when Windows HIP prebuilt is not found before silently falling through to CPU - tests/sh/test_get_torch_index_url.sh: replace no-op test 18 (could not actually test /opt/rocm path) with 3 real tests for malformed amd-smi output (empty version, "N/A", trailing text) Additional changes: - Export IS_ROCM from studio/backend/utils/hardware/__init__.py so downstream code can import it directly from the package Test results: 23/23 shell, 68/68 Python, 205/206 full suite (1 pre-existing failure unrelated to this change).
85 lines
2 KiB
Python
85 lines
2 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""
|
|
Hardware detection and GPU utilities
|
|
"""
|
|
|
|
from .hardware import (
|
|
DeviceType,
|
|
DEVICE,
|
|
CHAT_ONLY,
|
|
IS_ROCM,
|
|
detect_hardware,
|
|
get_device,
|
|
is_apple_silicon,
|
|
clear_gpu_cache,
|
|
get_gpu_memory_info,
|
|
log_gpu_memory,
|
|
get_gpu_summary,
|
|
get_package_versions,
|
|
get_gpu_utilization,
|
|
get_visible_gpu_utilization,
|
|
get_backend_visible_gpu_info,
|
|
get_physical_gpu_count,
|
|
get_visible_gpu_count,
|
|
get_parent_visible_gpu_ids,
|
|
resolve_requested_gpu_ids,
|
|
estimate_fp16_model_size_bytes,
|
|
estimate_required_model_memory_gb,
|
|
auto_select_gpu_ids,
|
|
prepare_gpu_selection,
|
|
safe_num_proc,
|
|
safe_thread_num_proc,
|
|
dataset_map_num_proc,
|
|
get_device_map,
|
|
get_offloaded_device_map_entries,
|
|
raise_if_offloaded,
|
|
apply_gpu_ids,
|
|
)
|
|
|
|
from .vram_estimation import (
|
|
ModelArchConfig,
|
|
TrainingVramConfig,
|
|
VramBreakdown,
|
|
extract_arch_config,
|
|
estimate_training_vram,
|
|
)
|
|
|
|
__all__ = [
|
|
"DeviceType",
|
|
"DEVICE",
|
|
"CHAT_ONLY",
|
|
"IS_ROCM",
|
|
"detect_hardware",
|
|
"get_device",
|
|
"is_apple_silicon",
|
|
"clear_gpu_cache",
|
|
"get_gpu_memory_info",
|
|
"log_gpu_memory",
|
|
"get_gpu_summary",
|
|
"get_package_versions",
|
|
"get_gpu_utilization",
|
|
"get_visible_gpu_utilization",
|
|
"get_backend_visible_gpu_info",
|
|
"get_physical_gpu_count",
|
|
"get_visible_gpu_count",
|
|
"get_parent_visible_gpu_ids",
|
|
"resolve_requested_gpu_ids",
|
|
"estimate_fp16_model_size_bytes",
|
|
"estimate_required_model_memory_gb",
|
|
"auto_select_gpu_ids",
|
|
"prepare_gpu_selection",
|
|
"safe_num_proc",
|
|
"safe_thread_num_proc",
|
|
"dataset_map_num_proc",
|
|
"get_device_map",
|
|
"get_offloaded_device_map_entries",
|
|
"raise_if_offloaded",
|
|
"apply_gpu_ids",
|
|
"ModelArchConfig",
|
|
"TrainingVramConfig",
|
|
"VramBreakdown",
|
|
"extract_arch_config",
|
|
"estimate_training_vram",
|
|
]
|