def test_send_receive(tmpdir_factory):
# Run kotekan bufferRecv
tmpdir = tmpdir_factory.mktemp("writer")
write_buffer = runner.DumpVisBuffer(str(tmpdir))
# the plan is: wait 5s and then kill it
rest_commands = [("wait", 5, None), ("get", "kill", None)]
receiver = runner.KotekanStageTester(
"bufferRecv",
{},
None,
write_buffer,
params_kotekan,
rest_commands=rest_commands,
)
# TODO: network buffer processes should use in_buf and out_buf to please the test framework
receiver._stages["bufferRecv_test"]["buf"] = receiver._stages[
"bufferRecv_test"]["out_buf"]
# Run kotekan bufferRecv in another thread
with concurrent.futures.ThreadPoolExecutor() as executor:
future_receiver = executor.submit(receiver.run)
# Wait for it to start so the sender doesn't drop frames
time.sleep(1)
fakevis_buffer = runner.FakeVisBuffer(
num_frames=params_kotekan["total_frames"],
mode=params_kotekan["mode"],
freq_ids=params_kotekan["freq_ids"],
sleep_before=2,
wait=False,
)
sender = runner.KotekanStageTester("bufferSend", {}, fakevis_buffer,
None, params_kotekan)
# TODO: network buffer processes should use in_buf and out_buf to please the test framework
sender._stages["bufferSend_test"]["buf"] = sender._stages[
"bufferSend_test"]["in_buf"]
# run kotekan bufferSend
sender.run()
# wait for kotekan bufferRecv to finish
future_receiver.result(timeout=7)
assert sender.return_code == 0
assert receiver.return_code == 0
vis_data = write_buffer.load()
assert len(vis_data) == params_kotekan["total_frames"]
def written_data_base(outdir, stage_extra=None, root_extra=None):
fakevis_buffer = runner.FakeVisBuffer(
freq_ids=writer_params['freq'],
num_frames=writer_params['total_frames'],
cadence=writer_params['cadence'])
root_params = writer_params.copy()
root_params['root_path'] = outdir
if root_extra is not None:
root_params.update(root_extra)
stage_params = {
'node_mode': False,
}
if stage_extra is not None:
stage_params.update(stage_extra)
test = runner.KotekanStageTester('visWriter', stage_params, fakevis_buffer,
None, root_params)
test.run()
import glob
files = sorted(glob.glob(outdir + '/20??????T??????Z_*_corr/*.h5'))
return [h5py.File(fname, 'r') for fname in files]
def run_baseband(tdir_factory, params=None, rest_commands=None):
p = dict(default_params)
tmpdir = tdir_factory.mktemp("baseband")
p['base_dir'] = str(tmpdir) + '/'
if params:
p.update(params)
fake_buffer = runner.FakeNetworkBuffer(
stage_name=DATAGEN_PNAME,
num_frames=p['total_frames'],
type=p['type'],
)
test = runner.KotekanStageTester(
'basebandReadout',
{},
fake_buffer,
None,
p,
rest_commands=rest_commands,
)
test.run()
dump_files = glob.glob(str(tmpdir) + '/*.h5')
return dump_files
def written_data(tmpdir_factory):
tmpdir = str(tmpdir_factory.mktemp("writer"))
fakevis_buffer = runner.FakeVisBuffer(
freq_ids=writer_params['freq'],
num_frames=writer_params['total_frames'],
cadence=writer_params['cadence']
)
params = writer_params.copy()
params['root_path'] = tmpdir
test = runner.KotekanStageTester(
'visWriter',
{'node_mode': False, 'file_type': 'raw'},
fakevis_buffer,
None,
params
)
test.run()
import glob
files = sorted(glob.glob(tmpdir + '/20??????T??????Z_*_corr/*.meta'))
yield [visbuffer.VisRaw(fname) for fname in files]
def vis_data(tmpdir_factory):
tmpdir_l = tmpdir_factory.mktemp("freqsplit_lower")
tmpdir_h = tmpdir_factory.mktemp("freqsplit_higher")
fakevis_buffer = runner.FakeVisBuffer(
num_frames=params['total_frames'],
mode=params['mode'],
freq_ids=params['freq_ids'],
wait=False
)
dump_buffer_l = runner.DumpVisBuffer(
str(tmpdir_l))
dump_buffer_h = runner.DumpVisBuffer(
str(tmpdir_h))
test = runner.KotekanStageTester(
'freqSplit', {},
fakevis_buffer,
[dump_buffer_l, dump_buffer_h],
params
)
test.run()
yield [dump_buffer_l.load(), dump_buffer_h.load()]
def run_baseband(tdir_factory,
params=None,
rest_commands=None,
expect_a_failure=False):
p = dict(default_params)
tmpdir = tdir_factory.mktemp("baseband")
if params:
p.update(params)
fake_buffer = runner.FakeNetworkBuffer(stage_name=DATAGEN_PNAME,
num_frames=p["total_frames"],
type=p["type"])
write_buffer = runner.DumpBasebandBuffer(
str(tmpdir),
num_frames="buffer_depth * 4",
frame_size="num_elements * samples_per_data_set / 4",
)
test = runner.KotekanStageTester(
"basebandReadout",
{},
fake_buffer,
write_buffer,
p,
rest_commands=rest_commands,
expect_failure=expect_a_failure,
)
test.run()
return write_buffer.load()
def write_data(tmpdir_factory):
tmpdir_l = tmpdir_factory.mktemp("freqsplit_write_lower")
tmpdir_h = tmpdir_factory.mktemp("freqsplit_write_higher")
fakevis_buffer = runner.FakeVisBuffer(
num_frames=params['total_frames'],
mode=params['mode'],
freq_ids=params['freq_ids'],
wait=False
)
write_buffer_l = runner.VisWriterBuffer(
str(tmpdir_l), 'raw')
write_buffer_h = runner.VisWriterBuffer(
str(tmpdir_h), 'raw')
test = runner.KotekanStageTester(
'freqSplit', {},
fakevis_buffer,
[write_buffer_l, write_buffer_h],
params
)
test.run()
return [write_buffer_l.load(), write_buffer_h.load()]
def written_data_base(outdir, stage_extra=None, root_extra=None):
fakevis_buffer = runner.FakeVisBuffer(
freq_ids=writer_params["freq"],
num_frames=writer_params["total_frames"],
cadence=writer_params["cadence"],
)
root_params = writer_params.copy()
root_params["root_path"] = outdir
if root_extra is not None:
root_params.update(root_extra)
stage_params = {"node_mode": False}
if stage_extra is not None:
stage_params.update(stage_extra)
test = runner.KotekanStageTester(
"VisWriter", stage_params, fakevis_buffer, None, root_params
)
test.run()
import glob
files = sorted(glob.glob(outdir + "/20??????T??????Z_*_corr/*.h5"))
return [h5py.File(fname, "r") for fname in files]
def run_flagging(tmpdir_factory, cmds):
""" Receive Flags """
tmpdir = tmpdir_factory.mktemp("receiveflags")
fakevis_buffer = runner.FakeVisBuffer(
num_frames=global_params["total_frames"],
mode=global_params["fakevis_mode"],
cadence=global_params["cadence"],
wait=global_params["wait"],
)
out_dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester(
"ReceiveFlags",
params,
buffers_in=fakevis_buffer,
buffers_out=out_dump_buffer,
global_config=global_params,
rest_commands=cmds,
expect_failure=True,
)
test.run()
return out_dump_buffer.load()
def run_baseband(tdir_factory,
params=None,
rest_commands=None,
expect_a_failure=False):
p = dict(default_params)
tmpdir = tdir_factory.mktemp("baseband")
p["base_dir"] = str(tmpdir) + "/"
if params:
p.update(params)
fake_buffer = runner.FakeNetworkBuffer(stage_name=DATAGEN_PNAME,
num_frames=p["total_frames"],
type=p["type"])
test = runner.KotekanStageTester(
"basebandReadout",
{},
fake_buffer,
None,
p,
rest_commands=rest_commands,
expect_failure=expect_a_failure,
)
test.run()
dump_files = glob.glob(str(tmpdir) + "/*.h5")
return dump_files
def frb_post_process_data(tmpdir_factory, missing_frames=None, test_env=None):
"""Runs the Kotekan using the short 6144 samples_per_dataset, optionally dropping frames
This works out to just one packet per stream in a frame, and uses a single
GPU, reading the beamforming inputs from files named starting with
"fake_cpu_beamform".
If there are any frames to drop, the config will include a
TestDropFrames stage, and drop the specified frame from the input
to the FrbPostProcess.
"""
if test_env is None and not glob.glob(
path.join(path.dirname(__file__), "data",
"fake_cpu_beamform_*.dump")):
pytest.skip("No test data and not requesting full generation")
tmpdir = tmpdir_factory.mktemp("frb_post_process")
if test_env == "run_slow_cpu_tests":
beamform_buffer = runner.FakeFrbBeamformBuffer(num_frames=7)
elif test_env == "run_amd_gpu_tests":
beamform_buffer = runner.FakeFrbBeamformBuffer(num_frames=7, cpu=False)
elif test_env is None:
beamform_buffer = runner.ReadRawBeamformBuffer(
path.join(path.dirname(__file__), "data"),
file_name="fake_cpu_beamform")
else:
pytest.skip("Unknown environment {}".format(test_env))
lost_samples_buffer = runner.FakeLostSamplesBuffer(num_frames=10)
dump_buffer = runner.DumpFrbPostProcessBuffer(str(tmpdir))
params = frb_params.copy()
params["num_gpus"] = 1
params["samples_per_data_set"] = 6144
if missing_frames:
drop_frames_buffer = runner.DropFramesBuffer(
beamform_buffer,
missing_frames,
frame_size=beamform_buffer.buffer_block[beamform_buffer.name]
["frame_size"],
)
input_buffers = [
beamform_buffer, drop_frames_buffer, lost_samples_buffer
]
params["in_buf_0"] = drop_frames_buffer.name
else:
input_buffers = [beamform_buffer, lost_samples_buffer]
params["in_buf_0"] = beamform_buffer.name
params["lost_samples_buf"] = lost_samples_buffer.name
test = runner.KotekanStageTester("frbPostProcess", frb_params,
input_buffers, dump_buffer, params)
test.run()
return dump_buffer
def apply_data(request, tmp_path_factory, gain_path, old_gains, new_gains,
cal_broker):
output_dir = str(tmp_path_factory.mktemp("output"))
global_params["gains_dir"] = str(gain_path)
global_params[
"read_from_file"] = True if request.param == "file" else False
# REST commands
cmds = [[
"post",
"gains",
{
"update_id": new_update_id,
"start_time": new_timestamp,
"transition_interval": transition_interval,
"new_state": new_state,
},
]]
fakevis_buffer = runner.FakeVisBuffer(
freq_ids=global_params["freq_ids"],
num_frames=global_params["total_frames"],
wait=global_params["wait"],
)
out_dump_buffer = runner.DumpVisBuffer(output_dir)
# Configuration for the direct dump of FakeVis's output
fakevis_dump_conf = {
"file_name": "fakevis_dump",
"file_ext": "dump",
"base_dir": output_dir,
}
host, port = cal_broker.server_address
global_params.update({"broker_host": host, "broker_port": port})
test = runner.KotekanStageTester(
"applyGains",
global_params,
buffers_in=fakevis_buffer,
buffers_out=out_dump_buffer,
global_config=global_params,
rest_commands=cmds,
parallel_stage_type="rawFileWrite",
parallel_stage_config=fakevis_dump_conf,
)
test.run()
in_dump = visbuffer.VisBuffer.load_files(
f"{output_dir}/*fakevis_dump*.dump")
return in_dump, out_dump_buffer.load()
def psr_post_process_data(tmpdir_factory, missing_frames=None, test_env=None):
"""Runs the Kotekan pulsarPostProcess stage, optionally dropping frames
If there are any frames to drop, the config will include a
TestDropFrames stage, and drop the specified frame from the input
to the pulsarPostProcess.
"""
if test_env is None and not glob.glob(
path.join(path.dirname(__file__), "data", "fake_psr_beamform_gpu_*.dump")
):
pytest.skip("No test data and not requesting full generation")
tmpdir = tmpdir_factory.mktemp("psr_post_process")
lost_samples_buffer = runner.FakeLostSamplesBuffer(num_frames=10)
dump_buffer = runner.DumpPsrPostProcessBuffer(str(tmpdir))
input_buffers = [lost_samples_buffer]
stage_params = {"pulsar_out_buf": dump_buffer.name}
for i in range(4):
beamform_buffer = runner.ReadRawBeamformBuffer(
path.join(path.dirname(__file__), "data"), file_name="fake_psr_beamform_gpu"
)
beamform_buffer.buffer_block[beamform_buffer.name][
"frame_size"
] = "samples_per_data_set * num_beams * num_pol * sizeof_float * 2"
input_buffers.append(beamform_buffer)
if missing_frames and missing_frames[i]:
drop_frames_buffer = runner.DropFramesBuffer(
beamform_buffer,
missing_frames[i],
frame_size=beamform_buffer.buffer_block[beamform_buffer.name][
"frame_size"
],
)
input_buffers.append(drop_frames_buffer)
stage_params[
"network_input_buffer_{:d}".format(i)
] = drop_frames_buffer.name
else:
stage_params["network_input_buffer_{:d}".format(i)] = beamform_buffer.name
params = global_params.copy()
test = runner.KotekanStageTester(
"pulsarPostProcess", stage_params, input_buffers, dump_buffer, params
)
test.run()
return dump_buffer
def pulsar_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("accumulate")
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester('visAccumulate', accumulate_params,
runner.FakeGPUBuffer(**pulsar_params),
dump_buffer, accumulate_params)
test.run()
yield dump_buffer.load()
def data(tmpdir_factory):
# keep all the data this test produces in a tmp directory
tmpdir = tmpdir_factory.mktemp("name_of_the_test_case")
# you can use FakeVisBuffer to produce fake data
fakevis_buffer = runner.FakeVisBuffer(num_frames=params["total_frames"], mode="gaussian")
# DumpVisBuffer can be used to dump data for testing
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
# KotekanStageTester is used to run kotekan with your config
test = runner.KotekanStageTester("stageUnderTest", {}, fakevis_buffer, dump_buffer, params)
def pulsar_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("pulsar")
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
dump_buffer_gated = runner.VisWriterBuffer(str(tmpdir), "raw")
# Insert an extra buffer for gated stream
dump_buffer.buffer_block.update(dump_buffer_gated.buffer_block)
dump_buffer.stage_block.update(dump_buffer_gated.stage_block)
acc_par = pulsar_params.copy()
acc_par.update({
"gating": {
"psr0": {
"mode": "pulsar",
"buf": dump_buffer_gated.name
}
},
"updatable_config": {
"psr0": "/updatable_config/psr0_config"
},
})
updatable_params = pulsar_params.copy()
updatable_params.update({
"updatable_config": {
"psr0_config": {
"kotekan_update_endpoint": "json",
"enabled": True,
"pulsar_name": "fakepsr",
"segment": 100.0,
"coeff": [pulsar_params["coeff"]],
"dm": pulsar_params["dm"],
"t_ref": [pulsar_params["t_ref"]],
"phase_ref": [pulsar_params["phase_ref"]],
"rot_freq": pulsar_params["rot_freq"],
"pulse_width": 1e-3,
}
}
})
test = runner.KotekanStageTester(
"visAccumulate",
acc_par,
runner.FakeGPUBuffer(**pulsar_params),
dump_buffer,
updatable_params,
)
test.run()
yield dump_buffer_gated.load()
def replace_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("replace")
fakevis_buffer = runner.FakeVisBuffer(
freq=replace_params["freq"], num_frames=replace_params["total_frames"])
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester("ReplaceVis", {}, fakevis_buffer,
dump_buffer, replace_params)
test.run()
yield dump_buffer.load()
def remove_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("remove")
fakevis_buffer = runner.FakeVisBuffer(
freq=remove_params['freq'], num_frames=remove_params['total_frames'])
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester('removeEv', {}, fakevis_buffer,
dump_buffer, remove_params)
test.run()
yield dump_buffer.load()
def lostsamples_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("lostsamples")
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester('visTransform', {'num_ev': 4}, [
runner.FakeGPUBuffer(mode='lostsamples',
freq=params['freq'],
num_frames=params['total_frames'])
for g in range(params['num_gpu_buffers'])
], dump_buffer, params)
test.run()
yield dump_buffer.load()
def lostsamples_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("lostsamples")
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester(
'visAccumulate', {'num_ev': 4},
runner.FakeGPUBuffer(mode='lostsamples',
freq=accumulate_params['freq'],
num_frames=accumulate_params['total_frames']),
dump_buffer, accumulate_params)
test.run()
yield dump_buffer.load()
def subset_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("subset")
fakevis_buffer = runner.FakeVisBuffer(
freq_ids=subset_params["freq_ids"],
num_frames=subset_params["total_frames"])
write_buffer = runner.VisWriterBuffer(str(tmpdir), "raw")
test = runner.KotekanStageTester("prodSubset", vis_params, fakevis_buffer,
write_buffer, subset_params)
test.run()
return write_buffer.load()
def gaussian_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("gaussian")
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester(
'visAccumulate', {'num_ev': 4},
runner.FakeGPUBuffer(mode='gaussian',
freq=gaussian_params['freq'],
num_frames=gaussian_params['total_frames']),
dump_buffer, gaussian_params)
test.run()
yield dump_buffer.load()
def pulsar_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("pulsar")
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
dump_buffer_gated = runner.VisWriterBuffer(str(tmpdir), "raw")
# Insert an extra buffer for gated stream
dump_buffer.buffer_block.update(dump_buffer_gated.buffer_block)
dump_buffer.stage_block.update(dump_buffer_gated.stage_block)
acc_par = pulsar_params.copy()
acc_par.update({
'gating': {
'psr0': {
'mode': 'pulsar',
'buf': dump_buffer_gated.name
},
},
'updatable_config': {
'psr0': "/updatable_config/psr0_config"
},
})
updatable_params = pulsar_params.copy()
updatable_params.update({
'updatable_config': {
'psr0_config': {
'kotekan_update_endpoint': 'json',
'enabled': True,
'pulsar_name': 'fakepsr',
'segment': 100.,
'coeff': [pulsar_params['coeff']],
'dm': pulsar_params['dm'],
't_ref': [pulsar_params['t_ref']],
'phase_ref': [pulsar_params['phase_ref']],
'rot_freq': pulsar_params['rot_freq'],
'pulse_width': 1e-3,
}
}
})
test = runner.KotekanStageTester('visAccumulate', acc_par,
runner.FakeGPUBuffer(**pulsar_params),
dump_buffer, updatable_params)
test.run()
yield dump_buffer_gated.load()
def write_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("freqsub_write")
fakevis_buffer = runner.FakeVisBuffer(num_frames=params['total_frames'],
mode=params['mode'],
freq_ids=params['freq_ids'],
wait=False)
write_buffer = runner.VisWriterBuffer(str(tmpdir), 'raw')
test = runner.KotekanStageTester('freqSubset', {}, fakevis_buffer,
write_buffer, params)
test.run()
return write_buffer.load()
def vis_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("freqsub")
fakevis_buffer = runner.FakeVisBuffer(num_frames=params['total_frames'],
mode=params['mode'],
freq_ids=params['freq_ids'],
wait=False)
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester('freqSubset', {}, fakevis_buffer,
dump_buffer, params)
test.run()
yield dump_buffer.load()
def run_kotekan(tmpdir_factory, nfreqs=1):
"""Starts Kotekan with a simulated baseband stream from `nfreq` frequencies being fed into a single `basebandWriter` stage.
The frequencies will have indexes [0:nfreq], and will be saved into raw baseband dump files in the `tmpdir_factory` subdirectory `baseband_raw_12345`, one frequency per file.
Returns:
--------
Sorted list of filenames of the saved baseband files.
"""
num_elements = global_params["num_elements"]
samples_per_data_set = global_params["samples_per_data_set"]
frame_list = []
for i in range(global_params["buffer_depth"]):
for freq_id in range(nfreqs):
fpga_seq = frame_size * i
frame_list.append(
baseband_buffer.BasebandBuffer.new_from_params(
event_id=12345,
freq_id=freq_id,
num_elements=num_elements,
frame_size=frame_size,
frame_data=generate_tpluse_data(fpga_seq, frame_size,
num_elements),
))
frame_list[-1].metadata.frame_fpga_seq = frame_size * i
frame_list[-1].metadata.valid_to = samples_per_data_set
frame_list[-1].metadata.valid_to -= 17
current_dir = str(tmpdir_factory.getbasetemp())
read_buffer = runner.ReadBasebandBuffer(current_dir, frame_list)
read_buffer.write()
test = runner.KotekanStageTester(
"BasebandWriter",
{"root_path": current_dir}, # stage_config
read_buffer, # buffers_in
None, # buffers_out is None
global_params, # global_config
expect_failure=True, # because rawFileRead runs out of files to read
)
test.run()
dump_files = sorted(
glob.glob(current_dir + "/baseband_raw_12345/baseband_12345_*.data"))
return dump_files
def vis_data(tmpdir_factory, comet_broker):
# keeping all the data this test produced here (probably do not need it)
# using FakeVisBuffer to produce fake data
fakevis_buffer = runner.FakeVisBuffer(**params_fakevis)
# pass comet port to kotekan
params["dataset_manager"]["ds_broker_port"] = comet_broker
# KotekanStageTester is used to run kotekan with my config
test = runner.KotekanStageTester(
stage_type="VisSharedMemWriter",
stage_config=params_writer_stage,
buffers_in=fakevis_buffer,
buffers_out=None,
global_config=params,
)
yield test
def vis_data(tmpdir_factory, request):
global num_frames
# keeping all the data this test produced here (probably do not need it)
# using FakeVisBuffer to produce fake data
fakevis_buffer = runner.FakeVisBuffer(**request.param)
num_frames = request.param["num_frames"]
# KotekanStageTester is used to run kotekan with my config
test = runner.KotekanStageTester(
stage_type="VisSharedMemWriter",
stage_config=params_writer_stage,
buffers_in=fakevis_buffer,
buffers_out=None,
global_config=global_params,
)
test.run()
def merge_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("merge")
fakevis_buffers = [
runner.FakeVisBuffer(freq_ids=[f],
num_frames=merge_params['total_frames'])
for f in merge_params['freq']
]
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester('bufferMerge', {}, fakevis_buffers,
dump_buffer, merge_params)
test.run()
yield dump_buffer.load()
def run_eigenvis(tdir_factory, params=None):
if not params:
params = default_params
tmpdir = tdir_factory.mktemp("eigenvis")
fakevis_buffer = runner.FakeVisBuffer(freq_ids=params["freq"],
num_frames=params["total_frames"],
mode=params["mode"])
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester("eigenVis", {}, fakevis_buffer,
dump_buffer, params)
test.run()
return dump_buffer.load()