marEx 2.0.0__py3-none-any.whl

marEx/helper.py

@@ -0,0 +1,251 @@
+"""
+HPC Dask Helper: Utilities for High-Performance Computing with Dask
+--------------------------------------------------------------------
+
+This module provides utilities for setting up and managing Dask clusters
+in HPC environments, with specific support for the DKRZ Levante Supercomputer.
+"""
+
+import os
+import re
+import subprocess
+import psutil
+from tempfile import TemporaryDirectory
+from getpass import getuser
+from pathlib import Path
+from typing import Dict, Optional, Union, Any
+
+import dask
+from dask_jobqueue import SLURMCluster
+from dask.distributed import Client, LocalCluster
+
+
+# Default configuration values
+DEFAULT_DASK_CONFIG = {
+    'array.slicing.split_large_chunks': False,
+    'distributed.comm.timeouts.connect': '120s',  # Increased from default
+    'distributed.comm.timeouts.tcp': '240s',      # Double the connection timeout
+    'distributed.comm.retry.count': 10,           # More retries before giving up
+}
+
+# DKRZ-specific paths and configuration
+DKRZ_SCRATCH_PATH = Path('/scratch') / getuser()[0] / getuser() / 'clients'
+DKRZ_LOG_PATH = Path('/home/b') / getuser() / '.log_trash'
+DKRZ_ACCOUNT = 'bk1377'
+
+# Memory configuration for different node types
+MEMORY_CONFIGS = {
+    256: {'client_memory': '250GB', 'constraint': '256', 'job_extra': ['--mem=0']},
+    512: {'client_memory': '500GB', 'constraint': '512', 'job_extra': ['--constraint=512G --mem=0']},
+    1024: {'client_memory': '1000GB', 'constraint': '1024', 'job_extra': ['--constraint=1024G --mem=0']}
+}
+
+
+def configure_dask(scratch_dir: Optional[Union[str, Path]] = None, 
+                  config: Optional[Dict[str, Any]] = None) -> TemporaryDirectory:
+    """
+    Configure Dask with appropriate settings for HPC environments.
+    
+    Parameters
+    ----------
+    scratch_dir : str or Path, optional
+        Directory to use for temporary files.
+    config : dict, optional
+        Additional Dask configuration settings to apply.
+    
+    Returns
+    -------
+    TemporaryDirectory
+        Temporary directory object that should be kept alive while Dask is in use.
+    """
+    # Use provided scratch directory or default to DKRZ scratch
+    scratch_path = Path(scratch_dir) if scratch_dir else DKRZ_SCRATCH_PATH
+    
+    # Create temporary directory
+    if not scratch_path.exists():
+        scratch_path.mkdir(parents=True, exist_ok=True)
+    
+    temp_dir = TemporaryDirectory(dir=scratch_path)
+    
+    # Apply default configuration
+    dask.config.set(temporary_directory=temp_dir.name)
+    print(f"Dask Scratch: {repr(temp_dir.name)}")
+    
+    # Apply default settings
+    for key, value in DEFAULT_DASK_CONFIG.items():
+        dask.config.set({key: value})
+    
+    # Apply any additional configuration
+    if config:
+        dask.config.set(config)
+    
+    return temp_dir
+
+
+def get_cluster_info(client: Client) -> Dict[str, str]:
+    """
+    Get and print cluster connection information.
+    
+    Parameters
+    ----------
+    client : Client
+        Dask client connected to a cluster.
+    
+    Returns
+    -------
+    dict
+        Dictionary containing connection information.
+    """
+    # Get hostname and dashboard port
+    remote_node = subprocess.run(['hostname'], capture_output=True, text=True).stdout.strip().split('.')[0]
+    port = re.search(r':(\d+)/', client.dashboard_link).group(1)
+    
+    # Print connection information
+    print(f"Hostname: {remote_node}")
+    print(f"Forward Port: {remote_node}:{port}")
+    print(f"Dashboard Link: localhost:{port}/status")
+    
+    return {
+        'hostname': remote_node,
+        'port': port,
+        'dashboard_link': f"localhost:{port}/status"
+    }
+
+
+def start_local_cluster(n_workers: int = 4, threads_per_worker: int = 1, 
+                       scratch_dir: Optional[Union[str, Path]] = None,
+                       **kwargs) -> Client:
+    """
+    Start a local Dask cluster.
+    
+    Parameters
+    ----------
+    n_workers : int, default=4
+        Number of worker processes to start.
+    threads_per_worker : int, default=1
+        Number of threads per worker.
+    scratch_dir : str or Path, optional
+        Directory to use for temporary files.
+    **kwargs
+        Additional keyword arguments to pass to LocalCluster.
+    
+    Returns
+    -------
+    Client
+        Dask client connected to the local cluster.
+    """
+    # Configure Dask
+    temp_dir = configure_dask(scratch_dir)
+    
+    # Check system resources
+    physical_cores = psutil.cpu_count(logical=False)
+    logical_cores = psutil.cpu_count(logical=True)
+    memory = psutil.virtual_memory()
+    
+    # Warn if requested resources exceed available
+    total_threads = n_workers * threads_per_worker
+    if total_threads > physical_cores:
+        print(f"Warning: Requested {n_workers} workers with {threads_per_worker} threads each, but only {physical_cores} physical cores available.")
+        print("Hyper-threading can reduce performance for compute-intensive tasks!")
+    elif total_threads > logical_cores:
+        print(f"Warning: Requested {n_workers} workers with {threads_per_worker} threads each, but only {logical_cores} logical cores available.")
+        print(f"Reducing to {logical_cores // threads_per_worker} workers.")
+        n_workers = logical_cores // threads_per_worker
+    
+    print(f"Memory per Worker: {memory.total / n_workers / (1024**3):.2f} GB")
+
+    # Create cluster and client
+    cluster = LocalCluster(n_workers=n_workers, threads_per_worker=threads_per_worker, **kwargs)
+    client = Client(cluster)
+    
+    # Store temporary directory (so it isn't garbage collected)
+    client._temp_dir = temp_dir
+
+    # Print connection information
+    get_cluster_info(client)
+    
+    return client
+
+
+def start_distributed_cluster(n_workers: int, workers_per_node: int, 
+                             runtime: int = 9, node_memory: int = 256, 
+                             dashboard_address: int = 8889, queue: str = 'compute', 
+                             scratch_dir: Optional[Union[str, Path]] = None, 
+                             account: Optional[str] = None, **kwargs) -> Client:
+    """
+    Start a distributed Dask cluster on a SLURM-based supercomputer.
+    
+    Parameters
+    ----------
+    n_workers : int
+        Total number of workers to request.
+    workers_per_node : int
+        Number of workers per node.
+    runtime : int, default=9
+        Maximum runtime in minutes.
+    node_memory : int, default=256
+        Memory per node in GB (256, 512, or 1024).
+    dashboard_address : int, default=8889
+        Port for the Dask dashboard.
+    queue : str, default='compute'
+        SLURM queue to submit jobs to.
+    scratch_dir : str or Path, optional
+        Directory to use for temporary files.
+    account : str, optional
+        SLURM account to charge. Defaults to DKRZ_ACCOUNT.
+    **kwargs
+        Additional keyword arguments to pass to SLURMCluster.
+    
+    Returns
+    -------
+    Client
+        Dask client connected to the distributed cluster.
+    """
+    # Configure Dask
+    temp_dir = configure_dask(scratch_dir)
+    
+    # Use default account if none specified
+    if account is None:
+        account = DKRZ_ACCOUNT
+    
+    # Validate node_memory
+    if node_memory not in MEMORY_CONFIGS:
+        raise ValueError(f"Unsupported node_memory value: {node_memory}. Must be one of {list(MEMORY_CONFIGS.keys())}.")
+    
+    config = MEMORY_CONFIGS[node_memory]
+    
+    # Calculate runtime in hours and minutes
+    runtime_hrs = runtime // 60
+    runtime_mins = runtime % 60
+    
+    # Create SLURM cluster
+    cluster = SLURMCluster(
+        name='dask-cluster',
+        cores=workers_per_node,
+        memory=config['client_memory'],
+        processes=workers_per_node,  # One process per core
+        interface='ib0',
+        queue=queue,
+        account=account,
+        walltime=f'{runtime_hrs:02d}:{runtime_mins:02d}:00',
+        asynchronous=0,
+        job_extra_directives=config['job_extra'],
+        log_directory=DKRZ_LOG_PATH,
+        local_directory=temp_dir.name,
+        scheduler_options={'dashboard_address': f':{dashboard_address}'},
+        **kwargs
+    )
+
+    print(f"Memory per Worker: {node_memory / workers_per_node:.2f} GB")
+    
+    # Scale the cluster
+    cluster.scale(n_workers)
+    client = Client(cluster)
+    
+    # Store temporary directory (so it isn't garbage collected)
+    client._temp_dir = temp_dir
+    
+    # Print connection information
+    get_cluster_info(client)
+    
+    return client

marEx/plotX/__init__.py

@@ -0,0 +1,83 @@
+from .base import PlotterBase, PlotConfig
+from .gridded import GriddedPlotter
+from .unstructured import UnstructuredPlotter, clear_cache
+import xarray as xr
+import warnings
+
+# Global variables to store grid information
+_fpath_tgrid = None
+_fpath_ckdtree = None
+_grid_type = None
+
+def _detect_grid_type(xarray_obj):
+    """
+    Deduce grid type based on coordinate structure.
+    
+    Returns:
+        str: 'gridded' or 'unstructured'
+    """
+    has_lat_lon_coords = 'lat' in xarray_obj.coords and 'lon' in xarray_obj.coords
+    has_lat_lon_dims = 'lat' in xarray_obj.dims and 'lon' in xarray_obj.dims
+    
+    # For unstructured data, lat/lon are coordinates but not dimensions
+    return 'unstructured' if (has_lat_lon_coords and not has_lat_lon_dims) else 'gridded'
+
+def register_plotter(xarray_obj):
+    """
+    Determine the appropriate plotter to use based on the data structure.
+    This function is called automatically by xarray's accessor system.
+    
+    First checks if grid type was specified via specify_grid(),
+    then falls back to coordinate-based detection if needed.
+    
+    Returns:
+        Appropriate plotter instance for the data structure
+    """
+    global _grid_type
+    
+    # Determine grid type
+    detected_type = _detect_grid_type(xarray_obj)
+    
+    # If grid type was explicitly specified, check for consistency
+    if _grid_type is not None:
+        if _grid_type != detected_type:
+            warnings.warn(
+                f"Specified grid type '{_grid_type}' differs from detected type '{detected_type}' "
+                f"based on coordinate structure. Using specified type '{_grid_type}'."
+            )
+        final_type = _grid_type
+    else:
+        final_type = detected_type
+    
+    # Create appropriate plotter
+    plotter_class = UnstructuredPlotter if final_type.lower() == 'unstructured' else GriddedPlotter
+    plotter = plotter_class(xarray_obj)
+    
+    # Set grid path if available for unstructured grids
+    if final_type == 'unstructured' and _fpath_tgrid is not None and _fpath_ckdtree is not None:
+        plotter.specify_grid(fpath_tgrid=_fpath_tgrid, fpath_ckdtree=_fpath_ckdtree)
+    
+    return plotter
+
+def specify_grid(grid_type=None, fpath_tgrid=None, fpath_ckdtree=None):
+    """
+    Set the global grid specification that will be used by all plotters.
+    
+    Args:
+        grid_type: str, either 'gridded' or 'unstructured'.
+                  If specified, this will be used as the primary method
+                  to determine grid type.
+        fpath_tgrid: Path to the triangulation grid file
+        fpath_ckdtree: Path to the pre-computed KDTree indices directory
+    """
+    global _fpath_tgrid, _fpath_ckdtree, _grid_type
+    
+    if grid_type is not None and grid_type.lower() not in ['gridded', 'unstructured']:
+        raise ValueError("grid_type must be either 'gridded' or 'unstructured'")
+    
+    _fpath_tgrid = str(fpath_tgrid) if fpath_tgrid else None
+    _fpath_ckdtree = str(fpath_ckdtree) if fpath_ckdtree else None
+    _grid_type = grid_type.lower() if grid_type else None
+
+# Register the accessor
+xr.register_dataarray_accessor('plotX')(register_plotter)