mirror of
https://github.com/invoke-ai/InvokeAI.git
synced 2026-01-14 06:08:00 -05:00
128 lines
4.9 KiB
Python
128 lines
4.9 KiB
Python
import pytest
|
|
import torch
|
|
|
|
from tests.dangerously_run_function_in_subprocess import dangerously_run_function_in_subprocess
|
|
|
|
# These tests are a bit fiddly, because the depend on the import behaviour of torch. They use subprocesses to isolate
|
|
# the import behaviour of torch, and then check that the function behaves as expected. We have to hack in some logging
|
|
# to check that the tested function is behaving as expected.
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA device.")
|
|
def test_configure_torch_cuda_allocator_configures_backend():
|
|
"""Test that configure_torch_cuda_allocator() raises a RuntimeError if the configured backend does not match the
|
|
expected backend."""
|
|
|
|
def test_func():
|
|
import os
|
|
|
|
# Unset the environment variable if it is set so that we can test setting it
|
|
try:
|
|
del os.environ["PYTORCH_CUDA_ALLOC_CONF"]
|
|
except KeyError:
|
|
pass
|
|
|
|
from unittest.mock import MagicMock
|
|
|
|
from invokeai.app.util.torch_cuda_allocator import configure_torch_cuda_allocator
|
|
|
|
mock_logger = MagicMock()
|
|
|
|
# Set the PyTorch CUDA memory allocator to cudaMallocAsync
|
|
configure_torch_cuda_allocator("backend:cudaMallocAsync", logger=mock_logger)
|
|
|
|
# Verify that the PyTorch CUDA memory allocator was configured correctly
|
|
import torch
|
|
|
|
assert torch.cuda.get_allocator_backend() == "cudaMallocAsync"
|
|
|
|
# Verify that the logger was called with the correct message
|
|
mock_logger.info.assert_called_once()
|
|
args, _kwargs = mock_logger.info.call_args
|
|
logged_message = args[0]
|
|
print(logged_message)
|
|
|
|
stdout, _stderr, returncode = dangerously_run_function_in_subprocess(test_func)
|
|
assert returncode == 0
|
|
assert "PyTorch CUDA memory allocator: cudaMallocAsync" in stdout
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA device.")
|
|
def test_configure_torch_cuda_allocator_raises_if_torch_already_imported():
|
|
"""Test that configure_torch_cuda_allocator() raises a RuntimeError if torch was already imported."""
|
|
|
|
def test_func():
|
|
from unittest.mock import MagicMock
|
|
|
|
# Import torch before calling configure_torch_cuda_allocator()
|
|
import torch # noqa: F401
|
|
|
|
from invokeai.app.util.torch_cuda_allocator import configure_torch_cuda_allocator
|
|
|
|
try:
|
|
configure_torch_cuda_allocator("backend:cudaMallocAsync", logger=MagicMock())
|
|
except RuntimeError as e:
|
|
print(e)
|
|
|
|
stdout, _stderr, returncode = dangerously_run_function_in_subprocess(test_func)
|
|
assert returncode == 0
|
|
assert "configure_torch_cuda_allocator() must be called before importing torch." in stdout
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA device.")
|
|
def test_configure_torch_cuda_allocator_warns_if_env_var_is_set_differently():
|
|
"""Test that configure_torch_cuda_allocator() logs at WARNING level if PYTORCH_CUDA_ALLOC_CONF is set and doesn't
|
|
match the requested configuration."""
|
|
|
|
def test_func():
|
|
import os
|
|
|
|
# Explicitly set the environment variable
|
|
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "backend:native"
|
|
|
|
from unittest.mock import MagicMock
|
|
|
|
from invokeai.app.util.torch_cuda_allocator import configure_torch_cuda_allocator
|
|
|
|
mock_logger = MagicMock()
|
|
|
|
# Set the PyTorch CUDA memory allocator a different configuration
|
|
configure_torch_cuda_allocator("backend:cudaMallocAsync", logger=mock_logger)
|
|
|
|
# Verify that the logger was called with the correct message
|
|
mock_logger.warning.assert_called_once()
|
|
args, _kwargs = mock_logger.warning.call_args
|
|
logged_message = args[0]
|
|
print(logged_message)
|
|
|
|
stdout, _stderr, returncode = dangerously_run_function_in_subprocess(test_func)
|
|
assert returncode == 0
|
|
assert "Attempted to configure the PyTorch CUDA memory allocator with 'backend:cudaMallocAsync'" in stdout
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA device.")
|
|
def test_configure_torch_cuda_allocator_logs_if_env_var_is_already_set_correctly():
|
|
"""Test that configure_torch_cuda_allocator() logs at INFO level if PYTORCH_CUDA_ALLOC_CONF is set and matches the
|
|
requested configuration."""
|
|
|
|
def test_func():
|
|
import os
|
|
|
|
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "backend:native"
|
|
from unittest.mock import MagicMock
|
|
|
|
from invokeai.app.util.torch_cuda_allocator import configure_torch_cuda_allocator
|
|
|
|
mock_logger = MagicMock()
|
|
|
|
configure_torch_cuda_allocator("backend:native", logger=mock_logger)
|
|
|
|
mock_logger.info.assert_called_once()
|
|
args, _kwargs = mock_logger.info.call_args
|
|
logged_message = args[0]
|
|
print(logged_message)
|
|
|
|
stdout, _stderr, returncode = dangerously_run_function_in_subprocess(test_func)
|
|
assert returncode == 0
|
|
assert "PYTORCH_CUDA_ALLOC_CONF is already set to 'backend:native'" in stdout
|