def test_bad_heap_address_bits(self):
with pytest.raises(ValueError):
spead2.Flavour(4, 64, 43, 0) # Not multiple of 8
with pytest.raises(ValueError):
spead2.Flavour(4, 64, 0, 0)
with pytest.raises(ValueError):
spead2.Flavour(4, 64, 64, 0)
def test_equality(self):
flavour1 = spead2.Flavour(4, 64, 40, 2)
flavour1b = spead2.Flavour(4, 64, 40, 2)
flavour2 = spead2.Flavour(4, 64, 48, 2)
flavour3 = spead2.Flavour(4, 64, 40, 4)
assert_true(flavour1 == flavour1b)
assert_true(flavour1 != flavour2)
assert_false(flavour1 != flavour1b)
assert_false(flavour1 == flavour2)
assert_false(flavour1 == flavour3)
def test_equality(self):
flavour1 = spead2.Flavour(4, 64, 40, 2)
flavour1b = spead2.Flavour(4, 64, 40, 2)
flavour2 = spead2.Flavour(4, 64, 48, 2)
flavour3 = spead2.Flavour(4, 64, 40, 4)
assert flavour1 == flavour1b
assert flavour1 != flavour2
assert not (flavour1 != flavour1b)
assert not (flavour1 == flavour2)
assert not (flavour1 == flavour3)
def _create_streams(self):
"""Construct streams, item group and item descriptions."""
# Construct the SPEAD flavour description
log = logging.getLogger(__name__)
# Define the SPEAD flavour
flavour_config = self._config.get('spead_flavour')
flavour = spead2.Flavour(flavour_config.get('version', 4),
flavour_config.get('item_pointer_bits', 64),
flavour_config.get('heap_address_bits', 40),
flavour_config.get('bug_compat_mask', 0))
log.debug('Using SPEAD flavour = SPEAD-%d-%d v%d compat:%d',
flavour.item_pointer_bits, flavour.heap_address_bits,
flavour.version, flavour.bug_compat)
# Create SPEAD streams
streams_config = self._config.get('sender_node').get('streams')
descriptor = get_heap_descriptor(self._frame_shape)
streams = list()
for i, config in enumerate(streams_config):
log.debug('Configuring stream %d:', i)
log.debug(' Address = %s:%d', config.get('host'),
config.get('port'))
log.debug(' Threads = %d', config.get('threads', 1))
streams.append(self._create_stream(config, flavour, descriptor))
return streams
def __init__(self, spead_config, precision=None, oskar_settings=None):
logger.info("Creating OSKAR interferometer")
oskar.Interferometer.__init__(self, precision, oskar_settings)
self._streams = []
self._vis_pack = None
self._write_ms = spead_config['write_ms']
# Construct UDP streams and associated item groups.
logger.info("Creating SPEAD2 stream config")
stream_config = spead2.send.StreamConfig(
spead_config['stream_config']['max_packet_size'],
spead_config['stream_config']['rate'],
spead_config['stream_config']['burst_size'],
spead_config['stream_config']['max_heaps'])
stream = spead_config['stream']
threads = stream['threads'] if 'threads' in stream else 1
logger.info("Creating SPEAD2 Thread pool")
thread_pool = spead2.ThreadPool(threads=threads)
logger.info("Creating SPEAD stream for host {} on port {}".format(
stream['host'], stream['port']))
udp_stream = spead2.send.TcpStream(thread_pool, stream['host'],
stream['port'], stream_config)
item_group = spead2.send.ItemGroup(
flavour=spead2.Flavour(4, 64, 40, 0))
# Append udp_stream and item_group to the stream list as a tuple.
self._streams.append((udp_stream, item_group))
self.abort = False
def _create_streams(self):
"""Construct streams, item group and item descriptions."""
# Construct the SPEAD flavour description
parent = 'spead_flavour'
version = self._get_config([parent, 'version'], 4)
item_pointer_bits = self._get_config([parent, 'item_pointer_bits'], 64)
heap_address_bits = self._get_config([parent, 'heap_address_bits'], 40)
bug_compat_mask = self._get_config([parent, 'bug_compat_mask'], 0)
flavour = spead2.Flavour(version, item_pointer_bits, heap_address_bits,
bug_compat_mask)
# Construct UDP stream objects and associated heap_descriptor item
# groups.
streams = list()
for i, stream in enumerate(self._config['sender_node']['streams']):
host = stream['host']
port = stream['port']
threads = stream['threads'] if 'threads' in stream else 1
stream_config = spead2.send.StreamConfig(rate=0)
thread_pool = spead2.ThreadPool(threads=threads)
# Blocking send
stream = spead2.send.UdpStream(thread_pool, host, port,
stream_config)
item_group = spead2.send.ItemGroup(flavour=flavour)
# Append stream & item group the stream list.
streams.append((stream, item_group))
self._log.debug('Configuring stream {}:'.format(i))
self._log.debug(' Address = {}:{}'.format(host, port))
self._log.debug(' Flavour = SPEAD-{}-{} v{} compat:{}'.format(
flavour.item_pointer_bits, flavour.heap_address_bits,
flavour.version, flavour.bug_compat))
self._log.debug(' Threads = {}'.format(threads))
# Add items to the item group based on the heap descriptor.
for key, item in self._heap_descriptor.items():
item_id = item['id']
if isinstance(item_id, str):
item_id = int(item_id, 0)
name = key
desc = item['description']
item_shape = item['shape'] if 'shape' in item else tuple()
item_type = item['type'] if 'type' in item else None
item_format = item['format'] if 'format' in item else None
item_group.add_item(item_id,
name,
desc,
shape=item_shape,
dtype=item_type,
format=item_format)
self._log.debug('Adding item: {} {}'.format(item_id, name))
self._log.debug(' description = {}'.format(desc))
if item_type is not None:
self._log.debug(' type = {}'.format(item_type))
if item_format is not None:
self._log.debug(' format = {}'.format(item_format))
self._log.debug(' shape = {}'.format(item_shape))
self._streams = streams
def measure_connection_once(args, rate, num_heaps, required_heaps):
reader, writer = yield From(trollius.open_connection(args.host, args.port))
def write(s):
writer.write(s.encode('ascii'))
write(json.dumps({'cmd': 'start', 'args': vars(args)}) + '\n')
# Wait for "ready" response
response = yield From(reader.readline())
assert response == b'ready\n'
if args.send_affinity is not None and len(args.send_affinity) > 0:
spead2.ThreadPool.set_affinity(args.send_affinity[0])
thread_pool = spead2.ThreadPool(
1, args.send_affinity[1:] + args.send_affinity[:1])
else:
thread_pool = spead2.ThreadPool()
thread_pool = spead2.ThreadPool()
config = spead2.send.StreamConfig(max_packet_size=args.packet,
burst_size=args.burst,
rate=rate,
max_heaps=num_heaps + 1,
burst_rate_ratio=args.burst_rate_ratio)
host = args.host
if args.multicast is not None:
host = args.multicast
if 'send_ibv' in args and args.send_ibv is not None:
stream = spead2.send.trollius.UdpIbvStream(
thread_pool, host, args.port, config, args.send_ibv,
args.send_buffer, 1, args.send_ibv_vector, args.send_ibv_max_poll)
else:
stream = spead2.send.trollius.UdpStream(thread_pool, host, args.port,
config, args.send_buffer)
item_group = spead2.send.ItemGroup(
flavour=spead2.Flavour(4, 64, args.addr_bits, 0))
item_group.add_item(id=None,
name='Test item',
description='A test item with arbitrary value',
shape=(args.heap_size, ),
dtype=np.uint8,
value=np.zeros((args.heap_size, ), dtype=np.uint8))
start = timeit.default_timer()
transferred = yield From(send_stream(item_group, stream, num_heaps))
end = timeit.default_timer()
elapsed = end - start
actual_rate = transferred / elapsed
# Give receiver time to catch up with any queue
yield From(trollius.sleep(0.1))
write(json.dumps({'cmd': 'stop'}) + '\n')
# Read number of heaps received
response = yield From(reader.readline())
response = json.loads(response.decode('ascii'))
received_heaps = response['received_heaps']
yield From(trollius.sleep(0.5))
yield From(writer.drain())
writer.close()
logging.debug("Received %d/%d heaps in %f seconds, rate %.3f Gbps",
received_heaps, num_heaps, elapsed, actual_rate * 8e-9)
raise Return(received_heaps >= required_heaps, actual_rate)
def main():
"""Runs the test sender."""
stream_config = spead2.send.StreamConfig(
max_packet_size=16356, rate=1000e6, burst_size=10, max_heaps=1)
item_group = spead2.send.ItemGroup(flavour=spead2.Flavour(4, 64, 48, 0))
# Add item descriptors to the heap.
num_baselines = (512 * 513) // 2
dtype = [('TCI', 'i1'), ('FD', 'u1'), ('VIS', '<c8', 4)]
item_group.add_item(
id=0x6000, name='visibility_timestamp_count', description='',
shape=tuple(), format=None, dtype='<u4')
item_group.add_item(
id=0x6001, name='visibility_timestamp_fraction', description='',
shape=tuple(), format=None, dtype='<u4')
item_group.add_item(
id=0x6005, name='visibility_baseline_count', description='',
shape=tuple(), format=None, dtype='<u4')
item_group.add_item(
id=0x6008, name='scan_id', description='',
shape=tuple(), format=None, dtype='<u8')
item_group.add_item(
id=0x600A, name='correlator_output_data', description='',
shape=(num_baselines,), dtype=dtype)
# Create streams and send start-of-stream message.
streams = []
num_streams = 2
for i in range(num_streams):
stream = spead2.send.UdpStream(
thread_pool=spead2.ThreadPool(threads=1),
hostname='127.0.0.1', port=41000 + i, config=stream_config)
stream.send_heap(item_group.get_start())
streams.append(stream)
vis = numpy.zeros(shape=(num_baselines,), dtype=dtype)
num_heaps = 200
start_time = time.time()
for stream in streams:
# Update values in the heap.
item_group['visibility_timestamp_count'].value = 1
item_group['visibility_timestamp_fraction'].value = 0
item_group['visibility_baseline_count'].value = num_baselines
item_group['scan_id'].value = 100000000
item_group['correlator_output_data'].value = vis
# Iterate heaps.
for i in range(num_heaps):
# Send heap.
stream.send_heap(item_group.get_heap(descriptors='all', data='all'))
# Print time taken.
duration = time.time() - start_time
data_size = num_streams * num_heaps * (vis.nbytes / 1e6)
print("Sent %.3f MB in %.3f sec (%.3f MB/sec)" % (
data_size, duration, (data_size/duration)))
# Send end-of-stream message.
for stream in streams:
stream.send_heap(item_group.get_end())
def _create_streams(config, log):
"""Construct streams, item group and item descriptions."""
# Construct the SPEAD flavour description
# Flavour(version, item_pointer_bits, heap_address_bits, bug_compat_mask)
version = config['spead_flavour']['version']
item_pointer_bits = config['spead_flavour']['item_pointer_bits']
heap_address_bits = config['spead_flavour']['heap_address_bits']
bug_compat_mask = config['spead_flavour']['bug_compat_mask']
flavour = spead2.Flavour(version, item_pointer_bits, heap_address_bits,
bug_compat_mask)
# Construct UDP stream objects and associated heap item groups.
streams = list()
for i, stream in enumerate(config['sender_node']['streams']):
host = stream['host']
port = stream['port']
rate = stream['rate'] if 'rate' in stream else 0
threads = stream['threads'] if 'threads' in stream else 1
stream_config = spead2.send.StreamConfig(rate=rate)
thread_pool = spead2.ThreadPool(threads=threads)
stream = spead2.send.UdpStream(thread_pool, host, port, stream_config)
item_group = spead2.send.ItemGroup(flavour=flavour)
# Append stream & item group the stream list.
streams.append((stream, item_group))
log.debug('Configuring stream {}:'.format(i))
log.debug(' Address = {}:{}'.format(host, port))
log.debug(' Flavour = SPEAD-{}-{} v{} compat:{}'.format(
flavour.item_pointer_bits, flavour.heap_address_bits,
flavour.version, flavour.bug_compat))
log.debug(' Threads = {}'.format(threads))
log.debug(' Rate = {}'.format(rate))
# Add items to the item group based on the heap description.
for j, item in enumerate(__heap__):
item_id = item['id']
if isinstance(item_id, str):
item_id = int(item_id, 0)
name = item['name']
desc = item['description']
item_type = item['type'] if 'type' in item else None
item_format = item['format'] if 'format' in item else None
shape = item['shape']
item_group.add_item(item_id,
name,
desc,
shape=shape,
dtype=item_type,
format=item_format)
log.debug('Adding item {} : {} {}'.format(j, item_id, name))
log.debug(' description = {}'.format(desc))
if item_type is not None:
log.debug(' type = {}'.format(item_type))
if item_format is not None:
log.debug(' format = {}'.format(item_format))
log.debug(' shape = {}'.format(shape))
return streams
def main():
args = get_args()
logging.basicConfig(level=getattr(logging, args.log.upper()))
dtype = np.dtype(args.dtype)
elements = args.heap_size // (args.items * dtype.itemsize)
heap_size = elements * args.items * dtype.itemsize
if heap_size != args.heap_size:
logging.warn(
'Heap size is not an exact multiple: using %d instead of %d',
heap_size, args.heap_size)
bug_compat = spead2.BUG_COMPAT_PYSPEAD_0_5_2 if args.pyspead else 0
item_group = spead2.send.ItemGroup(descriptor_frequency=args.descriptors,
flavour=spead2.Flavour(
4, 64, args.addr_bits, bug_compat))
for i in range(args.items):
item_group.add_item(id=None,
name='Test item {}'.format(i),
description='A test item with arbitrary value',
shape=(elements, ),
dtype=dtype,
value=np.zeros((elements, ), dtype=dtype))
if args.affinity is not None and len(args.affinity) > 0:
spead2.ThreadPool.set_affinity(args.affinity[0])
thread_pool = spead2.ThreadPool(args.threads,
args.affinity[1:] + args.affinity[:1])
else:
thread_pool = spead2.ThreadPool(args.threads)
config = spead2.send.StreamConfig(max_packet_size=args.packet,
burst_size=args.burst,
rate=args.rate * 10**9 / 8,
burst_rate_ratio=args.burst_rate_ratio)
if 'ibv' in args and args.ibv is not None:
stream = spead2.send.trollius.UdpIbvStream(thread_pool, args.host,
args.port, config, args.ibv,
args.buffer, args.ttl or 1,
args.ibv_vector,
args.ibv_max_poll)
elif args.ttl is not None:
stream = spead2.send.trollius.UdpStream(thread_pool,
args.host,
args.port,
config,
args.buffer,
ttl=args.ttl)
else:
# This is handled as a separate case, because passing TTL is only
# valid for multicast addresses.
stream = spead2.send.trollius.UdpStream(thread_pool, args.host,
args.port, config, args.buffer)
try:
trollius.get_event_loop().run_until_complete(
run(item_group, stream, args))
except KeyboardInterrupt:
sys.exit(1)
def __init__(self,
item_group,
descriptor_frequency=None,
flavour=_spead2.Flavour()):
self._item_group = item_group
self._info = {} # Maps ID to _ItemInfo
self._descriptor_frequency = descriptor_frequency
# Counter for calls to add_to_heap. This is independent of the
# protocol-level heap count.
self._descriptor_cnt = 0
self._flavour = flavour
def __init__(self,
item_group,
descriptor_frequency=None,
flavour=_spead2.Flavour()):
# Workaround to avoid creating a self-reference when it's bundled into
# the same class.
self._item_group = item_group if item_group is not self else None
self._info = {} # Maps ID to _ItemInfo
self._descriptor_frequency = descriptor_frequency
# Counter for calls to add_to_heap. This is independent of the
# protocol-level heap count.
self._descriptor_cnt = 0
self._flavour = flavour
async def async_main():
args = get_args()
logging.basicConfig(level=getattr(logging, args.log.upper()))
dtype = np.dtype(args.dtype)
elements = args.heap_size // (args.items * dtype.itemsize)
heap_size = elements * args.items * dtype.itemsize
if heap_size != args.heap_size:
logging.warn('Heap size is not an exact multiple: using %d instead of %d',
heap_size, args.heap_size)
bug_compat = spead2.BUG_COMPAT_PYSPEAD_0_5_2 if args.pyspead else 0
item_group = spead2.send.ItemGroup(
descriptor_frequency=args.descriptors,
flavour=spead2.Flavour(4, 64, args.addr_bits, bug_compat))
for i in range(args.items):
item_group.add_item(id=None, name='Test item {}'.format(i),
description='A test item with arbitrary value',
shape=(elements,), dtype=dtype,
value=np.zeros((elements,), dtype=dtype))
if args.affinity is not None and len(args.affinity) > 0:
spead2.ThreadPool.set_affinity(args.affinity[0])
thread_pool = spead2.ThreadPool(args.threads, args.affinity[1:] + args.affinity[:1])
else:
thread_pool = spead2.ThreadPool(args.threads)
config = spead2.send.StreamConfig(
max_packet_size=args.packet,
burst_size=args.burst,
rate=args.rate * 10**9 / 8,
burst_rate_ratio=args.burst_rate_ratio,
max_heaps=args.max_heaps)
if args.tcp:
stream = await spead2.send.asyncio.TcpStream.connect(
thread_pool, args.host, args.port, config, args.buffer, args.bind)
elif 'ibv' in args and args.ibv:
stream = spead2.send.asyncio.UdpIbvStream(
thread_pool, args.host, args.port, config, args.bind,
args.buffer, args.ttl or 1, args.ibv_vector, args.ibv_max_poll)
else:
kwargs = {}
if args.ttl is not None:
kwargs['ttl'] = args.ttl
if args.bind:
kwargs['interface_address'] = args.bind
stream = spead2.send.asyncio.UdpStream(
thread_pool, args.host, args.port, config, args.buffer, **kwargs)
await run(item_group, stream, args)
def __init__(self, spead_config):
self._config = spead_config
self._log = logging.getLogger('sip.sender')
self._streams = []
self._buffer = []
# Construct UDP streams and associated item groups.
stream_config = spead2.send.StreamConfig(
self._config['stream_config']['max_packet_size'],
self._config['stream_config']['rate'],
self._config['stream_config']['burst_size'],
self._config['stream_config']['max_heaps'])
for i_stream in range(self._config['num_streams']):
host = self._config['destination_host']
port = self._config['destination_port_start'] + i_stream
# It's much faster to have a thread pool of one thread per stream!
thread_pool = spead2.ThreadPool(threads=1)
self._log.info('Creating SPEAD stream on %s:%i ...', host, port)
udp_stream = spead2.send.asyncio.UdpStream(thread_pool, host, port,
stream_config)
item_group = spead2.send.ItemGroup(
flavour=spead2.Flavour(4, 64, 48, 0))
self._streams.append((udp_stream, item_group))
def __init__(self, log, spead_config, precision=None, oskar_settings=None):
oskar.Interferometer.__init__(self, precision, oskar_settings)
self._log = log
self._streams = []
self._vis_pack = None
# Construct UDP streams and associated item groups.
stream_config = spead2.send.StreamConfig(
spead_config['stream_config']['max_packet_size'],
spead_config['stream_config']['rate'],
spead_config['stream_config']['burst_size'],
spead_config['stream_config']['max_heaps'])
for stream in spead_config['streams']:
threads = stream['threads'] if 'threads' in stream else 1
thread_pool = spead2.ThreadPool(threads=threads)
log.info("Creating SPEAD stream for host {} on port {} ...".format(
stream['host'], stream['port']))
udp_stream = spead2.send.UdpStream(thread_pool, stream['host'],
stream['port'], stream_config)
item_group = spead2.send.ItemGroup(
flavour=spead2.Flavour(4, 64, 40, 0))
# Append udp_stream and item_group to the stream list as a tuple.
self._streams.append((udp_stream, item_group))
def test_nonascii_description(self):
"""Description with non-ASCII characters must fail"""
with pytest.raises(UnicodeEncodeError):
item = spead2.Item(0x1000, 'name', '\u0200', (), np.int32)
item.to_raw(spead2.Flavour())
# later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more
# details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
import spead2
import spead2.send
import sys
import logging
import numpy as np
logging.basicConfig(level=logging.INFO)
thread_pool = spead2.ThreadPool()
stream = spead2.send.UdpStream(thread_pool, "127.0.0.1", 8888,
spead2.send.StreamConfig(rate=1e7))
del thread_pool
shape = (40, 50)
ig = spead2.send.ItemGroup(
flavour=spead2.Flavour(4, 64, 48, spead2.BUG_COMPAT_PYSPEAD_0_5_2))
item = ig.add_item(0x1234, 'foo', 'a foo item', shape=shape, dtype=np.int32)
item.value = np.zeros(shape, np.int32)
stream.send_heap(ig.get_heap())
stream.send_heap(ig.get_end())
class TestPassthroughLegacyReceive64_48(BaseTestPassthroughLegacyReceive):
spead = spead64_48
flavour = spead2.Flavour(4, 64, 48, spead2.BUG_COMPAT_PYSPEAD_0_5_2)
def main():
"""Runs the test sender."""
num_stations = 4
num_heaps = 1500
num_streams = 1
rate = 1.0e5 # Bytes/s
target = ('127.0.0.1' if len(sys.argv) < 2 else sys.argv[1])
target_port = int('41000' if len(sys.argv) < 3 else sys.argv[2])
antenna_file = ('' if len(sys.argv) < 4 else sys.argv[3])
if (len(antenna_file) > 1):
coords = read_coordinates(antenna_file)
num_stations = len(coords)
num_baselines = (num_stations * (num_stations + 1)) // 2
print(f'no. stations : {num_stations}')
print(f'no. baselines : {num_baselines}')
print(f'no. times (heaps) : {num_heaps}')
print(f'host : {target}')
print(f'port : {target_port}')
stream_config = spead2.send.StreamConfig(max_packet_size=16356,
rate=rate,
burst_size=10,
max_heaps=1)
item_group = spead2.send.ItemGroup(flavour=spead2.Flavour(4, 64, 48, 0))
# Add item descriptors to the heap.
dtype = [('TCI', 'i1'), ('FD', 'u1'), ('VIS', '<c8', 4)]
item_group.add_item(id=0x6000,
name='visibility_timestamp_count',
description='',
shape=tuple(),
format=None,
dtype='<u4')
item_group.add_item(id=0x6001,
name='visibility_timestamp_fraction',
description='',
shape=tuple(),
format=None,
dtype='<u4')
item_group.add_item(id=0x6002,
name='visibility_channel_id',
description='',
shape=tuple(),
format=None,
dtype='<u4')
item_group.add_item(id=0x6003,
name='visibility_channel_count',
description='',
shape=tuple(),
format=None,
dtype='<u4')
item_group.add_item(id=0x6004,
name='visibility_polarisation_id',
description='',
shape=tuple(),
format=None,
dtype='<u4')
item_group.add_item(id=0x6005,
name='visibility_baseline_count',
description='',
shape=tuple(),
format=None,
dtype='<u4')
item_group.add_item(id=0x6008,
name='scan_id',
description='',
shape=tuple(),
format=None,
dtype='<u8')
item_group.add_item(id=0x600A,
name='correlator_output_data',
description='',
shape=(num_baselines, ),
dtype=dtype)
# Create streams and send start-of-stream message.
streams = []
for i in range(num_streams):
port = target_port + i
print("Sending to {}:{}".format(target, port))
stream = spead2.send.UdpStream(
thread_pool=spead2.ThreadPool(threads=1),
hostname=target,
port=port,
config=stream_config)
stream.send_heap(item_group.get_start())
streams.append(stream)
vis = numpy.zeros(shape=(num_baselines, ), dtype=dtype)
vis_amps = numpy.arange(num_baselines * 4, dtype='c8').reshape(
(num_baselines, 4)) / 1000.0
start_time = time.time()
for stream in streams:
# Update values in the heap.
item_group['visibility_timestamp_count'].value = 1
item_group['visibility_timestamp_fraction'].value = 0
item_group['visibility_baseline_count'].value = num_baselines
item_group['visibility_channel_id'].value = 12345
item_group['visibility_channel_count'].value = 0
item_group['visibility_polarisation_id'].value = 0
item_group['scan_id'].value = 100000000
item_group['correlator_output_data'].value = vis
# Iterate heaps.
for i in range(num_heaps):
item_group['correlator_output_data'].value['VIS'] = vis_amps + i
# Send heap.
stream.send_heap(item_group.get_heap(descriptors='all',
data='all'))
# Print time taken.
duration = time.time() - start_time
data_size = num_streams * num_heaps * (vis.nbytes / 1e6)
print("Sent %.3f MB in %.3f sec (%.3f MB/sec)" % (data_size, duration,
(data_size / duration)))
# Send end-of-stream message.
for stream in streams:
stream.send_heap(item_group.get_end())
async def _test_stream(self, end: bool, write: bool) -> None:
n_heaps = 30 # number of heaps in time
n_spectra = self.spectra_per_heap * n_heaps
# Pick some heaps to drop, including an entire slice and
# an entire channel for one stats dump
drop = np.zeros((self.n_bengs, n_heaps), np.bool_)
drop[:, 4] = True
drop[2, 9] = True
drop[7, 24] = True
drop[10, 12:18] = True
if not write:
self.args.file_base = None
# Start a receiver to get the signal display stream.
# It needs a deep queue because we don't service it while it is
# running.
rx = spead2.recv.Stream(
spead2.ThreadPool(),
spead2.recv.StreamConfig(max_heaps=2, stop_on_stop_item=False),
spead2.recv.RingStreamConfig(heaps=100))
rx.add_udp_reader(self.args.stats.host,
self.args.stats.port,
interface_address='127.0.0.1')
# Start up the server
server = bf_ingest_server.CaptureServer(self.args, self.loop)
filename = await server.start_capture('1122334455')
# Send it a SPEAD stream. Use small packets to ensure that each heap is
# split into multiple packets, to check that the data scatter works.
config = spead2.send.StreamConfig(max_packet_size=256)
flavour = spead2.Flavour(4, 64, 48, 0)
ig = spead2.send.ItemGroup(flavour=flavour)
ig.add_item(name='timestamp',
id=0x1600,
description='Timestamp',
shape=(),
format=[('u', 48)])
ig.add_item(
name='frequency',
id=0x4103,
description='The frequency channel of the data in this HEAP.',
shape=(),
format=[('u', 48)])
ig.add_item(name='bf_raw',
id=0x5000,
description='Beamformer data',
shape=(self.channels_per_heap, self.spectra_per_heap, 2),
dtype=np.dtype(np.int8))
# To guarantee in-order delivery (and hence make the test
# reliable/reproducible), we send all the data for the channels of
# interest through a single TCP socket. Data for channels outside the
# subscribed range is discarded. Note that giving multiple TcpStream's
# the same socket is dangerous because individual write() calls could
# interleave; it's safe only because we use only blocking calls so
# there is no concurrency between the streams.
subscribed_streams = self.args.channels // self.channels_per_endpoint
subscribed_bengs = self.args.channels // self.channels_per_heap
expected_heaps = 0
streams = [] # type: List[Optional[spead2.send.TcpStream]]
primary_ep = self.tcp_endpoints[subscribed_streams.start]
sock = socket.socket()
sock.setblocking(False)
await self.loop.sock_connect(sock, (primary_ep.host, primary_ep.port))
for i in range(len(self.endpoints)):
if i not in subscribed_streams:
streams.append(None)
continue
stream = spead2.send.TcpStream(spead2.ThreadPool(), sock, config)
streams.append(stream)
stream.set_cnt_sequence(i, len(self.endpoints))
stream.send_heap(ig.get_heap(descriptors='all'))
stream.send_heap(ig.get_start())
expected_heaps += 2
sock.close()
ts = 1234567890
for i in range(n_heaps):
data = np.zeros((self.n_channels, self.spectra_per_heap, 2),
np.int8)
for channel in range(self.n_channels):
data[channel, :, 0] = channel % 255 - 128
for t in range(self.spectra_per_heap):
data[:, t, 1] = (i * self.spectra_per_heap + t) % 255 - 128
for j in range(self.n_bengs):
ig['timestamp'].value = ts
ig['frequency'].value = j * self.channels_per_heap
ig['bf_raw'].value = data[j * self.channels_per_heap:(j + 1) *
self.channels_per_heap, ...]
if not drop[j, i]:
heap = ig.get_heap()
# The receiver looks at inline items in each packet to place
# data correctly.
heap.repeat_pointers = True
if j in subscribed_bengs:
out_stream = streams[j // (self.n_bengs //
len(self.endpoints))]
assert out_stream is not None
out_stream.send_heap(heap)
expected_heaps += 1
ts += self.spectra_per_heap * self.ticks_between_spectra
if end:
for out_stream in streams:
if out_stream is not None:
out_stream.send_heap(ig.get_end())
# They're all pointing at the same receiver, which will
# shut down after the first stop heap
break
streams = []
if not end:
# We only want to stop the capture once all the heaps we expect
# have been received, but time out after 5 seconds to avoid
# hanging the test.
for i in range(100):
if server.counters['heaps'] >= expected_heaps:
break
else:
print('Only {} / {} heaps received so far'.format(
server.counters['heaps'], expected_heaps))
await asyncio.sleep(0.05)
else:
print('Giving up waiting for heaps')
await server.stop_capture(force=not end)
expected_data = np.zeros((self.n_channels, n_spectra, 2), np.int8)
expected_weight = np.ones((self.n_channels, n_spectra), np.int8)
for channel in range(self.n_channels):
expected_data[channel, :, 0] = channel % 255 - 128
for t in range(n_spectra):
expected_data[:, t, 1] = t % 255 - 128
for i in range(self.n_bengs):
for j in range(n_heaps):
if drop[i, j]:
channel0 = i * self.channels_per_heap
spectrum0 = j * self.spectra_per_heap
index = np.s_[channel0:channel0 + self.channels_per_heap,
spectrum0:spectrum0 + self.spectra_per_heap]
expected_data[index] = 0
expected_weight[index] = 0
expected_data = expected_data[self.args.channels.asslice()]
expected_weight = expected_weight[self.args.channels.asslice()]
# Validate the output
if write:
h5file = h5py.File(filename, 'r')
with contextlib.closing(h5file):
bf_raw = h5file['/Data/bf_raw']
np.testing.assert_equal(expected_data, bf_raw)
timestamps = h5file['/Data/timestamps']
expected = 1234567890 \
+ self.ticks_between_spectra * np.arange(self.spectra_per_heap * n_heaps)
np.testing.assert_equal(expected, timestamps)
flags = h5file['/Data/flags']
expected = np.where(drop, 8, 0).astype(np.uint8)
expected = expected[self.args.channels.start //
self.channels_per_heap:self.args.channels.
stop // self.channels_per_heap]
np.testing.assert_equal(expected, flags)
data_set = h5file['/Data']
assert_equal('i0_tied_array_channelised_voltage_0x',
data_set.attrs['stream_name'])
assert_equal(self.args.channels.start,
data_set.attrs['channel_offset'])
else:
assert_is_none(filename)
# Validate the signal display stream
rx.stop()
heaps = list(rx)
# Note: would need updating if n_heaps is not a multiple of heaps_per_stats
assert_equal(n_heaps // self.heaps_per_stats + 2, len(heaps))
assert_true(heaps[0].is_start_of_stream())
assert_true(heaps[-1].is_end_of_stream())
ig = spead2.send.ItemGroup()
spectrum = 0
spectra_per_stats = self.heaps_per_stats * self.spectra_per_heap
for rx_heap in heaps[1:-1]:
updated = ig.update(rx_heap)
rx_data = updated['sd_data'].value
rx_flags = updated['sd_flags'].value
rx_timestamp = updated['sd_timestamp'].value
# Check types and shapes
assert_equal((len(self.args.channels), 2, 2), rx_data.shape)
assert_equal(np.float32, rx_data.dtype)
assert_equal((len(self.args.channels), 2), rx_flags.shape)
assert_equal(np.uint8, rx_flags.dtype)
np.testing.assert_equal(0, rx_data[..., 1]) # Should be real only
rx_power = rx_data[:, 0, 0]
rx_saturated = rx_data[:, 1, 0]
# Check calculations
ts_unix = (spectrum + 0.5 * spectra_per_stats) * self.ticks_between_spectra \
/ self.adc_sample_rate + 111111111.0
np.testing.assert_allclose(ts_unix * 100.0, rx_timestamp)
index = np.s_[:, spectrum:spectrum + spectra_per_stats]
frame_data = expected_data[index]
frame_weight = expected_weight[index]
weight_sum = np.sum(frame_weight, axis=1)
power = np.sum(frame_data.astype(np.float64)**2,
axis=2) # Sum real+imag
saturated = (frame_data == -128) | (frame_data == 127)
saturated = np.logical_or.reduce(saturated,
axis=2) # Combine real+imag
saturated = saturated.astype(np.float64)
# Average over time. Can't use np.average because it complains if
# weights sum to zero instead of giving a NaN.
with np.errstate(divide='ignore', invalid='ignore'):
power = np.sum(power * frame_weight, axis=1) / weight_sum
saturated = np.sum(saturated * frame_weight,
axis=1) / weight_sum
power = np.where(weight_sum, power, 0)
saturated = np.where(weight_sum, saturated, 0)
np.testing.assert_allclose(power, rx_power)
np.testing.assert_allclose(saturated, rx_saturated)
flags = np.where(weight_sum, 0, DATA_LOST)
np.testing.assert_equal(flags, rx_flags[:, 0])
spectrum += spectra_per_stats
def test_nonascii_name(self):
"""Name with non-ASCII characters must fail"""
with assert_raises(UnicodeEncodeError):
item = spead2.Item(0x1000, '\u0200', 'description', (), np.int32)
item.to_raw(spead2.Flavour())
def test_attributes(self):
flavour = spead2.Flavour(4, 64, 40, 2)
assert_equal(4, flavour.version)
assert_equal(64, flavour.item_pointer_bits)
assert_equal(40, flavour.heap_address_bits)
assert_equal(2, flavour.bug_compat)
def test_attributes(self):
flavour = spead2.Flavour(4, 64, 40, 2)
assert flavour.version == 4
assert flavour.item_pointer_bits == 64
assert flavour.heap_address_bits == 40
assert flavour.bug_compat == 2
def __init__(self,
spead_config,
disconnect_tolerance=0,
mpi_comm=None,
ports=None):
self._streams = []
self._ports = []
self.mpi_comm = mpi_comm
self._num_stream = len(spead_config['streams'])
self._num_stream_disconnect = 0
self._disconnect_tolerance = disconnect_tolerance
# Relay attributes
self._relay_stream = None
self._descriptor = None
self._measurement_set = None
self._header = {}
self._baseline_exclude = set()
self._baseline_map = []
try:
if spead_config['as_relay'] == 1:
self.as_relay = True
else:
self.as_relay = False
self.use_adios2 = spead_config.get('use_adios2', False)
self._file_name = spead_config.get('output_ms', 'output.ms')
try:
os.mkdir(os.path.dirname(self._file_name))
except OSError:
pass
if self.as_relay:
# NOTE: Don't do baseline exclusion if its a relay as it shares a codebase with the MS writer
# which will exclude baselines multiple times causing array dimension mismatches.
# So as a relay just average and pass on all baselines.
# Construct TCP streams and associated item groups.
stream_config = spead2.send.StreamConfig(
spead_config['relay']['stream_config']['max_packet_size'],
spead_config['relay']['stream_config']['rate'],
spead_config['relay']['stream_config']['burst_size'],
spead_config['relay']['stream_config']['max_heaps'])
stream = spead_config['relay']['stream']
threads = stream['threads'] if 'threads' in stream else 1
thread_pool = spead2.ThreadPool(threads=threads)
logger.info(
"Relaying visibilities to host {} on port {}".format(
stream['host'], stream['port']))
tcp_stream = spead2.send.TcpStream(thread_pool, stream['host'],
stream['port'],
stream_config)
item_group = spead2.send.ItemGroup(
flavour=spead2.Flavour(4, 64, 40, 0))
# Append udp_stream and item_group to the stream list as a tuple.
self._relay_stream = (tcp_stream, item_group)
else:
baseline_file = spead_config['baseline_map_filename']
if baseline_file:
full_baseline_map = parse_baseline_file(baseline_file)
index = 0
for b in full_baseline_map:
if b[2] == 0:
self._baseline_map.append(b)
else:
self._baseline_exclude.add(index)
index += 1
logger.info(
'Baseline count: total=%d, used=%d, excluded=%d',
len(full_baseline_map), len(self._baseline_map),
len(self._baseline_exclude))
# Ports can be given explicitly or taken from the config file
if not ports:
ports = (c['port'] for c in spead_config['streams'])
for port in ports:
logger.info('Adding TCP stream reader: port %d', port)
try:
stream = spead2.recv.Stream(spead2.ThreadPool(), 0)
stream.stop_on_stop_item = False
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sock.bind(('', port))
sock.listen(1)
stream.add_tcp_reader(sock)
self._ports.append(port)
self._streams.append(stream)
except Exception as e:
logger.error(
'Failed adding TCP stream reader: error %s port %d',
str(e), port)
raise
except:
self.close(graceful=False)
raise
def test_bad_item_pointer_bits(self):
with pytest.raises(ValueError):
spead2.Flavour(4, 32, 24, 0)
import logging
import asyncio
import numpy as np
import spead2
import spead2.send
import spead2.send.asyncio
logging.basicConfig(level=logging.INFO)
thread_pool = spead2.ThreadPool()
stream = spead2.send.asyncio.UdpStream(
thread_pool, [("127.0.0.1", 8888)], spead2.send.StreamConfig(rate=1e7))
del thread_pool # Make sure this doesn't crash anything
shape = (40, 50)
ig = spead2.send.ItemGroup(flavour=spead2.Flavour(4, 64, 48, 0))
item = ig.add_item(0x1234, 'foo', 'a foo item', shape=shape, dtype=np.int32)
item.value = np.zeros(shape, np.int32)
futures = [
stream.async_send_heap(ig.get_heap()),
stream.async_send_heap(ig.get_end())
]
# Delete things to check that there are no refcounting bugs
del ig
del stream
asyncio.get_event_loop().run_until_complete(asyncio.wait(futures))
def test_bad_version(self):
with pytest.raises(ValueError):
spead2.Flavour(3, 64, 40, 0)
