639 lines
22 KiB
Python
639 lines
22 KiB
Python
import importlib.util
|
|
import sys
|
|
import types
|
|
import asyncio
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
|
|
def _register_real_transfer_codec(package_name, utils_pkg):
|
|
"""Load the real utils/transfer_codec.py into the synthetic test package."""
|
|
codec_path = Path(__file__).resolve().parents[1] / "utils" / "transfer_codec.py"
|
|
codec_spec = importlib.util.spec_from_file_location(f"{package_name}.utils.transfer_codec", codec_path)
|
|
codec_module = importlib.util.module_from_spec(codec_spec)
|
|
sys.modules[f"{package_name}.utils.transfer_codec"] = codec_module
|
|
codec_spec.loader.exec_module(codec_module)
|
|
utils_pkg.transfer_codec = codec_module
|
|
return codec_module
|
|
|
|
|
|
def _load_collector_module():
|
|
# Restore shared dependencies even when a different test imports this loader.
|
|
# Keep real modules imported during loading (notably extension modules) cached.
|
|
with pytest.MonkeyPatch.context() as monkeypatch:
|
|
return _load_collector_module_with_stubs(monkeypatch)
|
|
|
|
|
|
def _load_collector_module_with_stubs(monkeypatch):
|
|
module_path = Path(__file__).resolve().parents[1] / "nodes" / "collector.py"
|
|
package_name = "dist_collector_list_testpkg"
|
|
|
|
for mod_name in list(sys.modules):
|
|
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
|
del sys.modules[mod_name]
|
|
|
|
root_pkg = types.ModuleType(package_name)
|
|
root_pkg.__path__ = []
|
|
sys.modules[package_name] = root_pkg
|
|
|
|
nodes_pkg = types.ModuleType(f"{package_name}.nodes")
|
|
nodes_pkg.__path__ = []
|
|
sys.modules[f"{package_name}.nodes"] = nodes_pkg
|
|
|
|
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
|
utils_pkg.__path__ = []
|
|
sys.modules[f"{package_name}.utils"] = utils_pkg
|
|
_register_real_transfer_codec(package_name, utils_pkg)
|
|
|
|
class _Routes:
|
|
def post(self, _path):
|
|
return lambda fn: fn
|
|
|
|
def get(self, _path):
|
|
return lambda fn: fn
|
|
|
|
prompt_server = types.SimpleNamespace(
|
|
routes=_Routes(),
|
|
distributed_jobs_lock=None,
|
|
distributed_pending_jobs={},
|
|
)
|
|
server_module = types.ModuleType("server")
|
|
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
|
|
monkeypatch.setitem(sys.modules, "server", server_module)
|
|
|
|
comfy_module = types.ModuleType("comfy")
|
|
model_management = types.ModuleType("comfy.model_management")
|
|
|
|
class InterruptProcessingException(Exception):
|
|
pass
|
|
|
|
model_management.InterruptProcessingException = InterruptProcessingException
|
|
model_management.throw_exception_if_processing_interrupted = lambda: None
|
|
comfy_module.model_management = model_management
|
|
|
|
comfy_utils = types.ModuleType("comfy.utils")
|
|
|
|
class ProgressBar:
|
|
def __init__(self, _total):
|
|
self.total = _total
|
|
self.updates = []
|
|
|
|
def update(self, value):
|
|
self.updates.append(value)
|
|
|
|
comfy_utils.ProgressBar = ProgressBar
|
|
comfy_module.utils = comfy_utils
|
|
monkeypatch.setitem(sys.modules, "comfy", comfy_module)
|
|
monkeypatch.setitem(sys.modules, "comfy.model_management", model_management)
|
|
monkeypatch.setitem(sys.modules, "comfy.utils", comfy_utils)
|
|
|
|
aiohttp_module = types.ModuleType("aiohttp")
|
|
aiohttp_module.ClientTimeout = lambda total: types.SimpleNamespace(total=total)
|
|
monkeypatch.setitem(sys.modules, "aiohttp", aiohttp_module)
|
|
|
|
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
|
logging_module.debug_log = lambda *_args, **_kwargs: None
|
|
logging_module.log = lambda *_args, **_kwargs: None
|
|
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
|
|
|
config_module = types.ModuleType(f"{package_name}.utils.config")
|
|
config_module.get_worker_timeout_seconds = lambda: 0.1
|
|
config_module.load_config = lambda: {"workers": []}
|
|
config_module.is_master_delegate_only = lambda: False
|
|
sys.modules[f"{package_name}.utils.config"] = config_module
|
|
|
|
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
|
constants_module.HEARTBEAT_INTERVAL = 1.0
|
|
constants_module.TRANSFER_MAX_CHUNK_BYTES = 32 * 1024 * 1024
|
|
constants_module.TRANSFER_MAX_REQUEST_BYTES = 96 * 1024 * 1024
|
|
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
|
|
|
image_module = types.ModuleType(f"{package_name}.utils.image")
|
|
def _ensure_contiguous(tensor):
|
|
return tensor.contiguous() if hasattr(tensor, "contiguous") else tensor
|
|
|
|
image_module.ensure_contiguous = _ensure_contiguous
|
|
image_module.tensor_to_pil = lambda *_args, **_kwargs: None
|
|
image_module.pil_to_tensor = lambda value: value
|
|
sys.modules[f"{package_name}.utils.image"] = image_module
|
|
|
|
network_module = types.ModuleType(f"{package_name}.utils.network")
|
|
network_module.build_worker_url = lambda worker: "http://worker"
|
|
network_module.get_client_session = lambda: None
|
|
network_module.probe_worker = lambda *_args, **_kwargs: None
|
|
sys.modules[f"{package_name}.utils.network"] = network_module
|
|
|
|
audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload")
|
|
audio_payload_module.encode_audio_payload = lambda audio: audio
|
|
sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_module
|
|
|
|
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
|
|
async_helpers_module.run_async_in_server_loop = lambda coro: asyncio.run(coro)
|
|
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
|
|
|
spec = importlib.util.spec_from_file_location(f"{package_name}.nodes.collector", module_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
assert spec is not None and spec.loader is not None
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def test_collector_opts_into_comfyui_list_inputs():
|
|
collector = _load_collector_module().DistributedCollectorNode
|
|
|
|
assert collector.INPUT_IS_LIST is True
|
|
|
|
|
|
def test_collector_exposes_images_as_optional_input():
|
|
input_types = _load_collector_module().DistributedCollectorNode.INPUT_TYPES()
|
|
|
|
assert "images" not in input_types["required"]
|
|
assert input_types["optional"]["images"] == ("IMAGE",)
|
|
|
|
|
|
def test_audio_only_pass_through_returns_no_images_and_preserves_audio():
|
|
collector = _load_collector_module().DistributedCollectorNode()
|
|
audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
|
|
|
|
images, returned_audio = collector.run(images=None, audio=[audio])
|
|
|
|
assert images is None
|
|
assert returned_audio is audio
|
|
|
|
|
|
def test_collector_rejects_missing_images_and_audio():
|
|
collector = _load_collector_module().DistributedCollectorNode()
|
|
|
|
try:
|
|
collector.run(images=None, audio=None)
|
|
except ValueError as exc:
|
|
assert "image or audio" in str(exc).lower()
|
|
else:
|
|
raise AssertionError("Expected collector to reject a run with no media input")
|
|
|
|
|
|
def test_delegate_only_master_allows_no_local_media_input():
|
|
collector = _load_collector_module().DistributedCollectorNode()
|
|
|
|
images, audio = collector.run(
|
|
images=None,
|
|
audio=None,
|
|
multi_job_id=["delegate-audio-job"],
|
|
delegate_only=[True],
|
|
enabled_worker_ids=["[]"],
|
|
)
|
|
|
|
assert images is None
|
|
assert tuple(audio["waveform"].shape) == (1, 2, 1)
|
|
|
|
|
|
def test_pass_through_collapses_comfyui_image_list_to_batch_and_unwraps_hidden_inputs():
|
|
collector = _load_collector_module().DistributedCollectorNode()
|
|
first = torch.zeros(1, 2, 2, 3)
|
|
second = torch.ones(1, 2, 2, 3)
|
|
|
|
images, audio = collector.run(
|
|
images=[first, second],
|
|
load_balance=[False],
|
|
audio=[None],
|
|
multi_job_id=[""],
|
|
is_worker=[False],
|
|
master_url=[""],
|
|
enabled_worker_ids=["[]"],
|
|
worker_batch_size=[1],
|
|
worker_id=[""],
|
|
pass_through=[False],
|
|
delegate_only=[False],
|
|
)
|
|
|
|
assert tuple(images.shape) == (2, 2, 2, 3)
|
|
assert torch.equal(images[0:1], first)
|
|
assert torch.equal(images[1:2], second)
|
|
assert tuple(audio["waveform"].shape) == (1, 2, 1)
|
|
|
|
|
|
def test_worker_list_input_sends_one_completion_sequence_with_last_only_on_final_item():
|
|
module = _load_collector_module()
|
|
collector = module.DistributedCollectorNode()
|
|
first = torch.zeros(1, 2, 2, 3)
|
|
second = torch.ones(1, 2, 2, 3)
|
|
posted_payloads = []
|
|
|
|
class _FakeImage:
|
|
def save(self, fp, format=None, compress_level=None):
|
|
fp.write(b"png-bytes")
|
|
|
|
class _FakeResponse:
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
def raise_for_status(self):
|
|
return None
|
|
|
|
class _FakeSession:
|
|
def post(self, url, json, timeout):
|
|
posted_payloads.append(json)
|
|
return _FakeResponse()
|
|
|
|
async def _fake_get_client_session():
|
|
return _FakeSession()
|
|
|
|
module.tensor_to_pil = lambda *_args, **_kwargs: _FakeImage()
|
|
module.get_client_session = _fake_get_client_session
|
|
module.encode_audio_payload = lambda _audio: None
|
|
|
|
images, audio = collector.run(
|
|
images=[first, second],
|
|
load_balance=[False],
|
|
audio=[None],
|
|
multi_job_id=["job-list-1"],
|
|
is_worker=[True],
|
|
master_url=["http://master"],
|
|
enabled_worker_ids=["[]"],
|
|
worker_batch_size=[1],
|
|
worker_id=["worker-a"],
|
|
pass_through=[False],
|
|
delegate_only=[False],
|
|
)
|
|
|
|
assert tuple(images.shape) == (2, 2, 2, 3)
|
|
assert tuple(audio["waveform"].shape) == (1, 2, 1)
|
|
assert len(posted_payloads) == 2
|
|
assert [payload["batch_idx"] for payload in posted_payloads] == [0, 1]
|
|
assert [payload["is_last"] for payload in posted_payloads] == [False, True]
|
|
assert {payload["job_id"] for payload in posted_payloads} == {"job-list-1"}
|
|
assert {payload["worker_id"] for payload in posted_payloads} == {"worker-a"}
|
|
|
|
|
|
def test_audio_only_worker_sends_one_completion_without_image():
|
|
module = _load_collector_module()
|
|
collector = module.DistributedCollectorNode()
|
|
audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
|
|
posted = []
|
|
|
|
class _FakeResponse:
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
def raise_for_status(self):
|
|
return None
|
|
|
|
class _FakeSession:
|
|
def post(self, url, json, timeout):
|
|
posted.append((url, json, timeout.total))
|
|
return _FakeResponse()
|
|
|
|
async def _fake_get_client_session():
|
|
return _FakeSession()
|
|
|
|
module.get_client_session = _fake_get_client_session
|
|
module.encode_audio_payload = lambda value: {"encoded": value is audio}
|
|
|
|
images, returned_audio = collector.run(
|
|
images=None,
|
|
audio=[audio],
|
|
multi_job_id=["audio-job"],
|
|
is_worker=[True],
|
|
master_url=["http://master"],
|
|
worker_id=["worker-a"],
|
|
)
|
|
|
|
assert images is None
|
|
assert returned_audio is audio
|
|
assert len(posted) == 1
|
|
assert posted[0][0] == "http://master/distributed/job_complete"
|
|
assert posted[0][2] == 600
|
|
assert posted[0][1] == {
|
|
"job_id": "audio-job",
|
|
"worker_id": "worker-a",
|
|
"batch_idx": 0,
|
|
"audio": {"encoded": True},
|
|
"is_last": True,
|
|
}
|
|
|
|
|
|
def test_audio_only_master_combines_local_and_worker_audio():
|
|
module = _load_collector_module()
|
|
collector = module.DistributedCollectorNode()
|
|
master_audio = {"waveform": torch.ones(1, 2, 2), "sample_rate": 48000}
|
|
worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
|
|
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
|
queue = asyncio.Queue()
|
|
queue.put_nowait(
|
|
{
|
|
"worker_id": "worker-a",
|
|
"image_index": 0,
|
|
"tensor": None,
|
|
"audio": worker_audio,
|
|
"is_last": True,
|
|
}
|
|
)
|
|
module.prompt_server.distributed_pending_jobs = {"audio-job": queue}
|
|
|
|
images, combined_audio = asyncio.run(
|
|
collector.execute(
|
|
images=None,
|
|
audio=master_audio,
|
|
multi_job_id="audio-job",
|
|
enabled_worker_ids='["worker-a"]',
|
|
)
|
|
)
|
|
|
|
assert images is None
|
|
assert combined_audio["sample_rate"] == 48000
|
|
assert tuple(combined_audio["waveform"].shape) == (1, 2, 5)
|
|
assert torch.equal(combined_audio["waveform"][..., :2], master_audio["waveform"])
|
|
assert torch.equal(combined_audio["waveform"][..., 2:], worker_audio["waveform"])
|
|
|
|
|
|
def test_delegate_only_audio_collects_worker_audio_without_placeholder_image():
|
|
module = _load_collector_module()
|
|
collector = module.DistributedCollectorNode()
|
|
worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
|
|
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
|
queue = asyncio.Queue()
|
|
queue.put_nowait(
|
|
{
|
|
"worker_id": "worker-a",
|
|
"image_index": 0,
|
|
"tensor": None,
|
|
"audio": worker_audio,
|
|
"is_last": True,
|
|
}
|
|
)
|
|
module.prompt_server.distributed_pending_jobs = {"delegate-audio-job": queue}
|
|
|
|
images, combined_audio = asyncio.run(
|
|
collector.execute(
|
|
images=None,
|
|
audio=None,
|
|
multi_job_id="delegate-audio-job",
|
|
enabled_worker_ids='["worker-a"]',
|
|
delegate_only=True,
|
|
)
|
|
)
|
|
|
|
assert images is None
|
|
assert combined_audio["sample_rate"] == 48000
|
|
assert torch.equal(combined_audio["waveform"], worker_audio["waveform"])
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Compressed worker -> master transfer
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
class _RecordingResponse:
|
|
def __init__(self, status=200, text=""):
|
|
self.status = status
|
|
self._text = text
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
async def text(self):
|
|
return self._text
|
|
|
|
def raise_for_status(self):
|
|
if self.status >= 400:
|
|
raise RuntimeError(f"HTTP {self.status}")
|
|
|
|
|
|
class _RecordingSession:
|
|
def __init__(self, statuses=None, text="Invalid JSON payload"):
|
|
self.calls = []
|
|
self._statuses = list(statuses or [])
|
|
self._text = text
|
|
|
|
def post(self, url, timeout, json=None, data=None):
|
|
self.calls.append({"url": url, "json": json, "data": data, "timeout": timeout.total})
|
|
status = self._statuses.pop(0) if self._statuses else 200
|
|
return _RecordingResponse(status=status, text=self._text)
|
|
|
|
|
|
def _install_fake_transport(module, session):
|
|
async def _fake_get_client_session():
|
|
return session
|
|
|
|
module.get_client_session = _fake_get_client_session
|
|
# Record the multipart content instead of building a real aiohttp.FormData.
|
|
module.DistributedCollectorNode._build_multipart = staticmethod(
|
|
lambda envelope, parts: {"envelope": envelope, "parts": list(parts)}
|
|
)
|
|
|
|
|
|
def _run_worker(collector, images, audio=None, **kwargs):
|
|
return collector.run(
|
|
images=[images],
|
|
audio=[audio],
|
|
multi_job_id=["job-c"],
|
|
is_worker=[True],
|
|
master_url=["http://master"],
|
|
worker_id=["worker-a"],
|
|
**{key: [value] for key, value in kwargs.items()},
|
|
)
|
|
|
|
|
|
def test_collector_exposes_transfer_widgets_with_legacy_default():
|
|
module = _load_collector_module()
|
|
optional = module.DistributedCollectorNode.INPUT_TYPES()["optional"]
|
|
|
|
formats, options = optional["transfer_format"]
|
|
assert "legacy_png" in formats and "h264" in formats and "ffv1" in formats
|
|
assert options["default"] == "legacy_png"
|
|
assert optional["transfer_quality"][0] == "INT"
|
|
|
|
|
|
def test_worker_sends_video_batch_as_one_multipart_chunk_with_audio_on_last():
|
|
module = _load_collector_module()
|
|
collector = module.DistributedCollectorNode()
|
|
session = _RecordingSession()
|
|
_install_fake_transport(module, session)
|
|
module.encode_audio_payload = lambda value: None if value is None else {"encoded": True}
|
|
images = torch.rand(6, 32, 48, 3)
|
|
|
|
_run_worker(
|
|
collector,
|
|
images,
|
|
audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
|
|
transfer_format="ffv1",
|
|
transfer_quality=90,
|
|
)
|
|
|
|
assert len(session.calls) == 1
|
|
call = session.calls[0]
|
|
assert call["url"] == "http://master/distributed/job_complete"
|
|
assert call["json"] is None
|
|
envelope = call["data"]["envelope"]
|
|
assert envelope["format"] == "ffv1"
|
|
assert (envelope["start_index"], envelope["count"]) == (0, 6)
|
|
assert (envelope["width"], envelope["height"]) == (48, 32)
|
|
assert envelope["is_last"] is True
|
|
assert envelope["audio"] == {"encoded": True}
|
|
assert envelope["job_id"] == "job-c" and envelope["worker_id"] == "worker-a"
|
|
assert len(call["data"]["parts"]) == 1
|
|
|
|
|
|
def test_worker_marks_only_final_chunk_as_last():
|
|
module = _load_collector_module()
|
|
collector = module.DistributedCollectorNode()
|
|
session = _RecordingSession()
|
|
_install_fake_transport(module, session)
|
|
module.TRANSFER_MAX_CHUNK_BYTES = 1 # force one frame per request
|
|
|
|
_run_worker(collector, torch.rand(3, 8, 8, 3), transfer_format="png")
|
|
|
|
envelopes = [call["data"]["envelope"] for call in session.calls]
|
|
assert [(e["start_index"], e["count"], e["is_last"]) for e in envelopes] == [
|
|
(0, 1, False),
|
|
(1, 1, False),
|
|
(2, 1, True),
|
|
]
|
|
assert "audio" not in envelopes[-1]
|
|
|
|
|
|
def test_worker_falls_back_to_legacy_png_when_master_rejects_multipart():
|
|
from PIL import Image
|
|
import numpy as np
|
|
|
|
module = _load_collector_module()
|
|
collector = module.DistributedCollectorNode()
|
|
session = _RecordingSession(statuses=[400])
|
|
_install_fake_transport(module, session)
|
|
module.tensor_to_pil = lambda tensor, idx: Image.fromarray(
|
|
(255 * tensor[idx].cpu().numpy()).astype(np.uint8)
|
|
)
|
|
|
|
_run_worker(collector, torch.rand(2, 8, 8, 3), transfer_format="h264")
|
|
|
|
assert session.calls[0]["data"] is not None # compressed attempt
|
|
legacy = [call["json"] for call in session.calls[1:]]
|
|
assert [payload["batch_idx"] for payload in legacy] == [0, 1]
|
|
assert all(payload["image"].startswith("data:image/png;base64,") for payload in legacy)
|
|
assert [payload["is_last"] for payload in legacy] == [False, True]
|
|
|
|
|
|
def test_worker_uses_legacy_path_for_unknown_transfer_format():
|
|
from PIL import Image
|
|
import numpy as np
|
|
|
|
module = _load_collector_module()
|
|
collector = module.DistributedCollectorNode()
|
|
session = _RecordingSession()
|
|
_install_fake_transport(module, session)
|
|
module.tensor_to_pil = lambda tensor, idx: Image.fromarray(
|
|
(255 * tensor[idx].cpu().numpy()).astype(np.uint8)
|
|
)
|
|
|
|
_run_worker(collector, torch.rand(1, 8, 8, 3), transfer_format="gif")
|
|
|
|
assert len(session.calls) == 1
|
|
assert session.calls[0]["json"]["image"].startswith("data:image/png;base64,")
|
|
|
|
|
|
def test_worker_falls_back_when_video_encoder_missing():
|
|
module = _load_collector_module()
|
|
collector = module.DistributedCollectorNode()
|
|
session = _RecordingSession()
|
|
_install_fake_transport(module, session)
|
|
codec = module.transfer_codec
|
|
original = codec.video_codec_available
|
|
codec.video_codec_available = lambda fmt: fmt == "h264" and original("h264")
|
|
try:
|
|
_run_worker(collector, torch.rand(2, 16, 16, 3), transfer_format="av1")
|
|
finally:
|
|
codec.video_codec_available = original
|
|
|
|
assert session.calls[0]["data"]["envelope"]["format"] == "h264"
|
|
|
|
|
|
def test_worker_does_not_fall_back_on_non_legacy_rejection():
|
|
import pytest
|
|
|
|
module = _load_collector_module()
|
|
collector = module.DistributedCollectorNode()
|
|
session = _RecordingSession(statuses=[400], text="Invalid multipart payload: Request Entity Too Large")
|
|
_install_fake_transport(module, session)
|
|
|
|
with pytest.raises(RuntimeError, match="MAX_REQUEST_BYTES"):
|
|
_run_worker(collector, torch.rand(2, 8, 8, 3), transfer_format="png")
|
|
|
|
# One compressed attempt, no legacy resend of the batch.
|
|
assert len(session.calls) == 1
|
|
assert session.calls[0]["json"] is None
|
|
|
|
|
|
def test_worker_sends_nothing_when_a_frame_cannot_fit():
|
|
import numpy as np
|
|
import pytest
|
|
|
|
module = _load_collector_module()
|
|
collector = module.DistributedCollectorNode()
|
|
session = _RecordingSession()
|
|
_install_fake_transport(module, session)
|
|
reserve = module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
|
module.TRANSFER_MAX_REQUEST_BYTES = reserve + 32 * 1024
|
|
module.TRANSFER_MAX_CHUNK_BYTES = 32 * 1024
|
|
# Frame 0 is tiny; frame 2 is noise that cannot fit in 32 KiB as PNG.
|
|
images = torch.zeros(3, 128, 128, 3)
|
|
images[2] = torch.from_numpy(np.random.default_rng(1).random((128, 128, 3), dtype=np.float32))
|
|
|
|
with pytest.raises(module.transfer_codec.FrameTooLargeError):
|
|
_run_worker(collector, images, transfer_format="png")
|
|
|
|
assert session.calls == []
|
|
|
|
|
|
def test_worker_sends_audio_separately_when_it_would_overflow_last_request():
|
|
module = _load_collector_module()
|
|
collector = module.DistributedCollectorNode()
|
|
session = _RecordingSession()
|
|
_install_fake_transport(module, session)
|
|
reserve = module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
|
module.TRANSFER_MAX_REQUEST_BYTES = reserve + 64 * 1024
|
|
module.TRANSFER_MAX_CHUNK_BYTES = 64 * 1024
|
|
module.encode_audio_payload = lambda value: None if value is None else {"data": "x" * (128 * 1024)}
|
|
|
|
_run_worker(
|
|
collector,
|
|
torch.zeros(2, 16, 16, 3),
|
|
audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
|
|
transfer_format="png",
|
|
)
|
|
|
|
multipart = [call["data"]["envelope"] for call in session.calls if call["data"] is not None]
|
|
assert [e["is_last"] for e in multipart] == [False]
|
|
assert all("audio" not in e for e in multipart)
|
|
final = session.calls[-1]["json"]
|
|
assert final["is_last"] is True and final["audio"]["data"].startswith("x")
|
|
assert "image" not in final
|
|
|
|
|
|
def test_worker_attaches_audio_to_last_chunk_when_it_fits():
|
|
module = _load_collector_module()
|
|
collector = module.DistributedCollectorNode()
|
|
session = _RecordingSession()
|
|
_install_fake_transport(module, session)
|
|
module.encode_audio_payload = lambda value: None if value is None else {"data": "small"}
|
|
|
|
_run_worker(
|
|
collector,
|
|
torch.zeros(2, 16, 16, 3),
|
|
audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
|
|
transfer_format="png",
|
|
)
|
|
|
|
assert len(session.calls) == 1
|
|
envelope = session.calls[0]["data"]["envelope"]
|
|
assert envelope["is_last"] is True and envelope["audio"] == {"data": "small"}
|