Enhance GeoZarr conversion: add support for Dask cluster setup for parallel processing, update chunking strategy for datasets, and ensure compatibility with existing S3 handling.

Author: emmanuelmathot

.vscode/launch.json

@@ -44,6 +44,7 @@
                 "--min-dimension", "256",
                 "--tile-width", "256",
                 "--max-retries", "2",
+                "--dask-cluster",
                 "--verbose"
             ],
             "cwd": "${workspaceFolder}",

eopf_geozarr/cli.py

@@ -8,13 +8,56 @@
 import argparse
 import sys
 from pathlib import Path
+from typing import Optional
 
 import xarray as xr
 
 from . import create_geozarr_dataset
 from .conversion import is_s3_path, validate_s3_access, get_s3_credentials_info
 
 
+def setup_dask_cluster(enable_dask: bool, verbose: bool = False) -> Optional[object]:
+    """
+    Set up a dask cluster for parallel processing.
+
+    Parameters
+    ----------
+    enable_dask : bool
+        Whether to enable dask cluster
+    verbose : bool, default False
+        Enable verbose output
+
+    Returns
+    -------
+    dask.distributed.Client or None
+        Dask client if enabled, None otherwise
+    """
+    if not enable_dask:
+        return None
+    
+    try:
+        from dask.distributed import Client
+        
+        # Set up local cluster
+        client = Client()  # set up local cluster
+        
+        if verbose:
+            print(f"🚀 Dask cluster started: {client}")
+            print(f"   Dashboard: {client.dashboard_link}")
+            print(f"   Workers: {len(client.scheduler_info()['workers'])}")
+        else:
+            print("🚀 Dask cluster started for parallel processing")
+        
+        return client
+        
+    except ImportError:
+        print("❌ Error: dask.distributed not available. Install with: pip install 'dask[distributed]'")
+        sys.exit(1)
+    except Exception as e:
+        print(f"❌ Error starting dask cluster: {e}")
+        sys.exit(1)
+
+
 def convert_command(args: argparse.Namespace) -> None:
     """
     Convert EOPF dataset to GeoZarr compliant format.
@@ -24,64 +67,70 @@ def convert_command(args: argparse.Namespace) -> None:
     args : argparse.Namespace
         Command line arguments
     """
-    # Validate input path (handle both local paths and URLs)
-    input_path_str = args.input_path
-    if input_path_str.startswith(("http://", "https://", "s3://", "gs://")):
-        # URL - no local validation needed
-        input_path = input_path_str
-    else:
-        # Local path - validate existence
-        input_path = Path(input_path_str)
-        if not input_path.exists():
-            print(f"Error: Input path {input_path} does not exist")
-            sys.exit(1)
-        input_path = str(input_path)
-
-    # Handle output path validation
-    output_path_str = args.output_path
-    if is_s3_path(output_path_str):
-        # S3 path - validate S3 access
-        print("🔍 Validating S3 access...")
-        success, error_msg = validate_s3_access(output_path_str)
-        if not success:
-            print(f"❌ Error: Cannot access S3 path {output_path_str}")
-            print(f"   Reason: {error_msg}")
-            print("\n💡 S3 Configuration Help:")
-            print("   Make sure you have S3 credentials configured:")
-            print("   - Set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables")
-            print("   - Set AWS_DEFAULT_REGION (default: us-east-1)")
-            print("   - For custom S3 providers (e.g., OVH Cloud), set AWS_S3_ENDPOINT")
-            print("   - Or configure AWS CLI with 'aws configure'")
-            print("   - Or use IAM roles if running on EC2")
-            
-            if args.verbose:
-                creds_info = get_s3_credentials_info()
-                print(f"\n🔧 Current AWS configuration:")
-                for key, value in creds_info.items():
-                    print(f"   {key}: {value or 'Not set'}")
+    # Set up dask cluster if requested
+    dask_client = setup_dask_cluster(
+        enable_dask=getattr(args, 'dask_cluster', False),
+        verbose=args.verbose
+    )
+    
+    try:
+        # Validate input path (handle both local paths and URLs)
+        input_path_str = args.input_path
+        if input_path_str.startswith(("http://", "https://", "s3://", "gs://")):
+            # URL - no local validation needed
+            input_path = input_path_str
+        else:
+            # Local path - validate existence
+            input_path = Path(input_path_str)
+            if not input_path.exists():
+                print(f"Error: Input path {input_path} does not exist")
+                sys.exit(1)
+            input_path = str(input_path)
+
+        # Handle output path validation
+        output_path_str = args.output_path
+        if is_s3_path(output_path_str):
+            # S3 path - validate S3 access
+            print("🔍 Validating S3 access...")
+            success, error_msg = validate_s3_access(output_path_str)
+            if not success:
+                print(f"❌ Error: Cannot access S3 path {output_path_str}")
+                print(f"   Reason: {error_msg}")
+                print("\n💡 S3 Configuration Help:")
+                print("   Make sure you have S3 credentials configured:")
+                print("   - Set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables")
+                print("   - Set AWS_DEFAULT_REGION (default: us-east-1)")
+                print("   - For custom S3 providers (e.g., OVH Cloud), set AWS_S3_ENDPOINT")
+                print("   - Or configure AWS CLI with 'aws configure'")
+                print("   - Or use IAM roles if running on EC2")
+                
+                if args.verbose:
+                    creds_info = get_s3_credentials_info()
+                    print(f"\n🔧 Current AWS configuration:")
+                    for key, value in creds_info.items():
+                        print(f"   {key}: {value or 'Not set'}")
+                
+                sys.exit(1)
 
-            sys.exit(1)
-        
-        print("✅ S3 access validated successfully")
-        output_path = output_path_str
-    else:
-        # Local path - create directory if it doesn't exist
-        output_path = Path(output_path_str)
-        output_path.parent.mkdir(parents=True, exist_ok=True)
-        output_path = str(output_path)
-
-    if args.verbose:
-        print(f"Loading EOPF dataset from: {input_path}")
-        print(f"Groups to convert: {args.groups}")
-        print(f"Output path: {output_path}")
-        print(f"Spatial chunk size: {args.spatial_chunk}")
-        print(f"Min dimension: {args.min_dimension}")
-        print(f"Tile width: {args.tile_width}")
+            print("✅ S3 access validated successfully")
+            output_path = output_path_str
+        else:
+            # Local path - create directory if it doesn't exist
+            output_path = Path(output_path_str)
+            output_path.parent.mkdir(parents=True, exist_ok=True)
+            output_path = str(output_path)
+
+        if args.verbose:
+            print(f"Loading EOPF dataset from: {input_path}")
+            print(f"Groups to convert: {args.groups}")
+            print(f"Output path: {output_path}")
+            print(f"Spatial chunk size: {args.spatial_chunk}")
+            print(f"Min dimension: {args.min_dimension}")
+            print(f"Tile width: {args.tile_width}")
 
-    try:
         # Load the EOPF DataTree
         print("Loading EOPF dataset...")
-        dt = xr.open_datatree(str(input_path), engine="zarr")
+        dt = xr.open_datatree(str(input_path), engine="zarr", chunks="auto")
 
         if args.verbose:
             print(f"Loaded DataTree with {len(dt.children)} groups")
@@ -117,6 +166,16 @@ def convert_command(args: argparse.Namespace) -> None:
 
             traceback.print_exc()
         sys.exit(1)
+    finally:
+        # Clean up dask client if it was created
+        if dask_client is not None:
+            try:
+                dask_client.close()
+                if args.verbose:
+                    print("🔄 Dask cluster closed")
+            except Exception as e:
+                if args.verbose:
+                    print(f"Warning: Error closing dask cluster: {e}")
 
 
 def info_command(args: argparse.Namespace) -> None:
@@ -143,7 +202,7 @@ def info_command(args: argparse.Namespace) -> None:
 
     try:
         print(f"Loading dataset from: {input_path}")
-        dt = xr.open_datatree(input_path, engine="zarr")
+        dt = xr.open_datatree(input_path, engine="zarr", chunks="auto")
 
         print("\nDataset Information:")
         print("==================")
@@ -185,7 +244,7 @@ def validate_command(args: argparse.Namespace) -> None:
 
     try:
         print(f"Validating GeoZarr compliance for: {input_path}")
-        dt = xr.open_datatree(input_path, engine="zarr")
+        dt = xr.open_datatree(input_path, engine="zarr", chunks="auto")
 
         compliance_issues = []
         total_variables = 0
@@ -306,6 +365,11 @@ def create_parser() -> argparse.ArgumentParser:
         help="Maximum number of retries for network operations (default: 3)",
     )
     convert_parser.add_argument("--verbose", action="store_true", help="Enable verbose output")
+    convert_parser.add_argument(
+        "--dask-cluster",
+        action="store_true",
+        help="Start a local dask cluster for parallel processing of chunks"
+    )
     convert_parser.set_defaults(func=convert_command)
 
     # Info command

eopf_geozarr/conversion/geozarr.py

@@ -269,8 +269,17 @@ def recursive_copy(
     if dt_node.data_vars:
         # Set up encoding for variables
         for var in ds.data_vars:
+            # create chunks to the size of the data
+            data_shape = ds[var].shape
+            if len(data_shape) >= 2:
+                chunking = (
+                    1,
+                    data_shape[-2],
+                    data_shape[-1],
+                )
             encoding[var] = {
                 "compressors": [compressor],
+                "chunks": chunking
             }
         for coord in ds.coords:
             encoding[coord] = {
@@ -527,11 +536,11 @@ def write_geozarr_group(
             if s3_utils.s3_path_exists(native_dataset_path):
                 store = s3_utils.create_s3_store(native_dataset_path)
                 storage_options = s3_utils.get_s3_storage_options(native_dataset_path)
-                existing_native_dataset = xr.open_zarr(store, zarr_format=3, storage_options=storage_options)
+                existing_native_dataset = xr.open_zarr(store, zarr_format=3, storage_options=storage_options, chunks="auto")
                 print(f"Found existing native dataset at {native_dataset_path}")
         else:
             if os.path.exists(native_dataset_path):
-                existing_native_dataset = xr.open_zarr(native_dataset_path, zarr_format=3)
+                existing_native_dataset = xr.open_zarr(native_dataset_path, zarr_format=3, chunks="auto")
                 print(f"Found existing native dataset at {native_dataset_path}")
     except Exception as e:
         print(f"Warning: Could not open existing native dataset at {native_dataset_path}: {e}")
@@ -583,10 +592,10 @@ def write_geozarr_group(
     if s3_utils.is_s3_path(output_path):
         store = s3_utils.create_s3_store(group_path)
         storage_options = s3_utils.get_s3_storage_options(output_path)
-        ds = xr.open_dataset(store, engine="zarr", zarr_format=3, decode_coords="all", storage_options=storage_options).compute()
+        ds = xr.open_dataset(store, engine="zarr", zarr_format=3, decode_coords="all", storage_options=storage_options, chunks="auto").compute()
     else:
         ds = xr.open_dataset(
-            group_path, engine="zarr", zarr_format=3, decode_coords="all"
+            group_path, engine="zarr", zarr_format=3, decode_coords="all", chunks="auto"
         ).compute()
 
     return ds
@@ -1292,9 +1301,9 @@ def write_dataset_band_by_band_with_validation(
     if s3_utils.is_s3_path(output_path):
         store = s3_utils.create_s3_store(output_path)
         storage_options = s3_utils.get_s3_storage_options(output_path)
-        ds = xr.open_dataset(store, engine="zarr", zarr_format=3, decode_coords="all", storage_options=storage_options).compute()
+        ds = xr.open_dataset(store, engine="zarr", zarr_format=3, decode_coords="all", storage_options=storage_options, chunks="auto").compute()
     else:
-        ds = xr.open_dataset(output_path, engine="zarr", zarr_format=3, decode_coords="all").compute()
+        ds = xr.open_dataset(output_path, engine="zarr", zarr_format=3, decode_coords="all", chunks="auto").compute()
 
     # Report results
     if failed_vars: