Merge pull request #49 from CMCC-Foundation/version-2025.06

Version 2025.06

Author: gtramonte

.github/workflows/build_on_release.yaml

@@ -1,9 +1,10 @@
 name: Build geolake docker images for geolake components and push to the repository
 
 on:
-  push:
-    tags:        
-      - 'v*'
+  release:
+    types: [published]
+  workflow_dispatch:
+
 jobs:
   build:
     runs-on: ubuntu-latest

drivers/Dockerfile

@@ -1,6 +1,6 @@
 ARG REGISTRY=rg.fr-par.scw.cloud/geokube
 #ARG TAG=2025.03.25.10.56
-ARG TAG=v0.2.7.10
+ARG TAG=2025.06
 FROM $REGISTRY/geokube:$TAG
 
 ADD . /opt/intake_geokube

drivers/intake_geokube/netcdf_with_ancillary.py

@@ -0,0 +1,93 @@
+"""geokube driver for intake."""
+import logging
+from typing import Mapping, Optional
+from .base import GeokubeSource
+from geokube import open_dataset, open_datacube
+from geokube.core.datacube import DataCube
+import pickle
+import os
+import xarray as xr
+import numpy as np
+import glob
+
+_PROJECTION = {"grid_mapping_name": "latitude_longitude"}
+
+class NetCDFAncillarySource(GeokubeSource):
+    name = "geokube_netcdf_ancillary"
+
+    def add_projection(self, dset: xr.Dataset, **kwargs) -> xr.Dataset:
+        """Add projection information to the dataset"""
+        coords = dset.coords
+        coords["crs"] = xr.DataArray(data=np.array(1), attrs=_PROJECTION)
+        for var in dset.data_vars.values():
+            enc = var.encoding
+            enc["grid_mapping"] = "crs"
+        return dset
+
+    def __init__(
+        self,
+        path: str,
+        ancillary_path: str,
+        pattern: str = None,
+        field_id: str = None,
+        delay_read_cubes: bool = False,
+        metadata_caching: bool = False,
+        metadata_cache_path: str = None,
+        storage_options: dict = None,
+        xarray_kwargs: dict = None,
+        metadata=None,
+        mapping: Optional[Mapping[str, Mapping[str, str]]] = None,
+        load_files_on_persistance: Optional[bool] = True,
+        **kwargs
+    ):
+        self._kube = None
+        self.path = path
+        self.ancillary_path = ancillary_path
+        self.pattern = pattern
+        self.field_id = field_id
+        self.delay_read_cubes = delay_read_cubes
+        self.metadata_caching = metadata_caching
+        self.metadata_cache_path = metadata_cache_path
+        self.storage_options = storage_options
+        self.mapping = mapping
+        self.xarray_kwargs = {} if xarray_kwargs is None else xarray_kwargs
+        self.load_files_on_persistance = load_files_on_persistance
+        #        self.xarray_kwargs.update({'engine' : 'netcdf'})
+        super(NetCDFAncillarySource, self).__init__(metadata=metadata, **kwargs)
+
+    def _open_dataset(self):
+
+        if self.metadata_caching:
+            cached = None
+            if os.path.exists(self.metadata_cache_path):
+                with open(self.metadata_cache_path, "rb") as f:
+                    cached = pickle.load(f)
+                self._kube = cached
+                return self._kube
+
+        afilepaths = glob.glob(self.ancillary_path)
+        filepaths = glob.glob(self.path)
+        ancillary = xr.open_mfdataset(afilepaths, compat='override')
+        ds = xr.open_mfdataset(filepaths, **self.xarray_kwargs)
+        finalds = xr.merge([ancillary, ds])
+
+        finalds.xgrid.attrs['standard_name'] = 'projection_grid_x_centers'
+        finalds.ygrid.attrs['standard_name'] = 'projection_grid_y_centers'
+
+        finalds2 = self.add_projection(finalds)
+        finalds3 = finalds2.assign_coords(tdim=np.arange(finalds2.sizes['tdim']))
+        time = finalds3.time.values
+        finalds4 = finalds3.assign_coords(time=("tdim", time)).swap_dims({"tdim": "time"})
+        finalds5 = finalds4.sortby("time")
+
+        for var in finalds5.data_vars.values():
+            if "grid_mapping" in var.attrs:
+                del var.attrs["grid_mapping"]
+
+        self._kube = DataCube.from_xarray(finalds5, mapping=self.mapping)
+
+        if self.metadata_caching:
+            with open(self.metadata_cache_path, "wb") as f:
+                pickle.dump(self._kube, f)
+
+        return self._kube

drivers/setup.py

@@ -18,7 +18,8 @@
         "intake.drivers": [
             "geokube_netcdf = intake_geokube.netcdf:NetCDFSource",
             "cmcc_wrf_geokube = intake_geokube.wrf:CMCCWRFSource",
-            "cmcc_sentinel_geokube = intake_geokube.sentinel:CMCCSentinelSource"
+            "cmcc_sentinel_geokube = intake_geokube.sentinel:CMCCSentinelSource",
+            "geokube_netcdf_ancillary = intake_geokube.netcdf_with_ancillary:NetCDFAncillarySource",
         ]
     },
     classifiers=[

executor/app/main.py

@@ -1,4 +1,5 @@
 import os
+import signal
 import tempfile
 import time
 import datetime
@@ -14,6 +15,7 @@
 from geokube.core.datacube import DataCube
 from geokube.core.dataset import Dataset
 from geokube.core.field import Field
+from pika.exceptions import ChannelClosed
 
 from datastore.datastore import Datastore
 from workflow import Workflow
@@ -277,6 +279,8 @@ def __init__(self, broker, store_path, dask_cluster_opts):
         self._channel = broker_conn.channel()
         self._db = DBManager()
         self.dask_cluster_opts = dask_cluster_opts
+        self.to_exit = False
+        self.processing = False
 
     def create_dask_cluster(self, dask_cluster_opts: dict = None):
         if dask_cluster_opts is None:
@@ -307,21 +311,21 @@ def create_dask_cluster(self, dask_cluster_opts: dict = None):
 
     def maybe_restart_cluster(self, status: RequestStatus):
         if status is RequestStatus.TIMEOUT:
-            self._LOG.info("recreating the cluster due to timeout")
+            self._LOG.info("recreating the cluster due to timeout", extra={"track_id": "N/A"})
             self._dask_client.cluster.close()
             self.create_dask_cluster()
         if self._dask_client.cluster.status is Status.failed:
-            self._LOG.info("attempt to restart the cluster...")
+            self._LOG.info("attempt to restart the cluster...", extra={"track_id": "N/A"})
             try:
                 asyncio.run(self._nanny.restart())
             except Exception as err:
                 self._LOG.error(
-                    "couldn't restart the cluster due to an error: %s", err
+                    "couldn't restart the cluster due to an error: %s", err, extra={"track_id": "N/A"}
                 )
-                self._LOG.info("closing the cluster")
+                self._LOG.info("closing the cluster", extra={"track_id": "N/A"})
                 self._dask_client.cluster.close()
-        if self._dask_client.cluster.status is Status.closed:
-            self._LOG.info("recreating the cluster")
+        if self._dask_client.cluster.status is Status.closed and not self.to_exit:
+            self._LOG.info("recreating the cluster", extra={"track_id": "N/A"})
             self.create_dask_cluster()
 
     def ack_message(self, channel, delivery_tag):
@@ -330,9 +334,13 @@ def ack_message(self, channel, delivery_tag):
         """
         if channel.is_open:
             channel.basic_ack(delivery_tag)
+            self.processing = False
+            if self.to_exit:
+                channel.stop_consuming()
         else:
             self._LOG.info(
-                "cannot acknowledge the message. channel is closed!"
+                "cannot acknowledge the message. channel is closed!",
+                extra={"track_id": "N/A"},
             )
             pass
 
@@ -391,6 +399,7 @@ def retry_until_timeout(
         return location_path, status, fail_reason
 
     def handle_message(self, connection, channel, delivery_tag, body):
+        self.processing = True
         message: Message = Message(body)
         self._LOG.debug(
             "executing query: `%s`",
@@ -468,8 +477,21 @@ def subscribe(self, etype):
         )
 
     def listen(self):
-        while True:
-            self._channel.start_consuming()
+        while not self.to_exit:
+            try:
+                self._channel.start_consuming()
+            except ChannelClosed as cc:
+                self._LOG.error("Channel closed exiting...", extra={"track_id": "N/A"})
+        self._LOG.info(f'Shutting down Dask...', extra={"track_id": "N/A"})
+        self._dask_client.shutdown()
+        self._LOG.info(f'Exiting...', extra={"track_id": "N/A"})
+        exit(0)
+
+    def stop_listening(self, signo, frame):
+        self._LOG.info(f'received signal {signo}:', extra={"track_id": "N/A"})
+        self.to_exit = True
+        if not self.processing:
+            self._channel.stop_consuming()
 
     def get_size(self, location_path):
         if location_path and os.path.exists(location_path):
@@ -503,5 +525,9 @@ def get_size(self, location_path):
 
         executor.subscribe(etype)
 
+    print('registering signal handlers')
+    signal.signal(signal.SIGTERM, executor.stop_listening)
+    signal.signal(signal.SIGINT, executor.stop_listening)
+
     print("waiting for requests ...")
     executor.listen()