203 lines
7.0 KiB
Python
203 lines
7.0 KiB
Python
import asyncio
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from test_collector_list_inputs import _load_collector_module
|
|
|
|
|
|
class _FakeManager:
|
|
def __init__(self, *, pending=False, owned=True):
|
|
self.pending = pending
|
|
self.owned = owned
|
|
self.pending_calls = []
|
|
self.owned_calls = []
|
|
self.cancel_calls = []
|
|
|
|
def is_pending(self, collector_job_id, participant_id):
|
|
self.pending_calls.append((collector_job_id, participant_id))
|
|
return self.pending
|
|
|
|
def owns(self, collector_job_id, participant_id):
|
|
self.owned_calls.append((collector_job_id, participant_id))
|
|
return self.owned
|
|
|
|
async def cancel(self, worker_id=None, execution_id=None, collector_job_id=None):
|
|
self.cancel_calls.append({
|
|
"worker_id": worker_id,
|
|
"execution_id": execution_id,
|
|
"collector_job_id": collector_job_id,
|
|
})
|
|
return 1
|
|
|
|
|
|
def _set_up_collector(module, job_id, queue, manager=None):
|
|
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
|
module.prompt_server.distributed_pending_jobs = {job_id: queue}
|
|
if manager is not None:
|
|
module.prompt_server.distributed_serverless = manager
|
|
return module.DistributedCollectorNode()
|
|
|
|
|
|
def test_serverless_pending_waits_past_normal_worker_timeout():
|
|
module = _load_collector_module()
|
|
job_id = "collector-pending"
|
|
participant = "serverless:pool:1"
|
|
queue = asyncio.Queue()
|
|
manager = _FakeManager(pending=True)
|
|
collector = _set_up_collector(module, job_id, queue, manager)
|
|
module.get_worker_timeout_seconds = lambda: 0.01
|
|
module.HEARTBEAT_INTERVAL = 0.1
|
|
probe_calls = []
|
|
module.probe_worker = lambda *args, **kwargs: probe_calls.append(args)
|
|
|
|
async def deliver_after_provider_work():
|
|
await asyncio.sleep(0.14)
|
|
queue.put_nowait({
|
|
"worker_id": participant,
|
|
"image_index": 0,
|
|
"tensor": torch.ones(1, 2, 2, 2),
|
|
"is_last": True,
|
|
})
|
|
|
|
async def run_collection():
|
|
producer = asyncio.create_task(deliver_after_provider_work())
|
|
result = await collector.execute(
|
|
images=None,
|
|
audio=None,
|
|
multi_job_id=job_id,
|
|
enabled_worker_ids=f'["{participant}"]',
|
|
delegate_only=True,
|
|
)
|
|
await producer
|
|
return result
|
|
|
|
images, _audio = asyncio.run(run_collection())
|
|
|
|
assert tuple(images.shape) == (1, 2, 2, 2)
|
|
assert manager.pending_calls
|
|
assert probe_calls == []
|
|
assert job_id not in module.prompt_server.distributed_pending_jobs
|
|
|
|
|
|
def test_serverless_queue_error_raises_and_cancels_only_collector_binding():
|
|
module = _load_collector_module()
|
|
job_id = "collector-error"
|
|
queue = asyncio.Queue()
|
|
queue.put_nowait({
|
|
"worker_id": "serverless:pool:1",
|
|
"serverless_error": "Provider execution failed safely",
|
|
"is_last": True,
|
|
})
|
|
manager = _FakeManager()
|
|
collector = _set_up_collector(module, job_id, queue, manager)
|
|
|
|
with pytest.raises(RuntimeError, match="serverless.*Provider execution failed safely"):
|
|
asyncio.run(collector.execute(
|
|
images=None,
|
|
audio=None,
|
|
multi_job_id=job_id,
|
|
enabled_worker_ids='["serverless:pool:1"]',
|
|
delegate_only=True,
|
|
))
|
|
|
|
assert manager.cancel_calls == [{
|
|
"worker_id": None,
|
|
"execution_id": None,
|
|
"collector_job_id": job_id,
|
|
}]
|
|
assert job_id not in module.prompt_server.distributed_pending_jobs
|
|
|
|
|
|
def test_owned_serverless_slot_without_pending_result_fails_instead_of_partial_success():
|
|
module = _load_collector_module()
|
|
job_id = "collector-stopped"
|
|
queue = asyncio.Queue()
|
|
manager = _FakeManager(pending=False, owned=True)
|
|
collector = _set_up_collector(module, job_id, queue, manager)
|
|
module.get_worker_timeout_seconds = lambda: 0.01
|
|
module.HEARTBEAT_INTERVAL = 0.1
|
|
module.probe_worker = lambda *_args, **_kwargs: pytest.fail("serverless slots are not worker URLs")
|
|
|
|
with pytest.raises(RuntimeError, match="serverless.*stopped.*result|serverless.*result.*stopped"):
|
|
asyncio.run(collector.execute(
|
|
images=None,
|
|
audio=None,
|
|
multi_job_id=job_id,
|
|
enabled_worker_ids='["serverless:pool:1"]',
|
|
delegate_only=True,
|
|
))
|
|
|
|
assert manager.owned_calls
|
|
assert manager.cancel_calls[0]["collector_job_id"] == job_id
|
|
assert job_id not in module.prompt_server.distributed_pending_jobs
|
|
|
|
|
|
def test_interruption_cancels_serverless_job_and_removes_collector_queue():
|
|
module = _load_collector_module()
|
|
job_id = "collector-interrupted"
|
|
queue = asyncio.Queue()
|
|
manager = _FakeManager(pending=True)
|
|
collector = _set_up_collector(module, job_id, queue, manager)
|
|
interrupted = module.comfy.model_management.InterruptProcessingException("stop")
|
|
module.comfy.model_management.throw_exception_if_processing_interrupted = lambda: (_ for _ in ()).throw(interrupted)
|
|
|
|
with pytest.raises(module.comfy.model_management.InterruptProcessingException, match="stop"):
|
|
asyncio.run(collector.execute(
|
|
images=None,
|
|
audio=None,
|
|
multi_job_id=job_id,
|
|
enabled_worker_ids='["serverless:pool:1"]',
|
|
delegate_only=True,
|
|
))
|
|
|
|
assert manager.cancel_calls[0]["collector_job_id"] == job_id
|
|
assert job_id not in module.prompt_server.distributed_pending_jobs
|
|
|
|
|
|
def test_serverless_tensor_shape_mismatch_raises_instead_of_silent_fallback():
|
|
module = _load_collector_module()
|
|
job_id = "collector-shape-mismatch"
|
|
queue = asyncio.Queue()
|
|
queue.put_nowait({
|
|
"worker_id": "serverless:pool:1",
|
|
"image_index": 0,
|
|
"tensor": torch.ones(1, 3, 2, 2),
|
|
"is_last": True,
|
|
})
|
|
manager = _FakeManager()
|
|
collector = _set_up_collector(module, job_id, queue, manager)
|
|
local_images = torch.zeros(1, 2, 2, 2)
|
|
|
|
with pytest.raises(RuntimeError, match="serverless.*shape|shape.*serverless"):
|
|
asyncio.run(collector.execute(
|
|
images=local_images,
|
|
audio=None,
|
|
multi_job_id=job_id,
|
|
enabled_worker_ids='["serverless:pool:1"]',
|
|
))
|
|
|
|
assert manager.cancel_calls[0]["collector_job_id"] == job_id
|
|
assert job_id not in module.prompt_server.distributed_pending_jobs
|
|
|
|
|
|
def test_regular_worker_timeout_keeps_partial_behavior_without_creating_manager():
|
|
module = _load_collector_module()
|
|
job_id = "collector-regular-timeout"
|
|
queue = asyncio.Queue()
|
|
collector = _set_up_collector(module, job_id, queue)
|
|
module.get_worker_timeout_seconds = lambda: 0.01
|
|
module.HEARTBEAT_INTERVAL = 0.1
|
|
local_images = torch.zeros(1, 2, 2, 2)
|
|
collector._get_serverless_manager = lambda *_args, **_kwargs: pytest.fail("regular collection must not create manager")
|
|
|
|
images, _audio = asyncio.run(collector.execute(
|
|
images=local_images,
|
|
audio=None,
|
|
multi_job_id=job_id,
|
|
enabled_worker_ids='["worker-a"]',
|
|
))
|
|
|
|
assert torch.equal(images, local_images)
|
|
assert job_id not in module.prompt_server.distributed_pending_jobs
|