684 lines
32 KiB
Python
684 lines
32 KiB
Python
"""CPU-only image-list compilation and lifecycle tests; providers are mocked."""
|
|
import asyncio
|
|
import base64
|
|
import copy
|
|
import importlib
|
|
import json
|
|
import sys
|
|
import tempfile
|
|
import types
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
PACKAGE = "_image_list_serverless_tests"
|
|
for suffix in ("", ".api", ".api.orchestration", ".image_list", ".serverless_bridge"):
|
|
module = types.ModuleType(PACKAGE + suffix)
|
|
module.__path__ = [str(ROOT.joinpath(*suffix.strip(".").split(".")))] if suffix else [str(ROOT)]
|
|
sys.modules[module.__name__] = module
|
|
adapter = importlib.import_module(PACKAGE + ".image_list.serverless")
|
|
runtime = adapter.runtime
|
|
manager_module = importlib.import_module(PACKAGE + ".api.orchestration.serverless_runtime")
|
|
handler = importlib.import_module(PACKAGE + ".serverless_bridge.handler")
|
|
codec, contract = adapter.codec, adapter.contract
|
|
|
|
|
|
def node(kind, **inputs):
|
|
return {"class_type": kind, "inputs": inputs}
|
|
|
|
|
|
def pool(wid="pool", jobs=1, enabled=True, **overrides):
|
|
return dict(id=wid, type="serverless", endpoint_id="ep", timeout_seconds=30,
|
|
jobs_per_run=jobs, enabled=enabled, media_profile=contract.PROFILE, **overrides)
|
|
|
|
|
|
def source():
|
|
return {
|
|
"upstream": node("LoadImage", image="never-read.png"),
|
|
"ns/start": node(runtime.START_CLASS, images=["upstream", 0]),
|
|
"model": node("CheckpointLoaderSimple", ckpt_name="fixed.safetensors"),
|
|
"process": node("Process", image=["ns/start", 0], mask=["ns/start", 1],
|
|
index=["ns/start", 2], text=["ns/start", 4], integer=["ns/start", 5],
|
|
floating=["ns/start", 6], boolean=["ns/start", 7], model=["model", 0]),
|
|
"ns/end": node(runtime.END_CLASS, context=["ns/start", 3], images=["process", 0],
|
|
masks=["process", 1], texts=["ns/start", 4], worker_result_token="drop-me"),
|
|
"downstream": node("SaveImage", images=["ns/end", 0]),
|
|
}
|
|
|
|
|
|
def region(prompt=None):
|
|
return {"start_id": "ns/start", "end_id": "ns/end",
|
|
"node_ids": set(prompt or source()) - {"upstream", "downstream"}}
|
|
|
|
|
|
def item(index=0, dtype=torch.float32):
|
|
return {"index": index, "image": torch.arange(24, dtype=dtype).reshape(1, 2, 3, 4) / 24,
|
|
"mask": torch.arange(6, dtype=dtype).reshape(1, 2, 3) / 6,
|
|
"values": {"value_1": " \n雪", "value_2": 123, "value_3": 0.25, "value_4": False}}
|
|
|
|
|
|
async def prepare(prompt=None, value=None, worker=None):
|
|
prompt = prompt or source()
|
|
return await adapter.prepare_serverless_item(prompt, region(prompt), value or item(), worker or pool(),
|
|
request_id="request:1", job_id="queue:1")
|
|
|
|
|
|
async def decoded_inputs(job):
|
|
return {artifact["id"]: codec.decode_payload(artifact["kind"], artifact["descriptor"],
|
|
await contract.resolve_artifact(artifact)) for artifact in job["media"]["artifacts"]}
|
|
|
|
|
|
def test_selection_is_exact_no_expansion_or_profile_fallback():
|
|
a, b = pool("a", jobs=16), pool("b", enabled=False)
|
|
direct = {"id": "direct", "type": "remote", "enabled": True}
|
|
config = {"workers": [a, b, direct]}
|
|
assert [w["id"] for w in adapter.resolve_serverless_workers(config, None)] == ["a"]
|
|
resolved = adapter.resolve_serverless_workers(config, ["b", "direct"])
|
|
assert [w["id"] for w in resolved] == ["b"]
|
|
assert resolved[0]["pool_id"] == "b"
|
|
resolved[0]["jobs_per_run"] = 2
|
|
assert b["jobs_per_run"] == 1
|
|
assert adapter.resolve_serverless_workers(config, []) == []
|
|
legacy = pool("legacy", enabled=False)
|
|
legacy.pop("media_profile")
|
|
config["workers"].append(legacy)
|
|
assert [w["id"] for w in adapter.resolve_serverless_workers(config, None)] == ["a"]
|
|
with pytest.raises(ValueError, match="comfy-media-v1"):
|
|
adapter.resolve_serverless_workers(config, ["legacy"])
|
|
with pytest.raises(ValueError, match="Unknown"):
|
|
adapter.resolve_serverless_workers(config, ["a:1"])
|
|
with pytest.raises(ValueError, match="Invalid"):
|
|
adapter.resolve_serverless_workers(config, [" a"])
|
|
with pytest.raises(ValueError, match="Duplicate"):
|
|
adapter.resolve_serverless_workers({"workers": [a, a]}, ["a"])
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
|
|
def test_deterministic_graph_exact_codecs_and_source_immutable(dtype):
|
|
async def check():
|
|
prompt, value = source(), item(dtype=dtype)
|
|
original, saved_image = copy.deepcopy(prompt), value["image"].clone()
|
|
one = await prepare(prompt, value)
|
|
two = await prepare(dict(reversed(list(prompt.items()))), value)
|
|
assert one == two
|
|
assert prompt == original and torch.equal(value["image"], saved_image)
|
|
assert "upstream" not in one["workflow"] and "downstream" not in one["workflow"]
|
|
assert one["workflow"]["model"] == prompt["model"]
|
|
assert not any("distributed" in n["class_type"].lower() for n in one["workflow"].values())
|
|
proc = one["workflow"]["process"]["inputs"]
|
|
for name, expected in [("index", 0), ("integer", 123), ("floating", 0.25), ("boolean", False), ("text", " \n雪")]:
|
|
assert proc[name] == expected and type(proc[name]) is type(expected)
|
|
output = next(b for b in one["manifest"] if b["socket"] == "texts")
|
|
text_source = one["workflow"][output["node_id"]]["inputs"]["texts"][0]
|
|
assert one["workflow"][text_source] == node("PrimitiveString", value=" \n雪")
|
|
values = await decoded_inputs(one)
|
|
for ref in one["media"]["inputs"]:
|
|
result = values[ref["artifact_id"]]
|
|
assert result.dtype == dtype
|
|
assert torch.equal(result, value[ref["input_name"]])
|
|
assert len(values) == 2
|
|
handler._validate_request({key: one[key] for key in ("workflow", "images", "media")})
|
|
asyncio.run(check())
|
|
|
|
|
|
def test_direct_passthrough_exports_and_tensor_generic_values_deduplicate():
|
|
async def check():
|
|
prompt, value = source(), item()
|
|
prompt["process"]["inputs"]["reference"] = ["ns/start", 8]
|
|
prompt["process"]["inputs"]["reference_mask"] = ["ns/start", 9]
|
|
value["values"].update(value_5=value["image"], value_6=value["mask"])
|
|
prompt["ns/end"]["inputs"]["images"] = ["ns/start", 0]
|
|
job = await prepare(prompt, value)
|
|
assert "process" in job["workflow"] # mask still depends on processing
|
|
assert len(job["media"]["artifacts"]) == 2
|
|
refs = job["media"]["inputs"]
|
|
assert len(refs) == 5
|
|
image_binding = next(b for b in job["manifest"] if b["socket"] == "images")
|
|
assert any(r["node_id"] == image_binding["node_id"] and r["input_name"] == "images" for r in refs)
|
|
decoded = await decoded_inputs(job)
|
|
for ref in refs:
|
|
tensor = decoded[ref["artifact_id"]]
|
|
assert torch.equal(tensor, value["mask"] if "mask" in ref["input_name"] else value["image"])
|
|
asyncio.run(check())
|
|
|
|
|
|
@pytest.mark.parametrize("bad", [None, object(), torch.ones(2), torch.ones(2, 3)])
|
|
def test_connected_unsupported_generic_values_fail(bad):
|
|
value = item()
|
|
value["values"]["value_1"] = bad
|
|
with pytest.raises((ValueError, TypeError)):
|
|
asyncio.run(prepare(value=value))
|
|
|
|
|
|
@pytest.mark.parametrize("bad", [True, -1, "0", 0.0])
|
|
def test_start_slot_identity_not_coerced(bad):
|
|
prompt = source()
|
|
prompt["process"]["inputs"]["image"] = ["ns/start", bad]
|
|
with pytest.raises(ValueError, match="link"):
|
|
asyncio.run(prepare(prompt))
|
|
|
|
|
|
@pytest.mark.parametrize("bad", ["../bad", "ns//bad", "ns/./bad", "bad ", 42])
|
|
def test_node_identity_not_normalized_even_on_pruned_branch(bad):
|
|
prompt = source()
|
|
prompt[bad] = node("PrimitiveString", value="unused")
|
|
with pytest.raises(ValueError):
|
|
asyncio.run(prepare(prompt))
|
|
|
|
|
|
def test_placeholder_namespace_collision_and_empty_text():
|
|
async def check():
|
|
prompt, value = source(), item()
|
|
prompt["image-list:input:0"] = node("PrimitiveInt", value=5)
|
|
prompt["process"]["inputs"]["fixed"] = ["image-list:input:0", 0]
|
|
value["values"]["value_1"] = ""
|
|
job = await prepare(prompt, value)
|
|
assert job["workflow"]["image-list:input:0"] == prompt["image-list:input:0"]
|
|
assert job["workflow"]["process"]["inputs"]["image"] == ["image-list:input:0:1", 0]
|
|
assert any(n == node("PrimitiveString", value="") for n in job["workflow"].values())
|
|
asyncio.run(check())
|
|
|
|
|
|
def test_no_connected_none_and_payload_budget_before_submission():
|
|
value = item()
|
|
value["mask"] = None
|
|
with pytest.raises(ValueError, match="empty"):
|
|
asyncio.run(prepare(value=value))
|
|
client = importlib.import_module(PACKAGE + ".api.orchestration.runpod_client")
|
|
with patch.object(client, "MAX_REQUEST_BYTES", 10):
|
|
with pytest.raises(ValueError, match="Serialized"):
|
|
asyncio.run(prepare())
|
|
|
|
|
|
def test_fixed_file_inputs_are_prepared_on_private_copy_and_merge(tmp_path, monkeypatch):
|
|
input_dir = tmp_path / "input"
|
|
input_dir.mkdir()
|
|
(tmp_path / "output").mkdir()
|
|
(tmp_path / "temp").mkdir()
|
|
payload = b"bounded original file bytes"
|
|
(input_dir / "fixed.png").write_bytes(payload)
|
|
folders = types.SimpleNamespace(get_input_directory=lambda: str(input_dir),
|
|
get_output_directory=lambda: str(tmp_path / "output"), get_temp_directory=lambda: str(tmp_path / "temp"),
|
|
get_annotated_filepath=lambda name: str(input_dir / name))
|
|
monkeypatch.setitem(sys.modules, "folder_paths", folders)
|
|
prompt = source()
|
|
prompt["fixed"] = node("LoadImage", image="fixed.png")
|
|
prompt["process"]["inputs"]["fixed"] = ["fixed", 0]
|
|
original = copy.deepcopy(prompt)
|
|
job = asyncio.run(prepare(prompt))
|
|
assert prompt == original
|
|
assert len(job["images"]) == 1
|
|
assert base64.b64decode(job["images"][0]["image"]) == payload
|
|
assert job["workflow"]["fixed"]["inputs"]["image"] == job["images"][0]["name"]
|
|
assert len(job["media"]["artifacts"]) == 2
|
|
|
|
|
|
def test_merge_rejects_conflicting_artifact_ids_and_combined_caps():
|
|
async def check():
|
|
job = await prepare()
|
|
artifact = job["media"]["artifacts"][0]
|
|
assert adapter._merge_artifacts([artifact], [copy.deepcopy(artifact)]) == [artifact]
|
|
bad = copy.deepcopy(artifact)
|
|
bad["sha256"] = "0" * 64
|
|
with pytest.raises(ValueError, match="identity"):
|
|
adapter._merge_artifacts([artifact], [bad])
|
|
with patch.object(contract, "MAX_TOTAL_BYTES", 1):
|
|
with pytest.raises(ValueError, match="limits"):
|
|
adapter._merge_artifacts([artifact])
|
|
asyncio.run(check())
|
|
|
|
|
|
def test_handler_injects_exact_native_values_and_prunes_placeholders():
|
|
async def check():
|
|
value = item(dtype=torch.float64)
|
|
prompt = source()
|
|
prompt["ns/end"]["inputs"].update(images=["ns/start", 0], masks=["ns/start", 1])
|
|
job = await prepare(prompt, value)
|
|
workflow, media, legacy, targets = handler._validate_request(
|
|
{key: job[key] for key in ("workflow", "images", "media")})
|
|
buffers = {}
|
|
def write(payload):
|
|
name = f"owned-{len(buffers)}.bin"
|
|
buffers[name] = payload
|
|
return name
|
|
class Submitted(Exception):
|
|
pass
|
|
async def submit(native_graph):
|
|
assert all(n["class_type"] in {"MediaExportV1", "MediaImportV1", "PrimitiveString"}
|
|
for n in native_graph.values())
|
|
for binding in job["manifest"]:
|
|
link = native_graph[binding["node_id"]]["inputs"][binding["socket"]]
|
|
native = native_graph[link[0]]
|
|
if binding["kind"] == "STRING":
|
|
assert native == node("PrimitiveString", value=value["values"]["value_1"])
|
|
continue
|
|
assert native["class_type"] == "MediaImportV1"
|
|
metadata = json.loads(native["inputs"]["artifact_json"])
|
|
decoded = codec.decode_payload(metadata["kind"], metadata["descriptor"],
|
|
buffers[native["inputs"]["filename"]])
|
|
assert link[1] == {"IMAGE": 0, "MASK": 2}[binding["kind"]]
|
|
assert torch.equal(decoded, value["image" if binding["kind"] == "IMAGE" else "mask"])
|
|
# Stop at the mocked HTTP boundary; no endpoint/process/provider exists.
|
|
raise Submitted()
|
|
execution = types.SimpleNamespace(media=media, targets=targets, remaining=lambda: 30,
|
|
stage=types.SimpleNamespace(write=write), submit_and_wait=submit)
|
|
with pytest.raises(Submitted):
|
|
await handler._run(execution, workflow, legacy, None, None)
|
|
asyncio.run(check())
|
|
|
|
|
|
def test_fixed_typed_inputs_merge_with_start_refs_without_losing_consumers():
|
|
async def check():
|
|
prompt = source()
|
|
prompt["process"]["inputs"]["fixed_audio"] = 0
|
|
audio = {"waveform": torch.ones(1, 2, 3), "sample_rate": 48000}
|
|
descriptor, payload = codec.encode_payload("AUDIO", audio)
|
|
artifact = contract.artifact_from_payload("fixed-audio", "AUDIO", descriptor, payload)
|
|
fixed = AsyncMock(return_value=([], [artifact], [
|
|
{"node_id": "process", "input_name": "fixed_audio", "artifact_id": "fixed-audio"}]))
|
|
with patch.object(adapter, "prepare_media_inputs", fixed):
|
|
job = await prepare(prompt)
|
|
assert len(job["media"]["artifacts"]) == 3
|
|
assert len(job["media"]["inputs"]) == 3
|
|
assert any(ref["input_name"] == "fixed_audio" for ref in job["media"]["inputs"])
|
|
decoded = await decoded_inputs(job)
|
|
assert decoded["fixed-audio"]["sample_rate"] == 48000
|
|
assert torch.equal(decoded["fixed-audio"]["waveform"], audio["waveform"])
|
|
asyncio.run(check())
|
|
|
|
|
|
class LifecycleTests(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.server = types.SimpleNamespace(distributed_pending_jobs={}, distributed_jobs_lock=asyncio.Lock())
|
|
self.manager = manager_module.ServerlessManager(self.server,
|
|
types.SimpleNamespace(get=lambda _wid: "fake-key"), Path(self.temp.name) / "jobs.json", AsyncMock())
|
|
self.manager.poll_seconds = 0
|
|
self.state = runtime.ImageListJob("context", "ns/start", "ns/end", source(),
|
|
[item(i) for i in range(40)], ["pool"], "", prompt_id="master-prompt")
|
|
self.gates, self.jobs, self.mutations = {}, {}, {}
|
|
self.started = asyncio.Event()
|
|
async def submit(workflow, images, timeout, *, media):
|
|
jid = "provider-" + str(len(self.jobs))
|
|
self.jobs[jid] = dict(workflow=workflow, images=images, media=media)
|
|
self.gates[jid] = asyncio.Event()
|
|
self.started.set()
|
|
return jid
|
|
async def status(jid):
|
|
await self.gates[jid].wait()
|
|
output = self.output(self.jobs[jid])
|
|
if jid in self.mutations:
|
|
self.mutations[jid](output)
|
|
return {"status": "COMPLETED", "output": output}
|
|
self.api = types.SimpleNamespace(submit=AsyncMock(side_effect=submit), status=AsyncMock(side_effect=status),
|
|
cancel=AsyncMock(return_value={"status": "CANCELLED"}))
|
|
self.patches = [patch.object(manager_module, "RunpodClient", return_value=self.api),
|
|
patch.object(adapter, "get_manager", return_value=self.manager)]
|
|
for p in self.patches:
|
|
p.start()
|
|
self.event = asyncio.Event()
|
|
self.controllers = []
|
|
self.controller = await self.controller_for()
|
|
|
|
async def controller_for(self, workers=None, state=None):
|
|
ctl = await adapter.create_serverless_slots(state or self.state, region(),
|
|
{"workers": workers or [pool()]}, self.event)
|
|
self.controllers.append(ctl)
|
|
return ctl
|
|
|
|
async def asyncTearDown(self):
|
|
for ctl in self.controllers:
|
|
await ctl.close()
|
|
await self.manager.cancel()
|
|
for p in reversed(self.patches):
|
|
p.stop()
|
|
self.temp.cleanup()
|
|
|
|
def output(self, job):
|
|
media = job["media"]
|
|
output = dict(version=1, profile=contract.PROFILE, request_id=media["request_id"],
|
|
bindings=[], artifacts=[], items=[])
|
|
index = int(media["request_id"].rsplit(":", 1)[1])
|
|
data = item(index)
|
|
values = {"images": data["image"], "masks": data["mask"], "texts": str(index)}
|
|
for n, binding in enumerate(media["outputs"]):
|
|
descriptor, payload = codec.encode_payload(binding["kind"], values[binding["socket"]])
|
|
aid = f"result:{n}"
|
|
output["bindings"].append(dict(binding, count=1))
|
|
output["artifacts"].append(contract.artifact_from_payload(aid, binding["kind"], descriptor, payload))
|
|
output["items"].insert(0, dict(binding_id=binding["binding_id"], item_index=0, artifact_id=aid))
|
|
return {"media": output}
|
|
|
|
async def dispatch(self, index=0, ctl=None, slot=None):
|
|
ctl = ctl or self.controller
|
|
slot = slot or ctl.slots[index]
|
|
handle = await ctl.dispatch(slot, index)
|
|
return slot, handle
|
|
|
|
async def finish(self, handle):
|
|
# Submission has run by the time dispatch's shield completes, but yield
|
|
# explicitly rather than assuming provider status has already been entered.
|
|
while len(self.jobs) <= handle["index"]:
|
|
await asyncio.sleep(0)
|
|
jid = next(j for j, job in self.jobs.items()
|
|
if job["media"]["request_id"] == handle["execution_id"])
|
|
self.gates[jid].set()
|
|
await handle["task"]
|
|
|
|
async def test_no_pool_does_not_initialize_manager_or_credentials(self):
|
|
state = copy.copy(self.state)
|
|
state.worker_ids = []
|
|
with patch.object(adapter, "get_manager", side_effect=AssertionError("must stay lazy")):
|
|
ctl = await adapter.create_serverless_slots(state, region(), {"workers": [pool()]}, self.event)
|
|
assert ctl.slots == []
|
|
await ctl.close()
|
|
|
|
async def test_slots_are_bounded_not_jobs_per_run_and_cannot_collide(self):
|
|
assert len(self.controller.slots) == manager_module.PER_WORKER_JOB_LIMIT == 16
|
|
ctl = await self.controller_for([pool(jobs=16)])
|
|
assert ctl.slots == self.controller.slots
|
|
assert all(isinstance(slot, tuple) and len(slot) == 3 for slot in ctl.slots)
|
|
assert len(set(ctl.slots)) == 16 and "pool" not in ctl.slots
|
|
many = [pool(f"p{i}") for i in range(8)]
|
|
state = copy.copy(self.state)
|
|
state.worker_ids = [w["id"] for w in many]
|
|
ctl = await self.controller_for(many, state)
|
|
assert len(ctl.slots) == 64
|
|
assert all(sum(s[1] == w["id"] for s in ctl.slots) <= 16 for w in many)
|
|
|
|
async def test_reverse_completion_order_and_slot_reuse(self):
|
|
first, second = await self.dispatch(0), await self.dispatch(1)
|
|
inflight = dict([first, second])
|
|
assert self.controller.poll(inflight) is None
|
|
await self.finish(second[1])
|
|
await asyncio.sleep(0)
|
|
assert self.event.is_set()
|
|
slot, index, result = self.controller.poll(inflight)
|
|
assert (slot, index) == (second[0], 1) and result["text"] == "1"
|
|
assert torch.equal(result["image"], item(1)["image"])
|
|
assert second[1]["job_id"] not in self.server.distributed_pending_jobs
|
|
del inflight[slot]
|
|
third = await self.dispatch(2, slot=slot)
|
|
assert third[1]["execution_id"] != second[1]["execution_id"]
|
|
await self.finish(first[1])
|
|
assert self.controller.poll(inflight)[1] == 0
|
|
await self.finish(third[1])
|
|
assert self.controller.poll(dict([third]))[1] == 2
|
|
|
|
async def test_many_items_submit_once_each_through_bounded_reusable_slots(self):
|
|
ctl = await self.controller_for([pool(jobs=16)])
|
|
first_wave = asyncio.Event()
|
|
active = 0
|
|
peak = 0
|
|
original_submit = self.api.submit.side_effect
|
|
async def submit(*args, **kwargs):
|
|
jid = await original_submit(*args, **kwargs)
|
|
if len(self.jobs) == 16:
|
|
first_wave.set()
|
|
return jid
|
|
async def status(jid):
|
|
nonlocal active, peak
|
|
active += 1
|
|
peak = max(peak, active)
|
|
try:
|
|
await first_wave.wait()
|
|
await asyncio.sleep(0.001)
|
|
return {"status": "COMPLETED", "output": self.output(self.jobs[jid])}
|
|
finally:
|
|
active -= 1
|
|
self.api.submit.side_effect = submit
|
|
self.api.status.side_effect = status
|
|
async def next_result(inflight):
|
|
while True:
|
|
self.event.clear()
|
|
completed = ctl.poll(inflight)
|
|
if completed is not None:
|
|
return completed
|
|
await self.event.wait()
|
|
async def cleanup(*_):
|
|
pass
|
|
results = await asyncio.wait_for(runtime.dynamic_schedule(
|
|
len(self.state.items), ctl.slots, ctl.dispatch, next_result, cleanup), 10)
|
|
assert self.api.submit.await_count == len(self.state.items)
|
|
assert [row["text"] for row in results] == [str(i) for i in range(len(self.state.items))]
|
|
assert peak == 16
|
|
assert not self.server.distributed_pending_jobs
|
|
assert all(task.done() for task in self.manager.tasks.values())
|
|
|
|
async def test_waits_for_task_drain_not_provider_terminal_or_queue_message(self):
|
|
entered, release = asyncio.Event(), asyncio.Event()
|
|
original = self.manager._notify_capacity
|
|
async def delayed():
|
|
entered.set()
|
|
await release.wait()
|
|
await original()
|
|
with patch.object(self.manager, "_notify_capacity", delayed):
|
|
slot, handle = await self.dispatch()
|
|
self.gates["provider-0"].set()
|
|
await entered.wait()
|
|
assert self.server.distributed_pending_jobs[handle["job_id"]].qsize() == 1
|
|
assert self.controller.poll({slot: handle}) is None
|
|
with pytest.raises(RuntimeError, match="occupied"):
|
|
await self.controller.dispatch(slot, 1)
|
|
release.set()
|
|
await handle["task"]
|
|
assert self.controller.poll({slot: handle})[1] == 0
|
|
|
|
async def test_unknown_reservations_survive_and_prevent_provider_admission(self):
|
|
for i in range(16):
|
|
self.manager.ledger.update(f"unknown-{i}", worker_id="pool", status="UNKNOWN", job_id=None)
|
|
slot, handle = await self.dispatch()
|
|
await asyncio.sleep(0)
|
|
self.api.submit.assert_not_awaited()
|
|
assert self.controller.poll({slot: handle}) is None
|
|
await self.controller.close()
|
|
assert sum(r["status"] == "UNKNOWN" for r in self.manager.ledger.all()) == 16
|
|
assert not self.server.distributed_pending_jobs
|
|
assert self.manager.pool_is_active("pool")
|
|
|
|
async def test_local_pending_cap_checked_before_journal(self):
|
|
waiting = [asyncio.create_task(asyncio.Event().wait()) for _ in range(64)]
|
|
self.manager.tasks.update({f"unrelated-{i}": t for i, t in enumerate(waiting)})
|
|
self.manager.runs.update({f"unrelated-{i}": {"execution_id": "unrelated"} for i in range(64)})
|
|
try:
|
|
with pytest.raises(ValueError, match="limit 64"):
|
|
await self.dispatch()
|
|
assert not self.manager.ledger.all()
|
|
self.api.submit.assert_not_awaited()
|
|
assert not self.server.distributed_pending_jobs
|
|
finally:
|
|
for i, task in enumerate(waiting):
|
|
self.manager.tasks.pop(f"unrelated-{i}")
|
|
self.manager.runs.pop(f"unrelated-{i}")
|
|
task.cancel()
|
|
await asyncio.gather(*waiting, return_exceptions=True)
|
|
|
|
async def test_missing_corrupt_or_mismatched_results_are_atomic_errors(self):
|
|
mutations = [lambda o: o["media"]["items"].pop(),
|
|
lambda o: o["media"]["artifacts"][-1].update(sha256="0" * 64),
|
|
lambda o: o["media"].update(request_id="wrong-request"),
|
|
lambda o: o["media"]["bindings"][-1].update(participant_id="wrong-worker")]
|
|
for index, mutation in enumerate(mutations):
|
|
slot, handle = await self.dispatch(index)
|
|
self.mutations[f"provider-{index}"] = mutation
|
|
await self.finish(handle)
|
|
queue = self.server.distributed_pending_jobs[handle["job_id"]]
|
|
assert queue.qsize() == 1
|
|
with pytest.raises(RuntimeError, match="failed"):
|
|
self.controller.poll({slot: handle})
|
|
|
|
async def test_terminal_queue_identity_completeness_and_multiplicity(self):
|
|
for index, kind in enumerate(("missing", "duplicate", "worker", "mask", "text", "multiple", "empty")):
|
|
slot, handle = await self.dispatch(index)
|
|
await self.finish(handle)
|
|
queue = self.server.distributed_pending_jobs[handle["job_id"]]
|
|
message = queue.get_nowait()
|
|
if kind == "duplicate":
|
|
queue.put_nowait(message)
|
|
if kind == "worker":
|
|
message["worker_id"] = "unrelated"
|
|
if kind in ("mask", "text"):
|
|
message["media"].pop("masks" if kind == "mask" else "texts")
|
|
if kind == "multiple":
|
|
message["media"]["images"].append(item()["image"])
|
|
if kind == "empty":
|
|
message["media"]["texts"] = [None]
|
|
if kind != "missing":
|
|
queue.put_nowait(message)
|
|
with pytest.raises(RuntimeError):
|
|
self.controller.poll({slot: handle})
|
|
|
|
async def test_wrong_frame_count_channels_and_unpaired_mask_fail(self):
|
|
for index, kind in enumerate(("batch", "channels", "grid")):
|
|
slot, handle = await self.dispatch(index)
|
|
await self.finish(handle)
|
|
queue = self.server.distributed_pending_jobs[handle["job_id"]]
|
|
message = queue.get_nowait()
|
|
if kind == "batch":
|
|
message["media"]["images"] = [torch.ones(2, 2, 3, 4)]
|
|
elif kind == "channels":
|
|
message["media"]["images"] = [torch.ones(1, 2, 3, 1)]
|
|
else:
|
|
message["media"]["masks"] = [torch.ones(1, 5, 5)]
|
|
queue.put_nowait(message)
|
|
with pytest.raises(ValueError):
|
|
self.controller.poll({slot: handle})
|
|
|
|
async def test_close_cancels_only_owned_execution_and_exact_queues(self):
|
|
state = copy.copy(self.state)
|
|
state.context_id = "unrelated"
|
|
unrelated = await self.controller_for(state=state)
|
|
other_slot, other = await self.dispatch(0, unrelated)
|
|
slot, handle = await self.dispatch(0)
|
|
queue = self.server.distributed_pending_jobs[handle["job_id"]]
|
|
replacement = asyncio.Queue()
|
|
self.server.distributed_pending_jobs[handle["job_id"]] = replacement
|
|
await self.controller.close()
|
|
assert other["task"].done() is False
|
|
assert self.server.distributed_pending_jobs[handle["job_id"]] is replacement
|
|
assert other["job_id"] in self.server.distributed_pending_jobs
|
|
assert handle["task"].done()
|
|
assert len(self.api.cancel.await_args_list) == 1
|
|
assert queue.empty()
|
|
self.server.distributed_pending_jobs.pop(handle["job_id"])
|
|
# A late manager callback skips a retired destination instead of recreating it.
|
|
run = next(r for r in self.manager.runs.values() if r["execution_id"] == handle["execution_id"])
|
|
await self.manager._deliver(run, error="late")
|
|
assert handle["job_id"] not in self.server.distributed_pending_jobs
|
|
|
|
async def test_cancel_dispatch_after_start_loses_handle_but_not_ownership(self):
|
|
entered = asyncio.Event()
|
|
original = self.manager.start
|
|
async def delayed(*args):
|
|
ids = await original(*args)
|
|
entered.set()
|
|
await asyncio.Event().wait()
|
|
return ids
|
|
with patch.object(self.manager, "start", delayed):
|
|
task = asyncio.create_task(self.controller.dispatch(self.controller.slots[0], 0))
|
|
await entered.wait()
|
|
await self.started.wait()
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
await self.controller.close()
|
|
assert all(t.done() for t in self.manager.tasks.values())
|
|
assert not self.server.distributed_pending_jobs
|
|
self.api.cancel.assert_awaited_once_with("provider-0")
|
|
|
|
async def test_close_during_preparation_prevents_late_start(self):
|
|
entered = asyncio.Event()
|
|
prepared = await prepare()
|
|
async def resistant(*args, **kwargs):
|
|
entered.set()
|
|
try:
|
|
await asyncio.Event().wait()
|
|
except asyncio.CancelledError:
|
|
return prepared
|
|
with patch.object(adapter, "prepare_serverless_item", resistant):
|
|
task = asyncio.create_task(self.controller.dispatch(self.controller.slots[0], 0))
|
|
await entered.wait()
|
|
await self.controller.close()
|
|
with pytest.raises(RuntimeError, match="closed"):
|
|
await task
|
|
assert not self.manager.ledger.all()
|
|
assert not self.server.distributed_pending_jobs
|
|
self.api.submit.assert_not_awaited()
|
|
|
|
async def test_cancel_during_provider_submission_preserves_unknown(self):
|
|
entered = asyncio.Event()
|
|
async def ambiguous(*args, **kwargs):
|
|
entered.set()
|
|
await asyncio.Event().wait()
|
|
self.api.submit.side_effect = ambiguous
|
|
slot, handle = await self.dispatch()
|
|
await entered.wait()
|
|
await self.controller.close()
|
|
assert next(iter(self.manager.ledger.all()))["status"] == "UNKNOWN"
|
|
assert self.manager.pool_is_active("pool")
|
|
self.api.cancel.assert_not_awaited()
|
|
assert not self.server.distributed_pending_jobs
|
|
|
|
async def test_cancelled_wait_closes_owned_work(self):
|
|
entered = asyncio.Event()
|
|
async def job():
|
|
try:
|
|
await self.dispatch()
|
|
entered.set()
|
|
await asyncio.Event().wait()
|
|
finally:
|
|
await self.controller.close()
|
|
task = asyncio.create_task(job())
|
|
await entered.wait()
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
assert all(t.done() for t in self.manager.tasks.values())
|
|
assert not self.server.distributed_pending_jobs
|
|
|
|
async def test_repeated_cancellation_does_not_abandon_close(self):
|
|
slot, handle = await self.dispatch()
|
|
entered, release = asyncio.Event(), asyncio.Event()
|
|
original = self.manager.cancel
|
|
async def delayed(**kwargs):
|
|
entered.set()
|
|
await release.wait()
|
|
return await original(**kwargs)
|
|
with patch.object(self.manager, "cancel", delayed):
|
|
closing = asyncio.create_task(self.controller.close())
|
|
await entered.wait()
|
|
closing.cancel()
|
|
await asyncio.sleep(0)
|
|
closing.cancel()
|
|
await asyncio.sleep(0)
|
|
assert not closing.done()
|
|
assert not self.server.distributed_pending_jobs
|
|
release.set()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await closing
|
|
assert handle["task"].done()
|
|
assert self.controller._close_task.done()
|
|
|
|
async def test_serialization_failure_never_journals_or_submits(self):
|
|
with patch.object(adapter, "serialize_request", side_effect=ValueError("budget")):
|
|
with pytest.raises(ValueError, match="budget"):
|
|
await self.dispatch()
|
|
assert not self.manager.ledger.all()
|
|
assert not self.server.distributed_pending_jobs
|
|
self.api.submit.assert_not_awaited()
|
|
|
|
async def test_duplicate_index_and_mismatched_handle_fail(self):
|
|
slot, handle = await self.dispatch()
|
|
with pytest.raises(ValueError, match="duplicate"):
|
|
await self.controller.dispatch(self.controller.slots[1], 0)
|
|
with pytest.raises(RuntimeError, match="Mismatched"):
|
|
self.controller.poll({slot: dict(handle)})
|
|
await self.finish(handle)
|
|
self.controller.poll({slot: handle})
|
|
with pytest.raises(RuntimeError, match="duplicate"):
|
|
self.controller.poll({slot: handle})
|