ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-03-31 11:08:36 +00:00
commit 14823265d2
3 changed files with 71 additions and 62 deletions
Download patch file Download diff file Expand all files Collapse all files

19
studio/backend/core/training/worker.py
View file

@ -263,8 +263,11 @@ def _install_package_wheel_first(
# Prefer uv for faster dependency resolution when available
if shutil.which("uv"):
pypi_cmd = [
"uv", "pip", "install",
"--python", sys.executable,
"uv",
"pip",
"install",
"--python",
sys.executable,
"--no-build-isolation",
"--no-deps",
f"{pypi_name}=={pypi_version}",
@ -294,7 +297,9 @@ def _install_package_wheel_first(
)
except _sp.TimeoutExpired:
logger.error(
"%s installation timed out after %ds", display_name, timeout,
"%s installation timed out after %ds",
display_name,
timeout,
)
_send_status(
event_queue,
@ -308,7 +313,9 @@ def _install_package_wheel_first(
error_lines = (result.stdout or "").strip().splitlines()
snippet = "\n".join(error_lines[-5:]) if error_lines else "(no output)"
logger.error(
"Failed to compile %s for ROCm:\n%s", display_name, result.stdout,
"Failed to compile %s for ROCm:\n%s",
display_name,
result.stdout,
)
_send_status(
event_queue,
@ -318,7 +325,9 @@ def _install_package_wheel_first(
)
else:
logger.error(
"Failed to install %s from PyPI:\n%s", display_name, result.stdout,
"Failed to install %s from PyPI:\n%s",
display_name,
result.stdout,
)
return

30
studio/backend/utils/hardware/amd.py
View file

@ -22,9 +22,9 @@ def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[dict]:
try:
result = subprocess.run(
["amd-smi", *args, "--json"],
capture_output=True,
text=True,
timeout=timeout,
capture_output = True,
text = True,
timeout = timeout,
)
except (OSError, subprocess.TimeoutExpired) as e:
logger.warning("amd-smi query failed: %s", e)
@ -62,7 +62,9 @@ def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]:
# amd-smi metric output structure varies by version; try common paths
usage = gpu_data.get("usage", gpu_data.get("gpu_activity", {}))
if isinstance(usage, dict):
gpu_util = _parse_numeric(usage.get("gfx_activity", usage.get("gpu_use_percent")))
gpu_util = _parse_numeric(
usage.get("gfx_activity", usage.get("gpu_use_percent"))
)
else:
gpu_util = _parse_numeric(usage)
@ -70,8 +72,13 @@ def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]:
temp_data = gpu_data.get("temperature", {})
if isinstance(temp_data, dict):
temp = _parse_numeric(
temp_data.get("edge", temp_data.get("temperature_edge",
temp_data.get("hotspot", temp_data.get("temperature_hotspot"))))
temp_data.get(
"edge",
temp_data.get(
"temperature_edge",
temp_data.get("hotspot", temp_data.get("temperature_hotspot")),
),
)
)
else:
temp = _parse_numeric(temp_data)
@ -80,9 +87,10 @@ def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]:
power_data = gpu_data.get("power", {})
if isinstance(power_data, dict):
power_draw = _parse_numeric(
power_data.get("current_socket_power",
power_data.get("average_socket_power",
power_data.get("socket_power")))
power_data.get(
"current_socket_power",
power_data.get("average_socket_power", power_data.get("socket_power")),
)
)
power_limit = _parse_numeric(
power_data.get("power_cap", power_data.get("max_power_limit"))
@ -120,7 +128,9 @@ def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]:
# Build the standardized dict (same shape as nvidia._build_gpu_metrics)
vram_used_gb = round(vram_used_mb / 1024, 2) if vram_used_mb is not None else None
vram_total_gb = round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None
vram_total_gb = (
round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None
)
vram_util = (
round((vram_used_mb / vram_total_mb) * 100, 1)
if vram_used_mb is not None and vram_total_mb and vram_total_mb > 0

84
tests/studio/install/test_rocm_support.py
View file

@ -992,7 +992,7 @@ class TestWorkerRocmMambaSsm:
# Mock all the imports worker.py needs
sys.modules["structlog"] = MagicMock()
sys.modules["loggers"] = MagicMock()
sys.modules["loggers"].get_logger = MagicMock(return_value=MagicMock())
sys.modules["loggers"].get_logger = MagicMock(return_value = MagicMock())
sys.modules["utils"] = MagicMock()
sys.modules["utils.hardware"] = MagicMock()
@ -1009,11 +1009,11 @@ class TestWorkerRocmMambaSsm:
"cxx11abi": "TRUE",
}
result = worker_mod._direct_wheel_url(
filename_prefix="causal_conv1d",
package_version="1.6.1",
release_tag="v1.6.1.post4",
release_base_url="https://github.com/Dao-AILab/causal-conv1d/releases/download",
env=env_rocm,
filename_prefix = "causal_conv1d",
package_version = "1.6.1",
release_tag = "v1.6.1.post4",
release_base_url = "https://github.com/Dao-AILab/causal-conv1d/releases/download",
env = env_rocm,
)
assert result is None
@ -1049,16 +1049,12 @@ class TestAmdGpuMonitoring:
def test_amd_py_exists(self):
"""amd.py should exist in the hardware directory."""
amd_path = (
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
)
amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
assert amd_path.exists()
def test_amd_py_has_required_functions(self):
"""amd.py should export the same function signatures as nvidia.py."""
amd_path = (
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
)
amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
source = amd_path.read_text()
assert "def get_physical_gpu_count" in source
assert "def get_primary_gpu_utilization" in source
@ -1066,15 +1062,13 @@ class TestAmdGpuMonitoring:
def test_amd_smi_json_parsing(self):
"""Verify _extract_gpu_metrics parses amd-smi JSON correctly."""
amd_path = (
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
)
amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
_amd_spec = importlib.util.spec_from_file_location("test_amd", amd_path)
assert _amd_spec is not None and _amd_spec.loader is not None
amd_mod = importlib.util.module_from_spec(_amd_spec)
sys.modules["loggers"] = MagicMock()
sys.modules["loggers"].get_logger = MagicMock(return_value=MagicMock())
sys.modules["loggers"].get_logger = MagicMock(return_value = MagicMock())
try:
_amd_spec.loader.exec_module(amd_mod)
@ -1106,32 +1100,34 @@ class TestAmdGpuMonitoring:
def test_amd_primary_gpu_with_mock(self):
"""get_primary_gpu_utilization returns correct dict with mocked amd-smi."""
amd_path = (
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
)
amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
_amd_spec = importlib.util.spec_from_file_location("test_amd2", amd_path)
assert _amd_spec is not None and _amd_spec.loader is not None
amd_mod = importlib.util.module_from_spec(_amd_spec)
sys.modules["loggers"] = MagicMock()
sys.modules["loggers"].get_logger = MagicMock(return_value=MagicMock())
sys.modules["loggers"].get_logger = MagicMock(return_value = MagicMock())
try:
_amd_spec.loader.exec_module(amd_mod)
except Exception:
pytest.skip("Could not load amd module")
mock_json = json.dumps([{
"usage": {"gfx_activity": "50"},
"temperature": {"edge": "65"},
"power": {"current_socket_power": "150", "power_cap": "250"},
"vram": {"vram_used": 4096, "vram_total": 16384},
}])
mock_json = json.dumps(
[
{
"usage": {"gfx_activity": "50"},
"temperature": {"edge": "65"},
"power": {"current_socket_power": "150", "power_cap": "250"},
"vram": {"vram_used": 4096, "vram_total": 16384},
}
]
)
mock_result = MagicMock()
mock_result.returncode = 0
mock_result.stdout = mock_json
with patch.object(subprocess, "run", return_value=mock_result):
with patch.object(subprocess, "run", return_value = mock_result):
result = amd_mod.get_primary_gpu_utilization()
assert result["available"] is True
assert result["gpu_utilization_pct"] == 50.0
@ -1139,38 +1135,32 @@ class TestAmdGpuMonitoring:
def test_amd_smi_not_found_returns_unavailable(self):
"""get_primary_gpu_utilization returns available=False when amd-smi is missing."""
amd_path = (
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
)
amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
_amd_spec = importlib.util.spec_from_file_location("test_amd3", amd_path)
assert _amd_spec is not None and _amd_spec.loader is not None
amd_mod = importlib.util.module_from_spec(_amd_spec)
sys.modules["loggers"] = MagicMock()
sys.modules["loggers"].get_logger = MagicMock(return_value=MagicMock())
sys.modules["loggers"].get_logger = MagicMock(return_value = MagicMock())
try:
_amd_spec.loader.exec_module(amd_mod)
except Exception:
pytest.skip("Could not load amd module")
with patch.object(
subprocess, "run", side_effect=OSError("amd-smi not found")
):
with patch.object(subprocess, "run", side_effect = OSError("amd-smi not found")):
result = amd_mod.get_primary_gpu_utilization()
assert result["available"] is False
def test_amd_timeout_returns_unavailable(self):
"""get_primary_gpu_utilization handles timeout gracefully."""
amd_path = (
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
)
amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
_amd_spec = importlib.util.spec_from_file_location("test_amd4", amd_path)
assert _amd_spec is not None and _amd_spec.loader is not None
amd_mod = importlib.util.module_from_spec(_amd_spec)
sys.modules["loggers"] = MagicMock()
sys.modules["loggers"].get_logger = MagicMock(return_value=MagicMock())
sys.modules["loggers"].get_logger = MagicMock(return_value = MagicMock())
try:
_amd_spec.loader.exec_module(amd_mod)
@ -1180,7 +1170,7 @@ class TestAmdGpuMonitoring:
with patch.object(
subprocess,
"run",
side_effect=subprocess.TimeoutExpired("amd-smi", 5),
side_effect = subprocess.TimeoutExpired("amd-smi", 5),
):
result = amd_mod.get_primary_gpu_utilization()
assert result["available"] is False
@ -1210,7 +1200,7 @@ class TestHardwareAmdBranching:
source = hw_path.read_text()
# Find the get_gpu_utilization function
func_start = source.find("def get_gpu_utilization")
func_body = source[func_start:source.find("\ndef ", func_start + 1)]
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert "IS_ROCM" in func_body
assert "amd.get_primary_gpu_utilization" in func_body
@ -1221,7 +1211,7 @@ class TestHardwareAmdBranching:
)
source = hw_path.read_text()
func_start = source.find("def get_visible_gpu_utilization")
func_body = source[func_start:source.find("\ndef ", func_start + 1)]
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert "IS_ROCM" in func_body
assert "amd.get_visible_gpu_utilization" in func_body
@ -1232,7 +1222,7 @@ class TestHardwareAmdBranching:
)
source = hw_path.read_text()
func_start = source.find("def get_physical_gpu_count")
func_body = source[func_start:source.find("\ndef ", func_start + 1)]
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert "IS_ROCM" in func_body
assert "amd.get_physical_gpu_count" in func_body
@ -1275,7 +1265,7 @@ class TestIsRdnaExpansion:
utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py"
source = utils_path.read_text()
func_start = source.find("def is_rdna()")
func_body = source[func_start:source.find("\ndef ", func_start + 1)]
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert "gfx1030" in func_body
assert "gfx1031" in func_body
assert "gfx1032" in func_body
@ -1285,7 +1275,7 @@ class TestIsRdnaExpansion:
utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py"
source = utils_path.read_text()
func_start = source.find("def is_rdna()")
func_body = source[func_start:source.find("\ndef ", func_start + 1)]
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert "gfx1100" in func_body
assert "gfx1101" in func_body
assert "gfx1102" in func_body
@ -1296,7 +1286,7 @@ class TestIsRdnaExpansion:
utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py"
source = utils_path.read_text()
func_start = source.find("def is_rdna()")
func_body = source[func_start:source.find("\ndef ", func_start + 1)]
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert "gfx1150" in func_body
assert "gfx1151" in func_body
assert "gfx1152" in func_body
@ -1306,7 +1296,7 @@ class TestIsRdnaExpansion:
utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py"
source = utils_path.read_text()
func_start = source.find("def is_rdna()")
func_body = source[func_start:source.find("\ndef ", func_start + 1)]
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert "gfx1200" in func_body
assert "gfx1201" in func_body
@ -1315,7 +1305,7 @@ class TestIsRdnaExpansion:
utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py"
source = utils_path.read_text()
func_start = source.find("def is_cdna()")
func_body = source[func_start:source.find("\ndef ", func_start + 1)]
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert "gfx940" in func_body
assert "gfx941" in func_body
assert "gfx942" in func_body