diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md index f0201b14d9ec9..1588d482108c8 100644 --- a/src/lightning/pytorch/CHANGELOG.md +++ b/src/lightning/pytorch/CHANGELOG.md @@ -22,6 +22,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - Added support for remote storage (fsspec URLs) when saving and loading distributed checkpoints with `ModelParallelStrategy` ([#21797](https://github.com/Lightning-AI/pytorch-lightning/issues/21797)) +- Added classmethod `Accelerator.device_name` so that users could see device information during setup. ([#21112](https://github.com/Lightning-AI/pytorch-lightning/pull/21112)) + ### Changed - diff --git a/src/lightning/pytorch/accelerators/accelerator.py b/src/lightning/pytorch/accelerators/accelerator.py index 9238071178a80..bed925b1c013c 100644 --- a/src/lightning/pytorch/accelerators/accelerator.py +++ b/src/lightning/pytorch/accelerators/accelerator.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from abc import ABC -from typing import Any +from typing import Any, Optional import lightning.pytorch as pl from lightning.fabric.accelerators.accelerator import Accelerator as _Accelerator @@ -45,3 +45,8 @@ def get_device_stats(self, device: _DEVICE) -> dict[str, Any]: """ raise NotImplementedError + + @classmethod + def device_name(cls, device: Optional[_DEVICE] = None) -> str: + """Get the device name for a given device.""" + return str(cls.is_available()) diff --git a/src/lightning/pytorch/accelerators/cuda.py b/src/lightning/pytorch/accelerators/cuda.py index 63a0d8adba8ea..9e28084c98e2d 100644 --- a/src/lightning/pytorch/accelerators/cuda.py +++ b/src/lightning/pytorch/accelerators/cuda.py @@ -118,6 +118,13 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No description=cls.__name__, ) + @classmethod + @override + def device_name(cls, device: Optional[_DEVICE] = None) -> str: + if not cls.is_available(): + return "" + return torch.cuda.get_device_name(device) + def get_nvidia_gpu_stats(device: _DEVICE) -> dict[str, float]: # pragma: no-cover """Get GPU stats including memory, fan speed, and temperature from nvidia-smi. diff --git a/src/lightning/pytorch/accelerators/mps.py b/src/lightning/pytorch/accelerators/mps.py index a999411f5ed8b..9c35b71681671 100644 --- a/src/lightning/pytorch/accelerators/mps.py +++ b/src/lightning/pytorch/accelerators/mps.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import subprocess from typing import Any, Optional, Union import torch @@ -92,6 +93,27 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No description=cls.__name__, ) + @classmethod + @override + def device_name(cls, device: Optional[_DEVICE] = None) -> str: + if not cls.is_available(): + return "" + return _get_mps_device_name() + + +def _get_mps_device_name() -> str: + try: + result = subprocess.run( + ["sysctl", "-n", "machdep.cpu.brand_string"], + capture_output=True, + text=True, + check=True, + ) + result_str = result.stdout.strip() + except subprocess.SubprocessError: + result_str = "True (mps)" + return result_str + # device metrics _VM_PERCENT = "M1_vm_percent" diff --git a/src/lightning/pytorch/accelerators/xla.py b/src/lightning/pytorch/accelerators/xla.py index 700a29c2e3e7b..5fd94d0fb464f 100644 --- a/src/lightning/pytorch/accelerators/xla.py +++ b/src/lightning/pytorch/accelerators/xla.py @@ -11,11 +11,12 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any +from typing import Any, Optional from typing_extensions import override from lightning.fabric.accelerators.registry import _AcceleratorRegistry +from lightning.fabric.accelerators.xla import _XLA_GREATER_EQUAL_2_1 from lightning.fabric.accelerators.xla import XLAAccelerator as FabricXLAAccelerator from lightning.fabric.utilities.types import _DEVICE from lightning.pytorch.accelerators.accelerator import Accelerator @@ -62,3 +63,25 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No cls, description=cls.__name__, ) + + @classmethod + @override + def device_name(cls, device: Optional[_DEVICE] = None) -> str: + is_available = cls.is_available() + if not is_available: + return "" + + if _XLA_GREATER_EQUAL_2_1: + from torch_xla._internal import tpu + else: + from torch_xla.experimental import tpu + import torch_xla.core.xla_env_vars as xenv + from requests.exceptions import HTTPError + + try: + ret = tpu.get_tpu_env()[xenv.ACCELERATOR_TYPE] + except HTTPError: + # Fallback to "True" if HTTPError is raised during retrieving device information + ret = str(is_available) + + return ret diff --git a/src/lightning/pytorch/trainer/setup.py b/src/lightning/pytorch/trainer/setup.py index 4f522c7c008bc..0ac22307d5a79 100644 --- a/src/lightning/pytorch/trainer/setup.py +++ b/src/lightning/pytorch/trainer/setup.py @@ -151,20 +151,23 @@ def _init_profiler(trainer: "pl.Trainer", profiler: Optional[Union[Profiler, str def _log_device_info(trainer: "pl.Trainer") -> None: if CUDAAccelerator.is_available(): - gpu_available = True - gpu_type = " (cuda)" + if isinstance(trainer.accelerator, CUDAAccelerator): + device_name = ", ".join(list({CUDAAccelerator.device_name(d) for d in trainer.device_ids})) + else: + device_name = CUDAAccelerator.device_name() elif MPSAccelerator.is_available(): - gpu_available = True - gpu_type = " (mps)" + device_name = MPSAccelerator.device_name() else: - gpu_available = False - gpu_type = "" + device_name = str(False) - gpu_used = isinstance(trainer.accelerator, (CUDAAccelerator, MPSAccelerator)) - rank_zero_info(f"GPU available: {gpu_available}{gpu_type}, used: {gpu_used}") + gpu_used = trainer.num_devices if isinstance(trainer.accelerator, (CUDAAccelerator, MPSAccelerator)) else 0 + rank_zero_info(f"GPU available: {device_name}, using: {gpu_used} devices.") num_tpu_cores = trainer.num_devices if isinstance(trainer.accelerator, XLAAccelerator) else 0 - rank_zero_info(f"TPU available: {XLAAccelerator.is_available()}, using: {num_tpu_cores} TPU cores") + rank_zero_info( + f"TPU available: {XLAAccelerator.device_name() if XLAAccelerator.is_available() else str(False)}, " + f"using: {num_tpu_cores} TPU cores" + ) if ( CUDAAccelerator.is_available() diff --git a/tests/tests_pytorch/accelerators/test_gpu.py b/tests/tests_pytorch/accelerators/test_gpu.py index 5a71887e17eec..02dd4bfdf2d97 100644 --- a/tests/tests_pytorch/accelerators/test_gpu.py +++ b/tests/tests_pytorch/accelerators/test_gpu.py @@ -68,3 +68,20 @@ def test_gpu_availability(): def test_warning_if_gpus_not_used(cuda_count_1): with pytest.warns(UserWarning, match="GPU available but not used"): Trainer(accelerator="cpu") + + +@RunIf(min_cuda_gpus=1) +def test_gpu_device_name(): + for i in range(torch.cuda.device_count()): + assert torch.cuda.get_device_name(i) == CUDAAccelerator.device_name(i) + + with torch.device("cuda:0"): + assert torch.cuda.get_device_name(0) == CUDAAccelerator.device_name() + + if torch.cuda.device_count() > 1: + with torch.device("cuda:1"): + assert torch.cuda.get_device_name(1) == CUDAAccelerator.device_name(1) + + +def test_gpu_device_name_no_gpu(cuda_count_0): + assert CUDAAccelerator.device_name() == "" diff --git a/tests/tests_pytorch/accelerators/test_mps.py b/tests/tests_pytorch/accelerators/test_mps.py index c0a28840f0ef6..242b21bde7fb7 100644 --- a/tests/tests_pytorch/accelerators/test_mps.py +++ b/tests/tests_pytorch/accelerators/test_mps.py @@ -13,6 +13,8 @@ # limitations under the License. from collections import namedtuple +from subprocess import SubprocessError +from unittest import mock import pytest import torch @@ -20,6 +22,7 @@ import tests_pytorch.helpers.pipelines as tpipes from lightning.pytorch import Trainer from lightning.pytorch.accelerators import MPSAccelerator +from lightning.pytorch.accelerators.mps import _get_mps_device_name from lightning.pytorch.demos.boring_classes import BoringModel from tests_pytorch.helpers.runif import RunIf @@ -39,6 +42,16 @@ def test_mps_availability(): assert MPSAccelerator.is_available() +@RunIf(mps=True) +def test_mps_device_name(): + assert MPSAccelerator.device_name() == "True (mps)" + + +def test_mps_device_name_not_available(): + with mock.patch("torch.backends.mps.is_available", return_value=False): + assert MPSAccelerator.device_name() == "" + + def test_warning_if_mps_not_used(mps_count_1): with pytest.warns(UserWarning, match="GPU available but not used"): Trainer(accelerator="cpu") @@ -140,3 +153,21 @@ def to(self, *args, **kwargs): batch = trainer.strategy.batch_to_device(CustomBatchType(), torch.device("mps")) assert batch.a.type() == "torch.mps.FloatTensor" + + +@mock.patch("lightning.pytorch.accelerators.mps.subprocess.run") +def test_get_mps_device_name(mock_run): + mock_stdout = mock.MagicMock() + mock_stdout.configure_mock(stdout="Apple M1 Pro\n") + + mock_run.return_value = mock_stdout + device_name = _get_mps_device_name() + assert device_name == "Apple M1 Pro" + + +@mock.patch( + "lightning.pytorch.accelerators.mps.subprocess.run", + side_effect=SubprocessError("test"), +) +def test_get_mps_device_name_exception(mock_run): + assert _get_mps_device_name() == "True (mps)" diff --git a/tests/tests_pytorch/accelerators/test_xla.py b/tests/tests_pytorch/accelerators/test_xla.py index 99364dd719c26..371b7b3a87e26 100644 --- a/tests/tests_pytorch/accelerators/test_xla.py +++ b/tests/tests_pytorch/accelerators/test_xla.py @@ -300,6 +300,42 @@ def test_warning_if_tpus_not_used(tpu_available): Trainer(accelerator="cpu") +@RunIf(tpu=True) +def test_tpu_device_name(): + from lightning.fabric.accelerators.xla import _XLA_GREATER_EQUAL_2_1 + + if _XLA_GREATER_EQUAL_2_1: + from torch_xla._internal import tpu + else: + from torch_xla.experimental import tpu + import torch_xla.core.xla_env_vars as xenv + + assert XLAAccelerator.device_name() == tpu.get_tpu_env()[xenv.ACCELERATOR_TYPE] + + +def test_tpu_device_name_exception(tpu_available, monkeypatch): + from requests.exceptions import HTTPError + + monkeypatch.delattr( + XLAAccelerator, + "device_name", + raising=False, + ) + + mock.patch( + "torch_xla._internal.tpu", + "get_tpu_env", + side_effect=HTTPError("Could not fetch TPU device name"), + ) + mock.patch( + "torch_xla.experimental.tpu", + "get_tpu_env", + side_effect=HTTPError("Could not fetch TPU device name"), + ) + + assert XLAAccelerator.device_name() == "True" + + @pytest.mark.parametrize( ("devices", "expected_device_ids"), [ diff --git a/tests/tests_pytorch/conftest.py b/tests/tests_pytorch/conftest.py index da48878c7f670..7ae17956571df 100644 --- a/tests/tests_pytorch/conftest.py +++ b/tests/tests_pytorch/conftest.py @@ -183,6 +183,7 @@ def thread_police_duuu_daaa_duuu_daaa(): def mock_cuda_count(monkeypatch, n: int) -> None: monkeypatch.setattr(lightning.fabric.accelerators.cuda, "num_cuda_devices", lambda: n) monkeypatch.setattr(lightning.pytorch.accelerators.cuda, "num_cuda_devices", lambda: n) + monkeypatch.setattr(torch.cuda, "get_device_name", lambda _: "Mocked CUDA Device") @pytest.fixture @@ -208,6 +209,9 @@ def cuda_count_4(monkeypatch): def mock_mps_count(monkeypatch, n: int) -> None: monkeypatch.setattr(lightning.fabric.accelerators.mps, "_get_all_available_mps_gpus", lambda: [0] if n > 0 else []) monkeypatch.setattr(lightning.fabric.accelerators.mps.MPSAccelerator, "is_available", lambda *_: n > 0) + monkeypatch.setattr( + lightning.pytorch.accelerators.mps, "_get_mps_device_name", lambda: "Mocked MPS Device" if n > 0 else "" + ) @pytest.fixture @@ -245,6 +249,11 @@ def mock_tpu_available(monkeypatch: pytest.MonkeyPatch, value: bool = True) -> N monkeypatch.setitem(sys.modules, "torch_xla", Mock()) monkeypatch.setitem(sys.modules, "torch_xla.core.xla_model", Mock()) monkeypatch.setitem(sys.modules, "torch_xla.experimental", Mock()) + monkeypatch.setattr( + lightning.pytorch.accelerators.xla.XLAAccelerator, + "device_name", + lambda *_: "Mocked TPU Device", + ) @pytest.fixture diff --git a/tests/tests_pytorch/plugins/test_cluster_integration.py b/tests/tests_pytorch/plugins/test_cluster_integration.py index 08bd1707b5cfd..f74b771199aaa 100644 --- a/tests/tests_pytorch/plugins/test_cluster_integration.py +++ b/tests/tests_pytorch/plugins/test_cluster_integration.py @@ -66,7 +66,7 @@ def test_ranks_available_manual_strategy_selection(_, strategy_cls): """Test that the rank information is readily available after Trainer initialization.""" num_nodes = 2 for cluster, variables, expected in environment_combinations(): - with mock.patch.dict(os.environ, variables): + with mock.patch.dict(os.environ, variables), mock.patch("torch.cuda.get_device_name", return_value="GPU"): strategy = strategy_cls( parallel_devices=[torch.device("cuda", 1), torch.device("cuda", 2)], cluster_environment=cluster )