92 lines
3.5 KiB
Python
92 lines
3.5 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
|
|
|
|
"""Tests for _nvidia_classify_spark_unified_memory (Spark OOM-guard classifier).
|
|
|
|
Two paths: (1) ``is_integrated`` property (authoritative on native Linux),
|
|
(2) name-token match -- needed because WSL2 GPU paravirtualization masks
|
|
``is_integrated`` to 0 and renames the device (N1X reports ``JMJWOA-Generic-GPU``;
|
|
verified live). Mirrors test_rocm_oom_guard.py, which the NVIDIA guard models.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from core.training.worker import _nvidia_classify_spark_unified_memory
|
|
|
|
|
|
def _props(**kwargs) -> SimpleNamespace:
|
|
"""Fake device-properties object with the given attributes."""
|
|
return SimpleNamespace(**kwargs)
|
|
|
|
|
|
# ── Path 1: is_integrated property ───────────────────────────────────────────
|
|
|
|
|
|
class TestIsIntegratedProperty:
|
|
"""``is_integrated`` truthy means unified memory, regardless of name."""
|
|
|
|
def test_integrated_native_spark(self) -> None:
|
|
props = _props(is_integrated = 1, name = "NVIDIA GB10")
|
|
marker, is_unified = _nvidia_classify_spark_unified_memory(props)
|
|
assert marker == "is_integrated"
|
|
assert is_unified is True
|
|
|
|
def test_integrated_wins_even_with_unknown_name(self) -> None:
|
|
props = _props(is_integrated = 1, name = "Some Future Unified Part")
|
|
marker, is_unified = _nvidia_classify_spark_unified_memory(props)
|
|
assert marker == "is_integrated"
|
|
assert is_unified is True
|
|
|
|
|
|
# ── Path 2: device-name token fallback (WSL masks is_integrated) ────────────
|
|
|
|
|
|
class TestDeviceNameTokenFallback:
|
|
"""is_integrated == 0 (or absent) -> classify by Spark name tokens."""
|
|
|
|
@pytest.mark.parametrize(
|
|
"name, expected_marker",
|
|
[
|
|
("JMJWOA-Generic-GPU", "JMJWOA"), # N1X under WSL2 (verified live)
|
|
("NVIDIA GB10", "GB10"), # native DGX Spark
|
|
("NVIDIA GB110", "GB110"), # "GB10" is not a substring of "GB110"
|
|
("NVIDIA DGX Spark", "DGX SPARK"),
|
|
("nvidia n1x prototype", "N1X"), # case-insensitive
|
|
],
|
|
)
|
|
def test_spark_names_unified(self, name: str, expected_marker: str) -> None:
|
|
props = _props(is_integrated = 0, name = name)
|
|
marker, is_unified = _nvidia_classify_spark_unified_memory(props)
|
|
assert is_unified is True
|
|
assert marker == expected_marker
|
|
|
|
@pytest.mark.parametrize(
|
|
"name",
|
|
[
|
|
"NVIDIA GeForce RTX 4090",
|
|
"NVIDIA H100 80GB HBM3",
|
|
"NVIDIA RTX 6000 Ada Generation",
|
|
"Tesla T4",
|
|
],
|
|
)
|
|
def test_discrete_names_not_unified(self, name: str) -> None:
|
|
props = _props(is_integrated = 0, name = name)
|
|
marker, is_unified = _nvidia_classify_spark_unified_memory(props)
|
|
assert is_unified is False
|
|
assert marker == ""
|
|
|
|
def test_missing_attrs_defaults_discrete(self) -> None:
|
|
"""No is_integrated, no name -> discrete (guard stays off)."""
|
|
marker, is_unified = _nvidia_classify_spark_unified_memory(_props())
|
|
assert is_unified is False
|
|
assert marker == ""
|
|
|
|
def test_none_name_defaults_discrete(self) -> None:
|
|
props = _props(is_integrated = 0, name = None)
|
|
marker, is_unified = _nvidia_classify_spark_unified_memory(props)
|
|
assert is_unified is False
|
|
assert marker == ""
|