"""Singleton-style helpers to start/reuse Dask clusters for COSEDA tasks."""
from coseda.logging_utils import log_print
from dask.distributed import Client, LocalCluster, TimeoutError
import psutil
import time
import tempfile
import os
import sys
from coseda.logging_utils import log_start
from coseda.logging_utils import get_logger
_LOGGER = get_logger(__name__)
[docs]
class DaskClientManager:
"""Create or reuse a LocalCluster, honoring optional override scheduler addresses."""
_client = None # Class-level variable to hold the singleton instance
temp_dir = tempfile.gettempdir()
_scheduler_file = os.path.join(temp_dir, 'dask_scheduler.json') # Scheduler file in temp directory
# Optional override for scheduler address (e.g., '127.0.0.1:8786')
_override_scheduler_address = os.getenv('DASK_SCHEDULER_ADDRESS', None)
[docs]
@classmethod
def set_scheduler_address(cls, address: str):
"""
Force all new Dask connections to use this scheduler address.
Useful when pointing the UI/CLI to a remote or pre-existing cluster
instead of spawning a LocalCluster.
"""
cls._override_scheduler_address = address
_LOGGER.info(f"Override scheduler address set to: {address}")
[docs]
@classmethod
def get_client(cls):
"""
Return the singleton Dask client, preferring reuse over new clusters.
Preference order:
1) explicit override address
2) existing scheduler file in temp dir
3) brand new LocalCluster
If an existing client reports zero workers, it is discarded and rebuilt.
"""
if cls._client is None:
# If an override address is specified, connect directly and return
if cls._override_scheduler_address:
try:
_LOGGER.info(f"Connecting to overridden Dask scheduler at {cls._override_scheduler_address}")
cls._client = Client(cls._override_scheduler_address, timeout=5)
_LOGGER.info("Connected to overridden Dask cluster.")
return cls._client
except Exception as e:
_LOGGER.warning(f"Failed to connect to override address {cls._override_scheduler_address}: {e}")
if os.path.exists(cls._scheduler_file):
# Try to connect to the existing cluster using the scheduler file
try:
_LOGGER.info(f"Attempting to connect to existing Dask cluster using scheduler file at {cls._scheduler_file}")
cls._client = Client(scheduler_file=cls._scheduler_file, timeout=5)
_LOGGER.info("Connected to existing Dask cluster.")
except (OSError, ValueError, TimeoutError, FileNotFoundError) as e:
_LOGGER.warning(f"Failed to connect to existing Dask cluster: {e}")
# Remove the stale scheduler file
try:
os.remove(cls._scheduler_file)
_LOGGER.info(f"Removed stale scheduler file at {cls._scheduler_file}")
except OSError as remove_error:
_LOGGER.warning(f"Error removing scheduler file: {remove_error}")
# Start a new cluster
cls._start_new_cluster()
else:
# No scheduler file exists, start a new cluster
cls._start_new_cluster()
else:
_LOGGER.info("Reusing existing Dask client.")
_, num_workers, _, _ = cls.get_client_info()
if num_workers == 0:
_LOGGER.warning("Existing cluster had 0 workers, setting up new cluster.")
cls._client.close()
cls._start_new_cluster()
_LOGGER.info("LocalCluster initialized.")
return cls._client
@classmethod
def _start_new_cluster(cls):
"""Start a LocalCluster with platform-aware defaults and attach a Client."""
_LOGGER.info("Starting a new Dask cluster...")
try:
total_cores = psutil.cpu_count(logical=True)
total_memory = psutil.virtual_memory().total / (1024 ** 3) # Convert bytes to GB
_LOGGER.info(f"Total cores: {total_cores}, Total memory: {total_memory:.2f} GB")
"""
n_workers = max(1, total_cores - 1) # Use all but one core
threads_per_worker = 1 # Optimal core utilization
log_print(f"Number of workers: {n_workers}, Threads per worker: {threads_per_worker}")
memory_per_worker = (total_memory / n_workers) * 0.8 # 80% of available memory per worker
log_print(f"Memory per worker: {memory_per_worker:.2f} GB")
"""
cluster_kwargs = {}
force_processes_env = os.getenv("COSEDA_DASK_PROCESSES")
if force_processes_env is not None:
# Any truthy value ("1", "true", etc.) forces process-based workers; falsy keeps threads
force_processes = force_processes_env.strip().lower() in ("1", "true", "yes", "on")
else:
# Default on macOS is now process-based workers to avoid HDF5 thread-safety issues
force_processes = True
if sys.platform == "darwin":
# Keep macOS workers thread-based to avoid per-process RSS bloat.
mac_workers = os.getenv("COSEDA_DASK_N_WORKERS")
if mac_workers:
try:
mac_workers = int(mac_workers)
except ValueError:
mac_workers = None
if not mac_workers:
mac_workers = max(1, total_cores - 1)
cluster_kwargs.update(
processes=force_processes,
threads_per_worker=1,
n_workers=mac_workers,
)
# Allow an opt-in memory limit override for macOS.
mac_mem_limit = os.getenv("COSEDA_DASK_MEMORY_LIMIT")
if mac_mem_limit:
cluster_kwargs["memory_limit"] = mac_mem_limit
_LOGGER.info(f"Detected macOS; initializing LocalCluster with {cluster_kwargs} (defaulting to processes unless COSEDA_DASK_PROCESSES=0)")
else:
_LOGGER.info("Initializing LocalCluster with default (process-based) settings...")
cluster = LocalCluster(**cluster_kwargs)
_LOGGER.info("LocalCluster initialized.")
cls._client = Client(cluster)
_LOGGER.info("Started a new Dask cluster.")
except Exception as e:
_LOGGER.error(f"Failed to start a new Dask cluster: {e}")
[docs]
@classmethod
def get_client_info(cls):
"""Return dashboard URL, worker count, threads per worker, and memory per worker."""
if cls._client is None:
_LOGGER.info("No Dask client is currently running.")
return None, None, None, None
else:
client = cls._client
dashboard_address = client.dashboard_link
# Get number of workers and threads per worker
workers_info = client.scheduler_info()['workers']
num_workers = len(workers_info)
threads_per_worker_set = set(worker_info['nthreads'] for worker_info in workers_info.values())
if len(threads_per_worker_set) == 1:
threads_per_worker = threads_per_worker_set.pop()
else:
threads_per_worker = 'Variable' # Workers have different numbers of threads
# Get memory limit per worker
memory_limits = set(worker_info['memory_limit'] for worker_info in workers_info.values())
if len(memory_limits) == 1:
memory_per_worker = memory_limits.pop() / (1024 ** 3) # Convert bytes to GB
memory_per_worker = f"{memory_per_worker:.1f}GB"
else:
memory_per_worker = 'Variable' # Workers have different memory limits
return dashboard_address, num_workers, threads_per_worker, memory_per_worker
[docs]
@classmethod
def log_client_info(cls, logfile_path):
"""
Log current worker counts, thread counts, and memory per worker to a logfile.
If no client is running, writes a short note instead of raising.
"""
if cls._client is None:
# Log that no Dask client is currently running
log_start(logfile_path, "No Dask client is currently running.")
return
else:
client = cls._client
# Get client information
dashboard_address = client.dashboard_link
# Get number of workers and threads per worker
workers_info = client.scheduler_info()['workers']
num_workers = len(workers_info)
threads_per_worker_set = set(worker_info['nthreads'] for worker_info in workers_info.values())
if len(threads_per_worker_set) == 1:
threads_per_worker = threads_per_worker_set.pop()
else:
threads_per_worker = 'Variable' # Workers have different numbers of threads
# Get memory limit per worker
memory_limits = set(worker_info['memory_limit'] for worker_info in workers_info.values())
if len(memory_limits) == 1:
memory_per_worker = memory_limits.pop() / (1024 ** 3) # Convert bytes to GB
memory_per_worker = f"{memory_per_worker:.1f}GB"
else:
memory_per_worker = 'Variable' # Workers have different memory limits
# Prepare the log message
log_message = (
f"Dask Client Information: "
f"Number of Workers: {num_workers}, "
f"Threads per Worker: {threads_per_worker}, "
f"Memory per Worker: {memory_per_worker}"
)
# Use log_start to log the information
log_start(logfile_path, log_message)
[docs]
@classmethod
def close_client(cls):
"""Close the active client handle (does not kill external clusters)."""
if cls._client is not None:
cls._client.close()
cls._client = None
# Optionally, you can log or print that you closed the client
else:
# Optionally, you can log or print that there's no client to close
pass
[docs]
@classmethod
def kill_current_cluster(cls):
"""
Gracefully shutdown the current Dask cluster via the client API and clean up.
Removes any persisted scheduler file so the next run does not attempt to
reuse a stale cluster.
"""
if cls._client is None:
_LOGGER.info("No Dask client to shutdown.")
return
try:
_LOGGER.info("Requesting scheduler to shutdown via client.shutdown()...")
cls._client.shutdown()
except Exception as e:
_LOGGER.warning(f"Error during shutdown request: {e}")
finally:
try:
cls._client.close()
except Exception:
pass
cls._client = None
# Remove scheduler file if it exists
try:
if os.path.exists(cls._scheduler_file):
os.remove(cls._scheduler_file)
_LOGGER.info(f"Removed scheduler file at {cls._scheduler_file}")
except Exception as e:
_LOGGER.warning(f"Error removing scheduler file: {e}")
[docs]
@classmethod
def kill_all_dask_clusters(cls):
"""Terminate all local Dask worker/scheduler processes (last resort)."""
cls.close_client()
time.sleep(2)
# Kill all running Dask clusters by terminating relevant processes
for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
try:
cmdline = proc.info['cmdline']
if cmdline and ('dask-worker' in cmdline or 'dask-scheduler' in cmdline):
proc.terminate()
# Optionally, you can log or print that you terminated a Dask process
except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
pass
time.sleep(2)