70 lines
2 KiB
Python
70 lines
2 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Unit tests for the single-GPU arbiter.
|
|
|
|
The real evictors (which tear down live backends) are replaced with recorders, so
|
|
these verify only the ownership/eviction sequencing — no torch, GPU, or subprocess.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
import core.inference.gpu_arbiter as arb
|
|
|
|
|
|
@pytest.fixture
|
|
def calls(monkeypatch):
|
|
recorded: list[str] = []
|
|
monkeypatch.setattr(arb, "_owner", None)
|
|
monkeypatch.setitem(arb._EVICTORS, arb.CHAT, lambda: recorded.append("evict-chat"))
|
|
monkeypatch.setitem(arb._EVICTORS, arb.DIFFUSION, lambda: recorded.append("evict-diffusion"))
|
|
return recorded
|
|
|
|
|
|
def test_first_acquire_evicts_nothing(calls):
|
|
arb.acquire_for(arb.CHAT)
|
|
assert calls == []
|
|
assert arb.current_owner() == arb.CHAT
|
|
|
|
|
|
def test_diffusion_load_evicts_chat(calls):
|
|
arb.acquire_for(arb.CHAT)
|
|
arb.acquire_for(arb.DIFFUSION)
|
|
assert calls == ["evict-chat"]
|
|
assert arb.current_owner() == arb.DIFFUSION
|
|
|
|
|
|
def test_chat_load_evicts_diffusion(calls):
|
|
arb.acquire_for(arb.DIFFUSION)
|
|
arb.acquire_for(arb.CHAT)
|
|
assert calls == ["evict-diffusion"]
|
|
assert arb.current_owner() == arb.CHAT
|
|
|
|
|
|
def test_reacquiring_same_owner_does_not_evict(calls):
|
|
arb.acquire_for(arb.CHAT)
|
|
arb.acquire_for(arb.CHAT)
|
|
assert calls == []
|
|
assert arb.current_owner() == arb.CHAT
|
|
|
|
|
|
def test_release_clears_owner(calls):
|
|
arb.acquire_for(arb.DIFFUSION)
|
|
arb.release(arb.DIFFUSION)
|
|
assert arb.current_owner() is None
|
|
# A subsequent chat load then has nothing to evict.
|
|
arb.acquire_for(arb.CHAT)
|
|
assert calls == []
|
|
|
|
|
|
def test_release_by_non_owner_is_noop(calls):
|
|
arb.acquire_for(arb.CHAT)
|
|
arb.release(arb.DIFFUSION)
|
|
assert arb.current_owner() == arb.CHAT
|
|
|
|
|
|
def test_unknown_owner_raises(calls):
|
|
with pytest.raises(ValueError):
|
|
arb.acquire_for("gpu")
|