github/InvokeAI
1
1
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI.git synced 2026-04-23 03:00:31 -04:00
Code Issues Packages Projects Releases Wiki Activity

tests: update tests for configure_torch_cuda_allocator

Browse Source
This commit is contained in:
psychedelicious
2025-03-05 08:55:57 +10:00
parent d5401e8443
commit d037d8f9aa
1 changed files with 21 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
tests/app/util/test_torch_cuda_allocator.py
View File

@@ -1,13 +1,31 @@
from unittest.mock import MagicMock
import pytest
import torch
from invokeai.app.util.torch_cuda_allocator import configure_torch_cuda_allocator
from tests.env_var_utils import set_env_var, unset_env_var
@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA device.")
def test_configure_torch_cuda_allocator_raises_if_torch_is_already_imported():
"""Test that configure_torch_cuda_allocator() raises a RuntimeError if torch is already imported."""
import torch # noqa: F401
with pytest.raises(RuntimeError, match="Failed to configure the PyTorch CUDA memory allocator."):
configure_torch_cuda_allocator("backend:cudaMallocAsync")
with unset_env_var("PYTORCH_CUDA_ALLOC_CONF"):
import torch # noqa: F401
mock_logger = MagicMock()
with pytest.raises(RuntimeError, match="Failed to configure the PyTorch CUDA memory allocator."):
configure_torch_cuda_allocator("backend:cudaMallocAsync", logger=mock_logger)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA device.")
def test_configure_torch_cuda_allocator_warns_if_env_var_is_already_set():
"""Test that configure_torch_cuda_allocator() logs a warning if PYTORCH_CUDA_ALLOC_CONF is already set."""
with set_env_var("PYTORCH_CUDA_ALLOC_CONF", "backend:native"):
mock_logger = MagicMock()
configure_torch_cuda_allocator("backend:cudaMallocAsync", logger=mock_logger)
mock_logger.warning.assert_called_once()
args, _kwargs = mock_logger.warning.call_args
assert "PYTORCH_CUDA_ALLOC_CONF is already set" in args[0]