sibi-dst 2025.9.9__py3-none-any.whl → 2025.9.10__py3-none-any.whl

sibi_dst/df_helper/_artifact_updater_async.py

@@ -6,10 +6,16 @@ import random
 import time
 from contextlib import ExitStack
 from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Type
+from typing import Any, Callable, Dict, List, Optional, Sequence, Type
 
 from sibi_dst.utils import ManagedResource
 
+try:
+    from dask.distributed import Client, LocalCluster
+except ImportError:
+    Client = None
+    LocalCluster = None
+
 
 @dataclass(slots=True)
 class _RetryCfg:
@@ -19,30 +25,59 @@ class _RetryCfg:
     jitter: float = 0.15
 
 
-_ORCHESTRATOR_KEYS = {
-    "retry_attempts",
-    "backoff_base",
-    "backoff_max",
-    "backoff_jitter",
-    "update_timeout_seconds",
-    "max_workers",
-    "priority_fn",
-    "artifact_class_kwargs",
-}
+# ---------------- Worker (safe for Dask pickling) ----------------
+def run_artifact_update(
+    cls: Type,
+    artifact_class_kwargs: Dict[str, Any],
+    retry: _RetryCfg,
+    period: str,
+    artifact_kwargs: Dict[str, Any],
+) -> Dict[str, Any]:
+    """Standalone worker — safe for Dask distributed execution."""
+    import logging
+
+    logger = logging.getLogger(cls.__name__)
 
+    start_wall = datetime.datetime.now()
+    attempt_count = 0
+    success = False
+    error_msg = None
+
+    for attempt in range(1, retry.attempts + 1):
+        attempt_count = attempt
+        try:
+            with ExitStack() as stack:
+                inst = cls(**artifact_class_kwargs)
+                inst = stack.enter_context(inst)
+                inst.update_parquet(period=period, **artifact_kwargs)
+            success = True
+            break
+        except Exception as e:
+            error_msg = str(e)
+            if attempt < retry.attempts:
+                delay = min(retry.backoff_base ** (attempt - 1), retry.backoff_max)
+                delay *= 1 + random.uniform(0, retry.jitter)
+                time.sleep(delay)
+
+    end_wall = datetime.datetime.now()
+    duration = (end_wall - start_wall).total_seconds()
 
-def _default_artifact_kwargs(resource: ManagedResource) -> Dict[str, Any]:
     return {
-        "logger": resource.logger,
-        "debug": resource.debug,
-        "fs": resource.fs,
-        "verbose": resource.verbose,
+        "artifact": cls.__name__,
+        "period": period,
+        "start": start_wall.isoformat(),
+        "end": end_wall.isoformat(),
+        "processing_time": duration,
+        "retries": attempt_count - 1 if success else attempt_count,
+        "success": success,
+        "error": error_msg,
     }
 
 
 class ArtifactUpdaterMultiWrapperAsync(ManagedResource):
     """
-    Backward-compatible async orchestrator with shutdown-aware scheduling.
+    Async/Threaded orchestrator.
+    Dask-enabled if a Client is passed (or created automatically).
     """
 
     def __init__(
@@ -57,13 +92,19 @@ class ArtifactUpdaterMultiWrapperAsync(ManagedResource):
         backoff_jitter: float = 0.15,
         priority_fn: Optional[Callable[[Type], int]] = None,
         artifact_class_kwargs: Optional[Dict[str, Any]] = None,
+        dask_client: Optional[Client] = None,
+        use_dask: bool = True,
         **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
+
         self.wrapped_classes = wrapped_classes
         self.max_workers = int(max_workers)
         self.update_timeout_seconds = int(update_timeout_seconds)
         self.priority_fn = priority_fn
+        self.use_dask = use_dask
+        self.client: Optional[Client] = dask_client
+        self._owns_client = False
 
         self._retry = _RetryCfg(
             attempts=int(retry_attempts),
@@ -72,29 +113,41 @@ class ArtifactUpdaterMultiWrapperAsync(ManagedResource):
             jitter=float(backoff_jitter),
         )
 
-        self.artifact_class_kwargs = {
-            **_default_artifact_kwargs(self),
-            **(artifact_class_kwargs or {}),
-        }
+        # Safe kwargs for artifacts
+        if self.use_dask:
+            self.artifact_class_kwargs = {
+                "debug": self.debug,
+                "verbose": self.verbose,
+                **(artifact_class_kwargs or {}),
+            }
+        else:
+            self.artifact_class_kwargs = {
+                "logger": self.logger,
+                "fs": self.fs,
+                "debug": self.debug,
+                "verbose": self.verbose,
+                **(artifact_class_kwargs or {}),
+            }
 
         self.completion_secs: Dict[str, float] = {}
         self.failed: List[str] = []
-
-        # NEW: async stop gate — tripped on cleanup/cancel
         self._stop = asyncio.Event()
 
-    # Trip stop gate on close paths
-    def _cleanup(self) -> None:
-        try:
-            loop = asyncio.get_running_loop()
-            loop.call_soon_threadsafe(self._stop.set)
-        except RuntimeError:
-            self._stop.set()
-
-    async def _acleanup(self) -> None:
-        self._stop.set()
+        if self.use_dask and Client is None:
+            raise RuntimeError("Dask is not installed, cannot use Dask mode")
+
+        # auto-start local client if requested
+        if self.use_dask and not self.client:
+            self.client = Client(
+                LocalCluster(
+                    n_workers=max_workers,
+                    threads_per_worker=1,
+                    dashboard_address=None,
+                )
+            )
+            self._owns_client = True
 
-    # ---- internals -----------------------------------------------------------
+    # ---- Internals ------------------------------------------------------------
 
     def _classes_for(self, period: str) -> List[Type]:
         try:
@@ -102,7 +155,7 @@ class ArtifactUpdaterMultiWrapperAsync(ManagedResource):
         except KeyError:
             raise ValueError(f"Unsupported period '{period}'.")
         if not classes:
-            raise ValueError(f"No artifact classes configured for period '{period}'.")
+            raise ValueError(f"No artifact classes configured for '{period}'.")
         if self.priority_fn:
             try:
                 classes.sort(key=self.priority_fn)
@@ -110,120 +163,106 @@ class ArtifactUpdaterMultiWrapperAsync(ManagedResource):
                 self.logger.warning(f"priority_fn failed; using listed order: {e}")
         return classes
 
-    @staticmethod
-    def _split_kwargs(raw: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict[str, Any]]:
-        orch: Dict[str, Any] = {}
-        art: Dict[str, Any] = {}
-        for k, v in raw.items():
-            if k in _ORCHESTRATOR_KEYS:
-                orch[k] = v
-            else:
-                art[k] = v
-        return orch, art
+    def _submit_one_dask(self, cls: Type, period: str, artifact_kwargs: Dict[str, Any]):
+        return self.client.submit(
+            run_artifact_update,
+            cls,
+            dict(self.artifact_class_kwargs),
+            self._retry,
+            period,
+            artifact_kwargs,
+            pure=False,
+        )
 
-    async def _run_one(self, cls: Type, period: str, sem: asyncio.Semaphore, artifact_kwargs: Dict[str, Any]) -> None:
+    async def _run_one_async(
+        self,
+        cls: Type,
+        period: str,
+        sem: asyncio.Semaphore,
+        artifact_kwargs: Dict[str, Any],
+    ) -> Dict[str, Any]:
+        """Async/threaded fallback execution."""
         name = cls.__name__
-        if self._stop.is_set() or self.closed:
-            raise asyncio.CancelledError()
-
-        self.logger.info(f"Running {name} with period '{period}'", extra={"artifact": name, "period": period})
-        async with sem:
-            loop = asyncio.get_running_loop()
-            start = loop.time()
-            for attempt in range(1, self._retry.attempts + 1):
-                if self._stop.is_set() or self.closed:
-                    raise asyncio.CancelledError()
-                try:
-                    def _sync_block() -> None:
-                        with ExitStack() as stack:
-                            inst = cls(**self.artifact_class_kwargs)
-                            inst = stack.enter_context(inst)
-                            inst.update_parquet(period=period, **artifact_kwargs)
-
-                    await asyncio.wait_for(asyncio.to_thread(_sync_block), timeout=self.update_timeout_seconds)
-                    dt_secs = loop.time() - start
-                    self.completion_secs[name] = dt_secs
-                    self.logger.info(f"✅ {name} ({period}) in {dt_secs:.2f}s")
-                    return
-
-                except asyncio.TimeoutError:
-                    self.logger.warning(f"Timeout in {name} attempt {attempt}/{self._retry.attempts}")
-                except asyncio.CancelledError:
-                    raise
-                except Exception as e:
-                    self.logger.error(
-                        f"{name} attempt {attempt}/{self._retry.attempts} failed: {e}",
-                        exc_info=self.debug,
-                    )
-
-                if attempt < self._retry.attempts and not self._stop.is_set():
-                    delay = min(self._retry.backoff_base ** (attempt - 1), self._retry.backoff_max)
-                    delay *= 1 + random.uniform(0, self._retry.jitter)
+        self.logger.info(f"▶️ Starting {name} for period '{period}'")
+        start_wall = datetime.datetime.now()
+
+        attempt_count = 0
+        success = False
+        error_msg = None
+
+        try:
+            async with sem:
+                for attempt in range(1, self._retry.attempts + 1):
+                    attempt_count = attempt
                     try:
-                        await asyncio.sleep(delay)
-                    except asyncio.CancelledError:
-                        raise
+                        def _sync_block():
+                            with ExitStack() as stack:
+                                inst = cls(**self.artifact_class_kwargs)
+                                inst = stack.enter_context(inst)
+                                inst.update_parquet(period=period, **artifact_kwargs)
+
+                        await asyncio.wait_for(
+                            asyncio.to_thread(_sync_block),
+                            timeout=self.update_timeout_seconds,
+                        )
+                        success = True
+                        break
+                    except Exception as e:
+                        error_msg = str(e)
+                        if attempt < self._retry.attempts and not self._stop.is_set():
+                            delay = min(
+                                self._retry.backoff_base ** (attempt - 1),
+                                self._retry.backoff_max,
+                            )
+                            delay *= 1 + random.uniform(0, self._retry.jitter)
+                            await asyncio.sleep(delay)
+        finally:
+            end_wall = datetime.datetime.now()
+            duration = (end_wall - start_wall).total_seconds()
 
-            self.failed.append(name)
-            self.logger.error(f"✖️  {name} permanently failed")
+            result = {
+                "artifact": name,
+                "period": period,
+                "start": start_wall.isoformat(),
+                "end": end_wall.isoformat(),
+                "processing_time": duration,
+                "retries": attempt_count - 1 if success else attempt_count,
+                "success": success,
+                "error": error_msg,
+            }
+
+            if success:
+                self.logger.info(f"✅ Artifact {name} succeeded", extra=result)
+                self.completion_secs[name] = duration
+            else:
+                self.logger.error(f"❌ Artifact {name} failed", extra=result)
+                self.failed.append(name)
 
-    # ---- public API ----------------------------------------------------------
+            return result
 
-    async def update_data(self, period: str, **kwargs: Any) -> None:
-        """
-        Backward-compatible:
-          - Accepts orchestrator knobs in kwargs (we consume them).
-          - Forwards only artifact-friendly kwargs to update_parquet.
-        """
-        _, artifact_kwargs = self._split_kwargs(kwargs)
+    # ---- Public API -----------------------------------------------------------
 
+    async def update_data(self, period: str, **kwargs: Any) -> List[Dict[str, Any]]:
         self.completion_secs.clear()
         self.failed.clear()
-
         classes = self._classes_for(period)
-        self.logger.info(
-            f"Starting update of {len(classes)} artifacts for period '{period}'",
-            extra={
-                "action_module_name": self.__class__.__name__,
-                "date_of_update": time.strftime('%Y-%m-%d'),
-                "start_time": time.strftime('%H:%M:%S'),
-                "period": period,
-            },
-        )
-
-        sem = asyncio.Semaphore(self.max_workers)
-        tasks = [asyncio.create_task(self._run_one(cls, period, sem, dict(artifact_kwargs))) for cls in classes]
 
         try:
-            for t in asyncio.as_completed(tasks):
-                if self._stop.is_set():
-                    break
-                await t
-        except (asyncio.CancelledError, KeyboardInterrupt):
-            self._stop.set()
-            for t in tasks:
-                t.cancel()
-            raise
+            if self.use_dask:
+                futures = [self._submit_one_dask(cls, period, kwargs) for cls in classes]
+                results = await asyncio.to_thread(lambda: self.client.gather(futures))
+            else:
+                sem = asyncio.Semaphore(self.max_workers)
+                tasks = [
+                    asyncio.create_task(self._run_one_async(cls, period, sem, kwargs))
+                    for cls in classes
+                ]
+                results = await asyncio.gather(*tasks)
+            return results
         finally:
-            # Drain/cancel everything deterministically
-            for t in tasks:
-                if not t.done():
-                    t.cancel()
-            await asyncio.gather(*tasks, return_exceptions=True)
-
-            self.logger.info(
-                f"Update completed for period: {period}",
-                extra={
-                    "action_module_name": self.__class__.__name__,
-                    "date_of_update": datetime.date.today().strftime('%Y-%m-%d'),
-                    "end_time": datetime.datetime.now().strftime('%H:%M:%S'),
-                    "period": period,
-                },
-            )
-            self.logger.info(
-                f"Artifacts processed: total={len(classes)}, "
-                f"completed={len(self.completion_secs)}, failed={len(self.failed)}"
-            )
+            # only shut down if we own the client
+            if self._owns_client:
+                self.close()
 
     def get_update_status(self) -> Dict[str, Any]:
         done = set(self.completion_secs)
@@ -235,4 +274,19 @@ class ArtifactUpdaterMultiWrapperAsync(ManagedResource):
             "failed": sorted(fail),
             "pending": sorted(all_names - done - fail),
             "completion_times": dict(self.completion_secs),
-        }
+        }
+
+    # ---- Lifecycle ------------------------------------------------------------
+
+    def _cleanup(self) -> None:
+        """Release any resources created by this wrapper."""
+        if self._owns_client and self.client is not None:
+            try:
+                cluster = getattr(self.client, "cluster", None)
+                self.client.close()
+                if cluster is not None:
+                    cluster.close()
+            finally:
+                self.client = None
+                self._owns_client = False
+