sibi-dst 0.3.58__py3-none-any.whl → 0.3.60__py3-none-any.whl

sibi_dst/df_helper/_artifact_updater_multi_wrapper.py

@@ -1,261 +1,239 @@
 import asyncio
 import logging
 import datetime
-import psutil
-import time
-from functools import total_ordering
-from collections import defaultdict
-from contextlib import asynccontextmanager
-import signal
-from sibi_dst.utils import Logger
-
-@total_ordering
-class PrioritizedItem:
-    def __init__(self, priority, artifact):
-        self.priority = priority
-        self.artifact = artifact
+import random
+from typing import Any, Callable, Dict, List, Optional, Type
 
-    def __lt__(self, other):
-        return self.priority < other.priority
+from sibi_dst.utils import Logger
 
-    def __eq__(self, other):
-        return self.priority == other.priority
 
 class ArtifactUpdaterMultiWrapper:
-    def __init__(self, wrapped_classes=None, debug=False, **kwargs):
-        self.wrapped_classes = wrapped_classes or {}
-        self.debug = debug
-        self.logger = kwargs.setdefault(
-            'logger', Logger.default_logger(logger_name=self.__class__.__name__)
+    """
+    Simplified wrapper that updates artifacts concurrently using an asyncio.Semaphore.
+
+    Features:
+    - Caps concurrency at max_workers via semaphore
+    - Optionally prioritises tasks via a priority function or static method on artifact classes
+    - Tracks per-artifact completion times
+    - Configurable retry/backoff strategy
+    - Optional metrics integration
+    - Thread-safe within a single asyncio loop
+
+    Usage:
+        wrapper = ArtifactUpdaterMultiWrapper(
+            wrapped_classes={
+                'mydata': [DataArtifactA, DataArtifactB],
+            },
+            max_workers=4,
+            retry_attempts=3,
+            update_timeout_seconds=600,
+            backoff_base=2,
+            backoff_max=60,
+            backoff_jitter=0.1,
+            priority_fn=None,  # or custom
+            metrics_client=None,
+            debug=True,
+            logger=None,
+            artifact_class_kwargs={
+                'fs': my_fs,
+                'parquet_storage_path': 's3://bucket/data',
+                'logger': my_logger,
+                'debug': True,
+            }
         )
-        self.logger.set_level(logging.DEBUG if debug else logging.INFO)
+        await wrapper.update_data('mydata', period='ytd', overwrite=True)
+    """
+    def __init__(
+        self,
+        wrapped_classes: Dict[str, List[Type]],
+        *,
+        max_workers: int = 3,
+        retry_attempts: int = 3,
+        update_timeout_seconds: int = 600,
+        backoff_base: int = 2,
+        backoff_max: Optional[int] = 60,
+        backoff_jitter: float = 0.1,
+        priority_fn: Optional[Callable[[Type], int]] = None,
+        metrics_client: Any = None,
+        debug: bool = False,
+        logger: Optional[logging.Logger] = None,
+        artifact_class_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        self.wrapped_classes = wrapped_classes
+        self.max_workers = max_workers
+        self.retry_attempts = retry_attempts
+        self.update_timeout_seconds = update_timeout_seconds
+        self.backoff_base = backoff_base
+        self.backoff_max = backoff_max
+        self.backoff_jitter = backoff_jitter
+        self.priority_fn = priority_fn
+        self.metrics_client = metrics_client
 
-        today = datetime.datetime.today()
-        self.parquet_start_date = kwargs.get(
-            'parquet_start_date',
-            datetime.date(today.year, 1, 1).strftime('%Y-%m-%d')
-        )
-        self.parquet_end_date = kwargs.get(
-            'parquet_end_date',
-            today.strftime('%Y-%m-%d')
+        self.debug = debug
+        self.logger = logger or Logger.default_logger(
+            logger_name=self.__class__.__name__,
+            log_level=Logger.DEBUG if debug else Logger.INFO
         )
 
-        # track pending/completed/failed artifacts
-        self.pending = set()
-        self.completed = set()
-        self.failed = set()
-
-        # concurrency primitives
-        self.locks = {}
-        self.locks_lock = asyncio.Lock()
-        self.worker_heartbeat = defaultdict(float)
-        self.workers_lock = asyncio.Lock()
-
-        # dynamic scaling config
-        self.min_workers = kwargs.get('min_workers', 1)
-        self.max_workers = kwargs.get('max_workers', 3)
-        self.memory_per_worker_gb = kwargs.get('memory_per_worker_gb', 1)
-        self.monitor_interval = kwargs.get('monitor_interval', 10)
-        self.retry_attempts = kwargs.get('retry_attempts', 3)
-        self.update_timeout_seconds = kwargs.get('update_timeout_seconds', 600)
-        self.lock_acquire_timeout_seconds = kwargs.get('lock_acquire_timeout_seconds', 10)
-
-    async def get_lock_for_artifact(self, artifact):
-        key = artifact.__class__.__name__
-        async with self.locks_lock:
-            if key not in self.locks:
-                self.locks[key] = asyncio.Lock()
-            return self.locks[key]
-
-    def get_artifacts(self, data_type):
+        # Default artifact init kwargs
+        today = datetime.datetime.today() + datetime.timedelta(days=1)
+        default_kwargs = {
+            'parquet_start_date': today.strftime('%Y-%m-%d'),
+            'parquet_end_date':   today.strftime('%Y-%m-%d'),
+            'logger':             self.logger,
+            'debug':              self.debug,
+        }
+        self.artifact_class_kwargs = artifact_class_kwargs or default_kwargs.copy()
+
+        # State
+        self.completion_times: Dict[str, float] = {}
+        self.failed: List[str] = []
+        self.original_classes: List[Type] = []
+
+    def get_artifact_classes(self, data_type: str) -> List[Type]:
+        """
+        Retrieve artifact classes by data type.
+        """
+        self.logger.info(f"Fetching artifact classes for '{data_type}'")
         if data_type not in self.wrapped_classes:
             raise ValueError(f"Unsupported data type: {data_type}")
-        artifacts = [cls(
-            parquet_start_date=self.parquet_start_date,
-            parquet_end_date=self.parquet_end_date,
-            logger=self.logger,
-            debug=self.debug
-        ) for cls in self.wrapped_classes[data_type]]
-        # seed pending set and clear others
-        self.pending = set(artifacts)
-        self.completed.clear()
-        self.failed.clear()
-        return artifacts
-
-    def estimate_complexity(self, artifact):
-        try:
-            return artifact.get_size_estimate()
-        except Exception:
-            return 1
-
-    def prioritize_tasks(self, artifacts):
-        queue = asyncio.PriorityQueue()
-        for art in artifacts:
-            queue.put_nowait(PrioritizedItem(self.estimate_complexity(art), art))
-        return queue
-
-    async def resource_monitor(self, queue, workers):
-        while not queue.empty():
+        classes = self.wrapped_classes[data_type]
+        self.logger.info(f"Found {len(classes)} artifact classes for '{data_type}'")
+        return classes
+
+    def estimate_priority(self, artifact_cls: Type) -> int:
+        """
+        Determine task priority for ordering. Lower values run first.
+        """
+        name = artifact_cls.__name__
+        if self.priority_fn:
             try:
-                avail = psutil.virtual_memory().available
-                max_by_mem = avail // (self.memory_per_worker_gb * 2**30)
-                optimal = max(self.min_workers,
-                              min(psutil.cpu_count(), max_by_mem, self.max_workers))
-                async with self.workers_lock:
-                    current = len(workers)
-                    if optimal > current:
-                        for _ in range(optimal - current):
-                            wid = len(workers)
-                            workers.append(asyncio.create_task(self.worker(queue, wid)))
-                            self.logger.info(f"Added worker {wid}")
-                    elif optimal < current:
-                        for _ in range(current - optimal):
-                            w = workers.pop()
-                            w.cancel()
-                            self.logger.info("Removed a worker")
-                await asyncio.sleep(self.monitor_interval)
-            except asyncio.CancelledError:
-                break
+                pr = self.priority_fn(artifact_cls)
+                self.logger.debug(f"priority_fn for {name}: {pr}")
+                return pr
             except Exception as e:
-                self.logger.error(f"Monitor error: {e}")
-                await asyncio.sleep(self.monitor_interval)
-
-    @asynccontextmanager
-    async def artifact_lock(self, artifact):
-        lock = await self.get_lock_for_artifact(artifact)
+                self.logger.warning(f"priority_fn error for {name}: {e}")
         try:
-            await asyncio.wait_for(lock.acquire(), timeout=self.lock_acquire_timeout_seconds)
-            yield
-        finally:
-            if lock.locked():
-                lock.release()
+            fs = self.artifact_class_kwargs.get('fs')
+            path = self.artifact_class_kwargs.get('parquet_storage_path')
+            pr=1
+            if hasattr(artifact_cls, 'get_size_estimate'):
+                pr = artifact_cls.get_size_estimate(fs, path)
+            self.logger.debug(f"Estimated priority for {name}: {pr}")
+            return pr
+        except Exception:
+            return 1
 
-    async def async_update_artifact(self, artifact, **kwargs):
-        for attempt in range(1, self.retry_attempts + 1):
-            lock = await self.get_lock_for_artifact(artifact)
+    async def _bounded_update(self, artifact_cls: Type, sem: asyncio.Semaphore, **update_kwargs) -> None:
+        """
+        Wrap update_artifact in a semaphore slot to limit concurrency.
+        """
+        async with sem:
+            name = artifact_cls.__name__
+            start = asyncio.get_event_loop().time()
+            self.logger.info(f"Starting update for {name}")
             try:
-                await asyncio.wait_for(lock.acquire(), timeout=self.lock_acquire_timeout_seconds)
-                try:
-                    self.logger.info(f"Updating {artifact.__class__.__name__} (attempt {attempt})")
-                    await asyncio.wait_for(
-                        asyncio.to_thread(artifact.update_parquet, **kwargs),
-                        timeout=self.update_timeout_seconds
-                    )
-                    # mark success
-                    async with self.workers_lock:
-                        self.pending.discard(artifact)
-                        self.completed.add(artifact)
-                    self.logger.info(
-                        f"✅ {artifact.__class__.__name__} done — "
-                        f"{len(self.completed)}/{len(self.completed) + len(self.pending) + len(self.failed)} completed, "
-                        f"{len(self.failed)} failed"
-                    )
-                    return
-                finally:
-                    if lock.locked():
-                        lock.release()
-            except asyncio.TimeoutError:
-                self.logger.warning(f"Timeout on {artifact.__class__.__name__}, attempt {attempt}")
-            except Exception as e:
-                self.logger.error(f"Error on {artifact}: {e}")
-            finally:
-                if lock.locked():
-                    lock.release()
-            await asyncio.sleep(2 ** (attempt - 1))
+                for attempt in range(1, self.retry_attempts + 1):
+                    try:
+                        artifact = await asyncio.to_thread(
+                            artifact_cls, **self.artifact_class_kwargs
+                        )
+                        await asyncio.wait_for(
+                            asyncio.to_thread(
+                                artifact.update_parquet, **update_kwargs
+                            ),
+                            timeout=self.update_timeout_seconds
+                        )
+                        duration = asyncio.get_event_loop().time() - start
+                        self.completion_times[name] = duration
+                        self.logger.info(f"✅ {name} updated in {duration:.2f}s (attempt {attempt})")
+                        if self.metrics_client:
+                            self.metrics_client.increment('task_succeeded')
+                        return
+                    except asyncio.TimeoutError:
+                        self.logger.warning(f"Timeout on {name}, attempt {attempt}")
+                    except Exception as e:
+                        self.logger.error(f"Error on {name} attempt {attempt}: {e}")
+
+                    delay = min(self.backoff_base ** (attempt - 1), self.backoff_max)
+                    delay *= 1 + random.uniform(0, self.backoff_jitter)
+                    self.logger.info(f"Sleeping {delay:.1f}s before retrying {name}")
+                    await asyncio.sleep(delay)
 
-        # all retries exhausted -> mark failure
-        async with self.workers_lock:
-            self.pending.discard(artifact)
-            self.failed.add(artifact)
-        self.logger.error(f"✖️  Permanently failed {artifact.__class__.__name__}")
-
-    async def worker(self, queue, worker_id, **kwargs):
-        while True:
-            try:
-                item = await queue.get()
-                art = item.artifact
-                self.worker_heartbeat[worker_id] = time.time()
-                await self.async_update_artifact(art, **kwargs)
             except asyncio.CancelledError:
-                self.logger.info(f"Worker {worker_id} stopped")
-                break
-            finally:
-                queue.task_done()
-
-    def calculate_initial_workers(self, count: int) -> int:
-        avail = psutil.virtual_memory().available
-        max_by_mem = avail // (self.memory_per_worker_gb * 2**30)
-        return max(self.min_workers,
-                   min(psutil.cpu_count(), max_by_mem, count, self.max_workers))
-
-    async def update_data(self, data_type, **kwargs):
-        self.logger.info(f"Starting update for {data_type}")
-        artifacts = self.get_artifacts(data_type)
-        queue = self.prioritize_tasks(artifacts)
-        init = self.calculate_initial_workers(len(artifacts))
-        tasks = [asyncio.create_task(self.worker(queue, i, **kwargs)) for i in range(init)]
-        monitor = asyncio.create_task(self.resource_monitor(queue, tasks))
-        await queue.join()
-        monitor.cancel()
-        for t in tasks:
-            t.cancel()
-        await asyncio.gather(*tasks, return_exceptions=True)
-        self.logger.info(self.format_results_table())
-        self.logger.info("All artifacts processed.")
-
-    def format_results_table(self):
-        results = self.get_update_status()
-        headers = ["Metric", "Value"]
-        rows = [
-            ["Total", results['total']],
-            ["Completed", results['completed']],
-            ["Pending", results['pending']],
-            ["Failed", results['failed']],
-            ["Pending Items", len(results['pending_items'])],
-            ["Failed Items", len(results['failed_items'])]
-        ]
-
-        # Find max lengths for alignment
-        max_metric = max(len(str(row[0])) for row in rows)
-        max_value = max(len(str(row[1])) for row in rows)
+                self.logger.warning(f"{name} update cancelled")
+                raise
+
+            # permanent failure
+            self.logger.error(f"✖️  {name} permanently failed after {self.retry_attempts} attempts")
+            if self.metrics_client:
+                self.metrics_client.increment('task_failed')
+            self.failed.append(name)
+
+    async def update_data(self, data_type: str, **kwargs: Any) -> None:
+        """
+        Entry point to update all artifacts of a given type concurrently.
+        """
+        self.logger.info(f"Starting update_data for '{data_type}' with kwargs={kwargs}")
+
+        # RESET STATE
+        self.completion_times.clear()
+        self.failed.clear()
+        self.original_classes = self.get_artifact_classes(data_type)
 
-        format_str = "{:<%d}  {:>%d}" % (max_metric, max_value)
+        # NON-DESTRUCTIVE SORTING
+        ordered = sorted(self.original_classes, key=self.estimate_priority)
 
-        table = [
-            "\n",
-            format_str.format(*headers),
-            "-" * (max_metric + max_value + 2)
+        sem = asyncio.Semaphore(self.max_workers)
+        tasks = [
+            asyncio.create_task(self._bounded_update(cls, sem, **kwargs))
+            for cls in ordered
         ]
 
-        for row in rows:
-            table.append(format_str.format(row[0], row[1]))
-
-        return "\n".join(table)
+        try:
+            for coro in asyncio.as_completed(tasks):
+                await coro
+        except asyncio.CancelledError:
+            self.logger.warning("update_data was cancelled—aborting remaining retries")
+            for t in tasks:
+                t.cancel()
+            raise
+        finally:
+            total = len(self.original_classes)
+            completed = len(self.completion_times)
+            failed = len(self.failed)
+            self.logger.info(f"All artifacts processed: total={total}, completed={completed}, failed={failed}")
+
+    def get_update_status(self) -> Dict[str, Any]:
+        """
+        Returns summary status including completion times.
+        """
+        total = len(self.original_classes)
+        completed = set(self.completion_times.keys())
+        failed = set(self.failed)
+        pending = {cls.__name__ for cls in self.original_classes} - completed - failed
 
-    def get_update_status(self):
-        total = len(self.pending) + len(self.completed) + len(self.failed)
         return {
-            "total": total,
-            "completed": len(self.completed),
-            "pending": len(self.pending),
-            "failed": len(self.failed),
-            "pending_items": [a.__class__.__name__ for a in self.pending],
-            "failed_items": [a.__class__.__name__ for a in self.failed]
+            'total': total,
+            'completed': list(completed),
+            'failed':    list(failed),
+            'pending':   list(pending),
+            'completion_times': self.completion_times,
         }
 
-# Top‑level driver
-# environment = None  # fill this in with your wrapped_classes dict
-#
-# async def main():
-#     wrapper = ArtifactUpdaterMultiWrapper(
-#         wrapped_classes=environment,
-#         debug=True
-#     )
-#     loop = asyncio.get_running_loop()
-#     for sig in (signal.SIGINT, signal.SIGTERM):
-#         loop.add_signal_handler(sig, lambda: asyncio.create_task(wrapper.shutdown()))
-#     await wrapper.update_data("your_data_type")
-#
-# if __name__ == "__main__":
-#     asyncio.run(main())
-
+    @staticmethod
+    def format_status_table(status: Dict[str, Any]) -> str:
+        """
+        Formats the status dict into a readable table.
+        """
+        lines = [
+            f"Total: {status['total']}",
+            f"Completed: {len(status['completed'])}  {status['completed']}",
+            f"Failed:    {len(status['failed'])}  {status['failed']}",
+            f"Pending:   {len(status['pending'])}  {status['pending']}",
+            "",
+            "Per-artifact timings:"
+        ]
+        for name, dur in status['completion_times'].items():
+            lines.append(f"  {name}: {dur:.2f}s")
+        return "\n".join(lines)

sibi_dst/df_helper/_df_helper.py

@@ -46,8 +46,7 @@ class DfHelper:
 
     :ivar df: The DataFrame currently being processed or loaded.
     :type df: Union[dd.DataFrame, pd.DataFrame]
-    :ivar backend_django: Configuration for interacting with Django database backends.
-    :type backend_connection: Optional[DjangoConnectionConfig]
+    :type backend_connection: Optional[DjangoConnectionConfig | SqlAlchemyConnectionConfig]
     :ivar _backend_query: Internal configuration for query handling.
     :type _backend_query: Optional[QueryConfig]
     :ivar _backend_params: Internal parameters configuration for DataFrame handling.
@@ -81,9 +80,10 @@ class DfHelper:
         self.debug = kwargs.setdefault("debug", False)
         self.logger = kwargs.get("logger", Logger.default_logger(logger_name=self.__class__.__name__))
         # Configure logger level
-        self.logger.set_level(logging.DEBUG if self.debug else logging.INFO)
+        self.logger.set_level(Logger.DEBUG if self.debug else Logger.INFO)
         self.logger.debug("Logger initialized in DEBUG mode.")
         self.parquet_storage_path = kwargs.setdefault("parquet_storage_path", None)
+        self.parquet_filename = kwargs.setdefault("parquet_filename", None)
         self.dt_field = kwargs.setdefault("dt_field", None)
         self.as_pandas = kwargs.setdefault("as_pandas", False)
         self.filesystem = kwargs.pop('filesystem', 'file')
@@ -429,7 +429,11 @@ class DfHelper:
                          the instance's attribute for storage path.
         :return: None
         """
+        if self.df.map_partitions(len).compute().sum() == 0:
+            self.logger.debug("Cannot save to parquet since DataFrame is empty")
+            return
         fs = kwargs.pop('fs', self.fs)
+        parquet_filename = parquet_filename or self.parquet_filename
         parquet_storage_path = kwargs.pop('parquet_storage_path', self.parquet_storage_path)
         ps = ParquetSaver(df_result=self.df, parquet_storage_path=parquet_storage_path, logger=self.logger, fs=fs)
         ps.save_to_parquet(parquet_filename)