166 lines
8.0 KiB
Python
166 lines
8.0 KiB
Python
"""Sparse End overrides through the standalone serverless bridge, CPU only."""
|
|
import asyncio
|
|
import copy
|
|
import math
|
|
import os
|
|
from pathlib import Path
|
|
import struct
|
|
import subprocess
|
|
import sys
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from unittest.mock import patch
|
|
|
|
from test_image_list_serverless import adapter, codec, contract, handler, item, prepare, source
|
|
|
|
|
|
@pytest.mark.skipif(not os.environ.get('COMFYUI_SOURCE_ROOT'), reason='set COMFYUI_SOURCE_ROOT for real CPU executor')
|
|
def test_real_cpu_standalone_bridge_calculated_overrides(tmp_path):
|
|
env = dict(os.environ, CUDA_VISIBLE_DEVICES='', HIP_VISIBLE_DEVICES='', ROCR_VISIBLE_DEVICES='',
|
|
PYTHONDONTWRITEBYTECODE='1')
|
|
env.pop('PYTHONPATH', None)
|
|
helper = Path(__file__).parent / 'helpers' / 'image_list_serverless_values_check.py'
|
|
result = subprocess.run([sys.executable, str(helper), os.environ['COMFYUI_SOURCE_ROOT'], str(tmp_path)],
|
|
env=env, text=True, capture_output=True, timeout=90)
|
|
assert result.returncode == 0, result.stdout + result.stderr
|
|
assert 'IMAGE_LIST_SERVERLESS_VALUES_CPU_OK' in result.stdout
|
|
|
|
|
|
def assert_exact(actual, expected):
|
|
assert type(actual) is type(expected)
|
|
if isinstance(expected, torch.Tensor):
|
|
assert actual.dtype == expected.dtype and actual.shape == expected.shape
|
|
assert torch.equal(actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
|
|
elif isinstance(expected, float):
|
|
assert struct.pack('!d', actual) == struct.pack('!d', expected)
|
|
else:
|
|
assert actual == expected
|
|
|
|
|
|
@pytest.mark.parametrize('value', [False, True, 2**63 + 123, -(2**100), -0.0, 1.0000000000000002,
|
|
'', ' \r\n雪 ', torch.ones(2, 2, 3, 4, dtype=torch.float64),
|
|
torch.ones(2, 2, 3, dtype=torch.bfloat16)])
|
|
def test_value_codec_exact_types(value):
|
|
descriptor, payload = codec.encode_payload('VALUE', value)
|
|
assert_exact(codec.decode_payload('VALUE', descriptor, payload), value)
|
|
|
|
|
|
@pytest.mark.parametrize('value', [None, {}, [], (1,), object(), float('inf'), float('nan'),
|
|
torch.ones(2), torch.ones(1, 2, 3, dtype=torch.int64),
|
|
torch.full((1, 2, 3), float('nan'))])
|
|
def test_value_codec_rejects_unsupported_or_nonfinite(value):
|
|
with pytest.raises((ValueError, TypeError)):
|
|
codec.encode_payload('VALUE', value)
|
|
|
|
|
|
def test_sparse_exports_never_create_collector_or_forward_unchanged_values():
|
|
async def check():
|
|
prompt, value = source(), item()
|
|
prompt['ns/end']['inputs'].update(value_2=['ns/start', 5], value_10=['process', 2])
|
|
original = copy.deepcopy(prompt)
|
|
job = await prepare(prompt, value)
|
|
assert prompt == original
|
|
bindings = {binding['socket']: binding for binding in job['manifest']}
|
|
assert set(bindings) == {'images', 'masks', 'texts', 'value_2', 'value_10'}
|
|
assert bindings['value_2']['kind'] == bindings['value_10']['kind'] == 'VALUE'
|
|
assert all('Collector' not in node['class_type'] for node in job['workflow'].values())
|
|
assert all(node['class_type'] not in (adapter.runtime.START_CLASS, adapter.runtime.END_CLASS)
|
|
for node in job['workflow'].values())
|
|
handler._validate_request({key: job[key] for key in ('workflow', 'images', 'media')})
|
|
asyncio.run(check())
|
|
|
|
|
|
@pytest.mark.parametrize('slot', ['value_0', 'value_11', 'value_01', 'values.value_1'])
|
|
def test_unknown_end_override_slot_rejected(slot):
|
|
prompt = source()
|
|
prompt['ns/end']['inputs'][slot] = ['ns/start', 4]
|
|
with pytest.raises(ValueError, match='slot'):
|
|
asyncio.run(prepare(prompt))
|
|
|
|
|
|
def test_value_codec_budgets_and_malformed_payload():
|
|
with pytest.raises(ValueError, match='byte limit'):
|
|
codec.encode_payload('VALUE', '雪', max_bytes=2)
|
|
descriptor, payload = codec.encode_payload('VALUE', True)
|
|
with pytest.raises(ValueError):
|
|
codec.decode_payload('VALUE', dict(descriptor, type='INT'), payload)
|
|
descriptor, payload = codec.encode_payload('VALUE', -0.0)
|
|
with pytest.raises(ValueError):
|
|
codec.decode_payload('VALUE', descriptor, struct.pack('<d', math.inf))
|
|
with pytest.raises(ValueError):
|
|
codec.decode_payload('VALUE', descriptor, payload + b'0')
|
|
with pytest.raises(ValueError, match='byte limit'):
|
|
codec.decode_payload('VALUE', descriptor, payload, max_bytes=7)
|
|
|
|
|
|
@pytest.mark.parametrize('fmt', ['exact', 'legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1'])
|
|
def test_end_image_policy_never_compresses_override_tensors(fmt):
|
|
prompt = source()
|
|
prompt['ns/end']['inputs'].update(value_10=['ns/start', 0], transfer_format=fmt, transfer_quality=73)
|
|
job = asyncio.run(prepare(prompt))
|
|
for binding in job['manifest']:
|
|
controls = job['workflow'][binding['node_id']]['inputs']
|
|
if binding['socket'] == 'images' and fmt != 'exact':
|
|
assert controls['transfer_format'] == fmt and controls['transfer_quality'] == 73
|
|
else:
|
|
assert 'transfer_format' not in controls and 'transfer_quality' not in controls
|
|
default = asyncio.run(prepare())
|
|
assert all('transfer_format' not in n['inputs'] and 'transfer_quality' not in n['inputs']
|
|
for n in default['workflow'].values())
|
|
|
|
|
|
def value_response():
|
|
binding = dict(binding_id='binding', collector_id='end', job_id='job', participant_id='pool',
|
|
node_id='export', socket='value_10', kind='VALUE')
|
|
descriptor, payload = codec.encode_payload('VALUE', 2**63 + 7)
|
|
artifact = contract.artifact_from_payload('artifact', 'VALUE', descriptor, payload)
|
|
return binding, dict(version=1, profile=contract.PROFILE, request_id='request',
|
|
bindings=[dict(binding, count=1)], artifacts=[artifact],
|
|
items=[dict(binding_id='binding', item_index=0, artifact_id='artifact')])
|
|
|
|
|
|
@pytest.mark.parametrize('bad', ['missing', 'extra', 'wrong_kind', 'multiple', 'extra_artifact', 'nonfinite'])
|
|
def test_override_response_rejected_atomically(bad):
|
|
binding, response = value_response()
|
|
if bad == 'missing':
|
|
response['bindings'] = []
|
|
elif bad == 'extra':
|
|
response['bindings'].append(dict(binding, binding_id='other', socket='value_9', count=1))
|
|
elif bad == 'wrong_kind':
|
|
response['artifacts'][0]['kind'] = 'STRING'
|
|
elif bad == 'multiple':
|
|
response['bindings'][0]['count'] = 2
|
|
elif bad == 'extra_artifact':
|
|
response['artifacts'].append(dict(response['artifacts'][0], id='extra'))
|
|
else:
|
|
descriptor, _ = codec.encode_payload('VALUE', 1.0)
|
|
response['artifacts'][0] = contract.artifact_from_payload(
|
|
'artifact', 'VALUE', descriptor, struct.pack('<d', float('nan')))
|
|
with pytest.raises(ValueError):
|
|
asyncio.run(contract.decode_response(response, request_id='request', bindings=[binding]))
|
|
|
|
|
|
def test_override_request_and_result_budgets():
|
|
prompt = source()
|
|
prompt['ns/end']['inputs']['value_10'] = '雪' * 100
|
|
with patch.object(contract, 'MAX_ARTIFACT_BYTES', 128):
|
|
with pytest.raises(ValueError, match='byte limit'):
|
|
asyncio.run(prepare(prompt))
|
|
binding, response = value_response()
|
|
with pytest.raises(ValueError, match='total byte limit'):
|
|
asyncio.run(contract.decode_response(response, request_id='request', bindings=[binding], max_total_bytes=1))
|
|
|
|
|
|
def test_collector_compiler_keeps_only_its_public_sockets():
|
|
from importlib import import_module
|
|
compiler = import_module(adapter.__package__.rsplit('.', 1)[0] + '.api.orchestration.serverless_media_prompt')
|
|
assert compiler.COLLECTOR_SOCKETS == ('images', 'audio')
|
|
graph = {'source': {'class_type': 'Source', 'inputs': {}},
|
|
'collector': {'class_type': 'DistributedCollector', 'inputs': {
|
|
'images': ['source', 0], 'value_10': ['source', 1], 'masks': ['source', 2],
|
|
'multi_job_id': 'job', 'worker_id': 'pool'}}}
|
|
_, bindings = compiler.compile_media_prompt(graph, 'request')
|
|
assert [binding['socket'] for binding in bindings] == ['images']
|