KOFiblto
8c9fcf5225
Deterministic mtime/path ordering, restore locking, duration and metadata validation, persistent queue entry repair and ordered create times, cancel/wipe recovery ordering, batch color sanitization, and concurrency-safe batch item indexes. Adds 61 Python tests and 3 Node DOM regression tests.
382 lines
15 KiB
Python
382 lines
15 KiB
Python
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
import test_helper # noqa: F401 (redirects SaturnNodes state to a temp dir)
|
|
|
|
import nodes.queue_control as queue_control
|
|
from nodes.queue_control import PersistentQueueManager
|
|
|
|
|
|
def make_entry(prompt_id, priority=0, item_pid=None, create_time=1, extra=None):
|
|
item = [
|
|
priority,
|
|
prompt_id if item_pid is None else item_pid,
|
|
{"prompt": {}},
|
|
{"client_id": "old-client", "create_time": create_time},
|
|
[],
|
|
]
|
|
entry = {"prompt_id": prompt_id, "item": item}
|
|
if extra:
|
|
entry.update(extra)
|
|
return entry
|
|
|
|
|
|
class FakePromptQueue:
|
|
def __init__(self, history=None):
|
|
self.history = {} if history is None else history
|
|
self.items = []
|
|
self.mutex = threading.RLock()
|
|
self.currently_running = {}
|
|
self.queue = []
|
|
self.fail_delete = False
|
|
self.fail_wipe = False
|
|
self.fail_ids = set()
|
|
|
|
def get_history(self):
|
|
return self.history
|
|
|
|
def put(self, item):
|
|
if item[1] in self.fail_ids:
|
|
raise RuntimeError(f"cannot queue {item[1]}")
|
|
self.items.append(item)
|
|
self.queue.append(item)
|
|
|
|
def task_done(self, item_id, history_result, status, process_item=None):
|
|
self.currently_running.pop(item_id, None)
|
|
|
|
def delete_queue_item(self, predicate):
|
|
if self.fail_delete:
|
|
raise RuntimeError("delete failed")
|
|
self.queue = [item for item in self.queue if not predicate(item)]
|
|
|
|
def wipe_queue(self):
|
|
if self.fail_wipe:
|
|
raise RuntimeError("wipe failed")
|
|
self.queue = []
|
|
|
|
|
|
class TestPersistentQueueRestore(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp_dir = tempfile.TemporaryDirectory()
|
|
self.temp_file = os.path.join(self.temp_dir.name, "persistent_queue.json")
|
|
self.patches = [
|
|
patch("nodes.queue_control.PERSISTENT_FILE", self.temp_file),
|
|
patch("nodes.queue_control.USER_DIR", self.temp_dir.name),
|
|
patch("nodes.queue_control.is_persistent_queue_enabled", return_value=True),
|
|
patch.object(
|
|
PersistentQueueManager,
|
|
"_schedule_save_locked",
|
|
new=lambda self: None,
|
|
),
|
|
]
|
|
for patcher in self.patches:
|
|
patcher.start()
|
|
self.manager = PersistentQueueManager()
|
|
|
|
def tearDown(self):
|
|
for patcher in reversed(self.patches):
|
|
patcher.stop()
|
|
self.temp_dir.cleanup()
|
|
|
|
def _server(self, queue, number=0):
|
|
return SimpleNamespace(
|
|
prompt_queue=queue,
|
|
number=number,
|
|
queue_updated=MagicMock(),
|
|
get_queue_status=MagicMock(return_value={}),
|
|
send_sync=MagicMock(),
|
|
)
|
|
|
|
def _write_payload(self, payload):
|
|
with open(self.temp_file, "w", encoding="utf-8") as handle:
|
|
json.dump(payload, handle)
|
|
|
|
def _read_payload(self):
|
|
with open(self.temp_file, "r", encoding="utf-8") as handle:
|
|
return json.load(handle)
|
|
|
|
def test_normalize_entry_uses_top_level_prompt_id(self):
|
|
raw = {"prompt_id": "top-level", "item": [0, "different", {"prompt": {}}]}
|
|
entry = PersistentQueueManager._normalize_entry(raw)
|
|
self.assertEqual(entry["prompt_id"], "top-level")
|
|
self.assertEqual(entry["item"][1], "different")
|
|
|
|
self.assertIsNone(PersistentQueueManager._normalize_entry(None))
|
|
self.assertIsNone(PersistentQueueManager._normalize_entry({"prompt_id": "x", "item": "bad"}))
|
|
self.assertIsNone(PersistentQueueManager._normalize_entry({"prompt_id": "x", "item": [0]}))
|
|
self.assertIsNone(PersistentQueueManager._normalize_entry({"item": [0, ""]}))
|
|
self.assertEqual(
|
|
PersistentQueueManager._normalize_entry({"item": [0, "from-item"]})["prompt_id"],
|
|
"from-item",
|
|
)
|
|
|
|
def test_normalize_items_repairs_and_dedupes_last_wins(self):
|
|
first = make_entry("p1", priority=1)
|
|
duplicate = make_entry("p1", priority=99)
|
|
malformed = {"prompt_id": "bad", "item": "nope"}
|
|
items, changed = PersistentQueueManager._normalize_items([first, malformed, duplicate])
|
|
self.assertTrue(changed)
|
|
self.assertEqual(len(items), 1)
|
|
self.assertEqual(items[0]["prompt_id"], "p1")
|
|
self.assertEqual(items[0]["item"][0], 99)
|
|
|
|
def test_restore_repairs_mismatched_item_prompt_id(self):
|
|
entry = make_entry("top-level", item_pid="wrong-id", priority=4)
|
|
self.manager.persistent_items = [entry]
|
|
queue = FakePromptQueue()
|
|
server = self._server(queue, number=2)
|
|
|
|
with patch.object(queue_control.PromptServer, "instance", server), \
|
|
patch.object(queue_control.pause_manager, "set_pause"):
|
|
self.manager.restore_queue("new-client")
|
|
|
|
self.assertEqual(len(queue.items), 1)
|
|
self.assertEqual(queue.items[0][1], "top-level")
|
|
self.assertEqual(self.manager.persistent_items[0]["item"][1], "top-level")
|
|
payload = self._read_payload()
|
|
self.assertEqual(payload["items"][0]["item"][1], "top-level")
|
|
|
|
def test_restore_drops_prompts_already_in_history(self):
|
|
self.manager.persistent_items = [make_entry("finished"), make_entry("pending", priority=2)]
|
|
queue = FakePromptQueue(history={"finished": {"status": "success"}})
|
|
server = self._server(queue)
|
|
|
|
with patch.object(queue_control.PromptServer, "instance", server), \
|
|
patch.object(queue_control.pause_manager, "set_pause"):
|
|
self.manager.restore_queue()
|
|
|
|
self.assertEqual([item[1] for item in queue.items], ["pending"])
|
|
self.assertEqual([x["prompt_id"] for x in self.manager.persistent_items], ["pending"])
|
|
|
|
def test_restore_drops_malformed_entries(self):
|
|
self.manager.persistent_items = [
|
|
{"prompt_id": "short", "item": [0, "short"]},
|
|
{"prompt_id": "bad-output", "item": [0, "bad-output", "not-a-dict"]},
|
|
{"prompt_id": "valid", "item": [0, "valid", {"prompt": {}}, {}, []]},
|
|
]
|
|
queue = FakePromptQueue()
|
|
server = self._server(queue)
|
|
|
|
with patch.object(queue_control.PromptServer, "instance", server), \
|
|
patch.object(queue_control.pause_manager, "set_pause"):
|
|
self.manager.restore_queue()
|
|
|
|
self.assertEqual([item[1] for item in queue.items], ["valid"])
|
|
self.assertEqual([x["prompt_id"] for x in self.manager.persistent_items], ["valid"])
|
|
|
|
def test_restore_assigns_ordered_unique_create_times_and_client_id(self):
|
|
self.manager.persistent_items = [
|
|
make_entry("p1", priority=10),
|
|
make_entry("p2", priority=5),
|
|
make_entry("p3", priority=7),
|
|
]
|
|
queue = FakePromptQueue()
|
|
server = self._server(queue, number=3)
|
|
|
|
with patch.object(queue_control.PromptServer, "instance", server), \
|
|
patch.object(queue_control.pause_manager, "set_pause"), \
|
|
patch("nodes.queue_control.time.time", return_value=2_000_000.0):
|
|
self.manager.restore_queue("client-42")
|
|
|
|
create_times = [item[3]["create_time"] for item in queue.items]
|
|
self.assertEqual(create_times, [1_999_999_998, 1_999_999_999, 2_000_000_000])
|
|
self.assertEqual(create_times, sorted(set(create_times)))
|
|
for item in queue.items:
|
|
self.assertEqual(item[3]["client_id"], "client-42")
|
|
self.assertEqual(server.number, 11, "server priority counter must exceed restored items")
|
|
|
|
def test_restore_retries_with_malformed_retry_counter(self):
|
|
entry = make_entry("flaky", extra={"retries": "not-a-number"})
|
|
self.manager.persistent_items = [entry]
|
|
queue = FakePromptQueue()
|
|
queue.fail_ids.add("flaky")
|
|
server = self._server(queue)
|
|
|
|
def env(key, default):
|
|
if key == "PERSISTENT_QUEUE_MAX_RETRIES":
|
|
return "2"
|
|
return default
|
|
|
|
with patch.object(queue_control.PromptServer, "instance", server), \
|
|
patch.object(queue_control.pause_manager, "set_pause"), \
|
|
patch("nodes.queue_control.get_env_setting", side_effect=env):
|
|
self.manager.restore_queue()
|
|
self.assertEqual(self.manager.persistent_items[0]["retries"], 1)
|
|
|
|
self.manager.has_claimed_once = False
|
|
self.manager.restore_queue()
|
|
self.assertEqual(self.manager.persistent_items[0]["retries"], 2)
|
|
|
|
self.manager.has_claimed_once = False
|
|
self.manager.restore_queue()
|
|
|
|
self.assertEqual(self.manager.persistent_items, [])
|
|
self.assertEqual(queue.items, [])
|
|
|
|
def test_restore_drops_after_first_failure_when_retries_disabled(self):
|
|
self.manager.persistent_items = [make_entry("flaky")]
|
|
queue = FakePromptQueue()
|
|
queue.fail_ids.add("flaky")
|
|
server = self._server(queue)
|
|
|
|
with patch.object(queue_control.PromptServer, "instance", server), \
|
|
patch.object(queue_control.pause_manager, "set_pause"):
|
|
self.manager.restore_queue()
|
|
|
|
self.assertEqual(self.manager.persistent_items, [])
|
|
|
|
def test_restore_tolerates_history_read_failure(self):
|
|
self.manager.persistent_items = [make_entry("p1")]
|
|
queue = FakePromptQueue()
|
|
|
|
def broken_history():
|
|
raise RuntimeError("history unavailable")
|
|
|
|
queue.get_history = broken_history
|
|
server = self._server(queue)
|
|
|
|
with patch.object(queue_control.PromptServer, "instance", server), \
|
|
patch.object(queue_control.pause_manager, "set_pause"):
|
|
self.manager.restore_queue()
|
|
|
|
self.assertEqual([item[1] for item in queue.items], ["p1"])
|
|
|
|
def test_restore_force_state_setting_is_case_insensitive(self):
|
|
for setting, expected in (("force paused", True), ("FORCE RUNNING", False)):
|
|
self.manager.persistent_items = [make_entry("p1")]
|
|
self.manager.has_claimed_once = False
|
|
queue = FakePromptQueue()
|
|
server = self._server(queue)
|
|
|
|
def env(key, default, value=setting):
|
|
if key == "PERSISTENT_QUEUE_RESTORED_STATE":
|
|
return value
|
|
return default
|
|
|
|
with patch.object(queue_control.PromptServer, "instance", server), \
|
|
patch.object(queue_control.pause_manager, "set_pause") as set_pause, \
|
|
patch("nodes.queue_control.get_env_setting", side_effect=env):
|
|
self.manager.restore_queue()
|
|
|
|
set_pause.assert_called_once_with(expected)
|
|
|
|
def test_restore_is_skipped_when_disabled_or_empty(self):
|
|
queue = FakePromptQueue()
|
|
server = self._server(queue)
|
|
|
|
with patch.object(queue_control.PromptServer, "instance", server), \
|
|
patch("nodes.queue_control.is_persistent_queue_enabled", return_value=False):
|
|
self.manager.persistent_items = [make_entry("p1")]
|
|
self.manager.restore_queue()
|
|
self.assertTrue(self.manager.has_claimed_once)
|
|
self.assertEqual(queue.items, [])
|
|
|
|
self.manager.has_claimed_once = False
|
|
self.manager.persistent_items = []
|
|
with patch.object(queue_control.PromptServer, "instance", server):
|
|
self.manager.restore_queue()
|
|
self.assertEqual(queue.items, [])
|
|
|
|
def test_newer_schema_file_is_not_overwritten_by_restore_or_add(self):
|
|
original = {
|
|
"schema_version": 99,
|
|
"items": [{"prompt_id": "future", "item": [0, "future", {}, {}, []], "future_field": True}],
|
|
}
|
|
self._write_payload(original)
|
|
self.manager.load_from_file()
|
|
self.assertTrue(self.manager._read_only_newer_schema)
|
|
self.assertEqual(self.manager.persistent_items, [])
|
|
|
|
self.manager.add_item((0, "new", {"prompt": {}}, {}, []))
|
|
self.manager.save_to_file(immediate=True)
|
|
self.assertEqual(self._read_payload(), original)
|
|
|
|
def test_load_from_file_rewrites_duplicate_entries(self):
|
|
duplicate = make_entry("p1", priority=99)
|
|
self._write_payload({
|
|
"schema_version": 1,
|
|
"items": [make_entry("p1", priority=1), duplicate],
|
|
})
|
|
manager = PersistentQueueManager()
|
|
self.assertEqual(len(manager.persistent_items), 1)
|
|
self.assertEqual(manager.persistent_items[0]["item"][0], 99)
|
|
self.assertEqual(len(self._read_payload()["items"]), 1)
|
|
|
|
def test_batch_info_requires_dict_and_ignores_finished_prompts(self):
|
|
self.manager.set_batch_info("p1", "not-a-dict")
|
|
self.manager.set_batch_info("", {"batchId": "empty"})
|
|
self.assertEqual(self.manager.batch_meta, {})
|
|
|
|
info = {"batchId": "b1", "color": "#10b981"}
|
|
self.manager.set_batch_info("p1", info)
|
|
self.assertEqual(self.manager.batch_meta["p1"], info)
|
|
|
|
self.manager.add_item((0, "p1", {"prompt": {}}, {}, []))
|
|
self.assertEqual(self.manager.persistent_items[0]["batch_info"], info)
|
|
|
|
self.manager.remove_item("p1")
|
|
self.manager.set_batch_info("p1", {"batchId": "late"})
|
|
self.assertNotIn("p1", self.manager.batch_meta)
|
|
|
|
def test_cancel_and_wipe_remove_persistence_only_after_success(self):
|
|
queue = FakePromptQueue()
|
|
server = self._server(queue)
|
|
self.manager._patched = False
|
|
|
|
with patch.object(queue_control.PromptServer, "instance", server):
|
|
self.manager.patch_server()
|
|
|
|
queue.put((0, "p1", {"prompt": {}}, {}, []))
|
|
self.assertEqual([x["prompt_id"] for x in self.manager.persistent_items], ["p1"])
|
|
|
|
queue.fail_delete = True
|
|
with self.assertRaises(RuntimeError):
|
|
queue.delete_queue_item(lambda item: True)
|
|
self.assertEqual([x["prompt_id"] for x in self.manager.persistent_items], ["p1"])
|
|
|
|
queue.fail_delete = False
|
|
queue.delete_queue_item(lambda item: True)
|
|
self.assertEqual(self.manager.persistent_items, [])
|
|
|
|
queue.put((0, "p2", {"prompt": {}}, {}, []))
|
|
self.assertEqual([x["prompt_id"] for x in self.manager.persistent_items], ["p2"])
|
|
queue.fail_wipe = True
|
|
with self.assertRaises(RuntimeError):
|
|
queue.wipe_queue()
|
|
self.assertEqual([x["prompt_id"] for x in self.manager.persistent_items], ["p2"])
|
|
|
|
queue.fail_wipe = False
|
|
queue.wipe_queue()
|
|
self.assertEqual(self.manager.persistent_items, [])
|
|
|
|
queue.put((0, "p3", {"prompt": {}}, {}, []))
|
|
queue.currently_running[7] = (0, "p3", {"prompt": {}}, {}, [])
|
|
queue.task_done(7, {}, "success")
|
|
self.assertEqual(self.manager.persistent_items, [])
|
|
|
|
def test_restore_syncs_client_id_to_all_entries(self):
|
|
self.manager.persistent_items = [make_entry("p1"), make_entry("p2", priority=2)]
|
|
queue = FakePromptQueue()
|
|
server = self._server(queue)
|
|
|
|
with patch.object(queue_control.PromptServer, "instance", server), \
|
|
patch.object(queue_control.pause_manager, "set_pause"):
|
|
self.manager.restore_queue("client-x")
|
|
|
|
payload = self._read_payload()
|
|
self.assertEqual(
|
|
[item["item"][3]["client_id"] for item in payload["items"]],
|
|
["client-x", "client-x"],
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|