AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-24 20:59:27 +03:00
Code Activity
Files
diffusers-serve
diffusers/tests/modular_pipelines/test_components_manager.py
Sayak Paul 102c24934f split up tests in modular pipelines (#14444) 
* notes from .ai

* ltx2

* cover the remaining tests

* fixes
2026-08-14 11:54:01 +02:00

298 lines
12 KiB
Python
Raw Permalink Blame History

# Copyright 2025 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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 gc
from unittest import mock
import torch
from diffusers import ComponentsManager
from diffusers.models import ModelMixin
from diffusers.utils import is_accelerate_available
from ..testing_utils import backend_empty_cache, require_accelerate, require_accelerator, torch_device
from .testing_utils import patch_free_memory
if is_accelerate_available():
from diffusers.modular_pipelines.components_manager import AutoOffloadStrategy
# The offload logic deals in bytes. We keep the test models tiny (a few KB of real
# parameters) and express every size as a multiple of this unit, then *simulate* the
# device's free memory at the same scale. This is what lets us exercise the offloading
# decisions deterministically instead of relying on the real free memory of the test
# hardware (an 80GB GPU never runs low on a handful of KB-sized models).
UNIT = 1024
class DummyModel(ModelMixin):
def __init__(self, footprint_bytes: int = UNIT):
super().__init__()
# A float32 parameter of `footprint_bytes // 4` elements weighs exactly
# `footprint_bytes`, so callers control the reported size directly.
self.weight = torch.nn.Parameter(torch.zeros(footprint_bytes // 4))
def forward(self, x):
return x + self.weight.sum()
class _FakeHook:
"""
Minimal stand-in for `UserCustomOffloadHook` in strategy-level unit tests.
`AutoOffloadStrategy` only reads `hook.model_id` and
`hook.model.get_memory_footprint()`, so we avoid attaching real accelerate hooks
(which would move modules around) and keep the logic test pure.
"""
def __init__(self, model_id: str, model: torch.nn.Module):
self.model_id = model_id
self.model = model
def _patch_cuda_mem_get_info(free_bytes: int, total_bytes: int = 80 * UNIT):
# Strategy unit tests use a `cuda:0` execution-device *descriptor* (which needs no
# real GPU), so they patch `torch.cuda.mem_get_info` directly.
return mock.patch.object(torch.cuda, "mem_get_info", return_value=(free_bytes, total_bytes))
@require_accelerate
class TestComponentsManager:
"""
Tests for `ComponentsManager` and its auto-offload strategy.
"""
# A `cuda:0` device descriptor is enough to drive the strategy's device-type and
# index logic; no real GPU is required because `mem_get_info` is mocked.
strategy_execution_device = torch.device("cuda:0")
def setup_method(self):
# Free VRAM before/after each test; the auto-offload integration tests below reason
# about device residency, so they must not inherit another test's allocations.
torch.compiler.reset()
gc.collect()
backend_empty_cache(torch_device)
def teardown_method(self):
torch.compiler.reset()
gc.collect()
backend_empty_cache(torch_device)
def get_dummy_model(self, footprint_bytes: int = UNIT) -> ModelMixin:
return DummyModel(footprint_bytes=footprint_bytes)
# ------------------------------------------------------------------
# AutoOffloadStrategy unit tests (hardware-independent)
# ------------------------------------------------------------------
def _select_offload(self, *, incoming_footprint, free_bytes, hook_sizes, memory_reserve_margin=UNIT):
strategy = AutoOffloadStrategy(memory_reserve_margin=memory_reserve_margin)
hooks = [_FakeHook(model_id, self.get_dummy_model(fp)) for model_id, fp in hook_sizes.items()]
incoming = self.get_dummy_model(incoming_footprint)
with _patch_cuda_mem_get_info(free_bytes):
selected = strategy(
hooks=hooks,
model_id="incoming",
model=incoming,
execution_device=self.strategy_execution_device,
)
return sorted(hook.model_id for hook in selected)
def test_strategy_no_offload_when_memory_is_sufficient(self):
# 70 units free, 1 reserved -> 69 usable, incoming needs 4: nothing to offload.
selected = self._select_offload(
incoming_footprint=4 * UNIT,
free_bytes=70 * UNIT,
hook_sizes={"a": 5 * UNIT, "b": 3 * UNIT},
)
assert selected == []
def test_strategy_offloads_minimal_single_model(self):
# usable = 4 - 1 = 3, incoming needs 6 -> must free 3.
# Smallest combination that frees >= 3 is "b" (exactly 3) on its own.
selected = self._select_offload(
incoming_footprint=6 * UNIT,
free_bytes=4 * UNIT,
hook_sizes={"a": 5 * UNIT, "b": 3 * UNIT, "c": 2 * UNIT},
)
assert selected == ["b"]
def test_strategy_offloads_smallest_sufficient_combination(self):
# usable = 4 - 1 = 3, incoming needs 8 -> must free 5.
# No single model frees 5 (max is 4), so the smallest sufficient combination of
# models is chosen: a (4) + c (1) = 5.
selected = self._select_offload(
incoming_footprint=8 * UNIT,
free_bytes=4 * UNIT,
hook_sizes={"a": 4 * UNIT, "b": 4 * UNIT, "c": 1 * UNIT},
)
assert selected == ["a", "c"]
def test_strategy_offloads_all_when_freeing_enough_is_impossible(self):
# incoming needs more than the sum of everything on device -> offload all.
selected = self._select_offload(
incoming_footprint=11 * UNIT,
free_bytes=1 * UNIT,
hook_sizes={"a": 5 * UNIT, "b": 3 * UNIT, "c": 2 * UNIT},
)
assert selected == ["a", "b", "c"]
def test_strategy_no_hooks_returns_empty(self):
selected = self._select_offload(
incoming_footprint=11 * UNIT,
free_bytes=0,
hook_sizes={},
)
assert selected == []
def test_strategy_memory_reserve_margin_changes_decision(self):
# Same device free memory and incoming model; only the reserve margin differs.
# A small margin leaves enough room; a large margin forces an offload. We check
# this both with a single resident model and with several, to confirm the margin
# participates in the selection regardless of how many candidates exist.
# Single candidate: free=5, incoming=3. margin 1 -> usable 4 (fits); margin 3 ->
# usable 2, must free 1 -> offload "a".
assert (
self._select_offload(
incoming_footprint=3 * UNIT,
free_bytes=5 * UNIT,
hook_sizes={"a": 2 * UNIT},
memory_reserve_margin=1 * UNIT,
)
== []
)
assert self._select_offload(
incoming_footprint=3 * UNIT,
free_bytes=5 * UNIT,
hook_sizes={"a": 2 * UNIT},
memory_reserve_margin=3 * UNIT,
) == ["a"]
# Multiple candidates: free=6, incoming=4. margin 1 -> usable 5 (fits); margin 3
# -> usable 3, must free 1 -> smallest sufficient model "c" (1) is offloaded.
multi_hooks = {"a": 3 * UNIT, "b": 2 * UNIT, "c": 1 * UNIT}
assert (
self._select_offload(
incoming_footprint=4 * UNIT,
free_bytes=6 * UNIT,
hook_sizes=multi_hooks,
memory_reserve_margin=1 * UNIT,
)
== []
)
assert self._select_offload(
incoming_footprint=4 * UNIT,
free_bytes=6 * UNIT,
hook_sizes=multi_hooks,
memory_reserve_margin=3 * UNIT,
) == ["c"]
# ------------------------------------------------------------------
# Registry tests (hardware-independent)
# ------------------------------------------------------------------
def test_add_and_get_one(self):
cm = ComponentsManager()
model = self.get_dummy_model()
component_id = cm.add("unet", model)
assert component_id in cm.components
assert cm.get_one(name="unet") is model
assert cm.get_one(component_id=component_id) is model
def test_add_same_component_twice_reuses_id(self):
cm = ComponentsManager()
model = self.get_dummy_model()
first_id = cm.add("unet", model)
second_id = cm.add("unet", model)
assert first_id == second_id
assert len(cm.components) == 1
def test_remove(self):
cm = ComponentsManager()
component_id = cm.add("unet", self.get_dummy_model())
cm.remove(component_id)
assert component_id not in cm.components
def test_get_model_info_reports_size(self):
cm = ComponentsManager()
model = self.get_dummy_model(footprint_bytes=2 * UNIT)
component_id = cm.add("unet", model)
info = cm.get_model_info(component_id, fields="size_gb")
assert info["size_gb"] == model.get_memory_footprint() / (1024**3)
# ------------------------------------------------------------------
# Auto-offload integration tests (require an accelerator)
# ------------------------------------------------------------------
@require_accelerator
def test_auto_offload_starts_with_all_components_on_cpu(self):
cm = ComponentsManager()
model = self.get_dummy_model(4 * UNIT)
cm.add("m1", model)
cm.enable_auto_cpu_offload(device=torch_device, memory_reserve_margin=UNIT)
try:
assert next(model.parameters()).device.type == "cpu"
finally:
cm.disable_auto_cpu_offload()
@require_accelerator
def test_auto_offload_evicts_resident_model_under_memory_pressure(self):
device_type = torch.device(torch_device).type
cm = ComponentsManager()
m1 = self.get_dummy_model(4 * UNIT)
m2 = self.get_dummy_model(4 * UNIT)
cm.add("m1", m1)
cm.add("m2", m2)
cm.enable_auto_cpu_offload(device=torch_device, memory_reserve_margin=UNIT)
try:
# Both components start offloaded on the CPU.
assert next(m1.parameters()).device.type == "cpu"
assert next(m2.parameters()).device.type == "cpu"
x = torch.randn(2, 4, device=torch_device)
# Ample free memory: running m1 just moves it onto the device, evicting
# nothing (m2 is not resident, so it is not even a candidate).
with patch_free_memory(70 * UNIT):
m1(x)
assert next(m1.parameters()).device.type == device_type
# Memory pressure: usable = 4 - 1 = 3 but m2 needs 4, so the only resident
# model (m1) must be evicted back to the CPU to make room for m2.
with patch_free_memory(4 * UNIT):
m2(x)
assert next(m2.parameters()).device.type == device_type
assert next(m1.parameters()).device.type == "cpu"
finally:
cm.disable_auto_cpu_offload()
@require_accelerator
def test_auto_offload_keeps_models_resident_when_memory_is_ample(self):
device_type = torch.device(torch_device).type
cm = ComponentsManager()
m1 = self.get_dummy_model(4 * UNIT)
m2 = self.get_dummy_model(4 * UNIT)
cm.add("m1", m1)
cm.add("m2", m2)
cm.enable_auto_cpu_offload(device=torch_device, memory_reserve_margin=UNIT)
try:
x = torch.randn(2, 4, device=torch_device)
with patch_free_memory(70 * UNIT):
m1(x)
m2(x)
# Both fit comfortably, so neither gets evicted.
assert next(m1.parameters()).device.type == device_type
assert next(m2.parameters()).device.type == device_type
finally:
cm.disable_auto_cpu_offload()
View Git Blame Copy Permalink