def open_ml_dataset_from_local_fs(
ctx: ServiceContext,
dataset_descriptor: DatasetDescriptor) -> MultiLevelDataset:
ds_id = dataset_descriptor.get('Identifier')
path = dataset_descriptor.get('Path')
if not path:
raise ServiceConfigError(
f"Missing 'path' entry in dataset descriptor {ds_id}")
if not os.path.isabs(path):
path = os.path.join(ctx.base_dir, path)
data_format = dataset_descriptor.get('Format', guess_cube_format(path))
if data_format == FORMAT_NAME_NETCDF4:
with measure_time(tag=f"opened local NetCDF dataset {path}"):
ds = assert_cube(xr.open_dataset(path))
return BaseMultiLevelDataset(ds)
if data_format == FORMAT_NAME_ZARR:
with measure_time(tag=f"opened local zarr dataset {path}"):
ds = assert_cube(xr.open_zarr(path))
return BaseMultiLevelDataset(ds)
if data_format == FORMAT_NAME_LEVELS:
with measure_time(tag=f"opened local levels dataset {path}"):
return FileStorageMultiLevelDataset(path)
raise ServiceConfigError(
f"Illegal data format {data_format!r} for dataset {ds_id}")
def test_it(self):
ds = _get_test_dataset()
ml_ds1 = BaseMultiLevelDataset(ds)
def input_ml_dataset_getter(ds_id):
if ds_id == "ml_ds1":
return ml_ds1
self.fail(f"unexpected ds_id={ds_id!r}")
ml_ds2 = ComputedMultiLevelDataset(os.path.join(os.path.dirname(__file__),
"..", "webapi", "res", "test", "script.py"),
"compute_dataset",
["ml_ds1"],
input_ml_dataset_getter,
input_parameters=dict(period='1W'),
ds_id="ml_ds2")
self.assertEqual(3, ml_ds2.num_levels)
self.assertEqual(TileGrid(3, 2, 1, 180, 180, (-180, -90, 180, 90), inv_y=False),
ml_ds2.tile_grid)
ds0 = ml_ds2.get_dataset(0)
self.assertEqual({'time': 3, 'lat': 720, 'lon': 1440, 'bnds': 2}, ds0.dims)
ds1 = ml_ds2.get_dataset(1)
self.assertEqual({'time': 3, 'lat': 360, 'lon': 720}, ds1.dims)
ds2 = ml_ds2.get_dataset(2)
self.assertEqual({'time': 3, 'lat': 180, 'lon': 360}, ds2.dims)
self.assertEqual([ds0, ds1, ds2], ml_ds2.datasets)
ml_ds1.close()
ml_ds2.close()
def test_it(self):
ml_ds_1 = BaseMultiLevelDataset(_get_test_dataset(('noise_1', 'noise_2')))
ml_ds_2 = BaseMultiLevelDataset(_get_test_dataset(('noise_3', 'noise_4')))
ml_ds_3 = BaseMultiLevelDataset(_get_test_dataset(('noise_5', 'noise_6')))
ml_ds = CombinedMultiLevelDataset([ml_ds_1, ml_ds_2, ml_ds_3])
self.assertEqual(3, ml_ds.num_levels)
self.assertEqual(TileGrid(3, 2, 1, 180, 180, (-180, -90, 180, 90), inv_y=False),
ml_ds.tile_grid)
expected_var_names = {'noise_1', 'noise_2',
'noise_3', 'noise_4',
'noise_5', 'noise_6'}
ds0 = ml_ds.get_dataset(0)
self.assertEqual({'time': 14, 'lat': 720, 'lon': 1440, 'bnds': 2}, ds0.dims)
self.assertEqual(expected_var_names, set(map(str, ds0.data_vars)))
self.assertTrue(all(v.dims == ('time', 'lat', 'lon') for v in ds0.data_vars.values()))
ds1 = ml_ds.get_dataset(1)
self.assertEqual({'time': 14, 'lat': 360, 'lon': 720}, ds1.dims)
self.assertEqual(expected_var_names, set(map(str, ds1.data_vars)))
self.assertTrue(all(v.dims == ('time', 'lat', 'lon') for v in ds1.data_vars.values()))
ds2 = ml_ds.get_dataset(2)
self.assertEqual({'time': 14, 'lat': 180, 'lon': 360}, ds2.dims)
self.assertEqual(expected_var_names, set(map(str, ds2.data_vars)))
self.assertTrue(all(v.dims == ('time', 'lat', 'lon') for v in ds2.data_vars.values()))
self.assertEqual([ds0, ds1, ds2], ml_ds.datasets)
ml_ds.close()
def test_new_ml_dataset_descriptor(self):
cube = new_cube(variables=dict(a=4.1, b=7.4))
ml_cube = BaseMultiLevelDataset(cube)
descriptor = new_data_descriptor('cube', ml_cube)
self.assertExpectedDescriptor(descriptor,
MultiLevelDatasetDescriptor,
'mldataset')
def test_it(self):
ds = _get_test_dataset()
ml_ds = BaseMultiLevelDataset(ds)
self.assertIsInstance(ml_ds.ds_id, str)
self.assertEqual(3, ml_ds.num_levels)
self.assertEqual(
TileGrid(3, 2, 1, 180, 180, (-180, -90, 180, 90), inv_y=False),
ml_ds.tile_grid)
ds0 = ml_ds.get_dataset(0)
self.assertIs(ds, ds0)
ds1 = ml_ds.get_dataset(1)
self.assertIsNot(ds, ds1)
self.assertEqual({'time': 14, 'lat': 360, 'lon': 720}, ds1.dims)
ds2 = ml_ds.get_dataset(2)
self.assertIsNot(ds, ds2)
self.assertEqual({'time': 14, 'lat': 180, 'lon': 360}, ds2.dims)
self.assertEqual([ds0, ds1, ds2], ml_ds.datasets)
ml_ds.close()
def test_it(self):
ds = _get_test_dataset()
ml_ds = BaseMultiLevelDataset(ds)
self.assertIsInstance(ml_ds.ds_id, str)
self.assertEqual(3, ml_ds.num_levels)
self.assertIsInstance(ml_ds.tile_grid, ImageTileGrid)
self.assertEqual((180, 180), ml_ds.tile_grid.tile_size)
self.assertEqual(3, ml_ds.tile_grid.num_levels)
ds0 = ml_ds.get_dataset(0)
self.assertIsNot(ds, ds0)
self.assertEqual({
'time': 14,
'lat': 720,
'lon': 1440,
'bnds': 2
}, ds0.dims)
ds1 = ml_ds.get_dataset(1)
self.assertIsNot(ds, ds1)
self.assertEqual({'time': 14, 'lat': 360, 'lon': 720}, ds1.dims)
ds2 = ml_ds.get_dataset(2)
self.assertIsNot(ds, ds2)
self.assertEqual({'time': 14, 'lat': 180, 'lon': 360}, ds2.dims)
self.assertEqual([ds0, ds1, ds2], ml_ds.datasets)
ml_ds.close()
def test_get_type_specifier(self):
self.assertIsNone(get_type_specifier(dict()))
self.assertEqual(get_type_specifier(new_cube()), TYPE_SPECIFIER_CUBE)
self.assertEqual(get_type_specifier(xr.Dataset()),
TYPE_SPECIFIER_DATASET)
self.assertEqual(
get_type_specifier(BaseMultiLevelDataset(xr.Dataset())),
TYPE_SPECIFIER_MULTILEVEL_DATASET)
self.assertEqual(get_type_specifier(gpd.GeoDataFrame()),
TYPE_SPECIFIER_GEODATAFRAME)
def _open_ml_dataset(self, dataset_config: DatasetConfigDict) \
-> MultiLevelDataset:
ds_id: str = dataset_config.get('Identifier')
store_instance_id = dataset_config.get('StoreInstanceId')
if store_instance_id:
data_store_pool = self.get_data_store_pool()
data_store = data_store_pool.get_store(store_instance_id)
data_id = dataset_config.get('Path')
open_params = dataset_config.get('StoreOpenParams') or {}
# Inject chunk_cache_capacity into open parameters
chunk_cache_capacity = self.get_dataset_chunk_cache_capacity(
dataset_config)
if chunk_cache_capacity \
and (data_id.endswith('.zarr')
or data_id.endswith('.levels')) \
and 'cache_size' not in open_params:
open_params['cache_size'] = chunk_cache_capacity
with self.measure_time(tag=f"opened dataset {ds_id!r}"
f" from data store"
f" {store_instance_id!r}"):
dataset = data_store.open_data(data_id, **open_params)
if isinstance(dataset, MultiLevelDataset):
ml_dataset = dataset
else:
cube, _, _ = decode_cube(dataset,
normalize=True,
force_non_empty=True,
force_geographic=True)
ml_dataset = BaseMultiLevelDataset(cube, ds_id=ds_id)
else:
fs_type = dataset_config.get('FileSystem')
if fs_type != 'memory':
raise ServiceConfigError(f"Invalid FileSystem {fs_type!r}"
f" in dataset configuration"
f" {ds_id!r}")
with self.measure_time(tag=f"opened dataset {ds_id!r}"
f" from {fs_type!r}"):
ml_dataset = _open_ml_dataset_from_python_code(
self, dataset_config)
augmentation = dataset_config.get('Augmentation')
if augmentation:
script_path = self.get_config_path(
augmentation,
f"'Augmentation' of dataset configuration {ds_id}")
input_parameters = augmentation.get('InputParameters')
callable_name = augmentation.get('Function', COMPUTE_VARIABLES)
ml_dataset = augment_ml_dataset(ml_dataset,
script_path,
callable_name,
self.get_ml_dataset,
self.set_ml_dataset,
input_parameters=input_parameters,
exception_type=ServiceConfigError)
return ml_dataset
def test_get_filename_ext(self):
import xarray as xr
import geopandas as gpd
from xcube.core.mldataset import BaseMultiLevelDataset
dataset = xr.Dataset()
self.assertEqual('.zarr', self.store._get_filename_ext(dataset))
frame = gpd.GeoDataFrame()
self.assertEqual('.geojson', self.store._get_filename_ext(frame))
mldataset = BaseMultiLevelDataset(base_dataset=dataset)
self.assertEqual('.levels', self.store._get_filename_ext(mldataset))
self.assertIsNone(self.store._get_filename_ext(None))
self.assertIsNone(
self.store._get_filename_ext(DataStoreError('A nonsense object')))
def _write_test_cube_pyramid(cls):
s3_kwargs = dict(key='test_fake_id', secret='test_fake_secret')
s3_client_kwargs = dict(endpoint_url=MOTO_SERVER_ENDPOINT_URL)
s3 = s3fs.S3FileSystem(**s3_kwargs, client_kwargs=s3_client_kwargs)
# Create bucket 'xcube-test', so it exists before we write a test pyramid
s3.mkdir('xcube-test')
# Create a test cube pyramid with just one variable "conc_chl"
zarr_path = os.path.join(os.path.dirname(__file__), '../../examples/serve/demo/cube-1-250-250.zarr')
base_dataset = xr.open_zarr(zarr_path)
base_dataset = cls._make_subset(base_dataset)
ml_dataset = BaseMultiLevelDataset(base_dataset)
# Write test cube pyramid
write_levels(ml_dataset,
'xcube-test/cube-1-250-250.levels',
s3_kwargs=s3_kwargs,
s3_client_kwargs=s3_client_kwargs)
def _write_test_cube_pyramid(cls):
# Create bucket 'xcube-test', so it exists before we write a test pyramid
s3_conn = boto3.client('s3')
s3_conn.create_bucket(Bucket='xcube-test', ACL='public-read')
# Create a test cube pyramid with just one variable "conc_chl"
zarr_path = os.path.join(
os.path.dirname(__file__),
'../../examples/serve/demo/cube-1-250-250.zarr')
base_dataset = xr.open_zarr(zarr_path)
base_dataset = xr.Dataset(dict(conc_chl=base_dataset.conc_chl))
ml_dataset = BaseMultiLevelDataset(base_dataset)
# Write test cube pyramid
write_levels(
ml_dataset,
'https://s3.amazonaws.com/xcube-test/cube-1-250-250.levels',
client_kwargs=dict(provider_access_key_id='test_fake_id',
provider_secret_access_key='test_fake_secret'))
def open_ml_dataset_from_object_storage(
ctx: ServiceContext,
dataset_descriptor: DatasetDescriptor) -> MultiLevelDataset:
ds_id = dataset_descriptor.get('Identifier')
path = dataset_descriptor.get('Path')
if not path:
raise ServiceConfigError(
f"Missing 'path' entry in dataset descriptor {ds_id}")
data_format = dataset_descriptor.get('Format', FORMAT_NAME_ZARR)
s3_client_kwargs = {}
if 'Endpoint' in dataset_descriptor:
s3_client_kwargs['endpoint_url'] = dataset_descriptor['Endpoint']
if 'Region' in dataset_descriptor:
s3_client_kwargs['region_name'] = dataset_descriptor['Region']
obs_file_system = s3fs.S3FileSystem(anon=True,
client_kwargs=s3_client_kwargs)
if data_format == FORMAT_NAME_ZARR:
store = s3fs.S3Map(root=path, s3=obs_file_system, check=False)
cached_store = zarr.LRUStoreCache(store, max_size=2**28)
with measure_time(tag=f"opened remote zarr dataset {path}"):
consolidated = obs_file_system.exists(f'{path}/.zmetadata')
ds = assert_cube(
xr.open_zarr(cached_store, consolidated=consolidated))
return BaseMultiLevelDataset(ds)
if data_format == FORMAT_NAME_LEVELS:
with measure_time(tag=f"opened remote levels dataset {path}"):
return ObjectStorageMultiLevelDataset(
ds_id,
obs_file_system,
path,
exception_type=ServiceConfigError)
def write_data(self,
data: Union[xr.Dataset, MultiLevelDataset],
data_id: str,
replace: bool = False,
**write_params) -> str:
assert_instance(data, (xr.Dataset, MultiLevelDataset), name='data')
assert_instance(data_id, str, name='data_id')
tile_size = write_params.pop('tile_size', None)
if isinstance(data, MultiLevelDataset):
ml_dataset = data
if tile_size:
warnings.warn('tile_size is ignored for multi-level datasets')
else:
base_dataset: xr.Dataset = data
if tile_size:
assert_instance(tile_size, int, name='tile_size')
gm = GridMapping.from_dataset(base_dataset)
x_name, y_name = gm.xy_dim_names
base_dataset = base_dataset.chunk({
x_name: tile_size,
y_name: tile_size
})
ml_dataset = BaseMultiLevelDataset(base_dataset)
fs, root, write_params = self.load_fs(write_params)
consolidated = write_params.pop('consolidated', True)
use_saved_levels = write_params.pop('use_saved_levels', False)
base_dataset_id = write_params.pop('base_dataset_id', None)
if use_saved_levels:
ml_dataset = BaseMultiLevelDataset(ml_dataset.get_dataset(0),
tile_grid=ml_dataset.tile_grid)
path_class = get_fs_path_class(fs)
data_path = path_class(data_id)
fs.mkdirs(str(data_path), exist_ok=replace)
for index in range(ml_dataset.num_levels):
level_dataset = ml_dataset.get_dataset(index)
if base_dataset_id and index == 0:
# Write file "0.link" instead of copying
# level zero dataset to "0.zarr".
# Compute a relative base dataset path first
base_dataset_path = path_class(root, base_dataset_id)
data_parent_path = data_path.parent
try:
base_dataset_path = base_dataset_path.relative_to(
data_parent_path)
except ValueError as e:
raise DataStoreError(
f'invalid base_dataset_id: {base_dataset_id}') from e
base_dataset_path = '..' / base_dataset_path
# Then write relative base dataset path into link file
link_path = data_path / f'{index}.link'
with fs.open(str(link_path), mode='w') as fp:
fp.write(f'{base_dataset_path}')
else:
# Write level "{index}.zarr"
level_path = data_path / f'{index}.zarr'
zarr_store = fs.get_mapper(str(level_path), create=True)
try:
level_dataset.to_zarr(zarr_store,
mode='w' if replace else None,
consolidated=consolidated,
**write_params)
except ValueError as e:
# TODO: remove already written data!
raise DataStoreError(f'Failed to write'
f' dataset {data_id}: {e}') from e
if use_saved_levels:
level_dataset = xr.open_zarr(zarr_store,
consolidated=consolidated)
ml_dataset.set_dataset(index, level_dataset)
return data_id