diff --git a/data_access_service/core/estimation_index.py b/data_access_service/core/estimation_index.py index f7962aa7..d512c576 100644 --- a/data_access_service/core/estimation_index.py +++ b/data_access_service/core/estimation_index.py @@ -206,6 +206,15 @@ def sidecar_extent_provider(api, output_format: str) -> ExtentProvider: def provider( uuid: str, key: str ) -> Tuple[Optional[pd.Timestamp], Optional[pd.Timestamp]]: + if key.endswith(".zarr"): + # A zarr extent must not go through get_temporal_extent: that + # caches a chunks="auto" open and builds the dask graph the + # estimate is trying to avoid. + from data_access_service.core.size_estimation import ( + zarr_extent_for_estimate, + ) + + return zarr_extent_for_estimate(api, uuid, key) if not key.endswith(".parquet") or output_format != OUTPUT_FORMAT_CSV: return api.get_temporal_extent(uuid, key) meta = require_sidecar(api, uuid, key) diff --git a/data_access_service/core/size_estimation.py b/data_access_service/core/size_estimation.py index 9749172d..1c0dc05a 100644 --- a/data_access_service/core/size_estimation.py +++ b/data_access_service/core/size_estimation.py @@ -48,6 +48,99 @@ log = logging.getLogger(__name__) +def zarr_extent_for_estimate(api, uuid: str, key: str): + """Min/max time of a zarr key, without caching a dask-backed open. + + Same day-alignment and future-date clamp as ``API.get_temporal_extent``. + The store is closed before this returns. + """ + ds = api._instance.get_dataset(key, chunks=None) + try: + try: + start_date, end_date = ds.get_temporal_extent() + except ValueError: + return None, None + return _clamp_extent(start_date, end_date) + finally: + _close_zarr(ds) + + +def _clamp_extent(start_date, end_date): + """Match ``API.get_temporal_extent`` so an estimate trims the same days.""" + if start_date is not None: + start_date = start_date.replace( + hour=0, minute=0, second=0, microsecond=0, nanosecond=0 + ) + if end_date is not None: + if end_date.tzinfo is None: + now_compare = pd.Timestamp.now() + else: + now_compare = pd.Timestamp.now(tz="UTC") + end_date = end_date.tz_convert("UTC") + end_date = end_date.replace( + hour=23, minute=59, second=59, microsecond=999999, nanosecond=999 + ) + if end_date > now_compare: + end_date = now_compare + return start_date, end_date + + +def _close_zarr(ds) -> None: + store = getattr(ds, "zarr_store", None) + if store is not None: + store.close() + + +def _estimate_zarr_key( + api, + uuid: str, + key: str, + resolved_subset_request: ResolvedSubsetRequest, + output_format: str, + cancellation: Optional[Cancellation], +) -> Optional[dict]: + """Size one zarr key from a chunks=None open that is not cached.""" + ds = api._instance.get_dataset(key, chunks=None) + try: + check_datasource_supports_format(api, ds, uuid, key, output_format) + if not resolved_subset_request.has_data: + return _empty_estimate(uuid, key, output_format) + raise_if_client_gone(cancellation) + + def _extent(_uuid: str, _key: str): + try: + start_date, end_date = ds.get_temporal_extent() + except ValueError: + return None, None + return _clamp_extent(start_date, end_date) + + date_start, date_end = trim_date_range_for_keys( + api, + uuid, + [key], + resolved_subset_request.start_date, + resolved_subset_request.end_date, + extent_provider=_extent, + ) + if date_start is None or date_end is None: + return _empty_estimate(uuid, key, output_format) + return _estimate_zarr_size( + api, + ds.zarr_store, + uuid, + key, + date_start, + date_end, + resolved_subset_request.bboxes, + resolved_subset_request.columns, + output_format, + resolved_subset_request.geometry, + cancellation, + ) + finally: + _close_zarr(ds) + + def _bytes_to_mb(num_bytes: int) -> float: """Bytes -> decimal MB (1 MB = 1,000,000 bytes), the way browsers report file sizes. Rounded to 1 decimal place.""" @@ -78,6 +171,14 @@ def estimate_single_key_size( :raises ClientGoneError: if the client disconnected """ uuid = resolved_subset_request.uuid + if key.endswith(".zarr"): + # Do not use get_datasource. That caches a chunks="auto" open, and + # the dask graph for a fine-grained satellite store is multiple GB. + # Downloads keep that open; only this estimate uses chunks=None. + return _estimate_zarr_key( + api, uuid, key, resolved_subset_request, output_format, cancellation + ) + ds = api.get_datasource(uuid, key) if ds is None: return None diff --git a/data_access_service/models/co_datasource/abstract_data_src.py b/data_access_service/models/co_datasource/abstract_data_src.py index f262674c..cd43ba86 100644 --- a/data_access_service/models/co_datasource/abstract_data_src.py +++ b/data_access_service/models/co_datasource/abstract_data_src.py @@ -29,12 +29,22 @@ def get_dataset_names(self) -> frozenset[str]: return frozenset(self.get_metadata_catalog()) @abstractmethod - def get_dataset(self, dataset_name_with_ext: str) -> DataSource: + def get_dataset( + self, dataset_name_with_ext: str, chunks: str | None = "auto" + ) -> DataSource: + """Open one dataset. + + ``chunks`` is forwarded to the zarr open. ``"auto"`` is the library + default and builds a dask graph. ``None`` leaves the store lazy with + no graph, which is what the size estimate uses. Parquet ignores it. + """ if dataset_name_with_ext not in self.get_dataset_names(): raise DatasetNotFoundError( dataset_name=dataset_name_with_ext, data_source_name=self.get_name() ) + if dataset_name_with_ext.endswith(".zarr"): + return self.get_data_src().get_dataset(dataset_name_with_ext, chunks=chunks) return self.get_data_src().get_dataset(dataset_name_with_ext) @abstractmethod diff --git a/data_access_service/models/co_datasource/aodn_data_src.py b/data_access_service/models/co_datasource/aodn_data_src.py index ad9ae353..f1bbc361 100644 --- a/data_access_service/models/co_datasource/aodn_data_src.py +++ b/data_access_service/models/co_datasource/aodn_data_src.py @@ -30,8 +30,12 @@ def get_dataset_names(self) -> frozenset[str]: self.__dataset_names = frozenset(self.__data_src.list_datasets()) return self.__dataset_names - def get_dataset(self, dataset_name_with_ext: str) -> DataSource: - return super().get_dataset(dataset_name_with_ext=dataset_name_with_ext) + def get_dataset( + self, dataset_name_with_ext: str, chunks: str | None = "auto" + ) -> DataSource: + return super().get_dataset( + dataset_name_with_ext=dataset_name_with_ext, chunks=chunks + ) def get_data_src(self) -> GetAodn: return self.__data_src diff --git a/data_access_service/models/co_datasource/co_data_registory.py b/data_access_service/models/co_datasource/co_data_registory.py index 23702429..c3ab3613 100644 --- a/data_access_service/models/co_datasource/co_data_registory.py +++ b/data_access_service/models/co_datasource/co_data_registory.py @@ -109,19 +109,29 @@ def get_metadata(self) -> Metadata: before_sleep=log_retry_attempt("get_dataset", log), reraise=True, ) - def get_dataset(self, dataset_name_with_ext: str) -> DataSource | None: - with self._datasets_lock: - cached = self._datasets.get(dataset_name_with_ext) - if cached is not None: - log.info("Reusing cached %s dataset", dataset_name_with_ext) - return cached + def get_dataset( + self, dataset_name_with_ext: str, chunks: str | None = "auto" + ) -> DataSource | None: + # Only the default dask open is cached. A chunks=None open is the + # size estimate: it must not replace the chunks="auto" store that + # downloads reuse, and it must not be kept after the estimate. + use_cache = chunks == "auto" + if use_cache: + with self._datasets_lock: + cached = self._datasets.get(dataset_name_with_ext) + if cached is not None: + log.info("Reusing cached %s dataset", dataset_name_with_ext) + return cached for data_src in self.data_source_list: try: log.info( f"Getting {dataset_name_with_ext} dataset from {data_src.get_name()}..." ) - dataset = data_src.get_dataset(dataset_name_with_ext) + if chunks == "auto": + dataset = data_src.get_dataset(dataset_name_with_ext) + else: + dataset = data_src.get_dataset(dataset_name_with_ext, chunks=chunks) except DatasetNotFoundError: # log the exception and continue to try the next data source log.info( @@ -131,6 +141,8 @@ def get_dataset(self, dataset_name_with_ext: str) -> DataSource | None: # Load happened outside the lock so a 4-minute S3 listing does # not block other datasets. setdefault keeps the first stored # instance if two threads both missed. + if not use_cache: + return dataset with self._datasets_lock: return self._datasets.setdefault(dataset_name_with_ext, dataset) diff --git a/data_access_service/models/co_datasource/csiro/csiro_data_src.py b/data_access_service/models/co_datasource/csiro/csiro_data_src.py index 68a1afdc..35d6a9f5 100644 --- a/data_access_service/models/co_datasource/csiro/csiro_data_src.py +++ b/data_access_service/models/co_datasource/csiro/csiro_data_src.py @@ -401,8 +401,12 @@ def get_metadata(self) -> Metadata: def get_metadata_catalog(self) -> dict: return self.__metadata_catalog - def get_dataset(self, dataset_name_with_ext: str) -> DataSource: - return super().get_dataset(dataset_name_with_ext=dataset_name_with_ext) + def get_dataset( + self, dataset_name_with_ext: str, chunks: str | None = "auto" + ) -> DataSource: + return super().get_dataset( + dataset_name_with_ext=dataset_name_with_ext, chunks=chunks + ) def get_name(self) -> str: return self.name diff --git a/tests/core/test_api_estimate_size.py b/tests/core/test_api_estimate_size.py index b84c1207..63afdc9a 100644 --- a/tests/core/test_api_estimate_size.py +++ b/tests/core/test_api_estimate_size.py @@ -121,6 +121,38 @@ def _api_with_zarr(dataset: xr.Dataset) -> tuple[API, MagicMock]: return api, mock_ds +def test_zarr_suffix_does_not_use_the_cached_dask_open(monkeypatch): + """A real .zarr key is sized from chunks=None and that open is closed.""" + from data_access_service.core import size_estimation + + dataset = _make_dataset() + opened = MagicMock(spec=ZarrDataSource) + opened.zarr_store = dataset + opened.get_temporal_extent.return_value = ( + pd.Timestamp("2020-01-01"), + pd.Timestamp("2020-01-05"), + ) + closed = [] + monkeypatch.setattr(size_estimation, "_close_zarr", lambda ds: closed.append(ds)) + api = MagicMock() + api.resolve_dim_names.return_value = ("LATITUDE", "LONGITUDE", "TIME") + api.map_column_names = MagicMock(side_effect=_identity_map) + api._instance.get_dataset.return_value = opened + api.get_datasource.side_effect = AssertionError("cached chunks=auto open") + + result = estimate_single_key_size( + api, + "satellite.zarr", + _resolved(keys=["satellite.zarr"]), + output_format="netcdf", + ) + + assert result["estimated_uncompressed_bytes"] == dataset.nbytes + api.get_datasource.assert_not_called() + api._instance.get_dataset.assert_called_with("satellite.zarr", chunks=None) + assert closed == [opened] + + def test_zarr_estimate_basic(): dataset = _make_dataset() api, _ = _api_with_zarr(dataset) diff --git a/tests/models/co_datasource/test_co_data_source.py b/tests/models/co_datasource/test_co_data_source.py index fcc11fbe..a0baf761 100644 --- a/tests/models/co_datasource/test_co_data_source.py +++ b/tests/models/co_datasource/test_co_data_source.py @@ -433,6 +433,19 @@ def test_get_metadata_raises_on_conflicting_dataset_names(self): with pytest.raises(Exception, match="Conflicted dataset names"): registry.get_metadata() + def test_chunks_none_is_not_cached(self): + registry, mock_aodn_src, _ = _make_registry( + aodn_catalog={"satellite.zarr": {}}, + ) + opened = MagicMock() + mock_aodn_src.get_dataset.return_value = opened + + assert registry.get_dataset("satellite.zarr", chunks=None) is opened + assert registry.get_dataset("satellite.zarr", chunks=None) is opened + assert mock_aodn_src.get_dataset.call_count == 2 + mock_aodn_src.get_dataset.assert_called_with("satellite.zarr", chunks=None) + assert "satellite.zarr" not in registry._datasets + def test_get_dataset_returns_from_first_matching_source(self): registry, mock_aodn_src, _ = _make_registry() expected = MagicMock()