71 lines
3.1 KiB
Python
71 lines
3.1 KiB
Python
"""End policy reaches cloud exports and initial worker compatibility preflight."""
|
|
import asyncio
|
|
import copy
|
|
import unittest
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
import test_image_list_serverless as cloud
|
|
import test_image_list_master_selection as selection
|
|
|
|
FORMATS = ('exact', 'legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1')
|
|
|
|
|
|
@pytest.mark.parametrize('fmt', FORMATS)
|
|
def test_cloud_returns_use_selected_image_format_but_inputs_remain_exact(fmt):
|
|
async def check():
|
|
prompt = cloud.source()
|
|
prompt['ns/end']['inputs'].update(transfer_format=fmt, transfer_quality=73)
|
|
original = copy.deepcopy(prompt)
|
|
value = cloud.item(dtype=torch.float64)
|
|
job = await cloud.prepare(prompt, value)
|
|
assert prompt == original
|
|
for binding in job['manifest']:
|
|
controls = job['workflow'][binding['node_id']]['inputs']
|
|
if binding['kind'] == 'IMAGE' and fmt != 'exact':
|
|
assert controls['transfer_format'] == fmt
|
|
assert controls['transfer_quality'] == 73
|
|
else:
|
|
assert 'transfer_format' not in controls and 'transfer_quality' not in controls
|
|
for artifact in job['media']['artifacts']:
|
|
assert artifact['descriptor']['encoding'] == 'tensor-le-v1'
|
|
inputs = await cloud.decoded_inputs(job)
|
|
for ref in job['media']['inputs']:
|
|
assert torch.equal(inputs[ref['artifact_id']], value[ref['input_name']])
|
|
cloud.handler._validate_request({k: job[k] for k in ('workflow', 'images', 'media')})
|
|
asyncio.run(check())
|
|
|
|
|
|
@pytest.mark.parametrize('fmt,quality', [('JPEG',90),('jpeg',True),('jpeg',101),('jpeg','90')])
|
|
def test_cloud_rejects_invalid_policy(fmt, quality):
|
|
prompt = cloud.source()
|
|
prompt['ns/end']['inputs'].update(transfer_format=fmt, transfer_quality=quality)
|
|
with pytest.raises(ValueError):
|
|
asyncio.run(cloud.prepare(prompt))
|
|
|
|
|
|
class TransferPreflightTests(unittest.IsolatedAsyncioTestCase):
|
|
asyncSetUp = selection.MasterSelectionTests.asyncSetUp
|
|
queue = selection.MasterSelectionTests.queue
|
|
|
|
async def test_queue_requests_image_codec_capability(self):
|
|
self.prompt['3']['inputs'].update(transfer_format='jpeg', transfer_quality=73)
|
|
await self.queue()
|
|
self.assertEqual(self.preflight.call_args.kwargs['required_transfer_formats'], ['jpeg'])
|
|
|
|
async def test_invalid_policy_fails_before_worker_probe_or_queue(self):
|
|
self.prompt['3']['inputs'].update(transfer_format='bad-codec')
|
|
with self.assertRaises(ValueError):
|
|
await self.queue()
|
|
self.q.select_active_workers.assert_not_awaited()
|
|
self.preflight.assert_not_awaited()
|
|
self.q.queue_prompt_payload.assert_not_awaited()
|
|
|
|
async def test_missing_capability_cannot_silently_fall_back_to_master(self):
|
|
self.prompt['3']['inputs'].update(transfer_format='jpeg')
|
|
self.preflight.side_effect = RuntimeError('upgrade worker for image return encoding')
|
|
with self.assertRaisesRegex(RuntimeError, 'upgrade'):
|
|
await self.queue()
|
|
self.q.queue_prompt_payload.assert_not_awaited()
|