chore: import upstream snapshot with attribution
docs / deploy (push) Has been cancelled
docs / changes (push) Has been cancelled
docs / check-and-build (push) Has been cancelled
build container image / cpu (push) Has been cancelled
build container image / cuda (push) Has been cancelled
build container image / rocm (push) Has been cancelled
frontend checks / frontend-checks (push) Has been cancelled
frontend tests / frontend-tests (push) Has been cancelled
lfs checks / lfs-check (push) Has been cancelled
python checks / python-checks (push) Has been cancelled
python tests / py3.12: macos-default (push) Has been cancelled
python tests / py3.11: windows-cpu (push) Has been cancelled
python tests / py3.12: windows-cpu (push) Has been cancelled
python tests / py3.11: linux-cpu (push) Has been cancelled
typegen checks / typegen-checks (push) Has been cancelled
uv lock checks / uv-lock-checks (push) Has been cancelled
openapi checks / openapi-checks (push) Has been cancelled
python tests / py3.11: macos-default (push) Has been cancelled
python tests / py3.12: linux-cpu (push) Has been cancelled
docs / deploy (push) Has been cancelled
docs / changes (push) Has been cancelled
docs / check-and-build (push) Has been cancelled
build container image / cpu (push) Has been cancelled
build container image / cuda (push) Has been cancelled
build container image / rocm (push) Has been cancelled
frontend checks / frontend-checks (push) Has been cancelled
frontend tests / frontend-tests (push) Has been cancelled
lfs checks / lfs-check (push) Has been cancelled
python checks / python-checks (push) Has been cancelled
python tests / py3.12: macos-default (push) Has been cancelled
python tests / py3.11: windows-cpu (push) Has been cancelled
python tests / py3.12: windows-cpu (push) Has been cancelled
python tests / py3.11: linux-cpu (push) Has been cancelled
typegen checks / typegen-checks (push) Has been cancelled
uv lock checks / uv-lock-checks (push) Has been cancelled
openapi checks / openapi-checks (push) Has been cancelled
python tests / py3.11: macos-default (push) Has been cancelled
python tests / py3.12: linux-cpu (push) Has been cancelled
This commit is contained in:
@@ -0,0 +1,293 @@
|
||||
"""Tests for session queue dequeue() ordering: FIFO and round-robin modes."""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
from pydantic_core import to_jsonable_python
|
||||
|
||||
from invokeai.app.services.config.config_default import InvokeAIAppConfig
|
||||
from invokeai.app.services.invoker import Invoker
|
||||
from invokeai.app.services.session_queue.session_queue_sqlite import (
|
||||
ROUND_ROBIN_DEQUEUE_QUERY,
|
||||
SqliteSessionQueue,
|
||||
)
|
||||
from invokeai.app.services.shared.graph import Graph, GraphExecutionState
|
||||
|
||||
_EMPTY_SESSION_JSON = json.dumps(to_jsonable_python(GraphExecutionState(graph=Graph()).model_dump()))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def session_queue_fifo(mock_invoker: Invoker) -> SqliteSessionQueue:
|
||||
"""Queue backed by a single-user (FIFO) invoker."""
|
||||
# Default config has multiuser=False, so FIFO is always used.
|
||||
db = mock_invoker.services.board_records._db
|
||||
queue = SqliteSessionQueue(db=db)
|
||||
queue.start(mock_invoker)
|
||||
return queue
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def session_queue_round_robin(mock_invoker: Invoker) -> SqliteSessionQueue:
|
||||
"""Queue backed by a multiuser invoker with round_robin mode."""
|
||||
mock_invoker.services.configuration = InvokeAIAppConfig(
|
||||
use_memory_db=True,
|
||||
node_cache_size=0,
|
||||
multiuser=True,
|
||||
session_queue_mode="round_robin",
|
||||
)
|
||||
db = mock_invoker.services.board_records._db
|
||||
queue = SqliteSessionQueue(db=db)
|
||||
queue.start(mock_invoker)
|
||||
return queue
|
||||
|
||||
|
||||
def _insert_queue_item(
|
||||
session_queue: SqliteSessionQueue,
|
||||
queue_id: str,
|
||||
user_id: str,
|
||||
priority: int = 0,
|
||||
) -> int:
|
||||
"""Directly insert a minimal queue item and return its item_id."""
|
||||
session_id = str(uuid.uuid4())
|
||||
batch_id = str(uuid.uuid4())
|
||||
with session_queue._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
INSERT INTO session_queue (queue_id, session, session_id, batch_id, field_values, priority, workflow, origin, destination, retried_from_item_id, user_id)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(queue_id, _EMPTY_SESSION_JSON, session_id, batch_id, None, priority, None, None, None, None, user_id),
|
||||
)
|
||||
return cursor.lastrowid # type: ignore[return-value]
|
||||
|
||||
|
||||
def _dequeue_user_ids(session_queue: SqliteSessionQueue, count: int) -> list[Optional[str]]:
|
||||
"""Dequeue `count` items and return the list of user_ids in dequeue order."""
|
||||
result = []
|
||||
for _ in range(count):
|
||||
item = session_queue.dequeue()
|
||||
result.append(item.user_id if item is not None else None)
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# FIFO tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_fifo_single_user_order(session_queue_fifo: SqliteSessionQueue) -> None:
|
||||
"""FIFO: items from a single user are dequeued in insertion order."""
|
||||
queue_id = "default"
|
||||
_insert_queue_item(session_queue_fifo, queue_id, "user_a")
|
||||
_insert_queue_item(session_queue_fifo, queue_id, "user_a")
|
||||
_insert_queue_item(session_queue_fifo, queue_id, "user_a")
|
||||
|
||||
user_ids = _dequeue_user_ids(session_queue_fifo, 3)
|
||||
assert user_ids == ["user_a", "user_a", "user_a"]
|
||||
|
||||
|
||||
def test_fifo_multi_user_preserves_insertion_order(session_queue_fifo: SqliteSessionQueue) -> None:
|
||||
"""FIFO: jobs from multiple users are dequeued in strict insertion order, not interleaved."""
|
||||
queue_id = "default"
|
||||
# Insert A1, A2, B1, C1, C2, A3 – FIFO should preserve this exact order.
|
||||
_insert_queue_item(session_queue_fifo, queue_id, "user_a")
|
||||
_insert_queue_item(session_queue_fifo, queue_id, "user_a")
|
||||
_insert_queue_item(session_queue_fifo, queue_id, "user_b")
|
||||
_insert_queue_item(session_queue_fifo, queue_id, "user_c")
|
||||
_insert_queue_item(session_queue_fifo, queue_id, "user_c")
|
||||
_insert_queue_item(session_queue_fifo, queue_id, "user_a")
|
||||
|
||||
user_ids = _dequeue_user_ids(session_queue_fifo, 6)
|
||||
assert user_ids == ["user_a", "user_a", "user_b", "user_c", "user_c", "user_a"]
|
||||
|
||||
|
||||
def test_fifo_priority_respected(session_queue_fifo: SqliteSessionQueue) -> None:
|
||||
"""FIFO: higher-priority items are dequeued before lower-priority ones."""
|
||||
queue_id = "default"
|
||||
_insert_queue_item(session_queue_fifo, queue_id, "user_a", priority=0)
|
||||
_insert_queue_item(session_queue_fifo, queue_id, "user_a", priority=10)
|
||||
|
||||
user_ids = _dequeue_user_ids(session_queue_fifo, 2)
|
||||
# Both are user_a; second inserted item has higher priority and should come first.
|
||||
assert user_ids == ["user_a", "user_a"]
|
||||
|
||||
|
||||
def test_fifo_returns_none_when_empty(session_queue_fifo: SqliteSessionQueue) -> None:
|
||||
"""FIFO: dequeue returns None when the queue is empty."""
|
||||
assert session_queue_fifo.dequeue() is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Round-robin tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_round_robin_interleaves_users(session_queue_round_robin: SqliteSessionQueue) -> None:
|
||||
"""Round-robin: jobs from multiple users are interleaved one per user per round.
|
||||
|
||||
Queue insertion order (matching the issue example):
|
||||
A job 1, A job 2, B job 1, C job 1, C job 2, A job 3
|
||||
|
||||
Expected dequeue order:
|
||||
A job 1, B job 1, C job 1, A job 2, C job 2, A job 3
|
||||
"""
|
||||
queue_id = "default"
|
||||
_insert_queue_item(session_queue_round_robin, queue_id, "user_a")
|
||||
_insert_queue_item(session_queue_round_robin, queue_id, "user_a")
|
||||
_insert_queue_item(session_queue_round_robin, queue_id, "user_b")
|
||||
_insert_queue_item(session_queue_round_robin, queue_id, "user_c")
|
||||
_insert_queue_item(session_queue_round_robin, queue_id, "user_c")
|
||||
_insert_queue_item(session_queue_round_robin, queue_id, "user_a")
|
||||
|
||||
user_ids = _dequeue_user_ids(session_queue_round_robin, 6)
|
||||
assert user_ids == ["user_a", "user_b", "user_c", "user_a", "user_c", "user_a"]
|
||||
|
||||
|
||||
def test_round_robin_single_user_behaves_like_fifo(session_queue_round_robin: SqliteSessionQueue) -> None:
|
||||
"""Round-robin with only one user produces the same order as FIFO."""
|
||||
queue_id = "default"
|
||||
_insert_queue_item(session_queue_round_robin, queue_id, "user_a")
|
||||
_insert_queue_item(session_queue_round_robin, queue_id, "user_a")
|
||||
_insert_queue_item(session_queue_round_robin, queue_id, "user_a")
|
||||
|
||||
user_ids = _dequeue_user_ids(session_queue_round_robin, 3)
|
||||
assert user_ids == ["user_a", "user_a", "user_a"]
|
||||
|
||||
|
||||
def test_round_robin_handles_user_joining_mid_queue(session_queue_round_robin: SqliteSessionQueue) -> None:
|
||||
"""Round-robin: a user who joins later is correctly interleaved."""
|
||||
queue_id = "default"
|
||||
_insert_queue_item(session_queue_round_robin, queue_id, "user_a")
|
||||
_insert_queue_item(session_queue_round_robin, queue_id, "user_a")
|
||||
_insert_queue_item(session_queue_round_robin, queue_id, "user_b")
|
||||
|
||||
user_ids = _dequeue_user_ids(session_queue_round_robin, 3)
|
||||
# Round 1: A (oldest rank-1 item), B (rank-1 item)
|
||||
# Round 2: A (rank-2 item)
|
||||
assert user_ids == ["user_a", "user_b", "user_a"]
|
||||
|
||||
|
||||
def test_round_robin_returns_none_when_empty(session_queue_round_robin: SqliteSessionQueue) -> None:
|
||||
"""Round-robin: dequeue returns None when the queue is empty."""
|
||||
assert session_queue_round_robin.dequeue() is None
|
||||
|
||||
|
||||
def test_round_robin_priority_within_user_respected(session_queue_round_robin: SqliteSessionQueue) -> None:
|
||||
"""Round-robin: within a single user's items, higher priority is dequeued first."""
|
||||
queue_id = "default"
|
||||
# Insert low-priority item first, then high-priority for same user.
|
||||
_insert_queue_item(session_queue_round_robin, queue_id, "user_a", priority=0)
|
||||
_insert_queue_item(session_queue_round_robin, queue_id, "user_a", priority=10)
|
||||
_insert_queue_item(session_queue_round_robin, queue_id, "user_b", priority=0)
|
||||
|
||||
# Round 1: user_a's best item (priority 10), user_b's only item.
|
||||
# Round 2: user_a's remaining item (priority 0).
|
||||
items = []
|
||||
for _ in range(3):
|
||||
item = session_queue_round_robin.dequeue()
|
||||
assert item is not None
|
||||
items.append((item.user_id, item.priority))
|
||||
|
||||
assert items[0] == ("user_a", 10)
|
||||
assert items[1] == ("user_b", 0)
|
||||
assert items[2] == ("user_a", 0)
|
||||
|
||||
|
||||
def _seed_completed_history(
|
||||
session_queue: SqliteSessionQueue,
|
||||
queue_id: str,
|
||||
user_id: str,
|
||||
count: int,
|
||||
) -> None:
|
||||
"""Insert `count` completed items (with started_at set) for a user, simulating retained history."""
|
||||
with session_queue._db.transaction() as cursor:
|
||||
for i in range(count):
|
||||
session_id = str(uuid.uuid4())
|
||||
batch_id = str(uuid.uuid4())
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
INSERT INTO session_queue
|
||||
(queue_id, session, session_id, batch_id, field_values, priority, workflow, origin, destination, retried_from_item_id, user_id, status, started_at, completed_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'completed', ?, ?)
|
||||
""",
|
||||
(
|
||||
queue_id,
|
||||
_EMPTY_SESSION_JSON,
|
||||
session_id,
|
||||
batch_id,
|
||||
None,
|
||||
0,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
user_id,
|
||||
# Monotonically increasing timestamps so MAX(started_at) is well-defined per user.
|
||||
f"2026-01-01 {i // 3600 % 24:02d}:{i // 60 % 60:02d}:{i % 60:02d}",
|
||||
f"2026-01-01 {i // 3600 % 24:02d}:{i // 60 % 60:02d}:{i % 60:02d}",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_round_robin_dequeue_does_not_scan_full_history(session_queue_round_robin: SqliteSessionQueue) -> None:
|
||||
"""Round-robin dequeue cost must scale with active users, not retained queue history.
|
||||
|
||||
Regression guard for the scaling concern: the per-user "last served" lookup must be an
|
||||
indexed seek (MAX(started_at) WHERE user_id = ?) rather than a GROUP BY / scan over every
|
||||
historical started row. `max_queue_history` is unbounded by default, so a plan that scans
|
||||
the full history makes each dequeue O(total history) instead of O(active users).
|
||||
|
||||
We seed a large completed history across several users plus a few pending items, then assert
|
||||
the dequeue query plan never scans the `session_queue` base table and resolves the
|
||||
last-served lookup via a seek on `idx_session_queue_user_started_at`.
|
||||
"""
|
||||
queue_id = "default"
|
||||
for u in ("user_a", "user_b", "user_c"):
|
||||
_seed_completed_history(session_queue_round_robin, queue_id, u, count=500)
|
||||
_insert_queue_item(session_queue_round_robin, queue_id, u)
|
||||
|
||||
with session_queue_round_robin._db.transaction() as cursor:
|
||||
plan_rows = cursor.execute("EXPLAIN QUERY PLAN " + ROUND_ROBIN_DEQUEUE_QUERY).fetchall()
|
||||
details = [row["detail"] for row in plan_rows]
|
||||
|
||||
# No step may scan the session_queue base table — that is the full-history scan we are
|
||||
# eliminating. (CTE result scans like "SCAN uni" / "SCAN (subquery-N)" are fine; those are
|
||||
# one row per pending user.)
|
||||
offending = [d for d in details if d.startswith("SCAN session_queue")]
|
||||
assert not offending, f"dequeue plan scans full queue history: {offending}\nfull plan: {details}"
|
||||
|
||||
# The last-served lookup must use the started_at index as a per-user seek.
|
||||
assert any("idx_session_queue_user_started_at" in d and "user_id=?" in d for d in details), (
|
||||
f"last-served lookup is not an indexed seek; plan: {details}"
|
||||
)
|
||||
|
||||
# And the dequeue must still return the least-recently-served user (correctness under history).
|
||||
# user_a's history ends earliest only if seeded first; all three were seeded equal counts with
|
||||
# identical timestamps, so item_id ASC tie-breaks to the first-inserted pending item (user_a).
|
||||
item = session_queue_round_robin.dequeue()
|
||||
assert item is not None
|
||||
assert item.user_id == "user_a"
|
||||
|
||||
|
||||
def test_round_robin_ignored_in_single_user_mode(mock_invoker: Invoker) -> None:
|
||||
"""When multiuser=False, round_robin config is ignored and FIFO is used."""
|
||||
mock_invoker.services.configuration = InvokeAIAppConfig(
|
||||
use_memory_db=True,
|
||||
node_cache_size=0,
|
||||
multiuser=False,
|
||||
session_queue_mode="round_robin",
|
||||
)
|
||||
db = mock_invoker.services.board_records._db
|
||||
queue = SqliteSessionQueue(db=db)
|
||||
queue.start(mock_invoker)
|
||||
|
||||
queue_id = "default"
|
||||
_insert_queue_item(queue, queue_id, "user_a")
|
||||
_insert_queue_item(queue, queue_id, "user_a")
|
||||
_insert_queue_item(queue, queue_id, "user_b")
|
||||
|
||||
# FIFO order: user_a, user_a, user_b
|
||||
user_ids = _dequeue_user_ids(queue, 3)
|
||||
assert user_ids == ["user_a", "user_a", "user_b"]
|
||||
Reference in New Issue
Block a user