iGavroche
a098c72d77
- Added optional in-memory cache to ROCmOptimizedCheckpointLoader to avoid redundant disk loads during repeated API runs. - Implemented cache eviction strategy to manage memory efficiently on gfx1151 architecture. - Updated INPUT_TYPES to include parameters for cache usage and forced reload. - Added integration and unit tests to verify cache behavior and functionality. - Documented caching mechanism in CHECKPOINT_CACHE_API.md and updated LOADER_NODES_GUIDE.md for clarity.
152 lines
6.6 KiB
Python
152 lines
6.6 KiB
Python
"""
|
|
Unit tests for ROCm Checkpoint Loader cache.
|
|
|
|
Tests cache hit/miss, force_reload bypass, use_cache behavior, and single-slot eviction.
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import Mock, patch
|
|
|
|
from rocm_nodes.core import checkpoint as checkpoint_module
|
|
from rocm_nodes.core.checkpoint import ROCMOptimizedCheckpointLoader
|
|
|
|
|
|
def _clear_checkpoint_cache():
|
|
"""Reset module-level cache so tests are isolated."""
|
|
checkpoint_module._checkpoint_cache.clear()
|
|
checkpoint_module._checkpoint_cache_key = None
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_cache():
|
|
"""Clear checkpoint cache before and after each test."""
|
|
_clear_checkpoint_cache()
|
|
yield
|
|
_clear_checkpoint_cache()
|
|
|
|
|
|
@pytest.fixture
|
|
def loader():
|
|
"""Create loader instance."""
|
|
return ROCMOptimizedCheckpointLoader()
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_model_clip_vae():
|
|
"""Fake (model, clip, vae) tuple returned by load_checkpoint_guess_config."""
|
|
return (Mock(name="model"), Mock(name="clip"), Mock(name="vae"))
|
|
|
|
|
|
class TestCheckpointLoaderInputTypes:
|
|
"""Test INPUT_TYPES and optional cache parameters."""
|
|
|
|
def test_required_ckpt_name(self, loader):
|
|
input_types = loader.INPUT_TYPES()
|
|
assert "ckpt_name" in input_types["required"]
|
|
|
|
def test_optional_use_cache_and_force_reload(self, loader):
|
|
input_types = loader.INPUT_TYPES()
|
|
optional = input_types["optional"]
|
|
assert "use_cache" in optional
|
|
assert optional["use_cache"][1]["default"] is True
|
|
assert "force_reload" in optional
|
|
assert optional["force_reload"][1]["default"] is False
|
|
|
|
|
|
class TestCheckpointLoaderCache:
|
|
"""Test cache hit, miss, force_reload, and eviction."""
|
|
|
|
@patch("rocm_nodes.core.checkpoint.gentle_memory_cleanup")
|
|
@patch("rocm_nodes.core.checkpoint.folder_paths.get_full_path_or_raise")
|
|
@patch("rocm_nodes.core.checkpoint.comfy.sd.load_checkpoint_guess_config")
|
|
def test_cache_miss_then_hit(
|
|
self, mock_load, mock_get_path, mock_cleanup, loader, mock_model_clip_vae
|
|
):
|
|
"""First load is cache miss (load called); second load with same ckpt is cache hit (load not called)."""
|
|
mock_get_path.return_value = "/fake/path/ckpt.safetensors"
|
|
mock_load.return_value = mock_model_clip_vae
|
|
|
|
out1 = loader.load_checkpoint("ckpt.safetensors", use_cache=True, force_reload=False)
|
|
assert out1 == mock_model_clip_vae
|
|
assert mock_load.call_count == 1
|
|
|
|
out2 = loader.load_checkpoint("ckpt.safetensors", use_cache=True, force_reload=False)
|
|
assert out2 == mock_model_clip_vae
|
|
assert mock_load.call_count == 1 # still 1: cache hit
|
|
assert out1[0] is out2[0] and out1[1] is out2[1] and out1[2] is out2[2] # same refs
|
|
|
|
@patch("rocm_nodes.core.checkpoint.gentle_memory_cleanup")
|
|
@patch("rocm_nodes.core.checkpoint.folder_paths.get_full_path_or_raise")
|
|
@patch("rocm_nodes.core.checkpoint.comfy.sd.load_checkpoint_guess_config")
|
|
def test_force_reload_bypasses_cache(
|
|
self, mock_load, mock_get_path, mock_cleanup, loader, mock_model_clip_vae
|
|
):
|
|
"""force_reload=True causes load from disk even when cache has the checkpoint."""
|
|
mock_get_path.return_value = "/fake/path/ckpt.safetensors"
|
|
mock_load.return_value = mock_model_clip_vae
|
|
|
|
loader.load_checkpoint("ckpt.safetensors", use_cache=True, force_reload=False)
|
|
assert mock_load.call_count == 1
|
|
|
|
loader.load_checkpoint("ckpt.safetensors", use_cache=True, force_reload=True)
|
|
assert mock_load.call_count == 2
|
|
|
|
@patch("rocm_nodes.core.checkpoint.gentle_memory_cleanup")
|
|
@patch("rocm_nodes.core.checkpoint.folder_paths.get_full_path_or_raise")
|
|
@patch("rocm_nodes.core.checkpoint.comfy.sd.load_checkpoint_guess_config")
|
|
def test_use_cache_false_still_loads_and_caches(
|
|
self, mock_load, mock_get_path, mock_cleanup, loader, mock_model_clip_vae
|
|
):
|
|
"""use_cache=False still loads and stores in cache; next call with use_cache=True can hit."""
|
|
mock_get_path.return_value = "/fake/path/ckpt.safetensors"
|
|
mock_load.return_value = mock_model_clip_vae
|
|
|
|
out1 = loader.load_checkpoint("ckpt.safetensors", use_cache=False, force_reload=False)
|
|
assert out1 == mock_model_clip_vae
|
|
assert mock_load.call_count == 1
|
|
|
|
# Cache was populated; with use_cache=True we get cache hit
|
|
out2 = loader.load_checkpoint("ckpt.safetensors", use_cache=True, force_reload=False)
|
|
assert out2 == mock_model_clip_vae
|
|
assert mock_load.call_count == 1
|
|
assert out1[0] is out2[0] and out1[1] is out2[1] and out1[2] is out2[2]
|
|
|
|
@patch("rocm_nodes.core.checkpoint.gentle_memory_cleanup")
|
|
@patch("rocm_nodes.core.checkpoint.folder_paths.get_full_path_or_raise")
|
|
@patch("rocm_nodes.core.checkpoint.comfy.sd.load_checkpoint_guess_config")
|
|
def test_eviction_when_loading_different_ckpt(
|
|
self, mock_load, mock_get_path, mock_cleanup, loader, mock_model_clip_vae
|
|
):
|
|
"""Loading a different ckpt_name evicts previous cache and calls gentle_memory_cleanup."""
|
|
mock_get_path.side_effect = lambda folder, name: f"/fake/{name}"
|
|
mock_load.return_value = mock_model_clip_vae
|
|
|
|
loader.load_checkpoint("ckpt_a.safetensors", use_cache=True, force_reload=False)
|
|
assert mock_load.call_count == 1
|
|
assert checkpoint_module._checkpoint_cache_key == "ckpt_a.safetensors"
|
|
assert "ckpt_a.safetensors" in checkpoint_module._checkpoint_cache
|
|
|
|
# Different checkpoint: evict and load
|
|
other_triple = (Mock(name="m2"), Mock(name="c2"), Mock(name="v2"))
|
|
mock_load.return_value = other_triple
|
|
loader.load_checkpoint("ckpt_b.safetensors", use_cache=True, force_reload=False)
|
|
|
|
assert mock_load.call_count == 2
|
|
assert mock_cleanup.call_count >= 1
|
|
assert checkpoint_module._checkpoint_cache_key == "ckpt_b.safetensors"
|
|
assert "ckpt_a.safetensors" not in checkpoint_module._checkpoint_cache
|
|
assert "ckpt_b.safetensors" in checkpoint_module._checkpoint_cache
|
|
|
|
@patch("rocm_nodes.core.checkpoint.gentle_memory_cleanup")
|
|
@patch("rocm_nodes.core.checkpoint.folder_paths.get_full_path_or_raise")
|
|
@patch("rocm_nodes.core.checkpoint.comfy.sd.load_checkpoint_guess_config")
|
|
def test_return_types_and_names(self, mock_load, mock_get_path, mock_cleanup, loader):
|
|
"""Loader returns MODEL, CLIP, VAE."""
|
|
mock_get_path.return_value = "/fake/path/ckpt.safetensors"
|
|
mock_load.return_value = (Mock(), Mock(), Mock())
|
|
|
|
assert loader.RETURN_TYPES == ("MODEL", "CLIP", "VAE")
|
|
assert loader.RETURN_NAMES == ("MODEL", "CLIP", "VAE")
|
|
result = loader.load_checkpoint("ckpt.safetensors")
|
|
assert len(result) == 3
|