sibi-dst 0.3.56__tar.gz → 0.3.58__tar.gz

{sibi_dst-0.3.56 → sibi_dst-0.3.58}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sibi-dst
-Version: 0.3.56
+Version: 0.3.58
 Summary: Data Science Toolkit
 Author: Luis Valverde
 Author-email: lvalverdeb@gmail.com

{sibi_dst-0.3.56 → sibi_dst-0.3.58}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sibi-dst"
-version = "0.3.56"
+version = "0.3.58"
 description = "Data Science Toolkit"
 authors = ["Luis Valverde <lvalverdeb@gmail.com>"]
 readme = "README.md"

sibi_dst-0.3.58/sibi_dst/df_helper/_artifact_updater_multi_wrapper.py

@@ -0,0 +1,261 @@
+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
+
+    def __lt__(self, other):
+        return self.priority < other.priority
+
+    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__)
+        )
+        self.logger.set_level(logging.DEBUG if debug else logging.INFO)
+
+        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')
+        )
+
+        # 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):
+        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():
+            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
+            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)
+        try:
+            await asyncio.wait_for(lock.acquire(), timeout=self.lock_acquire_timeout_seconds)
+            yield
+        finally:
+            if lock.locked():
+                lock.release()
+
+    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)
+            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))
+
+        # 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)
+
+        format_str = "{:<%d}  {:>%d}" % (max_metric, max_value)
+
+        table = [
+            "\n",
+            format_str.format(*headers),
+            "-" * (max_metric + max_value + 2)
+        ]
+
+        for row in rows:
+            table.append(format_str.format(row[0], row[1]))
+
+        return "\n".join(table)
+
+    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]
+        }
+
+# 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())
+

{sibi_dst-0.3.56 → sibi_dst-0.3.58}/sibi_dst/df_helper/_df_helper.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import asyncio
 import datetime
 import logging
@@ -6,10 +8,10 @@ from typing import Any, Dict, TypeVar
 from typing import Union, Optional
 
 import dask.dataframe as dd
-from dask import delayed, compute
+import fsspec
 import pandas as pd
+from dask import delayed, compute
 from pydantic import BaseModel
-import fsspec
 
 from sibi_dst.df_helper.core import QueryConfig, ParamsConfig, FilterHandler
 from sibi_dst.utils import Logger
@@ -45,7 +47,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_django: Optional[DjangoConnectionConfig]
+    :type backend_connection: Optional[DjangoConnectionConfig]
     :ivar _backend_query: Internal configuration for query handling.
     :type _backend_query: Optional[QueryConfig]
     :ivar _backend_params: Internal parameters configuration for DataFrame handling.
@@ -54,8 +56,6 @@ class DfHelper:
     :type backend_parquet: Optional[ParquetConfig]
     :ivar backend_http: Configuration for interacting with HTTP-based backends.
     :type backend_http: Optional[HttpConfig]
-    :ivar backend_sqlalchemy: Configuration for interacting with SQLAlchemy-based databases.
-    :type backend_sqlalchemy: Optional[SqlAlchemyConnectionConfig]
     :ivar parquet_filename: The filename for a Parquet file, if applicable.
     :type parquet_filename: str
     :ivar logger: Logger instance used for debugging and information logging.
@@ -64,12 +64,11 @@ class DfHelper:
     :type default_config: Dict
     """
     df: Union[dd.DataFrame, pd.DataFrame] = None
-    backend_django: Optional[DjangoConnectionConfig] = None
+    backend_db_connection: Optional[Union[DjangoConnectionConfig | SqlAlchemyConnectionConfig]] = None
     _backend_query: Optional[QueryConfig] = None
     _backend_params: Optional[ParamsConfig] = None
     backend_parquet: Optional[ParquetConfig] = None
     backend_http: Optional[HttpConfig] = None
-    backend_sqlalchemy: Optional[SqlAlchemyConnectionConfig] = None
     parquet_filename: str = None
     logger: Logger
     default_config: Dict = None
@@ -91,7 +90,7 @@ class DfHelper:
         self.filesystem_options = kwargs.pop('filesystem_options', {})
         kwargs.setdefault("live", True)
         kwargs.setdefault("logger", self.logger)
-        self.fs =kwargs.setdefault("fs", fsspec.filesystem('file'))
+        self.fs = kwargs.setdefault("fs", fsspec.filesystem('file'))
         self.__post_init(**kwargs)
 
     def __str__(self):
@@ -100,6 +99,34 @@ class DfHelper:
     def __call__(self, **options):
         return self.load(**options)
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.__cleanup()
+        return False
+
+    def __cleanup(self):
+        """
+        Clean up resources when exiting the context manager.
+        This method is called when the context manager exits.
+        """
+
+        if self.backend_db_connection:
+            if getattr(self.backend_db_connection, "dispose_idle_connections", None):
+                self.backend_db_connection.dispose_idle_connections()
+            if getattr(self.backend_db_connection, "close", None):
+                self.backend_db_connection.close()
+
+        self.backend_db_connection = None
+
+        if self.backend_parquet:
+            self.backend_parquet = None
+        if self.backend_http:
+            self.backend_http = None
+        self._backend_query = None
+        self._backend_params = None
+
     def __post_init(self, **kwargs):
         """
         Initializes backend-specific configurations based on the provided backend type and other
@@ -111,20 +138,19 @@ class DfHelper:
                        Additional parameters for specific backend types are extracted here.
         :return: None
         """
-        self.logger.debug(f"backend used: {self.backend}")
-        self.logger.debug(f"kwargs passed to backend plugins: {kwargs}")
+        # self.logger.debug(f"backend used: {self.backend}")
+        # self.logger.debug(f"kwargs passed to backend plugins: {kwargs}")
         self._backend_query = self.__get_config(QueryConfig, kwargs)
         self._backend_params = self.__get_config(ParamsConfig, kwargs)
         if self.backend == 'django_db':
-            self.backend_django = self.__get_config(DjangoConnectionConfig, kwargs)
+            self.backend_db_connection = self.__get_config(DjangoConnectionConfig, kwargs)
         elif self.backend == 'parquet':
             self.parquet_filename = kwargs.setdefault("parquet_filename", None)
             self.backend_parquet = ParquetConfig(**kwargs)
         elif self.backend == 'http':
             self.backend_http = HttpConfig(**kwargs)
         elif self.backend == 'sqlalchemy':
-            self.backend_sqlalchemy = self.__get_config(SqlAlchemyConnectionConfig, kwargs)
-
+            self.backend_db_connection = self.__get_config(SqlAlchemyConnectionConfig, kwargs)
 
     def __get_config(self, model: [T], kwargs: Dict[str, Any]) -> Union[T]:
         """
@@ -134,11 +160,13 @@ class DfHelper:
         :param kwargs: The dictionary of keyword arguments.
         :return: The initialized Pydantic model instance.
         """
+        kwargs.setdefault("debug", self.debug)
+        kwargs.setdefault("logger", self.logger)
         # Extract keys that the model can accept
         recognized_keys = set(model.__annotations__.keys())
         self.logger.debug(f"recognized keys: {recognized_keys}")
         model_kwargs = {k: kwargs.pop(k) for k in list(kwargs.keys()) if k in recognized_keys}
-        self.logger.debug(f"model_kwargs: {model_kwargs}")
+        # self.logger.debug(f"model_kwargs: {model_kwargs}")
         return model(**model_kwargs)
 
     def load_parallel(self, **options):
@@ -171,10 +199,10 @@ class DfHelper:
             `as_pandas` is set to True, or kept in its native backend format otherwise.
         """
         # this will be the universal method to load data from a df irrespective of the backend
-        df = self.__load(**options)
+        self.df = self.__load(**options)
         if self.as_pandas:
-            return df.compute()
-        return df
+            return self.df.compute()
+        return self.df
 
     def __load(self, **options):
         """
@@ -196,7 +224,7 @@ class DfHelper:
         """
         if self.backend == 'django_db':
             self._backend_params.parse_params(options)
-            return self.__load_from_db(**options)
+            return self.__load_from_django_db(**options)
         elif self.backend == 'sqlalchemy':
             self._backend_params.parse_params(options)
             return self.__load_from_sqlalchemy(**options)
@@ -227,7 +255,7 @@ class DfHelper:
         try:
             options.setdefault("debug", self.debug)
             db_loader = SqlAlchemyLoadFromDb(
-                self.backend_sqlalchemy,
+                self.backend_db_connection,
                 self._backend_query,
                 self._backend_params,
                 self.logger,
@@ -236,6 +264,7 @@ class DfHelper:
             self.df = db_loader.build_and_load()
             self.__process_loaded_data()
             self.__post_process_df()
+            self.backend_db_connection.close()
             self.logger.debug("Data successfully loaded from sqlalchemy database.")
         except Exception as e:
             self.logger.debug(f"Failed to load data from sqlalchemy database: {e}: options: {options}")
@@ -243,7 +272,7 @@ class DfHelper:
 
         return self.df
 
-    def __load_from_db(self, **options) -> Union[pd.DataFrame, dd.DataFrame]:
+    def __load_from_django_db(self, **options) -> Union[pd.DataFrame, dd.DataFrame]:
         """
         Loads data from a Django database using a specific backend query mechanism. Processes the loaded data
         and applies further post-processing before returning the dataframe. If the operation fails, an
@@ -258,7 +287,7 @@ class DfHelper:
         try:
             options.setdefault("debug", self.debug)
             db_loader = DjangoLoadFromDb(
-                self.backend_django,
+                self.backend_db_connection,
                 self._backend_query,
                 self._backend_params,
                 self.logger,
@@ -307,6 +336,7 @@ class DfHelper:
         :raises ValueError: If the lengths of `fieldnames` and `column_names` do not match,
             or if the specified `index_col` is not found in the DataFrame.
         """
+        self.logger.debug("Post-processing DataFrame.")
         df_params = self._backend_params.df_params
         fieldnames = df_params.get("fieldnames", None)
         index_col = df_params.get("index_col", None)
@@ -357,16 +387,16 @@ class DfHelper:
 
         :return: None
         """
-        self.logger.debug(f"Type of self.df: {type(self.df)}")
+        self.logger.debug(f"Processing loaded data...")
         if self.df.map_partitions(len).compute().sum() > 0:
             field_map = self._backend_params.field_map or {}
-            if isinstance(field_map, dict):
+            if isinstance(field_map, dict) and field_map != {}:
                 rename_mapping = {k: v for k, v in field_map.items() if k in self.df.columns}
                 missing_columns = [k for k in field_map.keys() if k not in self.df.columns]
 
                 if missing_columns:
                     self.logger.warning(
-                        f"The following columns in field_map are not in the DataFrame: {missing_columns}")
+                        f"The following columns in field_map are not in the DataFrame: {missing_columns}, field map: {field_map}")
 
                 def rename_columns(df, mapping):
                     return df.rename(columns=mapping)
@@ -376,6 +406,8 @@ class DfHelper:
                     self.df = self.df.map_partitions(rename_columns, mapping=rename_mapping)
 
             self.logger.debug("Processing of loaded data completed.")
+        else:
+            self.logger.debug("DataFrame is empty, skipping processing.")
 
     def save_to_parquet(self, parquet_filename: Optional[str] = None, **kwargs):
         """
@@ -536,14 +568,14 @@ class DfHelper:
 
         # Common logic for Django and SQLAlchemy
         if self.backend == 'django_db':
-            model_fields = {field.name: field for field in self.backend_django.model._meta.get_fields()}
+            model_fields = {field.name: field for field in self.backend_db_connection.model._meta.get_fields()}
             if mapped_field not in model_fields:
                 raise ValueError(f"Field '{dt_field}' does not exist in the Django model.")
             field_type = type(model_fields[mapped_field]).__name__
             is_date_field = field_type == 'DateField'
             is_datetime_field = field_type == 'DateTimeField'
         elif self.backend == 'sqlalchemy':
-            model = self.backend_sqlalchemy.model
+            model = self.backend_db_connection.model
             fields = [column.name for column in model.__table__.columns]
             if mapped_field not in fields:
                 raise ValueError(f"Field '{dt_field}' does not exist in the SQLAlchemy model.")

{sibi_dst-0.3.56 → sibi_dst-0.3.58}/sibi_dst/df_helper/_parquet_artifact.py

@@ -1,5 +1,6 @@
 import datetime
 import logging
+import threading
 from typing import Optional, Any, Dict
 
 import dask.dataframe as dd
@@ -78,6 +79,7 @@ class ParquetArtifact(DfHelper):
                             `parquet_filename`, `parquet_start_date`,
                             or `parquet_end_date`) are missing or not set properly.
         """
+        self._lock = threading.Lock()
         self.config = {
             **self.DEFAULT_CONFIG,
             **kwargs,
@@ -119,21 +121,36 @@ class ParquetArtifact(DfHelper):
         super().__init__(**self.config)
 
     def load(self, **kwargs):
-        self.df = super().load(**kwargs)
+        with self._lock:
+            self.df = super().load(**kwargs)
         return self.df
 
     def generate_parquet(self, **kwargs) -> None:
         """
         Generate a Parquet file using the configured DataWrapper class.
         """
-        params = self._prepare_params(kwargs)
-        dw = DataWrapper(self.data_wrapper_class, **params)
-        dw.process()
+        with self._lock:
+            params = self._prepare_params(kwargs)
+            dw = DataWrapper(self.data_wrapper_class, **params)
+            dw.process()
+
+    def __enter__(self):
+        if getattr(self, "_entered", False):
+            return self
+        self._entered = True
+        self.ensure_directory_exists(self.parquet_storage_path)
+        return self
 
     def __exit__(self, exc_type, exc_value, traceback):
-        # Ensure resources are cleaned up
-        if self.fs:
-            self.fs.close()
+        try:
+            if getattr(self, "_entered", False) and self.fs:
+                self.fs.close()
+        except Exception as e:
+            self.logger.warning(f"Error closing filesystem: {e}")
+        finally:
+            self._entered = False
+        # return False so exceptions aren’t suppressed
+        return False
 
     def update_parquet(self, period: str = 'today', **kwargs) -> None:
         """Update the Parquet file with data from a specific period."""
@@ -223,7 +240,8 @@ class ParquetArtifact(DfHelper):
 
     def ensure_directory_exists(self, path: str) -> None:
         """Ensure the directory exists in the specified filesystem."""
-        try:
-            self.fs.makedirs(path, exist_ok=True)
-        except Exception as e:
-            raise ValueError(f"Error creating directory {path} in filesystem {self.filesystem_type}: {e}")
+        with self._lock:
+            try:
+                self.fs.makedirs(path, exist_ok=True)
+            except Exception as e:
+                raise ValueError(f"Error creating directory {path} in filesystem {self.filesystem_type}: {e}")