def test_get_cubes_no_time_multiple(self, fake_get_region): """Test the get_cubes method - no time - multiple cubes""" # Generate random data cube1 = Cube.create_cube(self.resource, [self.x_dim, self.y_dim, self.z_dim]) cube1.random() cube1.morton_id = 32 cube2 = Cube.create_cube(self.resource, [self.x_dim, self.y_dim, self.z_dim]) cube2.random() cube2.morton_id = 33 cube3 = Cube.create_cube(self.resource, [self.x_dim, self.y_dim, self.z_dim]) cube3.random() cube3.morton_id = 36 db = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) # populate dummy data keys = self.write_test_cube(db, self.resource, 0, cube1, cache=True, s3=False) keys.extend(self.write_test_cube(db, self.resource, 0, cube2, cache=True, s3=False)) keys.extend(self.write_test_cube(db, self.resource, 0, cube3, cache=True, s3=False)) cube_read = db.get_cubes(self.resource, keys) np.testing.assert_array_equal(cube1.data, cube_read[0].data) np.testing.assert_array_equal(cube2.data, cube_read[1].data) np.testing.assert_array_equal(cube3.data, cube_read[2].data)
def test_get_cubes_missing_time_step(self, fake_get_region): """Test get_cubes() when not supplying keys for all time steps in a time range. """ EXTENTS = [self.x_dim, self.y_dim, self.z_dim] FIRST_T_RNG = (0, 4) cube1 = Cube.create_cube(self.resource, EXTENTS, FIRST_T_RNG) cube1.random() cube1.morton_id = 70 # Note, no data for time steps 4 and 5 provided. SECOND_T_RNG = (6, 9) cube2 = Cube.create_cube(self.resource, EXTENTS, SECOND_T_RNG) cube2.random() cube2.morton_id = 70 TOTAL_T_RNG = (0, 9) exp_cube = Cube.create_cube(self.resource, EXTENTS, TOTAL_T_RNG) exp_cube.zeros() exp_cube.morton_id = 70 exp_cube.overwrite(cube1.data, FIRST_T_RNG) exp_cube.overwrite(cube2.data, SECOND_T_RNG) db = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) # populate dummy data keys = self.write_test_cube(db, self.resource, 0, cube1, cache=True, s3=False) keys.extend(self.write_test_cube(db, self.resource, 0, cube2, cache=True, s3=False)) # Method under test. cube_read = db.get_cubes(self.resource, keys) np.testing.assert_array_equal(exp_cube.data, cube_read[0].data)
def test_cutout_iso_not_present(self): """Test write_cuboid and cutout methods with iso option, testing iso is stored in parallel""" data = self.data data["channel"]["base_resolution"] = 5 resource = BossResourceBasic(data) # Generate random data cube1 = Cube.create_cube(resource, [400, 400, 8]) cube1.random() cube1.morton_id = 0 cubez = Cube.create_cube(resource, [400, 400, 8]) cubez.zeros() cubez.morton_id = 0 # Write at 5, not iso, and verify sp = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) sp.write_cuboid(resource, (200, 600, 3), 5, cube1.data, iso=False) cube2 = sp.cutout(resource, (200, 600, 3), (400, 400, 8), 5, iso=False) np.testing.assert_array_equal(cube1.data, cube2.data) # Get at res 5 iso, which should be blank cube2 = sp.cutout(resource, (200, 600, 3), (400, 400, 8), 5, iso=True) np.testing.assert_array_equal(cubez.data, cube2.data)
def test_get_ids_in_region_multiple_cubes_and_x_partials(self): """ Region has some full cuboids and some partial cuboids along the x axis. """ cube_dim_tuple = (self.x_dim, self.y_dim, self.z_dim) cube1 = Cube.create_cube(self.resource, [self.x_dim, self.y_dim, self.z_dim]) cube1.zeros() cube1.data[0][0][40][105] = 55555 cube1.data[0][0][50][105] = 66666 pos1 = [7*self.x_dim, 5*self.y_dim, 2*self.z_dim] cube1.morton_id = XYZMorton(pos1) cube2 = Cube.create_cube(self.resource, [self.x_dim, self.y_dim, self.z_dim]) cube2.zeros() cube2.data[0][0][40][105] = 55555 cube2.data[0][0][50][105] = 77777 pos2 = [8*self.x_dim, 5*self.y_dim, 2*self.z_dim] cube2.morton_id = XYZMorton(pos2) cube3 = Cube.create_cube(self.resource, [self.x_dim, self.y_dim, self.z_dim]) cube3.zeros() cube3.data[0][0][0][105] = 88888 pos3 = [9*self.x_dim, 5*self.y_dim, 2*self.z_dim] cube3.morton_id = XYZMorton(pos3) sp = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) resolution = 0 sp.write_cuboid(self.resource, pos1, resolution, cube1.data, time_sample_start=0) sp.write_cuboid(self.resource, pos2, resolution, cube2.data, time_sample_start=0) sp.write_cuboid(self.resource, pos3, resolution, cube3.data, time_sample_start=0) # Make sure cube write complete and correct. actual_cube = sp.cutout(self.resource, pos1, cube_dim_tuple, resolution) np.testing.assert_array_equal(cube1.data, actual_cube.data) actual_cube = sp.cutout(self.resource, pos2, cube_dim_tuple, resolution) np.testing.assert_array_equal(cube2.data, actual_cube.data) actual_cube = sp.cutout(self.resource, pos3, cube_dim_tuple, resolution) np.testing.assert_array_equal(cube3.data, actual_cube.data) corner = (7*self.x_dim+100, 5*self.y_dim, 2*self.z_dim) extent = (2*self.x_dim+self.x_dim//2, self.y_dim, self.z_dim) t_range = [0, 1] version = 0 expected = ['55555', '66666', '77777', '88888'] # Method under test. actual = sp.get_ids_in_region( self.resource, resolution, corner, extent, t_range, version) self.assertIn('ids', actual) self.assertCountEqual(expected, actual['ids'])
def test_filtered_cutout(self): time_axis = [1] cube_dim = [self.x_dim, self.y_dim, self.z_dim] cube_dim_tuple = (self.x_dim, self.y_dim, self.z_dim) cube1 = Cube.create_cube(self.resource, cube_dim) cube1.data = np.ones(time_axis + [cube_dim[2], cube_dim[1], cube_dim[0]], dtype='uint64') cube1.morton_id = 0 corner = (0, 0, 0) expected = np.zeros(time_axis + [cube_dim[2], cube_dim[1], cube_dim[0]], dtype='uint64') # Will filter by these ids. id1 = 55555 id2 = 66666 cube1.data[0][0][40][0] = id1 cube1.data[0][0][50][0] = id2 expected[0][0][40][0] = id1 expected[0][0][50][0] = id2 sp = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) resolution = 0 sp.write_cuboid(self.resource, corner, resolution, cube1.data, time_sample_start=0) # Make sure cube written correctly. actual_cube = sp.cutout(self.resource, corner, cube_dim_tuple, resolution) np.testing.assert_array_equal(cube1.data, actual_cube.data) # Method under test. actual_filtered = sp.cutout(self.resource, corner, cube_dim_tuple, resolution, filter_ids=[id1, id2]) np.testing.assert_array_equal(expected, actual_filtered.data)
def test_get_ids_in_region_single_cube(self): """Test single cuboid using DynamoDB index.""" cube_dim_tuple = (self.x_dim, self.y_dim, self.z_dim) cube1 = Cube.create_cube(self.resource, [self.x_dim, self.y_dim, self.z_dim]) cube1.zeros() cube1.data[0][0][40][0] = 55555 cube1.data[0][0][50][0] = 66666000000000 pos1 = [2*self.x_dim, 3*self.y_dim, 2*self.z_dim] cube1.morton_id = XYZMorton(pos1) sp = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) resolution = 0 sp.write_cuboid(self.resource, pos1, resolution, cube1.data, time_sample_start=0) # Make sure cube write complete and correct. actual_cube = sp.cutout(self.resource, pos1, cube_dim_tuple, resolution) np.testing.assert_array_equal(cube1.data, actual_cube.data) corner = (2*self.x_dim, 3*self.y_dim, 2*self.z_dim) extent = (self.x_dim, self.y_dim, self.z_dim) t_range = [0, 1] version = 0 expected = ['55555', '66666000000000'] # Method under test. actual = sp.get_ids_in_region( self.resource, resolution, corner, extent, t_range, version) self.assertIn('ids', actual) self.assertCountEqual(expected, actual['ids'])
def test_cutout_no_time_single_aligned_miss(self): """Test the get_cubes method - no time - single - miss""" # Generate random data cube1 = Cube.create_cube(self.resource, [self.x_dim, self.y_dim, self.z_dim]) cube1.random() cube1.morton_id = 0 sp = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) sp.write_cuboid(self.resource, (1, 0, 0), 0, cube1.data) cube2 = sp.cutout(self.resource, (1, 0, 0), (self.x_dim, self.y_dim, self.z_dim), 0) # Make sure data is the same np.testing.assert_array_equal(cube1.data, cube2.data) # Delete everything in the cache sp.kvio.cache_client.flushdb() # Get the data again cube3 = sp.cutout(self.resource, (1, 0, 0), (self.x_dim, self.y_dim, self.z_dim), 0) # Make sure the data is the same np.testing.assert_array_equal(cube1.data, cube2.data) np.testing.assert_array_equal(cube1.data, cube3.data)
def test_page_in_multi_cuboids_y_dir(self): # Generate random data cube1 = Cube.create_cube(self.resource, [self.x_dim, self.y_dim * 2, self.z_dim]) cube1.random() cube1.morton_id = 0 sp = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) sp.write_cuboid(self.resource, (0, 0, 0), 0, cube1.data) cube2 = sp.cutout(self.resource, (0, 0, 0), (self.x_dim, self.y_dim * 2, self.z_dim), 0) np.testing.assert_array_equal(cube1.data, cube2.data) # Make sure data is the same np.testing.assert_array_equal(cube1.data, cube2.data) # Delete everything in the cache sp.kvio.cache_client.flushdb() # Force use of lambda function. sp.read_lambda_threshold = 0 # Get the data again, which should trigger lambda page in. cube3 = sp.cutout(self.resource, (0, 0, 0), (self.x_dim, self.y_dim * 2, self.z_dim), 0) # Make sure the data is the same np.testing.assert_array_equal(cube1.data, cube3.data)
def test_sqs_watcher_send_message(self): """Inject message into queue and test that SqsWatcher kicks off a lambda and writes cuboid to s3.""" # Generate random data cube1 = Cube.create_cube(self.resource, [512, 512, 16]) cube1.random() cube1.morton_id = 0 sp = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) base_write_cuboid_key = "WRITE-CUBOID&{}&{}".format(self.resource.get_lookup_key(), 0) morton_idx = ndlib.XYZMorton([0, 0, 0]) t = 0 write_cuboid_key = sp.kvio.insert_cube_in_write_buffer(base_write_cuboid_key, t, morton_idx, cube1.to_blosc_by_time_index(t)) # Put page out job on the queue sqs = boto3.client('sqs', region_name=get_region()) msg_data = {"config": self.config_data, "write_cuboid_key": write_cuboid_key, "lambda-name": "s3_flush", "resource": self.resource.to_dict()} response = sqs.send_message(QueueUrl=self.object_store_config["s3_flush_queue"], MessageBody=json.dumps(msg_data)) assert response['ResponseMetadata']['HTTPStatusCode'] == 200 watcher = SqsWatcher(self.lambda_data) # verify_queue() needs the be run multiple times to verify that the queue is not changing # only then does it send off a lambda message. time.sleep(5) watcher.verify_queue() time.sleep(5) lambdas_invoked = watcher.verify_queue() if lambdas_invoked < 1: time.sleep(5) watcher.verify_queue() time.sleep(15) client = boto3.client('sqs', region_name=get_region()) response = client.get_queue_attributes( QueueUrl=self.object_store_config["s3_flush_queue"], AttributeNames=[ 'ApproximateNumberOfMessages', 'ApproximateNumberOfMessagesNotVisible' ] ) https_status_code = response['ResponseMetadata']['HTTPStatusCode'] queue_count = int(response['Attributes']['ApproximateNumberOfMessages']) # test that the queue count is now 0 assert queue_count == 0 s3 = boto3.client('s3', region_name=get_region()) objects_list = s3.list_objects(Bucket=self.object_store_config['cuboid_bucket']) # tests that bucket has some Contents. assert "Contents" in objects_list.keys()
def test_factory(self): """Test the Cube factory in Cube""" data = get_image_dict() resource = BossResourceBasic(data) c = Cube.create_cube(resource, [30, 20, 13], [0, 15]) assert isinstance(c, ImageCube8) is True assert c.cube_size == [13, 20, 30] assert c.is_time_series is True assert c.time_range == [0, 15]
def test_get_tight_bounding_box_single_cuboid(self): """ Get the tight bounding box for an object that exists within a single cuboid. """ resolution = 0 [x_cube_dim, y_cube_dim, z_cube_dim] = CUBOIDSIZE[resolution] id = 33333 id_as_str = '33333' # Customize resource with so it writes to its own channel and uses a # coord frame large enough to encompass the data written. This is # important for proper loose bounding box calculations. data = get_anno_dict(boss_key='col1&exp1&ch50', lookup_key='1&1&50') data['coord_frame']['x_stop'] = 10000 data['coord_frame']['y_stop'] = 10000 data['coord_frame']['z_stop'] = 10000 resource = BossResourceBasic(data) time_sample = 0 version = 0 x_rng = [0, x_cube_dim] y_rng = [0, y_cube_dim] z_rng = [0, z_cube_dim] t_rng = [0, 1] cube_dim_tuple = (self.x_dim, self.y_dim, self.z_dim) cube1 = Cube.create_cube(resource, [self.x_dim, self.y_dim, self.z_dim]) cube1.zeros() cube1.data[0][14][500][104] = id cube1.data[0][15][501][105] = id cube1.data[0][15][502][104] = id cube1.data[0][14][503][105] = id pos1 = [10*self.x_dim, 15*self.y_dim, 2*self.z_dim] cube1.morton_id = XYZMorton(pos1) sp = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) sp.write_cuboid(resource, pos1, resolution, cube1.data, time_sample_start=0) # Make sure cube write complete and correct. actual_cube = sp.cutout(resource, pos1, cube_dim_tuple, resolution) np.testing.assert_array_equal(cube1.data, actual_cube.data) # Method under test. actual = sp.get_bounding_box(resource, resolution, id_as_str, bb_type='tight') expected = { 'x_range': [pos1[0]+104, pos1[0]+106], 'y_range': [pos1[1]+500, pos1[1]+504], 'z_range': [pos1[2]+14, pos1[2]+16], 't_range': t_rng } self.assertEqual(expected, actual)
def test_factory_no_time(self): """Test the Cube factory in Cube""" data = get_image_dict() data['channel']['datatype'] = 'uint16' resource = BossResourceBasic(data) c = Cube.create_cube(resource, [30, 20, 13]) assert isinstance(c, ImageCube16) is True assert c.cube_size == [13, 20, 30] assert c.is_time_series is False assert c.time_range == [0, 1]
def test_factory_no_time(self): """Test the Cube factory in Cube""" data = get_anno_dict() resource = BossResourceBasic(data) c = Cube.create_cube(resource, [30, 20, 13]) assert isinstance(c, AnnotateCube64) is True assert c.cube_size == [13, 20, 30] assert c.is_time_series is False assert c.time_range == [0, 1]
def test_write_cuboid_off_base_res(self, fake_get_region): """Test writing a cuboid to not the base resolution""" # Generate random data cube1 = Cube.create_cube(self.resource, [self.x_dim, self.y_dim, self.z_dim]) cube1.random() cube1.morton_id = 0 db = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) # populate dummy data with self.assertRaises(SpdbError): db.write_cuboid(self.resource, (0, 0, 0), 5, cube1.data, time_sample_start=0)
def test_get_ids_in_region_multiple_partial_cubes(self): """ Region cuboid aligned in x, but doesn't span full cuboids in the y and z. """ cube1 = Cube.create_cube(self.resource, [self.x_dim, self.y_dim, self.z_dim]) cube1.zeros() cube1.data[0][0][40][0] = 55555 cube1.data[0][0][50][0] = 66666 pos1 = [4*self.x_dim, 4*self.y_dim, 2*self.z_dim] cube1.morton_id = XYZMorton(pos1) cube2 = Cube.create_cube(self.resource, [self.x_dim, self.y_dim, self.z_dim]) cube2.zeros() cube2.data[0][0][40][0] = 55555 cube2.data[0][0][50][0] = 77777 pos2 = [5*self.x_dim, 4*self.y_dim, 2*self.z_dim] cube2.morton_id = XYZMorton(pos2) sp = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) resolution = 0 sp.write_cuboid(self.resource, pos1, resolution, cube1.data, time_sample_start=0) sp.write_cuboid(self.resource, pos2, resolution, cube2.data, time_sample_start=0) # Not verifying writes here because get_ids_in_region() should be doing # cutouts due to the region not containing full cuboids. corner = (4*self.x_dim, 4*self.y_dim, 2*self.z_dim) extent = (2*self.x_dim, 60, 10) t_range = [0, 1] version = 0 expected = ['55555', '66666', '77777'] # Method under test. actual = sp.get_ids_in_region( self.resource, resolution, corner, extent, t_range, version) self.assertIn('ids', actual) self.assertCountEqual(expected, actual['ids'])
def test_add_to_prefetch(self): cuboid_dims = CUBOIDSIZE[0] # Cuboid dimensions. x_dim = cuboid_dims[0] y_dim = cuboid_dims[1] z_dim = cuboid_dims[2] cube_above = Cube.create_cube(self.resource, [x_dim, y_dim, z_dim]) cube_above.random() # Write cuboid that are stacked vertically. self.sp.write_cuboid(self.resource, (0, 0, z_dim * 2), 0, cube_above.data) cube_above.morton_id = ndlib.XYZMorton([0, 0, z_dim * 2 // z_dim]) cube_above_cache_key = self.sp.kvio.generate_cached_cuboid_keys( self.resource, 0, [0], [cube_above.morton_id]) # Make sure cuboid saved. cube_act = self.sp.cutout(self.resource, (0, 0, z_dim * 2), (x_dim, y_dim, z_dim), 0) np.testing.assert_array_equal(cube_above.data, cube_act.data) # Clear cache so we can test prefetch. self.sp.kvio.cache_client.flushdb() # Also clear cache state before running test. self.sp.cache_state.status_client.flushdb() obj_keys = self.sp.objectio.cached_cuboid_to_object_keys( cube_above_cache_key) # Place a cuboid in the pretch queue. self.sp.cache_state.status_client.rpush('PRE-FETCH', obj_keys[0]) # This is the system under test. self.prefetch.process() # Wait for cube to be prefetched. i = 0 while not self.sp.kvio.cube_exists(cube_above_cache_key[0]) and i < 30: time.sleep(1) i += 1 # Confirm cuboid now in cache. self.assertTrue(self.sp.kvio.cube_exists(cube_above_cache_key[0])) cube_act = self.sp.cutout(self.resource, (0, 0, z_dim * 2), (x_dim, y_dim, z_dim), 0) np.testing.assert_array_equal(cube_above.data, cube_act.data)
def test_cutout_no_time_single_unaligned_hit(self): """Test the get_cubes method - no time - single - unaligned - hit""" # Generate random data cube1 = Cube.create_cube(self.resource, [self.x_dim, self.y_dim, self.z_dim]) cube1.random() cube1.morton_id = 0 sp = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) sp.write_cuboid(self.resource, (600, 0, 0), 0, cube1.data) cube2 = sp.cutout(self.resource, (600, 0, 0), (self.x_dim, self.y_dim, 16), 0) np.testing.assert_array_equal(cube1.data, cube2.data)
def test_cutout_time_offset_single_aligned_hit(self): """Test the get_cubes method - w/ time - single - hit""" # Generate random data cube1 = Cube.create_cube(self.resource, [self.x_dim, self.y_dim, self.z_dim], time_range=[0, 3]) cube1.random() cube1.morton_id = 0 sp = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) sp.write_cuboid(self.resource, (0, 0, 0), 0, cube1.data, time_sample_start=6) cube2 = sp.cutout(self.resource, (0, 0, 0), (self.x_dim, self.y_dim, self.z_dim), 0, time_sample_range=[6, 9]) np.testing.assert_array_equal(cube1.data, cube2.data)
def test_add_to_prefetch(self): cuboid_dims = CUBOIDSIZE[0] # Cuboid dimensions. x_dim = cuboid_dims[0] y_dim = cuboid_dims[1] z_dim = cuboid_dims[2] cube_above = Cube.create_cube(self.resource, [x_dim, y_dim, z_dim]) cube_above.random() # Write cuboid that are stacked vertically. self.sp.write_cuboid(self.resource, (0, 0, z_dim * 2), 0, cube_above.data) cube_above.morton_id = ndlib.XYZMorton([0, 0, z_dim * 2 // z_dim]) cube_above_cache_key = self.sp.kvio.generate_cached_cuboid_keys( self.resource, 0, [0], [cube_above.morton_id]) # Make sure cuboid saved. cube_act = self.sp.cutout(self.resource, (0, 0, z_dim * 2), (x_dim, y_dim, z_dim), 0) np.testing.assert_array_equal(cube_above.data, cube_act.data) # Clear cache so we can test prefetch. self.sp.kvio.cache_client.flushdb() # Also clear cache state before running test. self.sp.cache_state.status_client.flushdb() obj_keys = self.sp.objectio.cached_cuboid_to_object_keys( cube_above_cache_key) # Place a cuboid in the pretch queue. self.sp.cache_state.status_client.rpush('PRE-FETCH', obj_keys[0]) # This is the system under test. self.prefetch.process() # Wait for cube to be prefetched. i = 0 while not self.sp.kvio.cube_exists(cube_above_cache_key[0]) and i < 30: time.sleep(1) i += 1 # Confirm cuboid now in cache. self.assertTrue(self.sp.kvio.cube_exists(cube_above_cache_key[0])) cube_act = self.sp.cutout( self.resource, (0, 0, z_dim * 2), (x_dim, y_dim, z_dim), 0) np.testing.assert_array_equal(cube_above.data, cube_act.data)
def test_cutout_no_time_single_aligned_miss(self, fake_get_region): """Test the get_cubes method - no time - single""" # Generate random data cube1 = Cube.create_cube(self.resource, [self.x_dim, self.y_dim, self.z_dim]) cube1.random() cube1.morton_id = 0 db = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) # populate dummy data self.write_test_cube(db, self.resource, 0, cube1, cache=False, s3=True) cube2 = db.cutout(self.resource, (0, 0, 0), (self.x_dim, self.y_dim, self.z_dim), 0) np.testing.assert_array_equal(cube1.data, cube2.data)
def test_cutout_iso_below_fork(self): """Test write_cuboid and cutout methods with iso option, testing iso is equal below the res fork""" # Generate random data cube1 = Cube.create_cube(self.resource, [400, 400, 8]) cube1.random() cube1.morton_id = 0 cubez = Cube.create_cube(self.resource, [400, 400, 8]) cubez.zeros() cubez.morton_id = 0 # Write at 5, not iso, and verify sp = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) sp.write_cuboid(self.resource, (200, 600, 3), 0, cube1.data, iso=False) cube2 = sp.cutout(self.resource, (200, 600, 3), (400, 400, 8), 0, iso=False) np.testing.assert_array_equal(cube1.data, cube2.data) # Get at res 5 iso, which should be equal to non-iso call cube2 = sp.cutout(self.resource, (200, 600, 3), (400, 400, 8), 0, iso=True) np.testing.assert_array_equal(cube1.data, cube2.data)
def test_cutout_no_time_single_aligned_existing_hit(self): """Test the get_cubes method - no time - aligned - existing data - miss""" # Generate random data cube1 = Cube.create_cube(self.resource, [self.x_dim, self.y_dim, self.z_dim]) cube1.random() sp = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) sp.write_cuboid(self.resource, (0, 0, 0), 0, cube1.data) cube2 = sp.cutout(self.resource, (0, 0, 0), (self.x_dim, self.y_dim, self.z_dim), 0) np.testing.assert_array_equal(cube1.data, cube2.data) del cube1 del cube2 # now write to cuboid again cube3 = Cube.create_cube(self.resource, [self.x_dim, self.y_dim, self.z_dim]) cube3.random() sp.write_cuboid(self.resource, (0, 0, 0), 0, cube3.data) cube4 = sp.cutout(self.resource, (0, 0, 0), (self.x_dim, self.y_dim, self.z_dim), 0) np.testing.assert_array_equal(cube3.data, cube4.data)
def test_get_cubes_no_time_single(self, fake_get_region): """Test the get_cubes method - no time - single""" # Generate random data cube1 = Cube.create_cube(self.resource, [self.x_dim, self.y_dim, self.z_dim]) cube1.random() cube1.morton_id = 32 db = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) # populate dummy data keys = self.write_test_cube(db, self.resource, 0, cube1, cache=True, s3=False) cube2 = db.get_cubes(self.resource, keys) np.testing.assert_array_equal(cube1.data, cube2[0].data)
def test_filtered_cutout_bad_id_list(self): time_axis = [1] cube_dim = [self.x_dim, self.y_dim, self.z_dim] cube_dim_tuple = (self.x_dim, self.y_dim, self.z_dim) cube1 = Cube.create_cube(self.resource, cube_dim) cube1.data = np.ones(time_axis + [cube_dim[2], cube_dim[1], cube_dim[0]], dtype='uint64') cube1.morton_id = 0 corner = (6*self.x_dim, 6*self.y_dim, 2*self.z_dim) sp = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) resolution = 0 sp.write_cuboid(self.resource, corner, resolution, cube1.data, time_sample_start=0) # Method under test. with self.assertRaises(SpdbError): sp.cutout(self.resource, corner, cube_dim_tuple, resolution, filter_ids=['foo', 55555])
def test_cutout_no_time_multi_unaligned_hit_iso_below(self): """Test write_cuboid and cutout methods - no time - multi - unaligned - hit - isotropic, below iso fork""" # Generate random data cube1 = Cube.create_cube(self.resource, [400, 400, 8]) cube1.random() cube1.morton_id = 0 sp = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) sp.write_cuboid(self.resource, (200, 600, 3), 0, cube1.data, iso=True) cube2 = sp.cutout(self.resource, (200, 600, 3), (400, 400, 8), 0, iso=True) np.testing.assert_array_equal(cube1.data, cube2.data) # do it again...should be in cache cube2 = sp.cutout(self.resource, (200, 600, 3), (400, 400, 8), 0, iso=True) np.testing.assert_array_equal(cube1.data, cube2.data)
def test_cutout_no_time_multi_unaligned_hit(self): """Test the get_cubes method - no time - multi - unaligned - hit""" # Generate random data cube1 = Cube.create_cube(self.resource, [400, 400, 8]) cube1.random() cube1.morton_id = 0 sp = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) sp.write_cuboid(self.resource, (200, 600, 3), 0, cube1.data) cube2 = sp.cutout(self.resource, (200, 600, 3), (400, 400, 8), 0) np.testing.assert_array_equal(cube1.data, cube2.data) # do it again...shoudl be in cache cube2 = sp.cutout(self.resource, (200, 600, 3), (400, 400, 8), 0) np.testing.assert_array_equal(cube1.data, cube2.data)
def test_set_write_locked(self): # Cuboid dimensions. xy_dim = 128 z_dim = 16 lookup_key = self.data['lookup_key'] # Make sure this key isn't currently locked. self.dead_letter._sp.cache_state.set_project_lock(lookup_key, False) self.assertFalse( self.dead_letter._sp.cache_state.project_locked(lookup_key)) cube1 = Cube.create_cube(self.resource, [xy_dim, xy_dim, z_dim]) cube1.data = np.random.randint(1, 254, (1, z_dim, xy_dim, xy_dim)) # Ensure that the we are not in a page-out state by wiping the entire # cache state. self.dead_letter._sp.cache_state.status_client.flushdb() self.dead_letter._sp.write_cuboid(self.resource, (0, 0, 0), 0, cube1.data) try: with patch.object( self.dead_letter, 'send_alert', wraps=self.dead_letter.send_alert) as send_alert_spy: # Method under test. Returns True if it found a message. i = 0 while not self.dead_letter.check_queue() and i < 30: time.sleep(1) i += 1 self.assertTrue( self.dead_letter._sp.cache_state.project_locked( lookup_key)) # Ensure method that publishes to SNS topic called. send_alert_spy.assert_called_with(lookup_key, ANY) finally: # Make sure write lock is unset before terminating. self.dead_letter._sp.cache_state.set_project_lock( lookup_key, False)
def test_delayed_write_daemon_simple(self): """Test handling delayed writes""" sp = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) dwd = DelayedWriteDaemon("boss-delayedwrited-test.pid") # Create some a single delayed write cube1 = Cube.create_cube(self.resource, [512, 512, 16]) cube1.random() cube1.morton_id = 0 res = 0 time_sample = 0 write_cuboid_base = "WRITE-CUBOID&{}&{}".format(self.resource.get_lookup_key(), 0) write_cuboid_key = sp.kvio.insert_cube_in_write_buffer(write_cuboid_base, res, cube1.morton_id, cube1.to_blosc_by_time_index(time_sample)) sp.cache_state.add_to_delayed_write(write_cuboid_key, self.resource.get_lookup_key(), res, cube1.morton_id, time_sample, self.resource.to_json()) # Use Daemon To handle writes dwd.process(sp) time.sleep(30) # Make sure they went through cube2 = sp.cutout(self.resource, (0, 0, 0), (512, 512, 16), 0) np.testing.assert_array_equal(cube1.data, cube2.data) # Make sure delay key got deleted keys = sp.cache_state.get_all_delayed_write_keys() assert not keys
def test_cutout_no_time_multi_unaligned_hit_iso_above(self): """Test write_cuboid and cutout methods - no time - multi - unaligned - hit - isotropic, above iso fork""" data = self.data data["channel"]["base_resolution"] = 5 resource = BossResourceBasic(data) # Generate random data cube1 = Cube.create_cube(resource, [400, 400, 8]) cube1.random() cube1.morton_id = 0 sp = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) sp.write_cuboid(resource, (200, 600, 3), 5, cube1.data, iso=True) cube2 = sp.cutout(resource, (200, 600, 3), (400, 400, 8), 5, iso=True) np.testing.assert_array_equal(cube1.data, cube2.data) # do it again...should be in cache cube2 = sp.cutout(resource, (200, 600, 3), (400, 400, 8), 5, iso=True) np.testing.assert_array_equal(cube1.data, cube2.data)
def test_set_write_locked(self): # Cuboid dimensions. xy_dim = 128 z_dim = 16 lookup_key = self.data['lookup_key'] # Make sure this key isn't currently locked. self.dead_letter._sp.cache_state.set_project_lock(lookup_key, False) self.assertFalse(self.dead_letter._sp.cache_state.project_locked(lookup_key)) cube1 = Cube.create_cube(self.resource, [xy_dim, xy_dim, z_dim]) cube1.data = np.random.randint(1, 254, (1, z_dim, xy_dim, xy_dim)) # Ensure that the we are not in a page-out state by wiping the entire # cache state. self.dead_letter._sp.cache_state.status_client.flushdb() self.dead_letter._sp.write_cuboid( self.resource, (0, 0, 0), 0, cube1.data) try: with patch.object(self.dead_letter, 'send_alert', wraps=self.dead_letter.send_alert) as send_alert_spy: # Method under test. Returns True if it found a message. i = 0 while not self.dead_letter.check_queue() and i < 30: time.sleep(1) i += 1 self.assertTrue(self.dead_letter._sp.cache_state.project_locked(lookup_key)) # Ensure method that publishes to SNS topic called. send_alert_spy.assert_called_with(lookup_key, ANY) finally: # Make sure write lock is unset before terminating. self.dead_letter._sp.cache_state.set_project_lock(lookup_key, False)
def test_delayed_write_daemon_simple(self): """Test handling delayed writes""" sp = SpatialDB(self.kvio_config, self.state_config, self.object_store_config) dwd = DelayedWriteDaemon("boss-delayedwrited-test.pid") # Create some a single delayed write cube1 = Cube.create_cube(self.resource, [512, 512, 16]) cube1.random() cube1.morton_id = 0 res = 0 time_sample = 0 write_cuboid_base = "WRITE-CUBOID&{}&{}".format( self.resource.get_lookup_key(), 0) write_cuboid_key = sp.kvio.insert_cube_in_write_buffer( write_cuboid_base, res, cube1.morton_id, cube1.to_blosc_by_time_index(time_sample)) sp.cache_state.add_to_delayed_write(write_cuboid_key, self.resource.get_lookup_key(), res, cube1.morton_id, time_sample, self.resource.to_json()) # Use Daemon To handle writes dwd.process(sp) time.sleep(30) # Make sure they went through cube2 = sp.cutout(self.resource, (0, 0, 0), (512, 512, 16), 0) np.testing.assert_array_equal(cube1.data, cube2.data) # Make sure delay key got deleted keys = sp.cache_state.get_all_delayed_write_keys() assert not keys
def upload_data(self, session, args): """ Fill the coord frame with random data. Args: args (dict): This should be the dict returned by get_downsample_args(). """ cuboid_size = CUBOIDSIZE[0] x_dim = cuboid_size[0] y_dim = cuboid_size[1] z_dim = cuboid_size[2] resource = BossResourceBasic() resource.from_dict(self.get_image_dict()) resolution = 0 ts = 0 version = 0 # DP HACK: uploading all cubes will take longer than the actual downsample # just upload the first volume worth of cubes. # The downsample volume lambda will only read these cubes when # passed the 'test' argument. bucket = S3Bucket(session, args['s3_bucket']) print('Uploading test data', end='', flush=True) for cube in xyz_range(XYZ(0,0,0), XYZ(2,2,2)): key = AWSObjectStore.generate_object_key(resource, resolution, ts, cube.morton) key += "&0" # Add the version number #print('morton: {}'.format(cube.morton)) #print('key: {}'.format(key)) #print("{} -> {} -> {}".format(cube, cube.morton, key)) cube = Cube.create_cube(resource, [x_dim, y_dim, z_dim]) cube.random() data = cube.to_blosc() bucket.put(key, data) print('.', end='', flush=True) print(' Done uploading.')
def upload_data(self, args): """ Fill the coord frame with random data. Args: args (dict): This should be the dict returned by get_downsample_args(). """ cuboid_size = CUBOIDSIZE[0] x_dim = cuboid_size[0] y_dim = cuboid_size[1] z_dim = cuboid_size[2] resource = BossResourceBasic() resource.from_dict(self.get_image_dict()) resolution = 0 ts = 0 version = 0 # DP HACK: uploading all cubes will take longer than the actual downsample # just upload the first volume worth of cubes. # The downsample volume lambda will only read these cubes when # passed the 'test' argument. bucket = S3Bucket(self.bosslet_config.session, args['s3_bucket']) print('Uploading test data', end='', flush=True) for cube in xyz_range(XYZ(0,0,0), XYZ(2,2,2)): key = AWSObjectStore.generate_object_key(resource, resolution, ts, cube.morton) key += "&0" # Add the version number #print('morton: {}'.format(cube.morton)) #print('key: {}'.format(key)) #print("{} -> {} -> {}".format(cube, cube.morton, key)) cube = Cube.create_cube(resource, [x_dim, y_dim, z_dim]) cube.random() data = cube.to_blosc() bucket.put(key, data) print('.', end='', flush=True) print(' Done uploading.')
def handler(event, context): # Load settings SETTINGS = BossSettings.load() # Used as a guard against trying to delete the SQS message when lambda is # triggered by SQS. sqs_triggered = 'Records' in event and len(event['Records']) > 0 if sqs_triggered : # Lambda invoked by an SQS trigger. msg_data = json.loads(event['Records'][0]['body']) # Load the project info from the chunk key you are processing chunk_key = msg_data['chunk_key'] proj_info = BossIngestProj.fromSupercuboidKey(chunk_key) proj_info.job_id = msg_data['ingest_job'] else: # Standard async invoke of this lambda. # Load the project info from the chunk key you are processing proj_info = BossIngestProj.fromSupercuboidKey(event["chunk_key"]) proj_info.job_id = event["ingest_job"] # Get message from SQS ingest queue, try for ~2 seconds rx_cnt = 0 msg_data = None msg_id = None msg_rx_handle = None while rx_cnt < 6: ingest_queue = IngestQueue(proj_info) msg = [x for x in ingest_queue.receiveMessage()] if msg: msg = msg[0] print("MESSAGE: {}".format(msg)) print(len(msg)) msg_id = msg[0] msg_rx_handle = msg[1] msg_data = json.loads(msg[2]) print("MESSAGE DATA: {}".format(msg_data)) break else: rx_cnt += 1 print("No message found. Try {} of 6".format(rx_cnt)) time.sleep(1) if not msg_id: # No tiles ready to ingest. print("No ingest message available") return # Get the chunk key of the tiles to ingest. chunk_key = msg_data['chunk_key'] tile_error_queue = TileErrorQueue(proj_info) print("Ingesting Chunk {}".format(chunk_key)) tiles_in_chunk = int(chunk_key.split('&')[1]) # Setup SPDB instance sp = SpatialDB(msg_data['parameters']["KVIO_SETTINGS"], msg_data['parameters']["STATEIO_CONFIG"], msg_data['parameters']["OBJECTIO_CONFIG"]) # Get tile list from Tile Index Table tile_index_db = BossTileIndexDB(proj_info.project_name) # tile_index_result (dict): keys are S3 object keys of the tiles comprising the chunk. tile_index_result = tile_index_db.getCuboid(msg_data["chunk_key"], int(msg_data["ingest_job"])) if tile_index_result is None: # If chunk_key is gone, another lambda uploaded the cuboids and deleted the chunk_key afterwards. if not sqs_triggered: # Remove message so it's not redelivered. ingest_queue.deleteMessage(msg_id, msg_rx_handle) print("Aborting due to chunk key missing from tile index table") return # Sort the tile keys print("Tile Keys: {}".format(tile_index_result["tile_uploaded_map"])) tile_key_list = [x.rsplit("&", 2) for x in tile_index_result["tile_uploaded_map"].keys()] if len(tile_key_list) < tiles_in_chunk: print("Not a full set of 16 tiles. Assuming it has handled already, tiles: {}".format(len(tile_key_list))) if not sqs_triggered: ingest_queue.deleteMessage(msg_id, msg_rx_handle) return tile_key_list = sorted(tile_key_list, key=lambda x: int(x[1])) tile_key_list = ["&".join(x) for x in tile_key_list] print("Sorted Tile Keys: {}".format(tile_key_list)) # Augment Resource JSON data so it will instantiate properly that was pruned due to S3 metadata size limits resource_dict = msg_data['parameters']['resource'] _, exp_name, ch_name = resource_dict["boss_key"].split("&") resource_dict["channel"]["name"] = ch_name resource_dict["channel"]["description"] = "" resource_dict["channel"]["sources"] = [] resource_dict["channel"]["related"] = [] resource_dict["channel"]["default_time_sample"] = 0 resource_dict["channel"]["downsample_status"] = "NOT_DOWNSAMPLED" resource_dict["experiment"]["name"] = exp_name resource_dict["experiment"]["description"] = "" resource_dict["experiment"]["num_time_samples"] = 1 resource_dict["experiment"]["time_step"] = None resource_dict["experiment"]["time_step_unit"] = None resource_dict["coord_frame"]["name"] = "cf" resource_dict["coord_frame"]["name"] = "" resource_dict["coord_frame"]["x_start"] = 0 resource_dict["coord_frame"]["x_stop"] = 100000 resource_dict["coord_frame"]["y_start"] = 0 resource_dict["coord_frame"]["y_stop"] = 100000 resource_dict["coord_frame"]["z_start"] = 0 resource_dict["coord_frame"]["z_stop"] = 100000 resource_dict["coord_frame"]["voxel_unit"] = "nanometers" # Setup the resource resource = BossResourceBasic() resource.from_dict(resource_dict) dtype = resource.get_numpy_data_type() # read all tiles from bucket into a slab tile_bucket = TileBucket(proj_info.project_name) data = [] num_z_slices = 0 for tile_key in tile_key_list: try: image_data, message_id, receipt_handle, metadata = tile_bucket.getObjectByKey(tile_key) except KeyError: print('Key: {} not found in tile bucket, assuming redelivered SQS message and aborting.'.format( tile_key)) if not sqs_triggered: # Remove message so it's not redelivered. ingest_queue.deleteMessage(msg_id, msg_rx_handle) print("Aborting due to missing tile in bucket") return image_bytes = BytesIO(image_data) image_size = image_bytes.getbuffer().nbytes # Get tiles size from metadata, need to shape black tile if actual tile is corrupt. if 'x_size' in metadata: tile_size_x = metadata['x_size'] else: print('MetadataMissing: x_size not in tile metadata: using 1024.') tile_size_x = 1024 if 'y_size' in metadata: tile_size_y = metadata['y_size'] else: print('MetadataMissing: y_size not in tile metadata: using 1024.') tile_size_y = 1024 if image_size == 0: print('TileError: Zero length tile, using black instead: {}'.format(tile_key)) error_msg = 'Zero length tile' enqueue_tile_error(tile_error_queue, tile_key, chunk_key, error_msg) tile_img = np.zeros((tile_size_x, tile_size_y), dtype=dtype) else: try: tile_img = np.asarray(Image.open(image_bytes), dtype=dtype) except TypeError as te: print('TileError: Incomplete tile, using black instead (tile_size_in_bytes, tile_key): {}, {}' .format(image_size, tile_key)) error_msg = 'Incomplete tile' enqueue_tile_error(tile_error_queue, tile_key, chunk_key, error_msg) tile_img = np.zeros((tile_size_x, tile_size_y), dtype=dtype) except OSError as oe: print('TileError: OSError, using black instead (tile_size_in_bytes, tile_key): {}, {} ErrorMessage: {}' .format(image_size, tile_key, oe)) error_msg = 'OSError: {}'.format(oe) enqueue_tile_error(tile_error_queue, tile_key, chunk_key, error_msg) tile_img = np.zeros((tile_size_x, tile_size_y), dtype=dtype) data.append(tile_img) num_z_slices += 1 # Make 3D array of image data. It should be in XYZ at this point chunk_data = np.array(data) del data tile_dims = chunk_data.shape # Break into Cube instances print("Tile Dims: {}".format(tile_dims)) print("Num Z Slices: {}".format(num_z_slices)) num_x_cuboids = int(math.ceil(tile_dims[2] / CUBOIDSIZE[proj_info.resolution][0])) num_y_cuboids = int(math.ceil(tile_dims[1] / CUBOIDSIZE[proj_info.resolution][1])) print("Num X Cuboids: {}".format(num_x_cuboids)) print("Num Y Cuboids: {}".format(num_y_cuboids)) chunk_key_parts = BossUtil.decode_chunk_key(chunk_key) t_index = chunk_key_parts['t_index'] for x_idx in range(0, num_x_cuboids): for y_idx in range(0, num_y_cuboids): # TODO: check time series support cube = Cube.create_cube(resource, CUBOIDSIZE[proj_info.resolution]) cube.zeros() # Compute Morton ID # TODO: verify Morton indices correct! print(chunk_key_parts) morton_x_ind = x_idx + (chunk_key_parts["x_index"] * num_x_cuboids) morton_y_ind = y_idx + (chunk_key_parts["y_index"] * num_y_cuboids) print("Morton X: {}".format(morton_x_ind)) print("Morton Y: {}".format(morton_y_ind)) morton_index = XYZMorton([morton_x_ind, morton_y_ind, int(chunk_key_parts['z_index'])]) # Insert sub-region from chunk_data into cuboid x_start = x_idx * CUBOIDSIZE[proj_info.resolution][0] x_end = x_start + CUBOIDSIZE[proj_info.resolution][0] x_end = min(x_end, tile_dims[2]) y_start = y_idx * CUBOIDSIZE[proj_info.resolution][1] y_end = y_start + CUBOIDSIZE[proj_info.resolution][1] y_end = min(y_end, tile_dims[1]) z_end = CUBOIDSIZE[proj_info.resolution][2] # TODO: get sub-array w/o making a copy. print("Yrange: {}".format(y_end - y_start)) print("Xrange: {}".format(x_end - x_start)) print("X start: {}".format(x_start)) print("X stop: {}".format(x_end)) cube.data[0, 0:num_z_slices, 0:(y_end - y_start), 0:(x_end - x_start)] = chunk_data[0:num_z_slices, y_start:y_end, x_start:x_end] # Create object key object_key = sp.objectio.generate_object_key(resource, proj_info.resolution, t_index, morton_index) print("Object Key: {}".format(object_key)) # Put object in S3 sp.objectio.put_objects([object_key], [cube.to_blosc()]) # Add object to index sp.objectio.add_cuboid_to_index(object_key, ingest_job=int(msg_data["ingest_job"])) # Update id indices if this is an annotation channel # We no longer index during ingest. #if resource.data['channel']['type'] == 'annotation': # try: # sp.objectio.update_id_indices( # resource, proj_info.resolution, [object_key], [cube.data]) # except SpdbError as ex: # sns_client = boto3.client('sns') # topic_arn = msg_data['parameters']["OBJECTIO_CONFIG"]["prod_mailing_list"] # msg = 'During ingest:\n{}\nCollection: {}\nExperiment: {}\n Channel: {}\n'.format( # ex.message, # resource.data['collection']['name'], # resource.data['experiment']['name'], # resource.data['channel']['name']) # sns_client.publish( # TopicArn=topic_arn, # Subject='Object services misuse', # Message=msg) lambda_client = boto3.client('lambda', region_name=SETTINGS.REGION_NAME) names = AWSNames.create_from_lambda_name(context.function_name) delete_tiles_data = { 'tile_key_list': tile_key_list, 'region': SETTINGS.REGION_NAME, 'bucket': tile_bucket.bucket.name } # Delete tiles from tile bucket. lambda_client.invoke( FunctionName=names.delete_tile_objs_lambda, InvocationType='Event', Payload=json.dumps(delete_tiles_data).encode() ) delete_tile_entry_data = { 'tile_index': tile_index_db.table.name, 'region': SETTINGS.REGION_NAME, 'chunk_key': chunk_key, 'task_id': msg_data['ingest_job'] } # Delete entry from tile index. lambda_client.invoke( FunctionName=names.delete_tile_index_entry_lambda, InvocationType='Event', Payload=json.dumps(delete_tile_entry_data).encode() ) if not sqs_triggered: # Delete message since it was processed successfully ingest_queue.deleteMessage(msg_id, msg_rx_handle)
print("Tile Dims: {}".format(tile_dims)) print("Num Z Slices: {}".format(num_z_slices)) num_x_cuboids = int(math.ceil(tile_dims[2] / CUBOIDSIZE[proj_info.resolution][0])) num_y_cuboids = int(math.ceil(tile_dims[1] / CUBOIDSIZE[proj_info.resolution][1])) print("Num X Cuboids: {}".format(num_x_cuboids)) print("Num Y Cuboids: {}".format(num_y_cuboids)) # Cuboid List cuboids = [] chunk_key_parts = BossUtil.decode_chunk_key(chunk_key) t_index = chunk_key_parts['t_index'] for x_idx in range(0, num_x_cuboids): for y_idx in range(0, num_y_cuboids): # TODO: check time series support cube = Cube.create_cube(resource, CUBOIDSIZE[proj_info.resolution]) cube.zeros() # Compute Morton ID # TODO: verify Morton indices correct! print(chunk_key_parts) morton_x_ind = x_idx + (chunk_key_parts["x_index"] * num_x_cuboids) morton_y_ind = y_idx + (chunk_key_parts["y_index"] * num_y_cuboids) print("Morton X: {}".format(morton_x_ind)) print("Morton Y: {}".format(morton_y_ind)) morton_index = XYZMorton([morton_x_ind, morton_y_ind, int(chunk_key_parts['z_index'])]) # Insert sub-region from chunk_data into cuboid x_start = x_idx * CUBOIDSIZE[proj_info.resolution][0] x_end = x_start + CUBOIDSIZE[proj_info.resolution][0] x_end = min(x_end, tile_dims[2])
def test_add_to_prefetch(self): cuboid_dims = CUBOIDSIZE[0] # Cuboid dimensions. x_dim = cuboid_dims[0] y_dim = cuboid_dims[1] z_dim = cuboid_dims[2] cube = Cube.create_cube(self.resource, [x_dim, y_dim, z_dim]) cube.random() cube_above = Cube.create_cube(self.resource, [x_dim, y_dim, z_dim]) cube_above.random() cube_below = Cube.create_cube(self.resource, [x_dim, y_dim, z_dim]) cube_below.random() # Write 3 cuboids that are stacked vertically. self.sp.write_cuboid(self.resource, (0, 0, 0), 0, cube_below.data) self.sp.write_cuboid(self.resource, (0, 0, z_dim), 0, cube.data) self.sp.write_cuboid(self.resource, (0, 0, z_dim * 2), 0, cube_above.data) cube.morton_id = ndlib.XYZMorton([0, 0, z_dim // z_dim]) cube_below.morton_id = ndlib.XYZMorton([0, 0, 0]) cube_above.morton_id = ndlib.XYZMorton([0, 0, z_dim * 2 // z_dim]) print('mortons: {}, {}, {}'.format( cube_below.morton_id, cube.morton_id, cube_above.morton_id)) cube_below_cache_key, cube_cache_key, cube_above_cache_key = self.sp.kvio.generate_cached_cuboid_keys( self.resource, 0, [0], [cube_below.morton_id, cube.morton_id, cube_above.morton_id]) # Make sure cuboids saved. cube_act = self.sp.cutout(self.resource, (0, 0, 0), (x_dim, y_dim, z_dim), 0) np.testing.assert_array_equal(cube_below.data, cube_act.data) cube_act = self.sp.cutout(self.resource, (0, 0, z_dim), (x_dim, y_dim, z_dim), 0) np.testing.assert_array_equal(cube.data, cube_act.data) cube_act = self.sp.cutout(self.resource, (0, 0, z_dim * 2), (x_dim, y_dim, z_dim), 0) np.testing.assert_array_equal(cube_above.data, cube_act.data) # Clear cache so we can get a cache miss. self.sp.kvio.cache_client.flushdb() # Also clear CACHE-MISS before running testing. self.sp.cache_state.status_client.flushdb() # Get middle cube again. This should trigger a cache miss. cube_act = self.sp.cutout(self.resource, (0, 0, z_dim), (x_dim, y_dim, z_dim), 0) # Confirm there is a cache miss. misses = self.sp.cache_state.status_client.lrange('CACHE-MISS', 0, 10) print('misses:') print(misses) miss_actual = self.sp.cache_state.status_client.lindex('CACHE-MISS', 0) self.assertEqual(cube_cache_key, str(miss_actual, 'utf-8')) # This is the system under test. self.cache_miss.process() # Confirm PRE-FETCH has the object keys for the cube above and below. fetch_actual1 = self.sp.cache_state.status_client.lindex('PRE-FETCH', 0) fetch_actual2 = self.sp.cache_state.status_client.lindex('PRE-FETCH', 1) obj_keys = self.sp.objectio.cached_cuboid_to_object_keys( [cube_above_cache_key, cube_below_cache_key]) self.assertEqual(obj_keys[0], str(fetch_actual1, 'utf-8')) self.assertEqual(obj_keys[1], str(fetch_actual2, 'utf-8'))
if exist_keys: # Cuboid Exists # Get cuboid to flush from write buffer write_cuboid_bytes = sp.kvio.get_cube_from_write_buffer(write_cuboid_key) if write_cuboid_bytes is None: # Didn't get any data back. Assume another lambda already # served this request. Remove message and continue. print("No data returned from write buffer, ignoring and deleting message.") sqs_client.delete_message( QueueUrl=event["config"]["object_store_config"]["s3_flush_queue"], ReceiptHandle=rx_handle) # Increment run counter run_cnt += 1 continue new_cube = Cube.create_cube(resource, cube_dim) new_cube.morton_id = morton new_cube.from_blosc(write_cuboid_bytes) # Get existing cuboid from S3 existing_cube = Cube.create_cube(resource, cube_dim) existing_cube.morton_id = new_cube.morton_id existing_cube_bytes = sp.objectio.get_single_object(object_keys[0]) existing_cube.from_blosc(existing_cube_bytes, new_cube.time_range) # Merge cuboids existing_cube.overwrite(new_cube.data, new_cube.time_range) # Get bytes cuboid_bytes = existing_cube.to_blosc() uncompressed_cuboid_bytes = existing_cube.data