sibi-flux 2025.12.0__py3-none-any.whl

sibi_flux/dask_cluster/client_manager.py

@@ -0,0 +1,549 @@
+"""
+Dask Client Manager module.
+
+Manages Dask cluster lifecycles, registry paths, and client connections.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+import tempfile
+import time
+import psutil
+import multiprocessing
+import threading
+from pathlib import Path
+from contextlib import asynccontextmanager, contextmanager, suppress
+from typing import Optional, Dict, Any
+
+import shutil
+
+try:
+    from dask.distributed import Client, LocalCluster, get_client
+    HAS_DISTRIBUTED = True
+except ImportError:
+    Client = object
+    LocalCluster = object
+    get_client = None
+    HAS_DISTRIBUTED = False
+
+# Lazy import watchdog to allow remote workers (which might lack it) to import this module
+try:
+    from watchdog.observers import Observer
+    from watchdog.events import FileSystemEventHandler
+
+    HAS_WATCHDOG = True
+except ImportError:
+    Observer = object
+    FileSystemEventHandler = object
+    HAS_WATCHDOG = False
+
+try:
+    from sibi_flux.logger import Logger
+
+    _default_logger = Logger.default_logger(logger_name="dask_lifecycle_manager")
+except ImportError:
+    _default_logger = logging.getLogger(__name__)
+    if not _default_logger.handlers:
+        _default_logger.addHandler(logging.StreamHandler())
+    _default_logger.setLevel(logging.INFO)
+
+_cores = multiprocessing.cpu_count()
+DEFAULT_WORKERS = max(2, min(4, _cores))
+DEFAULT_THREADS = 2
+
+
+class RegistryEventHandler(FileSystemEventHandler):
+    """
+    Watchdog handler to trigger logic when the registry file changes.
+    Used to detect if the shared cluster configuration has been updated by another process.
+    """
+
+    def __init__(self, callback):
+        if not HAS_WATCHDOG:
+            return
+        self.callback = callback
+
+    def on_modified(self, event):
+        if not event.is_directory and event.src_path.endswith(
+            "shared_dask_cluster.json"
+        ):
+            self.callback()
+
+
+class DaskClientMixin:
+    """
+    Stabilised Dask lifecycle manager.
+
+    This class manages a shared Dask LocalCluster instance, enabling:
+    - Singleton cluster management across processes via a file-based registry.
+    - Automatic healing and reconnection if the cluster becomes unresponsive.
+    - Event-driven monitoring using `watchdog` to detect external registry updates.
+    """
+
+    REGISTRY_PATH = Path(tempfile.gettempdir()) / "shared_dask_cluster.json"
+    REGISTRY_LOCK_PATH = REGISTRY_PATH.parent / (REGISTRY_PATH.name + ".lock")
+    WATCHDOG_INTERVAL = 30
+    HEAL_COOLDOWN = 5.0
+
+    def __init__(self, **kwargs):
+        self.dask_client: Optional[Client] = None
+        self.own_dask_client = False
+        self.logger: Any = kwargs.get("logger") or _default_logger
+        self._watchdog_task: Optional[asyncio.Task] = None
+        self._watchdog_stop = asyncio.Event()
+        self._fs_observer: Optional[Observer] = None
+        self._init_params = {}
+        self._last_heal_time = 0.0
+        self._last_registry_write = 0.0
+
+    def _deploy_local_code(self, client: Client) -> None:
+        """
+        Zips the local 'sibi_flux' package and uploads it to the scheduler.
+        Dask distributed will propagate this to all workers.
+        """
+        try:
+            import sibi_flux
+
+            pkg_path = Path(sibi_flux.__file__).parent
+            src_root = pkg_path.parent
+
+            with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
+                archive_path = Path(tmp.name)
+
+            shutil.make_archive(
+                str(archive_path.with_suffix("")),  # trim extension, shutil adds it
+                "zip",
+                root_dir=str(src_root),
+                base_dir="sibi_flux",
+            )
+
+            self.logger.info(f"Deploying local code to cluster: {archive_path}")
+            client.upload_file(str(archive_path))
+
+            # Cleanup
+            archive_path.unlink(missing_ok=True)
+
+        except Exception as e:
+            self.logger.warning(f"Failed to deploy local code: {e}")
+
+    def _verify_connectivity(self, client: Client) -> bool:
+        """
+        Verifies that the client can actually run tasks on the cluster.
+        This catches issues like missing environments, version mismatches,
+        or deserialization errors that occur after connection.
+        """
+        try:
+            # 1. Check for workers
+            # Giving it a moment to discover workers if just connected
+            client.wait_for_workers(1, timeout=2.0)
+
+            def environment_probe():
+                return True
+
+            f = client.submit(environment_probe)
+            # Retry loop for race conditions (workers restart/receive file)
+            for attempt in range(3):
+                try:
+                    if f.result(timeout=5.0) is True:
+                        return True
+                except Exception:
+                    if attempt < 2:
+                        time.sleep(2.0)
+                        continue
+            return False
+        except Exception as e:
+            self.logger.warning(
+                f"Cluster verification failed (Environment mismatch?): {e}",
+                extra={"error": str(e)},
+            )
+            return False
+
+    def _remove_registry(self):
+        with suppress(FileNotFoundError):
+            self.REGISTRY_PATH.unlink()
+
+    def _cleanup_stale_registry(self):
+        """Removes the registry file if the owning process is no longer running."""
+        reg = self._read_registry()
+        if reg and reg.get("pid"):
+            if not psutil.pid_exists(reg["pid"]):
+                self.logger.warning(
+                    "Cleaning zombie registry", extra={"pid": reg["pid"]}
+                )
+                self._remove_registry()
+
+    def _retire_all_workers(self):
+        """Attempts to gracefully retire all workers."""
+        if self.dask_client:
+            with suppress(Exception):
+                self.dask_client.retire_workers()
+
+    def _has_inflight(self) -> bool:
+        return False
+
+    @contextmanager
+    def _registry_lock(self, timeout: float = 10.0):
+        """
+        Acquires an exclusive file lock to ensure atomic registry updates.
+        """
+        start_time = time.time()
+        while True:
+            try:
+                # Atomically create the lock file. Fails if exists.
+                self.REGISTRY_LOCK_PATH.touch(exist_ok=False)
+                try:
+                    yield
+                finally:
+                    with suppress(OSError):
+                        self.REGISTRY_LOCK_PATH.unlink()
+                return
+            except FileExistsError:
+                if time.time() - start_time > timeout:
+                    raise TimeoutError("Could not acquire registry lock")
+                time.sleep(0.05)
+
+    def _read_registry(self) -> Optional[Dict[str, Any]]:
+        """Reads and parses the registry file JSON."""
+        if not self.REGISTRY_PATH.exists():
+            return None
+        try:
+            with self.REGISTRY_PATH.open("r") as f:
+                return json.load(f)
+        except (json.JSONDecodeError, OSError):
+            return None
+
+    def _write_registry(self, data: Dict[str, Any]) -> None:
+        """
+        Writes data to the registry file atomically via a temporary file.
+        Updates the timestamp of the last write.
+        """
+        self._last_registry_write = time.time()
+        # Use a temporary file in the same directory to ensure atomic move support
+        tmp = self.REGISTRY_PATH.with_suffix(".tmp")
+        with tmp.open("w") as f:
+            json.dump(data, f)
+        tmp.replace(self.REGISTRY_PATH)
+
+    def _init_dask_client(self, **kwargs) -> None:
+        self._init_params = kwargs
+        if not HAS_DISTRIBUTED:
+             self.logger.info("Dask Distributed not installed. Skipping cluster initialization.")
+             return
+
+        if kwargs.get("dask_client"):
+            self.dask_client = kwargs["dask_client"]
+            return
+        if kwargs.get("reuse_context_client", True):
+            with suppress(ValueError):
+                self.dask_client = get_client()
+                return
+        self._connect_or_create()
+        if kwargs.get("watchdog", True):
+            self._start_watchdog()
+
+    def _connect_or_create(self) -> None:
+        """
+        Connects to an existing Dask cluster or creates a new one.
+        Handles locking, registry validation, and fallback to ephemeral clusters.
+        """
+        p = self._init_params
+        try:
+            # 1. Explicit argument
+            explicit_address = p.get("scheduler_address")
+            if explicit_address:
+                self.logger.info(
+                    f"Connecting to EXTERNAL Dask Scheduler (Explicit): {explicit_address}"
+                )
+                try:
+                    candidate_client = Client(
+                        address=explicit_address,
+                        timeout=p.get("timeout", 5),
+                        set_as_default=p.get("set_as_default", True),
+                    )
+                    if self._verify_connectivity(candidate_client):
+                        self.dask_client = candidate_client
+                        self.own_dask_client = False
+                        return
+                    else:
+                        candidate_client.close()
+                except Exception as e:
+                    self.logger.error(f"Failed to connect to specified scheduler: {e}")
+
+            # 2. Environment Variable
+            env_address = os.environ.get("DASK_SCHEDULER_ADDRESS")
+            if env_address:
+                self.logger.info(
+                    f"Connecting to EXTERNAL Dask Scheduler (Env): {env_address}"
+                )
+                try:
+                    candidate_client = Client(
+                        address=env_address,
+                        timeout=p.get("timeout", 5),
+                        set_as_default=p.get("set_as_default", True),
+                    )
+
+                    # Optional: Deploy code if requested
+                    if p.get("deploy_mode", False):
+                        self._deploy_local_code(candidate_client)
+
+                    if self._verify_connectivity(candidate_client):
+                        self.dask_client = candidate_client
+                        self.own_dask_client = False
+                        return
+                    else:
+                        self.logger.warning(
+                            "External cluster unusable. Falling back to LOCAL."
+                        )
+                        candidate_client.close()
+                except Exception as e:
+                    self.logger.error(
+                        f"Failed to connect to env-specified scheduler: {e}"
+                    )
+
+            try:
+                lock_context = self._registry_lock()
+            except TimeoutError:
+                self.logger.warning(
+                    "Lock acquisition timed out; falling back to ephemeral cluster."
+                )
+                self._bootstrap_new_cluster(p)
+                return
+
+            try:
+                with lock_context:
+                    self._cleanup_stale_registry()
+                    reg = self._read_registry()
+
+                    if reg and not reg.get("closing"):
+                        try:
+                            self.dask_client = Client(
+                                address=reg["address"],
+                                timeout=p.get("timeout", 30),
+                                set_as_default=p.get("set_as_default", True),
+                            )
+                            reg["refcount"] = int(reg.get("refcount", 0)) + 1
+                            self._write_registry(reg)
+                            self.own_dask_client = False
+                            self.logger.info(
+                                f"Connected to SHARED LOCAL Dask cluster at {reg['address']}\n"
+                                f"Dashboard: {reg.get('dashboard_link', 'Unknown')}"
+                            )
+                            return
+                        except Exception:
+                            self.logger.warning("Failed to connect to shared cluster.")
+
+                    self._bootstrap_new_cluster(p)
+            except TimeoutError:
+                self.logger.warning("Lock context entry timed out; falling back.")
+                self._bootstrap_new_cluster(p)
+
+        except Exception as e:
+            self.logger.error("Dask initialization failed", extra={"error": str(e)})
+
+    def _bootstrap_new_cluster(self, p: Dict[str, Any]) -> None:
+        """
+        Creates a new LocalCluster instance.
+
+        Configures the cluster based on input parameters or defaults.
+        Logs:
+            Silences 'distributed.shuffle' logs to ERROR level to reduce noise.
+        """
+
+        self.logger.info("Starting NEW LOCAL Dask cluster...")
+        cluster = LocalCluster(
+            n_workers=p.get("n_workers", DEFAULT_WORKERS),
+            threads_per_worker=p.get("threads_per_worker", DEFAULT_THREADS),
+            processes=p.get("processes", False),
+            dashboard_address=p.get("dashboard_address", ":0"),
+        )
+        self.dask_client = Client(cluster, set_as_default=p.get("set_as_default", True))
+        self.own_dask_client = True
+        self.logger.info(f"Local Cluster Dashboard: {cluster.dashboard_link}")
+        self._write_registry(
+            {
+                "address": cluster.scheduler_address,
+                "dashboard_link": cluster.dashboard_link,
+                "refcount": 1,
+                "closing": False,
+                "pid": os.getpid(),
+                "created_at": time.time(),
+            }
+        )
+
+    def _close_dask_client(self) -> None:
+        """
+        Closes the Dask client and cleans up resources.
+        Decrements refcount in registry and shuts down cluster if count hits zero.
+        """
+        if self._fs_observer:
+            self._fs_observer.stop()
+            self._fs_observer.join(timeout=2.0)
+            self._fs_observer = None
+
+        if self._watchdog_task:
+            self._watchdog_stop.set()
+
+        if not self.dask_client:
+            return
+
+        if not self.own_dask_client:
+            self.dask_client = None
+            return
+
+        try:
+            with self._registry_lock():
+                if self._has_inflight():
+                    return
+
+                reg = self._read_registry()
+                if reg:
+                    reg["refcount"] = max(0, int(reg.get("refcount", 1)) - 1)
+                    if reg["refcount"] > 0:
+                        self._write_registry(reg)
+                        self.dask_client.close()
+                        return
+                    reg["closing"] = True
+                    self._write_registry(reg)
+                    self._retire_all_workers()
+                    self._remove_registry()
+
+                cluster = getattr(self.dask_client, "cluster", None)
+                self.dask_client.close()
+                if cluster:
+                    cluster.close()
+                self.dask_client = None
+        except Exception as e:
+            self.logger.error("Shutdown error", extra={"error": str(e)})
+
+    def _start_watchdog(self) -> None:
+        if not HAS_WATCHDOG:
+            return
+
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            return
+
+        def on_fs_event():
+            if time.time() - self._last_registry_write < 1.5:
+                return
+            loop.call_soon_threadsafe(
+                lambda: asyncio.create_task(self._heal_if_disconnected())
+            )
+
+        handler = RegistryEventHandler(on_fs_event)
+        self._fs_observer = Observer()
+        self._fs_observer.schedule(
+            handler, str(self.REGISTRY_PATH.parent), recursive=False
+        )
+        self._fs_observer.start()
+
+        async def liveness_checker():
+            while not self._watchdog_stop.is_set():
+                await asyncio.sleep(self.WATCHDOG_INTERVAL)
+                await self._heal_if_disconnected()
+
+        self._watchdog_task = loop.create_task(liveness_checker())
+
+    async def _heal_if_disconnected(self) -> None:
+        now = time.time()
+        if now - self._last_heal_time < self.HEAL_COOLDOWN:
+            return
+        is_healthy = False
+        if self.dask_client:
+            try:
+                # Handle both sync and async clients
+                if getattr(self.dask_client, "asynchronous", False):
+                    await self.dask_client.scheduler_info()
+                else:
+                    # Sync client: scheduler_info() returns dict immediately
+                    self.dask_client.scheduler_info()
+                is_healthy = True
+            except Exception:
+                pass
+        if not is_healthy:
+            self._last_heal_time = now
+            self.logger.warning("Dask heartbeat lost. Healing.")
+            with suppress(Exception):
+                if self.dask_client:
+                    await self.dask_client.close(timeout=1)
+            self.dask_client = None
+            self._connect_or_create()
+
+
+_persistent_mixin: Optional[DaskClientMixin] = None
+_singleton_lock = threading.Lock()
+
+
+def get_persistent_client(**kwargs) -> Client:
+    """
+    Retrieves or creates a singleton Dask client instance.
+    This client persists until `force_close_persistent_client` is called.
+    """
+    if not HAS_DISTRIBUTED:
+        return None
+    global _persistent_mixin
+    with _singleton_lock:
+        if _persistent_mixin is None:
+            _persistent_mixin = DaskClientMixin(logger=kwargs.get("logger"))
+
+        # If client is None, initialize it
+        if _persistent_mixin.dask_client is None:
+            _persistent_mixin._init_dask_client(**kwargs)
+        else:
+            # Check if we need to upgrade from Local to External+Deploy
+            current_local = _persistent_mixin.own_dask_client
+            want_deploy = kwargs.get("deploy_mode", False)
+
+            # Update params for future healings
+            _persistent_mixin._init_params.update(kwargs)
+
+            if want_deploy and current_local:
+                _persistent_mixin.logger.info(
+                    "Re-initialising client to apply 'deploy_mode' configuration."
+                )
+                _persistent_mixin._close_dask_client()
+                # _init_dask_client will use the updated params
+                _persistent_mixin._init_dask_client(**kwargs)
+
+    return _persistent_mixin.dask_client  # type: ignore
+
+
+def force_close_persistent_client():
+    global _persistent_mixin
+    with _singleton_lock:
+        if _persistent_mixin:
+            _persistent_mixin._close_dask_client()
+            _persistent_mixin = None
+
+
+def shared_dask_session(*, async_mode: bool = True, **kwargs):
+    """
+    Context manager for a shared Dask session.
+    """
+    mixin = DaskClientMixin(logger=kwargs.get("logger"))
+    mixin._init_dask_client(**kwargs)
+    if async_mode:
+
+        @asynccontextmanager
+        async def _async_manager():
+            try:
+                yield mixin.dask_client
+            finally:
+                mixin._close_dask_client()
+
+        return _async_manager()
+    else:
+
+        @contextmanager
+        def _sync_manager():
+            try:
+                yield mixin.dask_client
+            finally:
+                mixin._close_dask_client()
+
+        return _sync_manager()