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.
79 lines
2.7 KiB
Python
79 lines
2.7 KiB
Python
"""
|
|
Integration tests for ROCm Checkpoint Loader cache.
|
|
|
|
Verifies that when the same checkpoint is requested twice (e.g. two API runs),
|
|
the second run does not reload from disk (load_checkpoint_guess_config called once).
|
|
"""
|
|
|
|
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():
|
|
checkpoint_module._checkpoint_cache.clear()
|
|
checkpoint_module._checkpoint_cache_key = None
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_cache():
|
|
_clear_checkpoint_cache()
|
|
yield
|
|
_clear_checkpoint_cache()
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_model_clip_vae():
|
|
return (Mock(name="model"), Mock(name="clip"), Mock(name="vae"))
|
|
|
|
|
|
@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_two_runs_same_checkpoint_only_loads_once(
|
|
mock_load, mock_get_path, mock_cleanup, mock_model_clip_vae
|
|
):
|
|
"""
|
|
Simulate two API runs with the same checkpoint: load_checkpoint_guess_config
|
|
must be called only once (second run uses cache).
|
|
"""
|
|
mock_get_path.return_value = "/fake/checkpoints/ltx-2-19b-dev-fp8.safetensors"
|
|
mock_load.return_value = mock_model_clip_vae
|
|
|
|
loader = ROCMOptimizedCheckpointLoader()
|
|
|
|
# First "API run"
|
|
out1 = loader.load_checkpoint(
|
|
"ltx-2-19b-dev-fp8.safetensors", use_cache=True, force_reload=False
|
|
)
|
|
assert len(out1) == 3
|
|
assert mock_load.call_count == 1
|
|
|
|
# Second "API run" with same workflow (same checkpoint)
|
|
out2 = loader.load_checkpoint(
|
|
"ltx-2-19b-dev-fp8.safetensors", use_cache=True, force_reload=False
|
|
)
|
|
assert len(out2) == 3
|
|
assert mock_load.call_count == 1 # no second load
|
|
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_force_reload_on_second_run_calls_load_again(
|
|
mock_load, mock_get_path, mock_cleanup, mock_model_clip_vae
|
|
):
|
|
"""Second run with force_reload=True triggers a reload (e.g. user requested refresh)."""
|
|
mock_get_path.return_value = "/fake/checkpoints/ckpt.safetensors"
|
|
mock_load.return_value = mock_model_clip_vae
|
|
|
|
loader = ROCMOptimizedCheckpointLoader()
|
|
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
|