def test_create_messages(self):
job_id = 8
num_tiles = 16
exp_upload_queue = 'upload-test-queue'
exp_ingest_queue = 'ingest-test-queue'
args = {
'job_id': job_id,
'project_info': [5, 3, 2], # Collection/Experiment/Channel ids
'z_chunk_size': 16,
'resolution': 0,
'upload_queue': exp_upload_queue,
'ingest_queue': exp_ingest_queue
}
backend = BossBackend(None)
x1, y1, z1 = (0, 0, 0)
x2, y2, z2 = (1024, 1024, 16)
res = 0
t = 2
chunk_key1 = backend.encode_chunk_key(num_tiles, args['project_info'],
res, x1, y1, z1, t)
chunk_key2 = backend.encode_chunk_key(num_tiles, args['project_info'],
res, x2, y2, z2, t)
with NamedTemporaryFile(mode='wt', suffix='.csv',
delete=False) as output_csv:
csv_file = output_csv.name
output_csv.write('{},{},{}'.format(chunk_key1, job_id, num_tiles))
output_csv.write('{},{},{}'.format(chunk_key2, job_id, num_tiles))
try:
actual = _create_messages(args, csv_file)
for i, raw_msg in enumerate(actual):
msg = json.loads(raw_msg)
self.assertEqual(job_id, msg['job_id'])
self.assertEqual(exp_upload_queue, msg['upload_queue_arn'])
self.assertEqual(exp_ingest_queue, msg['ingest_queue_arn'])
if i < 16:
self.assertEqual(chunk_key1, msg['chunk_key'])
exp_tile_key = backend.encode_tile_key(
args['project_info'], res, x1, y1, z1 + i, t)
self.assertEqual(exp_tile_key, msg['tile_key'])
elif i < 32:
self.assertEqual(chunk_key2, msg['chunk_key'])
exp_tile_key = backend.encode_tile_key(
args['project_info'], res, x2, y2, z2 + i - 16, t)
self.assertEqual(exp_tile_key, msg['tile_key'])
else:
self.fail('Too many messages returned')
finally:
os.remove(csv_file)
def test_decode_tile_key(self):
"""Test encoding an object key"""
b = BossBackend(self.example_config_data)
b.setup(self.api_token)
params = {"collection": 1,
"experiment": 2,
"channel": 3,
"resolution": 0,
"x_index": 5,
"y_index": 6,
"z_index": 1,
"t_index": 0,
}
proj = [str(params['collection']), str(params['experiment']), str(params['channel'])]
key = b.encode_tile_key(proj,
params['resolution'],
params['x_index'],
params['y_index'],
params['z_index'],
params['t_index'],
)
parts = b.decode_tile_key(key)
assert parts["collection"] == params['collection']
assert parts["experiment"] == params['experiment']
assert parts["channel"] == params['channel']
assert parts["resolution"] == params['resolution']
assert parts["x_index"] == params['x_index']
assert parts["y_index"] == params['y_index']
assert parts["z_index"] == params['z_index']
assert parts["t_index"] == params['t_index']
def test_encode_tile_key(self):
"""Test encoding an object key"""
b = BossBackend(self.example_config_data)
b.setup(self.api_token)
params = {"collection": 1,
"experiment": 2,
"channel": 3,
"resolution": 0,
"x_index": 5,
"y_index": 6,
"z_index": 1,
"t_index": 0,
}
proj = [str(params['collection']), str(params['experiment']), str(params['channel'])]
key = b.encode_tile_key(proj,
params['resolution'],
params['x_index'],
params['y_index'],
params['z_index'],
params['t_index'],
)
assert key == six.u("03ca58a12ec662954ac12e06517d4269&1&2&3&0&5&6&1&0")
def create_messages(args):
"""Create all of the tile messages to be enqueued
Args:
args (dict): Same arguments as populate_upload_queue()
Returns:
list: List of strings containing Json data
"""
tile_size = lambda v: args[v + "_tile_size"]
range_ = lambda v: range(args[v + '_start'], args[v + '_stop'], tile_size(v))
# DP NOTE: configuration is not actually used by encode_*_key method
backend = BossBackend(None)
msgs = []
for t in range_('t'):
for z in range_('z'):
for y in range_('y'):
for x in range_('x'):
chunk_x = int(x/tile_size('x'))
chunk_y = int(y/tile_size('y'))
chunk_z = int(z/tile_size('z'))
num_of_tiles = min(tile_size('z'), args['z_stop'] - z)
chunk_key = backend.encode_chunk_key(num_of_tiles,
args['project_info'],
args['resolution'],
chunk_x,
chunk_y,
chunk_z,
t)
for tile in range(z, z + num_of_tiles):
tile_key = backend.encode_tile_key(args['project_info'],
args['resolution'],
chunk_x,
chunk_y,
tile,
t)
msg = {
'job_id': args['job_id'],
'upload_queue_arn': args['upload_queue'],
'ingest_queue_arn': args['ingest_queue'],
'chunk_key': chunk_key,
'tile_key': tile_key,
}
yield json.dumps(msg)
def _create_messages(args, index_csv):
"""
Create all of the tile messages to be enqueued.
Args:
args (dict): Same arguments as patch_upload_queue().
index_csv (str): CSV file with tile info.
Generates:
(list): List of strings containing Json data
"""
# DP NOTE: configuration is not actually used by encode_*_key method
backend = BossBackend(None)
chunk_key_list = []
with open(index_csv, "rt") as data:
lines = data.readlines()
for line in lines:
parts = line.split(",")
chunk_key_list.append(parts[0])
msgs = []
for base_chunk_key in chunk_key_list:
parts = backend.decode_chunk_key(base_chunk_key)
chunk_x = parts['x_index']
chunk_y = parts['y_index']
chunk_z = parts['z_index']
t = parts['t_index']
num_of_tiles = parts['num_tiles']
chunk_key = base_chunk_key
z_start = chunk_z * args['z_chunk_size']
for tile in range(z_start, z_start + num_of_tiles):
tile_key = backend.encode_tile_key(args['project_info'],
args['resolution'], chunk_x,
chunk_y, tile, t)
msg = {
'job_id': args['job_id'],
'upload_queue_arn': args['upload_queue'],
'ingest_queue_arn': args['ingest_queue'],
'chunk_key': chunk_key,
'tile_key': tile_key,
}
yield json.dumps(msg)
def check_tiles(self, chunk_key, tiles):
"""
Check the chunk's tile map for missing tiles. If any are missing,
generate the proper stringified JSON for putting those missing tiles
back in the tile upload queue.
Args:
chunk_key (str): Identifies chunk of tiles.
tiles (): List of tiles uploaded for the chunk.
Yields:
(str): JSON string for sending to SQS tile upload queue.
"""
# Only using encode|decode_*_key methods, so don't need to provide a
# config.
ingest_backend = BossBackend(None)
chunk_key_parts = ingest_backend.decode_chunk_key(chunk_key)
chunk_x = chunk_key_parts['x_index']
chunk_y = chunk_key_parts['y_index']
chunk_z = chunk_key_parts['z_index']
t = chunk_key_parts['t_index']
num_tiles = chunk_key_parts['num_tiles']
z_start = chunk_z * self.job['z_chunk_size']
for tile_z in range(z_start, z_start + num_tiles):
# First arg is a list of [collection, experiment, channel] ids.
tile_key = ingest_backend.encode_tile_key(self._get_project_info(),
self.job['resolution'],
chunk_x, chunk_y, tile_z,
t)
if tile_key in tiles:
continue
msg = {
'job_id': self.job['task_id'],
'upload_queue_arn': self.job['upload_queue'],
'ingest_queue_arn': self.job['ingest_queue'],
'chunk_key': chunk_key,
'tile_key': tile_key
}
log.info(
f'Re-enqueuing tile: {tile_key} belonging to chunk: {chunk_key}'
)
yield json.dumps(msg)